| from msal import ConfidentialClientApplication | |
| from langchain_openai import AzureOpenAIEmbeddings, AzureChatOpenAI | |
| from langchain_groq import ChatGroq | |
| from langchain.vectorstores.azuresearch import AzureSearch | |
| import os | |
| class LLM: | |
| def __init__(self, llm): | |
| self.llm = llm | |
| self.callbacks = [] | |
| def stream(self, prompt, prompt_arguments): | |
| self.llm.streaming = True | |
| streamed_content = self.llm.stream(prompt.format_messages(**prompt_arguments)) | |
| output = "" | |
| for op in streamed_content: | |
| output += op.content | |
| yield output | |
| def get_prediction(self, prompt, prompt_arguments): | |
| self.llm.callbacks = self.callbacks | |
| return self.llm.predict_messages( | |
| prompt.format_messages(**prompt_arguments) | |
| ).content | |
| async def get_aprediction(self, prompt, prompt_arguments): | |
| self.llm.callbacks = self.callbacks | |
| prediction = await self.llm.apredict_messages( | |
| prompt.format_messages(**prompt_arguments) | |
| ) | |
| return prediction | |
| async def get_apredictions(self, prompts, prompts_arguments): | |
| self.llm.callbacks = self.callbacks | |
| predictions = [] | |
| for prompt_, prompt_args_ in zip(prompts.keys(), prompts_arguments): | |
| prediction = await self.llm.apredict_messages( | |
| prompts[prompt_].format_messages(**prompt_args_) | |
| ) | |
| predictions.append(prediction.content) | |
| return predictions | |
| def get_llm_api(groq_model_name): | |
| if os.getenv("EKI_OPENAI_LLM_DEPLOYMENT_NAME"): | |
| print("Using Azure OpenAI API") | |
| return LLM( | |
| AzureChatOpenAI( | |
| deployment_name=os.getenv("EKI_OPENAI_LLM_DEPLOYMENT_NAME"), | |
| openai_api_key=os.getenv("EKI_OPENAI_API_KEY"), | |
| azure_endpoint=os.getenv("EKI_OPENAI_LLM_API_ENDPOINT"), | |
| openai_api_version=os.getenv("EKI_OPENAI_API_VERSION"), | |
| streaming=True, | |
| temperature=0, | |
| max_tokens=2048, | |
| stop=["<|im_end|>"], | |
| ) | |
| ) | |
| else: | |
| print("Using GROQ API") | |
| return LLM( | |
| ChatGroq( | |
| model=groq_model_name, | |
| temperature=0, | |
| max_tokens=2048, | |
| ) | |
| ) | |
| def get_vectorstore_api(index_name): | |
| aoai_embeddings = AzureOpenAIEmbeddings( | |
| model="text-embedding-ada-002", | |
| azure_deployment=os.getenv("EKI_OPENAI_EMB_DEPLOYMENT_NAME"), | |
| api_key=os.getenv("EKI_OPENAI_API_KEY"), | |
| azure_endpoint=os.environ["EKI_OPENAI_EMB_API_ENDPOINT"], | |
| openai_api_version=os.getenv("EKI_OPENAI_API_VERSION"), | |
| ) | |
| vector_store: AzureSearch = AzureSearch( | |
| azure_search_endpoint=os.getenv("EKI_VECTOR_STORE_ADDRESS"), | |
| azure_search_key=os.getenv("EKI_VECTOR_STORE_PASSWORD"), | |
| index_name=index_name, | |
| embedding_function=aoai_embeddings.embed_query, | |
| ) | |
| return vector_store | |