# src/custom_llm.py

from langchain.llms.base import LLM
from typing import Optional, List, Any
from llm_setup import HuggingFaceLLM

class HuggingFaceLLMWrapper(LLM):
    def __init__(self, model_name: str, device: str = 'cuda'):
        self.llm = HuggingFaceLLM(model_name=model_name, device=device)

    @property
    def _llm_type(self) -> str:
        return "custom_huggingface"

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        response = self.llm.generate(prompt)
        if stop:
            for s in stop:
                response = response.split(s)[0]
        return response

    def _identifying_params(self) -> dict:
        return {"model_name": self.llm.model.name_or_path, "device": self.llm.device}