File size: 2,552 Bytes
9023ac5 d43f1a7 9023ac5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | 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
## --- LLMCall ---
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 |