| import os |
| from src.controller.customlogger import logging |
| from llama_index.llms.google_genai import GoogleGenAI |
| from llama_index.llms.openai import OpenAI |
| from llama_index.llms.huggingface import HuggingFaceLLM |
| from mistralai import Mistral |
| from llama_index.core.llms import ChatMessage |
| from dotenv import load_dotenv |
| load_dotenv() |
|
|
|
|
| class MistralWrapper: |
| def __init__(self): |
| self.client = Mistral(api_key=os.environ["MISTRAL_API_KEY"]) |
| self.model = "mistral-large-latest" |
|
|
| def complete(self, messages, **kwargs): |
| chat_response = self.client.chat.complete( |
| model= self.model, |
| messages = [{"role":"user", "content":messages}]) |
| return chat_response.choices[0].message.content |
|
|
| |
| class LLMCall: |
| def __init__(self, llm_type="OpenAi", config=None): |
| self.llm_type = llm_type |
| self.config = config or {} |
| self.api_key_google = os.getenv("GOOGLE_API_KEY") |
| self.client = self.get_llm() |
|
|
| def get_llm(self): |
| if self.llm_type == "OpenAi": |
| logging.info("Initializing OpenAI LLM") |
| return self._get_openai_llm() |
| elif self.llm_type == "Google": |
| logging.info("Initializing Google Gemini LLM") |
| return self._get_google_llm() |
| elif self.llm_type == "HuggingFace": |
| logging.info("Initializing HuggingFace LLM") |
| return self._get_huggingface_llm() |
| elif self.llm_type == "Mistral": |
| logging.info("Initializing HuggingFace LLM") |
| return self._get_mistral() |
| else: |
| raise ValueError(f"Unsupported LLM type: {self.llm_type}") |
|
|
| def _get_openai_llm(self): |
| return OpenAI( |
| model=self.config.get("model", "gpt-3.5-turbo-0613"), |
| api_key=self.config.get("api_key"), |
| temperature=self.config.get("temperature", 0.7), |
| max_tokens=self.config.get("max_tokens", 1024) |
| ) |
|
|
| def _get_google_llm(self): |
| return GoogleGenAI( |
| model="models/gemini-2.0-flash", |
| ) |
| |
| def _get_mistral(self): |
| return MistralWrapper() |
|
|
|
|
| def _get_huggingface_llm(self): |
| return HuggingFaceLLM( |
| model_name="HuggingFaceH4/zephyr-7b-beta", |
| tokenizer_name="HuggingFaceH4/zephyr-7b-beta", |
| context_window=self.config.get("context_window", 2048), |
| max_new_tokens=self.config.get("max_new_tokens", 256) |
| ) |
| |
| |
| def get_client(self): |
| return self.client |
