Merge branch 'feat/limit_or_rephrase' into main

app.py

@@ -4,7 +4,7 @@ import os
 import streamlit as st
 import streamlit.components.v1 as components
 from langchain.callbacks import get_openai_callback
-from langchain.chains import ConversationalRetrievalChain
+
 from langchain.chains.conversation.memory import ConversationBufferMemory
 from langchain.embeddings import GPT4AllEmbeddings
 from langchain.llms import OpenAI
@@ -14,6 +14,8 @@ from connection import connect
 from css import load_css
 from message import Message
 from vector_store import CustomVectorStore
+from conversation_retrieval_chain import CustomConversationalRetrievalChain
+
 
 st.set_page_config(layout="wide")
 
@@ -50,12 +52,13 @@ def initialize_session_state():
         memory = ConversationBufferMemory(
             output_key="answer", memory_key="chat_history", return_messages=True
         )
-        st.session_state.conversation = ConversationalRetrievalChain.from_llm(
+        st.session_state.conversation = CustomConversationalRetrievalChain.from_llm(
             llm=llm,
             retriever=retriever,
             verbose=True,
             memory=memory,
             return_source_documents=True,
+            max_tokens_limit=3700,
         )
 
 

conversation_retrieval_chain.py

@@ -0,0 +1,64 @@
+import inspect
+from typing import Any, Dict, Optional
+
+from langchain.chains.conversational_retrieval.base import (
+    ConversationalRetrievalChain,
+    _get_chat_history,
+)
+from langchain.callbacks.manager import CallbackManagerForChainRun
+
+
+class CustomConversationalRetrievalChain(ConversationalRetrievalChain):
+    def _handle_docs(self, docs):
+        if len(docs) == 0:
+            return False, "No documents found. Can you rephrase ?"
+        elif len(docs) == 1:
+            return False, "Only one document found. Can you rephrase ?"
+        elif len(docs) > 10:
+            return False, "Too many documents found. Can you specify your request ?"
+        return True, ""
+
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
+        question = inputs["question"]
+        get_chat_history = self.get_chat_history or _get_chat_history
+        chat_history_str = get_chat_history(inputs["chat_history"])
+
+        if chat_history_str:
+            callbacks = _run_manager.get_child()
+            new_question = self.question_generator.run(
+                question=question, chat_history=chat_history_str, callbacks=callbacks
+            )
+        else:
+            new_question = question
+        accepts_run_manager = (
+            "run_manager" in inspect.signature(self._get_docs).parameters
+        )
+        if accepts_run_manager:
+            docs = self._get_docs(new_question, inputs, run_manager=_run_manager)
+        else:
+            docs = self._get_docs(new_question, inputs)  # type: ignore[call-arg]
+        valid_docs, message = self._handle_docs(docs)
+        if not valid_docs:
+            return {
+                self.output_key: message,
+                "source_documents": docs,
+            }
+
+        new_inputs = inputs.copy()
+        if self.rephrase_question:
+            new_inputs["question"] = new_question
+        new_inputs["chat_history"] = chat_history_str
+        answer = self.combine_docs_chain.run(
+            input_documents=docs, callbacks=_run_manager.get_child(), **new_inputs
+        )
+        output: Dict[str, Any] = {self.output_key: answer}
+        if self.return_source_documents:
+            output["source_documents"] = docs
+        if self.return_generated_question:
+            output["generated_question"] = new_question
+        return output

models/distance.py

@@ -1,6 +1,13 @@
 import enum
 
 
+distance_strategy_limit = {
+    "l2": 1.05,
+    "cosine": 0.55,
+    "inner": 1.0,
+}
+
+
 class DistanceStrategy(str, enum.Enum):
     """Enumerator of the Distance strategies."""
 

requirements.txt

@@ -1,11 +1,12 @@
 gpt4all==1.0.12
 langchain==0.0.313
 openai==0.28.1
+pandas==2.1.1
+pgvector==0.2.3
+psycopg2-binary==2.9.9
+psycopg2==2.9.9
 streamlit==1.27.2
 streamlit-chat==0.1.1
 SQLAlchemy==2.0.22
 sqlite-vss==0.1.2
-pandas==2.1.1
-pgvector==0.2.3
-psycopg2-binary==2.9.9
-psycopg2==2.9.9
+tiktoken==0.5.1

vector_store.py

@@ -14,10 +14,10 @@ from sqlalchemy import delete, text
 from sqlalchemy.orm import Session
 
 from model import Article
-from models.distance import DistanceStrategy
+from models.distance import DistanceStrategy, distance_strategy_limit
 from utils import str_to_list
 
-DEFAULT_DISTANCE_STRATEGY = DistanceStrategy.EUCLIDEAN
+DEFAULT_DISTANCE_STRATEGY = DistanceStrategy.COSINE
 
 _LANGCHAIN_DEFAULT_COLLECTION_NAME = "langchain"
 
@@ -252,6 +252,8 @@ class CustomVectorStore(VectorStore):
         k: int = 4,
     ) -> List[Any]:
         """Query the collection."""
+
+        limit = distance_strategy_limit[self._distance_strategy]
         with Session(self._conn) as session:
             results = session.execute(
                 text(
@@ -272,10 +274,11 @@ class CustomVectorStore(VectorStore):
                     left join author on author.id = article_author.author_id  
                     where
                         abstract_en != '' and
-                        abstract_en != 'None'
+                        abstract_en != 'None' and
+                        abstract_embedding_en {self.distance_strategy} '{str(embedding)}' < {limit}
                     GROUP BY a.id
                     ORDER BY distance
-                    LIMIT {k};
+                    LIMIT 100;
                 """
                 )
             )