add chat history

.gitignore

@@ -1,4 +1,7 @@
 venv
 configs/.env
 .idea
-*__pycache__
+*__pycache__
+venv
+embeddings
+*.gradio

app.py

@@ -20,16 +20,22 @@ service.scrape_and_get_store_vector_retriever(config.URLS)
 # Initialize the LLMService with logger, prompt, and store vector retriever
 llm_svc = LLMService(logger, config.SYSTEM_PROMPT, store.get_chroma_instance().as_retriever())
 
-
-
-def respond(user_input, history):
-    print(f"{user_input}")
-
-    response = llm_svc.conversational_rag_chain().invoke(user_input)
+def respond(user_input):
+    response = llm_svc.conversational_rag_chain().invoke(
+        {"input": user_input},
+        config={"configurable": {"session_id": "abc"}},
+    )["answer"]
 
     return response
 
-
 if __name__ == '__main__':
     logging.info("Starting AIVIz Bot")
-    gr.ChatInterface(respond).launch(share=True)
+
+    # Using ChatInterface to create the chatbot interface
+    chat_interface = gr.ChatInterface(
+        fn=respond,
+        title="AISDb Bot",
+        description="LLM's are prone to hallucinations"
+    )
+
+    chat_interface.launch(share=True)

configs/.env

@@ -1 +1 @@
-GOOGLE_API_KEY="AIzaSyCWQsPEq-D3nJZFdMgsTlxDOweTzPKOTwI"
+GOOGLE_API_KEY=""

configs/config.py

@@ -57,13 +57,12 @@ EMBEDDINGS = HuggingFaceEmbeddings(
     model_kwargs={"device": "cpu"},
 )
 
-SYSTEM_PROMPT = """
-You are a chatbot to assist users asking about Automatic Identification systems (AIS) database from the context given to you. 
-Use this Context: {context}. The users are building great Machine learning models using this Database, 
-so assist them with code, definitions, summarization and so forth like a tutor. 
-Give scenario based answer that can clearly explain it to users and explain step by step.
-Based on this, now answer for this user's question: {question}."""
-
+SYSTEM_PROMPT = """You are an assistant for question-answering tasks. \
+Use the following pieces of retrieved context to answer the question. \
+Try to keep the answer concise, unless aksed by the user to be eloborated.\
+Analyse the question, and provide necessary python code help if necessary, as you will be mainly used for ML research.\
+If you don't know the answer, just say that you don't know. \
+Context:  {context}"""
 
 def set_envs():
     if "GOOGLE_API_KEY" not in os.environ:

llm_setup/llm_setup.py

@@ -1,13 +1,16 @@
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import (
     ChatPromptTemplate,
-    PromptTemplate,
-    HumanMessagePromptTemplate
+    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
 
 
@@ -15,20 +18,12 @@ def _initialize_llm() -> ChatGoogleGenerativeAI:
     """
     Initializes the LLM instance.
     """
-    llm = ChatGoogleGenerativeAI(model="gemini-pro")
+    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash-exp")
+
     return llm
 
 
 class LLMService:
-    """
-    Service for managing LLM interactions and conversational RAG chain.
-
-    Args:
-        logger: Logger instance for logging.
-        system_prompt: The prompt for the QA system.
-        web_retriever: A VectorStoreRetriever instance for retrieving web documents.
-    """
-
     def __init__(self, logger, system_prompt: str, web_retriever: VectorStoreRetriever):
         self._conversational_rag_chain = None
         self._logger = logger
@@ -39,27 +34,55 @@ class LLMService:
 
         self._initialize_conversational_rag_chain()
 
+        ### Statefully manage chat history ###
+        self.store = {}
+
     def _initialize_conversational_rag_chain(self):
         """
         Initializes the conversational RAG chain.
         """
-        # Initialize RAG (Retrieval-Augmented Generation) chain
-        prompt = ChatPromptTemplate(input_variables=['history','context', 'question'], messages=[HumanMessagePromptTemplate(
-            prompt=PromptTemplate(input_variables=['history','context', 'question'], template="{history}\nContext: {context}\nQuestion: {question}"))])
-
-        # Initialize conversational RAG chain
-        self._conversational_rag_chain = (
-                {"context": self._web_retriever | format_documents, "question": RunnablePassthrough(), "history":self.get_chat_history}
-                | prompt
-                | self.llm
-                | StrOutputParser()
+        ### 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}"),
+            ]
         )
 
-    def get_chat_history(self):
-        """
-        Retrieves the last 3 chat messages formatted as a string.
-        """
-        return "\n".join(self.chat_history) if self.chat_history else "No prior conversation."    
+
+
+        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:
+            if session_id not in self.store:
+                self.store[session_id] = ChatMessageHistory()
+            return self.store[session_id]
 
     def conversational_rag_chain(self):
         """
@@ -69,20 +92,6 @@ class LLMService:
             The conversational RAG chain instance.
         """
         return self._conversational_rag_chain
-
-    def update_chat_history(self, user_input: str, llm_response: str):
-        """
-        Updates the chat history with the latest question and response.
-        """
-        self.chat_history.append(f"User: {user_input}\nAI: {llm_response}")
-
-    def ask_question(self, question: str):
-        """
-        Processes a user question using the conversational RAG chain and updates history.
-        """
-        response = self._conversational_rag_chain.invoke(question)
-        self.update_chat_history(question, response)
-        return response
     
     def get_llm(self) -> ChatGoogleGenerativeAI:
         """

requirements.txt

@@ -1,154 +1,155 @@
+
 huggingface_hub
-aiofiles==23.2.1
-aiohappyeyeballs==2.4.0
-aiohttp==3.10.5
-aiosignal==1.3.1
-annotated-types==0.7.0
-anyio==4.4.0
-asgiref==3.8.1
-attrs==24.2.0
-backoff==2.2.1
-bcrypt==4.2.0
-build==1.2.2
-cachetools==5.5.0
-certifi==2024.8.30
-charset-normalizer==3.3.2
-chroma-hnswlib==0.7.6
-chromadb==0.5.7
-click==8.1.7
-colorama==0.4.6
-coloredlogs==15.0.1
-contourpy==1.3.0
-cycler==0.12.1
-dataclasses-json==0.6.7
-Deprecated==1.2.14
-fastapi==0.115.0
-ffmpy==0.4.0
-filelock==3.16.1
-flatbuffers==24.3.25
-fonttools==4.53.1
-frozenlist==1.4.1
-fsspec==2024.9.0
-google-ai-generativelanguage==0.6.6
-google-api-core==2.19.2
-google-api-python-client==2.146.0
-google-auth==2.34.0
-google-auth-httplib2==0.2.0
-google-generativeai==0.7.2
-googleapis-common-protos==1.65.0
-gradio==4.44.0
-gradio_client==1.3.0
-greenlet==3.1.0
-grpcio==1.66.1
-grpcio-status==1.62.3
-h11==0.14.0
-httpcore==1.0.5
-httplib2==0.22.0
-httptools==0.6.1
-httpx==0.27.2
-huggingface-hub==0.25.0
-humanfriendly==10.0
-idna==3.10
-importlib_metadata==8.4.0
-importlib_resources==6.4.5
-Jinja2==3.1.4
-joblib==1.4.2
-jsonpatch==1.33
-jsonpointer==3.0.0
-kiwisolver==1.4.7
-kubernetes==30.1.0
-langchain==0.3.0
-langchain-chroma==0.1.4
-langchain-community==0.3.0
-langchain-core==0.3.1
-langchain-google-genai==2.0.0
-langchain-huggingface==0.1.0
-langchain-text-splitters==0.3.0
-langsmith==0.1.121
-markdown-it-py==3.0.0
-MarkupSafe==2.1.5
-marshmallow==3.22.0
-matplotlib==3.9.2
-mdurl==0.1.2
-mmh3==5.0.0
-monotonic==1.6
-mpmath==1.3.0
-multidict==6.1.0
-mypy-extensions==1.0.0
-networkx==3.3
-numpy==1.26.4
-oauthlib==3.2.2
-onnxruntime==1.19.2
-opentelemetry-api==1.27.0
-opentelemetry-exporter-otlp-proto-common==1.27.0
-opentelemetry-exporter-otlp-proto-grpc==1.27.0
-opentelemetry-instrumentation==0.48b0
-opentelemetry-instrumentation-asgi==0.48b0
-opentelemetry-instrumentation-fastapi==0.48b0
-opentelemetry-proto==1.27.0
-opentelemetry-sdk==1.27.0
-opentelemetry-semantic-conventions==0.48b0
-opentelemetry-util-http==0.48b0
-orjson==3.10.7
-overrides==7.7.0
-packaging==24.1
-pandas==2.2.2
-pillow==10.4.0
-posthog==3.6.6
-proto-plus==1.24.0
-protobuf==4.25.4
-pyasn1==0.6.1
-pyasn1_modules==0.4.1
-pydantic==2.9.2
-pydantic-settings==2.5.2
-pydantic_core==2.23.4
-pydub==0.25.1
-Pygments==2.18.0
-pyparsing==3.1.4
-PyPika==0.48.9
-pyproject_hooks==1.1.0
-pyreadline3==3.5.4
-python-dateutil==2.9.0.post0
-python-dotenv==1.0.1
-python-multipart==0.0.9
-pytz==2024.2
-PyYAML==6.0.2
-regex==2024.9.11
-requests==2.32.3
-requests-oauthlib==2.0.0
-rich==13.8.1
-rsa==4.9
-ruff==0.6.5
-safetensors==0.4.5
-scikit-learn==1.5.2
-scipy==1.14.1
-semantic-version==2.10.0
-sentence-transformers==3.1.0
-setuptools==75.1.0
-shellingham==1.5.4
-six==1.16.0
-sniffio==1.3.1
-SQLAlchemy==2.0.35
-starlette==0.38.5
-sympy==1.13.2
-tenacity==8.5.0
-threadpoolctl==3.5.0
-tokenizers==0.19.1
-tomlkit==0.12.0
-torch==2.4.1
-tqdm==4.66.5
-transformers==4.44.2
-typer==0.12.5
-typing-inspect==0.9.0
-typing_extensions==4.12.2
-tzdata==2024.1
-uritemplate==4.1.1
-urllib3==2.2.3
-uvicorn==0.30.6
-watchfiles==0.24.0
-websocket-client==1.8.0
-websockets==12.0
-wrapt==1.16.0
-yarl==1.11.1
-zipp==3.20.2
+aiofiles
+aiohappyeyeballs
+aiohttp
+aiosignal
+annotated-types
+anyio
+asgiref
+attrs
+backoff
+bcrypt
+build
+cachetools
+certifi
+charset-normalizer
+chroma-hnswlib
+chromadb
+click
+colorama
+coloredlogs
+contourpy
+cycler
+dataclasses-json
+Deprecated
+fastapi
+ffmpy
+filelock
+flatbuffers
+fonttools
+frozenlist
+fsspec
+google-ai-generativelanguage
+google-api-core
+google-api-python-client
+google-auth
+google-auth-httplib2
+google-generativeai
+googleapis-common-protos
+gradio
+gradio_client
+greenlet
+grpcio
+grpcio-status
+h11
+httpcore
+httplib2
+httptools
+httpx
+huggingface-hub
+humanfriendly
+idna
+importlib_metadata
+importlib_resources
+Jinja2
+joblib
+jsonpatch
+jsonpointer
+kiwisolver
+kubernetes
+langchain
+langchain-chroma
+langchain-community
+langchain-core
+langchain-google-genai
+langchain-huggingface
+langchain-text-splitters
+langsmith
+markdown-it-py
+MarkupSafe
+marshmallow
+matplotlib
+mdurl
+mmh3
+monotonic
+mpmath
+multidict
+mypy-extensions
+networkx
+numpy
+oauthlib
+onnxruntime
+opentelemetry-api
+opentelemetry-exporter-otlp-proto-common
+opentelemetry-exporter-otlp-proto-grpc
+opentelemetry-instrumentation
+opentelemetry-instrumentation-asgi
+opentelemetry-instrumentation-fastapi
+opentelemetry-proto
+opentelemetry-sdk
+opentelemetry-semantic-conventions
+opentelemetry-util-http
+orjson
+overrides
+packaging
+pandas
+pillow
+posthog
+proto-plus
+protobuf
+pyasn1
+pyasn1_modules
+pydantic
+pydantic-settings
+pydantic_core
+pydub
+Pygments
+pyparsing
+PyPika
+pyproject_hooks
+pyreadline3
+python-dateutil
+python-dotenv
+python-multipart
+pytz
+PyYAML
+regex
+requests
+requests-oauthlib
+rich
+rsa
+ruff
+safetensors
+scikit-learn
+scipy
+semantic-version
+sentence-transformers
+setuptools
+shellingham
+six
+sniffio
+SQLAlchemy
+starlette
+sympy
+tenacity
+threadpoolctl
+tokenizers
+tomlkit
+torch
+tqdm
+transformers
+typer
+typing-inspect
+typing_extensions
+tzdata
+uritemplate
+urllib3
+uvicorn
+watchfiles
+websocket-client
+websockets
+wrapt
+yarl
+zipp
 bs4