from typing import Optional, List
from langchain.llms.base import LLM
from litellm import completion

class LiteLLMWrapper(LLM):
    model_name: str = "gpt-3.5-turbo"

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

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        response = completion(
            model=self.model_name,
            messages=[{"role": "user", "content": prompt}],
            stop=stop
        )
        return response.choices[0].message.content

# Usage in LangChain chains
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

llm = LiteLLMWrapper(model_name="claude-3-sonnet-20240229")
prompt = PromptTemplate(
    input_variables=["topic"],
    template="Write a brief summary about {topic}"
)

chain = LLMChain(llm=llm, prompt=prompt)
result = chain.run(topic="artificial intelligence")

        __
        __