Adding chat history to RAG app and refactor to better utilize LangChain

applications/rag/frontend/container/Dockerfile

@@ -19,4 +19,9 @@ WORKDIR /workspace/frontend
 
 RUN pip install -r requirements.txt
 
-CMD ["python", "main.py"]
+EXPOSE 8080
+
+ENV FLASK_APP=/workspace/frontend/main.py
+ENV PYTHONPATH=.
+# Run the application with Gunicorn
+CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:8080", "main:app"]

applications/rag/frontend/container/rai/__init__.py → applications/rag/frontend/container/__init__.py

@@ -11,5 +11,3 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-# This file is required to make Python treat the subfolder as a package

applications/rag/frontend/container/application/__init__.py

@@ -0,0 +1,24 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+from flask import Flask
+
+def create_app():
+    app = Flask(__name__, static_folder='static', template_folder='templates')
+    app.jinja_env.trim_blocks = True
+    app.jinja_env.lstrip_blocks = True
+    app.config['SECRET_KEY'] = os.environ.get("SECRET_KEY")
+
+    return app

applications/rag/frontend/container/application/models/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .vector_embeddings import VectorEmbeddings
+
+__all__ = ["VectorEmbeddings"]

applications/rag/frontend/container/application/models/vector_embeddings.py

@@ -0,0 +1,31 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+from sqlalchemy import Column, String, Text
+from sqlalchemy.orm import mapped_column, declarative_base
+from pgvector.sqlalchemy import Vector
+
+Base = declarative_base()
+
+VECTOR_EMBEDDINGS_TABLE_NAME = os.environ.get("EMBEDDINGS_TABLE_NAME", "")
+
+
+class VectorEmbeddings(Base):
+    __tablename__ = VECTOR_EMBEDDINGS_TABLE_NAME
+
+    id = Column(String(255), primary_key=True)
+    text = Column(Text)
+    text_embedding = mapped_column(Vector(384))

applications/rag/frontend/container/application/rag_langchain/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

applications/rag/frontend/container/application/rag_langchain/huggingface_inference_model.py

@@ -0,0 +1,131 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import logging
+from typing import Any, Dict, Iterator, List, Optional
+
+from langchain_core.callbacks.manager import CallbackManagerForLLMRun
+from langchain_core.language_models.llms import LLM
+from langchain_core.outputs import GenerationChunk
+
+from application.utils import post_request
+
+INFERENCE_ENDPOINT = os.environ.get("INFERENCE_ENDPOINT", "127.0.0.1:8081")
+
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+
+
+class HuggingFaceCustomChatModel(LLM):
+    """A custom chat model that calls to an HuggingFace TGI API and returns the generated
+    content based on the given message.
+
+    Example:
+
+        .. code-block:: python
+
+            model = HuggingFaceCustomChatModel()
+            result = model.invoke([HumanMessage(content="hello")])
+            result = model.batch([[HumanMessage(content="hello")],
+                                 [HumanMessage(content="world")]])
+    """
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Run the LLM on the given input.
+        Args:
+            prompt: The prompt to generate from.
+            stop: Stop words to use when generating. Model output is cut off at the
+                first occurrence of any of the stop substrings.
+                If stop tokens are not supported consider raising NotImplementedError.
+            run_manager: Callback manager for the run.
+            **kwargs: Arbitrary additional keyword arguments. These are usually passed
+                to the model provider API call.
+
+        Returns:
+            The model output as a string.
+        """
+        if stop is not None:
+            raise ValueError("stop kwargs are not permitted.")
+
+        api_endpoint = f"http://{INFERENCE_ENDPOINT}/generate"
+        body = {"inputs": prompt}
+        headers = {"Content-Type": "application/json"}
+        generated_output = post_request(api_endpoint, body, headers)
+        generated_text = generated_output.get("generated_text", "")
+
+        return generated_text
+
+    def _stream(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> Iterator[GenerationChunk]:
+        """Stream the LLM on the given prompt.
+
+        This method should be overridden by subclasses that support streaming.
+
+        If not implemented, the default behavior of calls to stream will be to
+        fallback to the non-streaming version of the model and return
+        the output as a single chunk.
+
+        Args:
+            prompt: The prompt to generate from.
+            stop: Stop words to use when generating. Model output is cut off at the
+                first occurrence of any of these substrings.
+            run_manager: Callback manager for the run.
+            **kwargs: Arbitrary additional keyword arguments. These are usually passed
+                to the model provider API call.
+
+        Returns:
+            An iterator of GenerationChunks.
+        """
+        api_endpoint = f"http://{INFERENCE_ENDPOINT}/generate_stream"
+        body = {"inputs": prompt}
+        headers = {"Content-Type": "application/json"}
+        logging.info("Calling external model")
+        generated_output = post_request(api_endpoint, body, headers)
+        generated_text = generated_output.get("generated_text", "")
+
+        for char in generated_text:
+            chunk = GenerationChunk(text=char)
+            if run_manager:
+                run_manager.on_llm_new_token(chunk.text, chunk=chunk)
+
+            yield chunk
+
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        """Return a dictionary of identifying parameters."""
+        return {
+            # The model name allows users to specify custom token counting
+            # rules in LLM monitoring applications (e.g., in LangSmith users
+            # can provide per token pricing for their model and monitor
+            # costs for the given LLM.)
+            "model_name": "HuggingFaceTGI",
+        }
+
+    @property
+    def _llm_type(self) -> str:
+        """Get the type of language model used by this chat model. Used for logging purposes only."""
+        return "custom"

applications/rag/frontend/container/application/rag_langchain/rag_chain.py

@@ -0,0 +1,119 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import logging
+
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.runnables import RunnableParallel, RunnableLambda
+from langchain_core.runnables.history import RunnableWithMessageHistory
+
+from langchain_community.embeddings.huggingface import HuggingFaceEmbeddings
+from langchain_google_cloud_sql_pg import PostgresChatMessageHistory
+
+
+from application.utils import (
+    create_sync_postgres_engine
+)
+from application.rag_langchain.huggingface_inference_model import (
+    HuggingFaceCustomChatModel,
+)
+from application.vector_storages import CloudSQLVectorStore
+
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+
+CHAT_HISTORY_TABLE_NAME = os.environ.get("CHAT_HISTORY_TABLE_NAME", "message_store")
+
+QUESTION = "input"
+HISTORY = "chat_history"
+CONTEXT = "context"
+
+SENTENCE_TRANSFORMER_MODEL = "intfloat/multilingual-e5-small"  # Transformer to use for converting text chunks to vector embeddings
+
+template_str = """Answer the question given by the user in no more than 2 sentences. 
+Use the provided context to improve upon your previous answers. Stick to the facts and be brief. Avoid conversational format.
+\n\n
+Context: {context} 
+"""
+
+prompt = ChatPromptTemplate.from_messages(
+    [
+        ("system", template_str),
+        MessagesPlaceholder("chat_history"),
+        ("human", "{input}"),
+    ]
+)
+
+engine = create_sync_postgres_engine()
+
+
+def get_chat_history(session_id: str) -> PostgresChatMessageHistory:
+    history = PostgresChatMessageHistory.create_sync(
+        engine, session_id=session_id, table_name=CHAT_HISTORY_TABLE_NAME
+    )
+
+    logging.info(
+        f"Retrieving history for session {session_id} with {len(history.messages)}"
+    )
+    return history
+
+
+def clear_chat_history(session_id: str):
+    history = PostgresChatMessageHistory.create_sync(
+        engine, session_id=session_id, table_name=CHAT_HISTORY_TABLE_NAME
+    )
+    history.clear()
+
+
+def create_chain() -> RunnableWithMessageHistory:
+    try:
+        model = HuggingFaceCustomChatModel()
+
+        langchain_embed = HuggingFaceEmbeddings(model_name=SENTENCE_TRANSFORMER_MODEL)
+        vector_store = CloudSQLVectorStore(langchain_embed, engine)
+
+        retriever = vector_store.as_retriever()
+
+        setup_and_retrieval = RunnableParallel(
+            {
+                "context": retriever,
+                QUESTION: RunnableLambda(lambda d: d[QUESTION]),
+                HISTORY: RunnableLambda(lambda d: d[HISTORY]),
+            }
+        )
+
+        chain = setup_and_retrieval | prompt | model
+        chain_with_history = RunnableWithMessageHistory(
+            chain,
+            get_chat_history,
+            input_messages_key=QUESTION,
+            history_messages_key=HISTORY,
+            output_messages_key="output",
+        )
+        return chain_with_history
+    except Exception as e:
+        logging.info(e)
+        raise e
+
+def take_chat_turn(
+    chain: RunnableWithMessageHistory, session_id: str, query_text: str
+) -> str:
+    try:
+        config = {"configurable": {"session_id": session_id}}
+        result = chain.invoke({"input": query_text}, config=config)
+        return result
+    except Exception as e:
+        logging.info(e)
+        raise e

applications/rag/frontend/container/cloud_sql/__init__.py → applications/rag/frontend/container/application/rai/__init__.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# This file is required to make Python treat the subfolder as a package
+# This file is required to make Python treat the subfolder as a package