feat: added the ability to log the history to DynamoDB

EurLexChat.py

@@ -1,3 +1,4 @@
+import boto3
 from langchain_community.vectorstores import Qdrant
 from langchain_core.runnables.history import RunnableWithMessageHistory
 from langchain_core.runnables.base import RunnableLambda
@@ -6,7 +7,7 @@ from langchain_core.tools import StructuredTool
 from langchain_core.utils.function_calling import convert_to_openai_tool
 from langchain_core.messages import AIMessage
 from typing import List
-from chat_utils import get_init_modules, SYSTEM_PROMPT, SYSTEM_PROMPT_LOOP, ContextInput, Answer
+from chat_utils import get_init_modules, SYSTEM_PROMPT, SYSTEM_PROMPT_LOOP, ContextInput, Answer, get_vectorDB_module
 from langchain_core.documents.base import Document
 
 
@@ -59,7 +60,7 @@ class EurLexChat:
             input_messages_key="question",
             history_messages_key="history",
         )
-        
+
         self.relevant_documents_pipeline = ( self.retriever | self._parse_documents )
 
 
@@ -96,6 +97,14 @@ class EurLexChat:
         if self.config["chatDB"]["class"] == 'FileChatMessageHistory':
             file_path = f"{kwargs['output_path']}/{session_id}.json"
             return self.chatDB_class(file_path=file_path)
+        elif self.config["chatDB"]["class"] == 'DynamoDBChatMessageHistory':
+            table_name = kwargs["table_name"]
+            session = boto3.Session(aws_access_key_id=kwargs["aws_access_key_id"],
+                                    aws_secret_access_key=kwargs["aws_secret_access_key"],
+                                    region_name='eu-west-1')
+            return self.chatDB_class(session_id=session_id,
+                                     table_name=table_name,
+                                     boto3_session=session)
         else:
             return self.chatDB_class(session_id=session_id, **kwargs)
 
@@ -270,3 +279,48 @@ class EurLexChat:
                 return Answer(answer=result.answer)        
         return Answer(answer=result.content)
 
+
+class EurLexChatAkn(EurLexChat):
+    def _parse_documents(self, docs: List[Document]) -> List[dict]:
+        """
+        Parse a list of documents into a standardized format.
+
+        Args:
+            docs (List[Document]): A list of documents to parse.
+
+        Returns:
+            List[dict]: A list of dictionaries, each containing parsed information from the input documents.
+        """
+
+        parsed_documents = []
+
+        for doc in docs:
+            parsed_documents.append({
+                'text': doc.page_content,
+                'source': doc.metadata["uri"],
+                '_id': doc.metadata["uri"] + doc.metadata["article_id"]
+            })
+        return parsed_documents
+
+    def get_relevant_docs(self, question: str, eurovoc: str = None) -> List[dict]:
+        """
+        Retrieve relevant documents based on a given question.
+
+        Args:
+            question (str): The question for which relevant documents are retrieved.
+            eurovoc (str): The Eurovoc to be used as filter
+
+        Returns:
+            List[dict]: A list of relevant documents.
+        """
+        if eurovoc:
+            retriever = get_vectorDB_module(
+                self.config['vectorDB'], self.embedder, metadata={'filter': {'eurovoc': ''}}
+            )
+            relevant_documents_pipeline_with_filter = (retriever | self._parse_documents)
+            docs = relevant_documents_pipeline_with_filter.invoke(
+                question
+            )
+        else:
+            docs = self.relevant_documents_pipeline.invoke(question)
+        return docs

app.py

@@ -1,18 +1,8 @@
 import gradio as gr
 from EurLexChat import EurLexChat
-import yaml
 import random
 import string
-import argparse
-import os
-
-openai_org_key = os.getenv("OPENAI_ORG_KEY")
-openai_key = os.getenv("OPENAI_KEY")
-ui_pwd = os.getenv("pwd")
-ui_user = os.getenv("user")
-qdrant_url=os.getenv("url")
-qdrant_key=os.getenv("qdrant_key")
-
+from config import CONFIG, UI_USER, UI_PWD
 
 def generate_random_string(length):
     # Generate a random string of the specified length 
@@ -25,32 +15,8 @@ class Documents():
     def __init__(self) -> None:
         self.documents = []
 
-parser = argparse.ArgumentParser(description="Chat-eur-lex ui")
-
-parser.add_argument('--config_path',
-                    dest='config_path',
-                    metavar='config_path',
-                    type=str,
-                    help='The path to the config file that contains all the settings for the chat engine' ,
-                    default='config.yaml')
-args = parser.parse_args()
-
-# Read config file
-with open(args.config_path, 'r') as file:
-    config = yaml.safe_load(file)
-
-config["embeddings"]["kwargs"]["openai_api_key"] = openai_key
-config["embeddings"]["kwargs"]["openai_organization"] = openai_org_key
-config["llm"]["kwargs"]["openai_api_key"] = openai_key
-config["llm"]["kwargs"]["openai_organization"] = openai_org_key
-config["vectorDB"]["kwargs"]["url"] = qdrant_url
-config["vectorDB"]["kwargs"]["api_key"] = qdrant_key
-
-
-
-
 
-chat = EurLexChat(config=config)
+chat = EurLexChat(config=CONFIG)
 docs = Documents()
 
 
@@ -113,7 +79,7 @@ with block:
             
         with gr.Column(scale=1, visible=False) as col:
             gr.Markdown("""<h3><center>Context documents</center></h3>""")
-            for i in range(config['vectorDB']['retriever_args']['search_kwargs']['k']):
+            for i in range(CONFIG['vectorDB']['retriever_args']['search_kwargs']['k']):
                 with gr.Accordion(label="", elem_id=f'accordion_{i}', open=False) as acc:
                     list_texts.append(gr.Textbox("", elem_id=f'text_{i}', show_label=False, lines=10))
                     btn = gr.Button(f"Remove document")
@@ -141,4 +107,4 @@ with block:
     for i, b in enumerate(delete_buttons):
         b.click(remove_doc, inputs=states[i], outputs=[*accordions, *list_texts])
 
-block.launch(debug=True, auth=(ui_user, ui_pwd))
+block.launch(debug=True, auth=(UI_USER, UI_PWD))

chat_utils.py

@@ -64,7 +64,7 @@ def get_init_modules(config):
     return embedder, llm, chatDB_class, retriever
 
 
-def get_vectorDB_module(db_config, embedder):
+def get_vectorDB_module(db_config, embedder, metadata=None):
     mod_chat = __import__("langchain_community.vectorstores",
                           fromlist=[db_config["class"]])
     vectorDB_class = getattr(mod_chat, db_config["class"])
@@ -85,10 +85,13 @@ def get_vectorDB_module(db_config, embedder):
 
         client = QdrantClient(**client_kwargs)
 
+        if metadata is None:
+            metadata = {}
         retriever = vectorDB_class(
             client, embeddings=embedder, **db_kwargs).as_retriever(
                 search_type=db_config["retriever_args"]["search_type"],
-                search_kwargs=db_config["retriever_args"]["search_kwargs"]
+                search_kwargs={**db_config["retriever_args"]["search_kwargs"], **metadata},
+                filter=metadata
         )
 
     else:

config.py

@@ -0,0 +1,35 @@
+import os
+import yaml
+
+
+# Read config file
+if os.path.exists('config.yaml'):
+    with open('config.yaml', 'r') as file:
+        CONFIG = yaml.safe_load(file)
+else:
+    raise FileNotFoundError('config.yml not found Aborting!')
+
+OPENAI_ORG_KEY = os.getenv("OPENAI_ORG_KEY", "")
+OPENAI_KEY = os.getenv("OPENAI_KEY", "")
+QDRANT_URL = os.getenv("url", CONFIG["vectorDB"]["kwargs"].get("url", ""))
+QDRANT_KEY = os.getenv("qdrant_key", CONFIG["vectorDB"]["kwargs"].get("api_key", ""))
+
+UI_USER = os.getenv("user", "admin")
+UI_PWD = os.getenv("pwd", "admin")
+
+CONFIG["embeddings"]["kwargs"]["openai_api_key"] = OPENAI_KEY
+CONFIG["embeddings"]["kwargs"]["openai_organization"] = OPENAI_ORG_KEY
+CONFIG["llm"]["kwargs"]["openai_api_key"] = OPENAI_KEY
+CONFIG["llm"]["kwargs"]["openai_organization"] = OPENAI_ORG_KEY
+CONFIG["vectorDB"]["kwargs"]["url"] = QDRANT_URL
+CONFIG["vectorDB"]["kwargs"]["api_key"] = QDRANT_KEY
+
+# if the history should be stored on AWS DynamoDB
+# otherwise it will be stored on local FS to the output_path defined in the config.yaml file
+if CONFIG['chatDB']['class'] == 'DynamoDBChatMessageHistory':
+    CHATDB_TABLE_NAME = os.getenv("CHATDB_TABLE_NAME", CONFIG["chatDB"]["kwargs"].get("table_name", "ChatEurlexHistory"))
+    AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID", CONFIG["chatDB"]["kwargs"].get("aws_access_key_id", ""))
+    AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY", CONFIG["chatDB"]["kwargs"].get("aws_secret_access_key", ""))
+    CONFIG["chatDB"]["kwargs"]["table_name"] = CHATDB_TABLE_NAME
+    CONFIG["chatDB"]["kwargs"]["aws_access_key_id"] = AWS_ACCESS_KEY_ID
+    CONFIG["chatDB"]["kwargs"]["aws_secret_access_key"] = AWS_SECRET_ACCESS_KEY

config.yaml

@@ -16,8 +16,6 @@ vectorDB:
 embeddings: 
     class: OpenAIEmbeddings
     kwargs:
-        openai_api_key: ""
-        openai_organization: ""
         model: text-embedding-ada-002
 
 
@@ -26,15 +24,15 @@ llm:
     use_context_function: True
     max_context_size: 6000
     kwargs:
-        openai_organization: ""
-        openai_api_key: ""
         model_name: gpt-4
         temperature: 0.8
 
 
 chatDB:
-    class: FileChatMessageHistory
+    class: DynamoDBChatMessageHistory
     kwargs:
-        output_path: ./output
+        table_name: ''
+        aws_access_key_id: ''
+        aws_secret_access_key: ''
 
 max_history_messages: 5

requirements.txt

@@ -4,4 +4,5 @@ tiktoken==0.6.0
 qdrant-client==1.7.3
 transformers==4.37.2
 openai==1.12.0
-gradio==4.18.0
+gradio==4.18.0
+boto3==1.34