Spaces:
Running
Running
| from langchain_core.output_parsers import StrOutputParser | |
| from langchain_core.prompts import ( | |
| ChatPromptTemplate, | |
| MessagesPlaceholder, | |
| ) | |
| from langchain.chains import create_history_aware_retriever, create_retrieval_chain | |
| from langchain.chains.combine_documents import create_stuff_documents_chain | |
| from langchain_core.runnables import RunnablePassthrough | |
| from langchain_core.vectorstores import VectorStoreRetriever | |
| from langchain_google_genai import ChatGoogleGenerativeAI | |
| from langchain_core.chat_history import BaseChatMessageHistory | |
| from langchain_community.chat_message_histories import ChatMessageHistory | |
| from langchain_core.runnables.history import RunnableWithMessageHistory | |
| from processing.documents import format_documents | |
| from caching.lfu import LFUCache | |
| def _initialize_llm(model_name) -> ChatGoogleGenerativeAI: | |
| """ | |
| Initializes the LLM instance. | |
| """ | |
| llm = ChatGoogleGenerativeAI(model= model_name) | |
| return llm | |
| class LLMService: | |
| def __init__(self, logger, system_prompt: str, web_retriever: VectorStoreRetriever,cache_capacity: int = 50, llm_model_name = "gemini-2.0-flash-thinking-exp-01-21"): | |
| self._conversational_rag_chain = None | |
| self._logger = logger | |
| self.system_prompt = system_prompt | |
| self._web_retriever = web_retriever | |
| self.llm = _initialize_llm(llm_model_name) | |
| self._initialize_conversational_rag_chain() | |
| ### Statefully manage chat history ### | |
| self.store = LFUCache(capacity=cache_capacity) | |
| def _initialize_conversational_rag_chain(self): | |
| """ | |
| Initializes the conversational RAG chain. | |
| """ | |
| ### Contextualize question ### | |
| contextualize_q_system_prompt = """Given a chat history and the latest user question \ | |
| which might reference context in the chat history, formulate a standalone question \ | |
| which can be understood without the chat history. Do NOT answer the question, \ | |
| just reformulate it if needed and otherwise return it as is.""" | |
| contextualize_q_prompt = ChatPromptTemplate.from_messages( | |
| [ | |
| ("system", contextualize_q_system_prompt), | |
| MessagesPlaceholder("chat_history"), | |
| ("human", "{input}"), | |
| ] | |
| ) | |
| history_aware_retriever = create_history_aware_retriever( | |
| self.llm, self._web_retriever, contextualize_q_prompt) | |
| qa_prompt = ChatPromptTemplate.from_messages( | |
| [ | |
| ("system", self.system_prompt), | |
| MessagesPlaceholder("chat_history"), | |
| ("human", "{input}"), | |
| ] | |
| ) | |
| question_answer_chain = create_stuff_documents_chain(self.llm, qa_prompt) | |
| rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain) | |
| self._conversational_rag_chain = RunnableWithMessageHistory( | |
| rag_chain, | |
| self._get_session_history, | |
| input_messages_key="input", | |
| history_messages_key="chat_history", | |
| output_messages_key="answer", | |
| ) | |
| def _get_session_history(self, session_id: str) -> BaseChatMessageHistory: | |
| history = self.store.get(session_id) | |
| if history is None: | |
| history = ChatMessageHistory() | |
| self.store.put(session_id, history) | |
| return history | |
| def conversational_rag_chain(self): | |
| """ | |
| Returns the initialized conversational RAG chain. | |
| Returns: | |
| The conversational RAG chain instance. | |
| """ | |
| return self._conversational_rag_chain | |
| def get_llm(self) -> ChatGoogleGenerativeAI: | |
| """ | |
| Returns the LLM instance. | |
| """ | |
| if self.llm is None: | |
| raise Exception("llm is not initialized") | |
| return self.llm | |