refactor: refactoe

app.py

@@ -10,401 +10,24 @@ from langchain_core.runnables import ConfigurableFieldSpec
 from langchain.schema import Document
 from langchain.prompts import PromptTemplate
 from langchain_core.output_parsers import JsonOutputParser
-from langchain.retrievers.multi_query import MultiQueryRetriever
 from langchain_core.output_parsers import StrOutputParser
 from langchain_pinecone import PineconeVectorStore
 from typing_extensions import TypedDict
 from typing import Dict, List
-from langgraph.graph import END, StateGraph
 import warnings
 
-from lib.embedding import EmbeddingBuilder, build_embedding
-from lib.model_builder import ModelBuilder, ModelBuilderV2
 from lib.gradio_custom_theme import DarkTheme
-warnings.filterwarnings('ignore')
-from dotenv import load_dotenv
-
-load_dotenv()
-store = {}
-
-def get_session_history(user_id: str, conversation_id: str) -> BaseChatMessageHistory:
-    if (user_id, conversation_id) not in store:
-        store[(user_id, conversation_id)] = ChatMessageHistory()
-    return store[(user_id, conversation_id)]
-
-
-def combine_vectors(vectors):
-    result = []
-    vec1_count = len(vectors["vector1"])
-    # vec2_count = len(vectors["vector2"])
-    for i in range(vec1_count):
-        if i < vec1_count:
-            result.append(vectors['vector1'][i])
-        # if i < vec2_count:
-        #     result.append(vectors['vector2'][i])
-    return result
-
-
-class GraphState(TypedDict):
-    """
-    Represents the state of our graph.
+from lib.graph import build_workflow
 
-    Attributes:
-        question: question
-        generation: LLM generation
-        web_search: whether to add search
-        documents: list of documents
-    """
-
-    question: str
-    generation: str
-    documents: List[Dict]
-    filtered_documets: List[Dict]
-    is_fuse: bool
-    count_regenerations: int
-
-class FAISSBuilder:
-    def __enter__(self):
-        # Initialize resources
-        with EmbeddingBuilder("intfloat/multilingual-e5-large") as rag_emb:
-            faiss_db = FAISS.load_local("data/faiss_nk_28_05", rag_emb, allow_dangerous_deserialization=True)
-            return faiss_db.as_retriever()
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        # Cleanup resources if necessary
-        pass
 
+warnings.filterwarnings("ignore")
+from dotenv import load_dotenv
 
-class PineConeBuilder:
-    def __init__(self, index_name, embedding_model):
-        self.index_name = index_name
-        self.embedding_model = embedding_model
+load_dotenv()
 
-    def __enter__(self):
-        with EmbeddingBuilder(self.embedding_model) as embeddings:
-            pc_db = PineconeVectorStore.from_existing_index(self.index_name, embeddings)
-            return pc_db.as_retriever()
-        
-    def __exit__(self, exc_type, exc_value, traceback):
-        # Cleanup resources if necessary
-        pass
 
 def deploy():
-
-    casual_prompt = PromptTemplate(
-    template="""Just answer a question as casual chatter. \n 
-        Here is the user question: {question} \n
-        Always reply in Russian. Leave clear answer, without any addition.
-        Chat history: 
-        {chat_history}
-        Answer:
-        """,
-        input_variables=["question", "chat_history"],
-    )
-    with ModelBuilderV2("openchat/openchat-7b", 0.7) as llm:
-        casual_llm = RunnableWithMessageHistory(
-            casual_prompt | llm,
-            get_session_history,
-            input_messages_key="question",
-            history_messages_key="chat_history",
-            history_factory_config=[
-                ConfigurableFieldSpec(
-                    id="user_id",
-                    annotation=str,
-                    name="User ID",
-                    description="Unique identifier for the user.",
-                    default="default_user",
-                    is_shared=True,
-                ),
-                ConfigurableFieldSpec(
-                    id="conversation_id",
-                    annotation=str,
-                    name="Conversation ID",
-                    description="Unique identifier for the conversation.",
-                    default="default_session",
-                    is_shared=True,
-                ),
-            ],
-        ) | StrOutputParser()
-
-
-    prompt = PromptTemplate(
-    template="""You are a grader assessing relevance of a retrieved documents to a user question. \n 
-        Here is the first retrieved document: \n\n {document_1} \n\n
-        Here is the second retrieved document: \n\n {document_2} \n\n
-        Here is the third retrieved document: \n\n {document_3} \n\n
-        Here is the user question: {question} \n
-        If the document contains keywords related to the user question, grade it as relevant. \n
-        It does not need to be a stringent test. The goal is to filter out erroneous retrievals. \n
-        For each! document give a score from 0 to 1, score to indicate whether the document is relevant to the question. \n
-        Provide the scores as a JSON list that contains an objects with single key 'score' and no premable or explanation.""",
-        input_variables=["question", "document_1", "document_2", "document_3"],
-    )
-    with ModelBuilderV2("cohere/command-r") as llm:
-        retrieval_grader_3_docs = prompt | llm | JsonOutputParser()
-    
-    template = """
-        SYSTEM: You are an assistant for question-answering tasks. 
-        Use the following pieces of retrieved context to answer the question.
-        Use previous messages then current message higly likely  
-        If you don't find the answer in the context, transform the question ans ask the user to specify his qusetion. 
-
-        Keep the answer concise.  
-        Print a most possible topic of conversation.
-        Always reply in Russian, all text must be in Russian!
-
-        Context: {context}
-
-        Previous messages: {chat_history}
-
-        Question: {question}
-
-        Answer:
-        
-        """
-    prompt = ChatPromptTemplate.from_template(template)
-    
-    with ModelBuilderV2("mistralai/mixtral-8x22b-instruct") as llm:
-    # Post-processing
-        def format_docs(docs):
-            return "\n\n".join(doc.page_content for doc in docs)
-
-        # Chain
-        rag_chain = RunnableWithMessageHistory(
-            prompt | llm,
-            get_session_history,
-            input_messages_key="question",
-            history_messages_key="chat_history",
-            history_factory_config=[
-                ConfigurableFieldSpec(
-                    id="user_id",
-                    annotation=str,
-                    name="User ID",
-                    description="Unique identifier for the user.",
-                    default="default_user",
-                    is_shared=True,
-                ),
-                ConfigurableFieldSpec(
-                    id="conversation_id",
-                    annotation=str,
-                    name="Conversation ID",
-                    description="Unique identifier for the conversation.",
-                    default="default_session",
-                    is_shared=True,
-                ),
-            ],
-        ) | StrOutputParser()
-
-    classification_conversation_template = """
-        SYSTEM: You are message classificator. You classify message into classes:
-        - TAX: tax payment or other similar jura topic;
-        - CASUAL: casual conversation, any generic question that can't fit in the conversation;
-        - SPAM: any rude, useless messages.
-        Here is a question: {question}
-        Provide the answer as a JSON object that contains an object with key 'message_type' and with key 'system_message' with short remark about message and no other premable or explanation.
-    """
-
-    with ModelBuilderV2("openchat/openchat-7b") as decider_llm:
-        message_classificator = PromptTemplate.from_template(classification_conversation_template) | decider_llm | JsonOutputParser()
-
-    def retrieve(state):
-        """
-        Retrieve documents
-
-        Args:
-            state (dict): The current graph state
-
-        Returns:
-            state (dict): New key added to state, documents, that contains retrieved documents
-        """
-        print("---RETRIEVE---")
-        question = state["question"]
-
-        # Retrieval
-        with FAISSBuilder() as faiss_retriever:
-            with ModelBuilderV2("openchat/openchat-7b") as mq_llm:
-                retriever = MultiQueryRetriever.from_llm(retriever=faiss_retriever, llm=mq_llm)
-                documents = retriever.get_relevant_documents(question)
-                return {"documents": documents, "question": question}
-
-
-    def start_point(state):
-        """
-        Start point, just return state
-
-        Args:
-            state (dict): The current graph state 
-
-        Returns:
-            state (dict): The current graph state
-        """
-        return state
-
-    def casual_chat(state):
-        """
-        Define type of message
-
-        Args:
-            state (dict): The current graph state
-
-        Returns:
-            state (dict): New key added to state, generation, that contains message with casual answer
-        """
-        question = state["question"]
-        print("---CASUAL CHAT---")
-        generation = casual_llm.invoke({"question": question}, config={"configurable": {"conversation_id": "default_session", "user_id": "deafault_user"}})
-        state['generation'] = generation
-
-        return state
-
-
-    def define_message_type(state):
-        """
-        Define type of message
-
-        Args:
-            state (dict): The current graph state
-
-        Returns:
-
-        """
-        print("---MESSAGE CLASSIFICATION---")
-        question = state["question"]
-        msg_type_obj = message_classificator.invoke({"question": question})
-        print(f"---MESSAGE TYPE: {msg_type_obj['message_type']} SYSTEM MESSAGE---\n {msg_type_obj['system_message']}")
-        msg_type = msg_type_obj['message_type']
-
-        if msg_type == "TAX":
-            return "retrieve"
-        # if msg_type == "":
-        return "casual_chat"
-        return "__end__"
-
-    def generate(state):
-        """
-        Generate answer
-
-        Args:
-            state (dict): The current graph state
-
-        Returns:
-            state (dict): New key added to state, generation, that contains LLM generation, based on documents
-        """
-        print("---GENERATE---")
-        question = state["question"]
-        documents = state["documents"]
-
-        # RAG generation
-        generation = rag_chain.invoke({"context": documents, "question": question}, config={"configurable": {"conversation_id": "default_session", "user_id": "deafault_user"}},)
-        return {"documents": documents, "question": question, "generation": generation}
-
-
-    def grade_documents(state):
-        """
-        Determines whether the retrieved documents are relevant to the question.
-
-        Args:
-            state (dict): The current graph state
-
-        Returns:
-            state (dict): Updates documents key with only filtered relevant documents
-        """
-
-        print("---CHECK DOCUMENT RELEVANCE TO QUESTION---")
-        question = state["question"]
-        documents = state["documents"]
-
-        # Score each doc
-        filtered_docs = []
-        count_docs = len(documents)
-        for ind_d in range(0, count_docs, 3):
-            d_1 = documents[ind_d] if ind_d < count_docs else None
-            d_2 = documents[ind_d + 1] if ind_d + 1 < count_docs else None
-            d_3 = documents[ind_d + 2] if ind_d + 2 < count_docs else None
-            scores = retrieval_grader_3_docs.invoke(
-                {"question": question, "document_1": d_1, "document_2": d_2, "document_3": d_3}
-            )
-            for j in range(len(scores)):
-                grade = scores[j]["score"]
-                if grade > 0.7:
-                    print(f"---GRADE: DOCUMENT RELEVANT--- GRADE: {grade}")
-                    filtered_docs.append(documents[ind_d + j])
-                else:
-                    print("---GRADE: DOCUMENT NOT RELEVANT---")
-        is_fuse = len(filtered_docs) / len(documents) <= 0.5
-            
-        return {"documents": filtered_docs, "question": question}
-    
-
-    def make_collapsable_source_message(doc: Dict):
-        file_path = doc.metadata.get("file_name", "")
-        file_name = file_path.replace(".pdf", "")
-        chapter_title = doc.metadata.get("chapter_title", None)
-        page_num = doc.metadata.get("first_page_num", None)
-        title = f"""
-            {file_name}
-            {f": {chapter_title} " if chapter_title is not None else ""}
-            {f"Стр. {page_num} " if page_num is not None else ""}
-        """.replace("\n", " ")
-        content = doc.page_content.replace("\n\n", "\n")
-
-        if page_num is None:
-            message = rf"""
-                <details>
-                <summary>{title}</summary>
-                {str(content)}
-                </details>
-            """
-        else:
-            base_url = "http://localhost:5000/sta"
-            url = f"{base_url}?file={file_path}&#page={page_num}&zoom=90&toolbar=0"
-            message = f"""
-                <a class="open_pdf" href='{url}' onclick="return openPdf('{url}')">{title}</a>
-            """
-
-        return message
-
-
-    def add_sources(state):
-        """
-        Determines whether the retrieved documents are relevant to the question.
-
-        Args:
-            state (dict): The current graph state
-
-        Returns:
-            state (dict): Add collapsable sources
-        """
-        question = state["question"]
-        documents = state["documents"]
-        generation = state["generation"]
-
-        sources_message = "<i></i>".join(map(make_collapsable_source_message, documents))
-        extended_generation_message = f"{generation} {sources_message}"
-        return {"documents": documents, "question": question, "generation": extended_generation_message}
-
-
-    ### Edges
-    workflow = StateGraph(GraphState)
-
-    # Define the nodes
-    workflow.add_node("start_point", start_point)
-    workflow.add_node("retrieve", retrieve)  # retrieve
-    workflow.add_node("grade_documents", grade_documents)  # grade documents
-    workflow.add_node("generate", generate)  # generate
-    workflow.add_node("casual_chat", casual_chat) # simple chat
-    workflow.add_node("add_sources", add_sources)
-    # Build graph
-    workflow.set_entry_point("start_point")
-    workflow.add_conditional_edges("start_point", define_message_type)
-    workflow.add_edge("retrieve", "grade_documents")
-    workflow.add_edge("grade_documents", "generate")
-    workflow.add_edge("generate", "add_sources")
-    workflow.add_edge("add_sources", END)
-    workflow.add_edge("casual_chat", END)
-
-    # Compile
-    app = workflow.compile()
+    app = build_workflow()
 
     pdf_open_js = """
     <script>
@@ -417,54 +40,48 @@ def deploy():
     </script>
     """
 
-    def print_source_documents(documents):
-        return "\n\n".join([f"Взято из файла: {doc.metadata['file_name']} \n Metadata: {doc.metadata}"  for doc in documents])
-
-
     dark_theme = DarkTheme()
     with gr.Blocks(head=pdf_open_js, fill_height=True, theme=dark_theme) as demo:
         with gr.Row():
             with gr.Column(scale=1):
-                chatbot_rag = gr.Chatbot(label=f"RAG: llama3 + документы", height=740, sanitize_html=False, show_copy_button=True)
-                chat_input = gr.MultimodalTextbox(interactive=True, file_types=None, placeholder="Введите сообщение...", show_label=False, scale=4)
+                chatbot_rag = gr.Chatbot(
+                    label=f"RAG: llama3 + документы",
+                    height=740,
+                    sanitize_html=False,
+                    show_copy_button=True,
+                )
+                chat_input = gr.MultimodalTextbox(
+                    interactive=True,
+                    file_types=None,
+                    placeholder="Введите сообщение...",
+                    show_label=False,
+                    scale=4,
+                )
             with gr.Column(scale=1.5):
-                pdf_output = gr.HTML("<iframe id='opener' width='100%' height='740px' src=''></iframe>")
+                pdf_output = gr.HTML(
+                    "<iframe id='opener' width='100%' height='740px' src=''></iframe>"
+                )
         # clear = gr.Button("Clear")
 
         def user_rag(history, message):
             if message["text"] is not None:
                 history.append((message["text"], None))
             return history, gr.update(value=None, interactive=False)
-    
 
         def bot_rag(history):
             result = app.invoke({"question": history[-1][0]})
-            form_answer = result["generation"].strip() 
+            form_answer = result["generation"].strip()
             history[-1][1] = form_answer
             return history
 
-        # def bot_llm(history):
-        #     result = casual_llm.invoke({"question": history[-1][0]}, config={"configurable": {"conversation_id": "default_session_standalone", "user_id": "deafault_user_standalone"}})
-        #     history[-1][1] = result.strip()
-        #     return history
-
-        chat_input.submit(user_rag,  [chatbot_rag, chat_input], [chatbot_rag, chat_input], queue=False).then(
-            bot_rag, chatbot_rag, chatbot_rag
-        ).then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])
-        # pdf_output.change(lambda x: gr.HTML(pdf_open_js), chatbot_rag, pdf_output, queue=False)
-
-        # chat_input.submit(user_llm, [chatbot_llm, chat_input], [chatbot_llm, chat_input], queue=False).then(
-        #     bot_llm, chatbot_llm, chatbot_llm
-        # ).then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])
-        # clear.click(lambda: None, None, chatbot_rag, queue=False)
-        # clear.click(lambda: None, None, chatbot_llm, queue=False)
+        chat_input.submit(
+            user_rag, [chatbot_rag, chat_input], [chatbot_rag, chat_input], queue=False
+        ).then(bot_rag, chatbot_rag, chatbot_rag).then(
+            lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input]
+        )
 
     demo.launch(share=True)
 
 
 if __name__ == "__main__":
-    # parser = argparse.ArgumentParser(description='Deploy llm chat')
-    # parser.add_argument('--model_name', metavar='M', type=str,
-    #                 help='model name as: openai/gpt-3.5-turbo')
-    
-    deploy()
+    deploy()

lib/embedding.py

@@ -2,13 +2,17 @@ from langchain_community.embeddings import HuggingFaceEmbeddings
 from dotenv import load_dotenv
 
 load_dotenv()
+
+
 def build_embedding(model_name: str):
-    embedding = HuggingFaceEmbeddings(model_name=model_name, \
-                                    #    model_kwargs={"device": "cuda"}, \
-                                        encode_kwargs={"normalize_embeddings": True})
+    embedding = HuggingFaceEmbeddings(
+        model_name=model_name,  #    model_kwargs={"device": "cuda"}, \
+        encode_kwargs={"normalize_embeddings": True},
+    )
     embedding.show_progress = True
     return embedding
 
+
 class EmbeddingBuilder:
     def __init__(self, model_name: str, device: str = "cpu"):
         self.model_name = model_name
@@ -16,12 +20,14 @@ class EmbeddingBuilder:
 
     def __enter__(self):
         # Initialize resources
-        embedding = HuggingFaceEmbeddings(model_name=self.model_name, \
-                                       model_kwargs={"device": self.device}, \
-                                        encode_kwargs={"normalize_embeddings": True})
+        embedding = HuggingFaceEmbeddings(
+            model_name=self.model_name,
+            model_kwargs={"device": self.device},
+            encode_kwargs={"normalize_embeddings": True},
+        )
         embedding.show_progress = True
         return embedding
 
     def __exit__(self, exc_type, exc_value, traceback):
         # Cleanup resources if necessary
-        pass
+        pass

lib/gradio_custom_theme.py

@@ -14,16 +14,12 @@ class DarkTheme(Base):
         spacing_size: sizes.Size | str = sizes.spacing_md,
         radius_size: sizes.Size | str = sizes.radius_md,
         text_size: sizes.Size | str = sizes.text_lg,
-        font: fonts.Font
-        | str
-        | Iterable[fonts.Font | str] = (
+        font: fonts.Font | str | Iterable[fonts.Font | str] = (
             fonts.GoogleFont("Quicksand"),
             "ui-sans-serif",
             "sans-serif",
         ),
-        font_mono: fonts.Font
-        | str
-        | Iterable[fonts.Font | str] = (
+        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
             fonts.GoogleFont("Roboto"),
             "ui-monospace",
             "monospace",

lib/graph.py

@@ -0,0 +1,279 @@
+from typing import Dict, List
+from typing_extensions import TypedDict
+
+from langchain_core.documents import Document
+
+from lib.model_builder import ModelBuilderV2
+from lib.vectorestores import FAISSBuilder
+from lib.model_builder import ModelBuilderV2
+from lib.vectorestores import FAISSBuilder
+from langchain.retrievers.multi_query import MultiQueryRetriever
+from lib.runnables import (
+    casual_llm,
+    retrieval_grader_3,
+    rag_chain,
+    message_classificator,
+)
+from langgraph.graph import END, StateGraph
+
+
+class GraphState(TypedDict):
+    """
+    Represents the state of our graph.
+
+    Attributes:
+        question: question
+        generation: LLM generation
+        web_search: whether to add search
+        documents: list of documents
+    """
+
+    question: str
+    generation: str
+    documents: List[Document]
+    filtered_documets: List[Document]
+    is_fuse: bool
+    count_regenerations: int
+
+
+def combine_vectors(vectors):
+    result = []
+    vec1_count = len(vectors["vector1"])
+    # vec2_count = len(vectors["vector2"])
+    for i in range(vec1_count):
+        if i < vec1_count:
+            result.append(vectors["vector1"][i])
+        # if i < vec2_count:
+        #     result.append(vectors['vector2'][i])
+    return result
+
+
+def retrieve(state):
+    """
+    Retrieve documents
+
+    Args:
+        state (dict): The current graph state
+
+    Returns:
+        state (dict): New key added to state, documents, that contains retrieved documents
+    """
+    print("---RETRIEVE---")
+    question = state["question"]
+
+    # Retrieval
+    with FAISSBuilder() as faiss_retriever:
+        with ModelBuilderV2("openchat/openchat-7b") as mq_llm:
+            retriever = MultiQueryRetriever.from_llm(
+                retriever=faiss_retriever, llm=mq_llm
+            )
+            documents = retriever.get_relevant_documents(question)
+            return {"documents": documents, "question": question}
+
+
+def start_point(state):
+    """
+    Start point, just return state
+
+    Args:
+        state (dict): The current graph state
+
+    Returns:
+        state (dict): The current graph state
+    """
+    return state
+
+
+def casual_chat(state):
+    """
+    Define type of message
+
+    Args:
+        state (dict): The current graph state
+
+    Returns:
+        state (dict): New key added to state, generation, that contains message with casual answer
+    """
+    question = state["question"]
+    print("---CASUAL CHAT---")
+    generation = casual_llm.invoke(
+        {"question": question},
+        config={
+            "configurable": {
+                "conversation_id": "default_session",
+                "user_id": "deafault_user",
+            }
+        },
+    )
+    state["generation"] = generation
+
+    return state
+
+
+def define_message_type(state):
+    """
+    Define type of message
+
+    Args:
+        state (dict): The current graph state
+
+    Returns:
+
+    """
+    print("---MESSAGE CLASSIFICATION---")
+    question = state["question"]
+    msg_type_obj = message_classificator.invoke({"question": question})
+    print(
+        f"---MESSAGE TYPE: {msg_type_obj['message_type']} SYSTEM MESSAGE---\n {msg_type_obj['system_message']}"
+    )
+    msg_type = msg_type_obj["message_type"]
+
+    if msg_type == "TAX":
+        return "retrieve"
+    # if msg_type == "":
+    return "casual_chat"
+    return "__end__"
+
+
+def generate(state):
+    """
+    Generate answer
+
+    Args:
+        state (dict): The current graph state
+
+    Returns:
+        state (dict): New key added to state, generation, that contains LLM generation, based on documents
+    """
+    print("---GENERATE---")
+    question = state["question"]
+    documents = state["documents"]
+
+    # RAG generation
+    generation = rag_chain.invoke(
+        {"context": documents, "question": question},
+        config={
+            "configurable": {
+                "conversation_id": "default_session",
+                "user_id": "deafault_user",
+            }
+        },
+    )
+    return {"documents": documents, "question": question, "generation": generation}
+
+
+def grade_documents(state):
+    """
+    Determines whether the retrieved documents are relevant to the question.
+
+    Args:
+        state (dict): The current graph state
+
+    Returns:
+        state (dict): Updates documents key with only filtered relevant documents
+    """
+
+    print("---CHECK DOCUMENT RELEVANCE TO QUESTION---")
+    question = state["question"]
+    documents = state["documents"]
+
+    # Score each doc
+    filtered_docs = []
+    count_docs = len(documents)
+    for ind_d in range(0, count_docs, 3):
+        d_1 = documents[ind_d] if ind_d < count_docs else None
+        d_2 = documents[ind_d + 1] if ind_d + 1 < count_docs else None
+        d_3 = documents[ind_d + 2] if ind_d + 2 < count_docs else None
+        scores = retrieval_grader_3.invoke(
+            {
+                "question": question,
+                "document_1": d_1,
+                "document_2": d_2,
+                "document_3": d_3,
+            }
+        )
+        for j in range(len(scores)):
+            grade = scores[j]["score"]
+            if grade > 0.7:
+                print(f"---GRADE: DOCUMENT RELEVANT--- GRADE: {grade}")
+                filtered_docs.append(documents[ind_d + j])
+            else:
+                print("---GRADE: DOCUMENT NOT RELEVANT---")
+    is_fuse = len(filtered_docs) / len(documents) <= 0.5
+
+    return {"documents": filtered_docs, "question": question}
+
+
+def make_collapsable_source_message(doc: Dict):
+    file_path = doc.metadata.get("file_name", "")
+    file_name = file_path.replace(".pdf", "")
+    chapter_title = doc.metadata.get("chapter_title", None)
+    page_num = doc.metadata.get("first_page_num", None)
+    title = f"""
+        {file_name}
+        {f": {chapter_title} " if chapter_title is not None else ""}
+        {f"Стр. {page_num} " if page_num is not None else ""}
+    """.replace(
+        "\n", " "
+    )
+    content = doc.page_content.replace("\n\n", "\n")
+
+    if page_num is None:
+        message = rf"""
+            <details>
+            <summary>{title}</summary>
+            {str(content)}
+            </details>
+        """
+    else:
+        base_url = "http://localhost:5000/sta"
+        url = f"{base_url}?file={file_path}&#page={page_num}&zoom=90&toolbar=0"
+        message = f"""
+            <a class="open_pdf" href='{url}' onclick="return openPdf('{url}')">{title}</a>
+        """
+
+    return message
+
+
+def add_sources(state):
+    """
+    Determines whether the retrieved documents are relevant to the question.
+
+    Args:
+        state (dict): The current graph state
+
+    Returns:
+        state (dict): Add collapsable sources
+    """
+    question = state["question"]
+    documents = state["documents"]
+    generation = state["generation"]
+
+    sources_message = "<i></i>".join(map(make_collapsable_source_message, documents))
+    extended_generation_message = f"{generation} {sources_message}"
+    return {
+        "documents": documents,
+        "question": question,
+        "generation": extended_generation_message,
+    }
+
+
+def build_workflow():
+    workflow = StateGraph(GraphState)
+
+    # Define the nodes
+    workflow.add_node("start_point", start_point)
+    workflow.add_node("retrieve", retrieve)  # retrieve
+    workflow.add_node("grade_documents", grade_documents)  # grade documents
+    workflow.add_node("generate", generate)  # generate
+    workflow.add_node("casual_chat", casual_chat)  # simple chat
+    workflow.add_node("add_sources", add_sources)
+    # Build graph
+    workflow.set_entry_point("start_point")
+    workflow.add_conditional_edges("start_point", define_message_type)
+    workflow.add_edge("retrieve", "grade_documents")
+    workflow.add_edge("grade_documents", "generate")
+    workflow.add_edge("generate", "add_sources")
+    workflow.add_edge("add_sources", END)
+    workflow.add_edge("casual_chat", END)
+    return workflow.compile()

lib/model_builder.py

@@ -3,16 +3,24 @@ from langchain_openai import ChatOpenAI
 from dotenv import load_dotenv
 
 load_dotenv()
-VSEGPT_KEY = os.getenv('VSEGPT_KEY')
-OPENAI_BASE = os.getenv('OPENAI_BASE')
+VSEGPT_KEY = os.getenv("VSEGPT_KEY")
+OPENAI_BASE = os.getenv("OPENAI_BASE")
+
 
 class ModelBuilder:
     def createVseGptModel(model, temperature):
-        return ChatOpenAI(temperature=temperature, model_name=model, \
-                          api_key=VSEGPT_KEY, base_url = OPENAI_BASE)
+        return ChatOpenAI(
+            temperature=temperature,
+            model_name=model,
+            api_key=VSEGPT_KEY,
+            base_url=OPENAI_BASE,
+        )
+
 
 class ModelBuilderV2:
-    def __init__(self, model_name: str, temperature=0, api_key=VSEGPT_KEY, base_url=OPENAI_BASE):
+    def __init__(
+        self, model_name: str, temperature=0, api_key=VSEGPT_KEY, base_url=OPENAI_BASE
+    ):
         self.model_name = model_name
         self.temperature = temperature
         self.api_key = api_key
@@ -20,9 +28,13 @@ class ModelBuilderV2:
 
     def __enter__(self):
         # Initialize resources
-        return ChatOpenAI(temperature=self.temperature, model_name=self.model_name, \
-                          api_key=VSEGPT_KEY, base_url = OPENAI_BASE)
+        return ChatOpenAI(
+            temperature=self.temperature,
+            model_name=self.model_name,
+            api_key=VSEGPT_KEY,
+            base_url=OPENAI_BASE,
+        )
 
     def __exit__(self, exc_type, exc_value, traceback):
         # Cleanup resources if necessary
-        pass
+        pass

lib/prompts.py

@@ -0,0 +1,61 @@
+from langchain_core.prompts import PromptTemplate
+
+
+casual_prompt = PromptTemplate(
+    template="""Just answer a question as casual chatter. \n 
+    Here is the user question: {question} \n
+    Always reply in Russian. Leave clear answer, without any addition.
+    Chat history: 
+    {chat_history}
+    Answer:
+    """,
+    input_variables=["question", "chat_history"],
+)
+
+grader_3_doc_prompt = PromptTemplate(
+    template="""You are a grader assessing relevance of a retrieved documents to a user question. \n 
+        Here is the first retrieved document: \n\n {document_1} \n\n
+        Here is the second retrieved document: \n\n {document_2} \n\n
+        Here is the third retrieved document: \n\n {document_3} \n\n
+        Here is the user question: {question} \n
+        If the document contains keywords related to the user question, grade it as relevant. \n
+        It does not need to be a stringent test. The goal is to filter out erroneous retrievals. \n
+        For each! document give a score from 0 to 1, score to indicate whether the document is relevant to the question. \n
+        Provide the scores as a JSON list that contains an objects with single key 'score' and no premable or explanation.""",
+    input_variables=["question", "document_1", "document_2", "document_3"],
+)
+
+rag_assistant_prompt = PromptTemplate(
+    template="""
+        SYSTEM: You are an assistant for question-answering tasks. 
+        Use the following pieces of retrieved context to answer the question.
+        Use previous messages then current message higly likely  
+        If you don't find the answer in the context, transform the question ans ask the user to specify his qusetion. 
+
+        Keep the answer concise.  
+        Print a most possible topic of conversation.
+        Always reply in Russian, all text must be in Russian!
+
+        Context: {context}
+
+        Previous messages: {chat_history}
+
+        Question: {question}
+
+        Answer:
+        
+        """,
+    input_variables=["context", "chat_history", "question"],
+)
+
+classificator_question_prompt = PromptTemplate(
+    template="""
+        SYSTEM: You are message classificator. You classify message into classes:
+        - TAX: tax payment or other similar jura topic;
+        - CASUAL: casual conversation, any generic question that can't fit in the conversation;
+        - SPAM: any rude, useless messages.
+        Here is a question: {question}
+        Provide the answer as a JSON object that contains an object with key 'message_type' and with key 'system_message' with short remark about message and no other premable or explanation.
+    """,
+    input_variables=["question"],
+)

lib/runnables.py

@@ -0,0 +1,114 @@
+import contextlib
+
+from lib.model_builder import ModelBuilderV2
+from lib.prompts import (
+    casual_prompt,
+    grader_3_doc_prompt,
+    rag_assistant_prompt,
+    classificator_question_prompt,
+)
+from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
+from langchain_core.runnables.history import RunnableWithMessageHistory
+from langchain_core.chat_history import (
+    BaseChatMessageHistory,
+    InMemoryChatMessageHistory,
+)
+from langchain_core.runnables import ConfigurableFieldSpec
+
+store = {}
+
+
+def get_session_history(user_id: str, conversation_id: str) -> BaseChatMessageHistory:
+    if (user_id, conversation_id) not in store:
+        store[(user_id, conversation_id)] = InMemoryChatMessageHistory()
+    return store[(user_id, conversation_id)]
+
+
+class ModelConfig:
+    def __init__(self, model_name, temperature=0.7):
+        self.model_name = model_name
+        self.temperature = temperature
+
+
+class ConfigField:
+    def __init__(self, id, annotation, name, description, default, is_shared):
+        self.id = id
+        self.annotation = annotation
+        self.name = name
+        self.description = description
+        self.default = default
+        self.is_shared = is_shared
+
+
+USER_ID_FIELD = ConfigurableFieldSpec(
+    id="user_id",
+    annotation=str,
+    name="User ID",
+    description="Unique identifier for the user.",
+    default="default_user",
+    is_shared=True,
+)
+
+CONVERSATION_ID_FIELD = ConfigurableFieldSpec(
+    id="conversation_id",
+    annotation=str,
+    name="Conversation ID",
+    description="Unique identifier for the conversation.",
+    default="default_session",
+    is_shared=True,
+)
+
+
+def create_runnable_with_history(
+    prompt, llm, input_messages_key, history_messages_key, history_factory_config
+):
+    return RunnableWithMessageHistory(
+        prompt | llm,
+        get_session_history,
+        input_messages_key=input_messages_key,
+        history_messages_key=history_messages_key,
+        history_factory_config=history_factory_config,
+    )
+
+
+@contextlib.contextmanager
+def create_model_builder(config):
+    with ModelBuilderV2(config.model_name, config.temperature) as llm:
+        yield llm
+    # try:
+    #     yield llm
+    # finally:
+    #     llm.release()  # Assuming ModelBuilderV2 has a release method to clear resources
+
+
+casual_config = ModelConfig("openchat/openchat-7b", 0.7)
+retrieval_config = ModelConfig("cohere/command-r")
+rag_config = ModelConfig("mistralai/mixtral-8x22b-instruct")
+classificator_msg_config = ModelConfig("openchat/openchat-7b")
+
+history_config = [USER_ID_FIELD, CONVERSATION_ID_FIELD]
+
+with create_model_builder(casual_config) as llm:
+    casual_llm = (
+        create_runnable_with_history(
+            casual_prompt, llm, "question", "chat_history", history_config
+        )
+        | StrOutputParser()
+    )
+
+with create_model_builder(retrieval_config) as llm:
+    retrieval_grader_3 = grader_3_doc_prompt | llm | JsonOutputParser()
+
+with create_model_builder(rag_config) as llm:
+    rag_chain = (
+        create_runnable_with_history(
+            rag_assistant_prompt, llm, "question", "chat_history", history_config
+        )
+        | StrOutputParser()
+    )
+
+
+with create_model_builder(classificator_msg_config) as decider_llm:
+    message_classificator = (
+        classificator_question_prompt | decider_llm | JsonOutputParser()
+    )

lib/vectorestores.py

@@ -0,0 +1,40 @@
+from langchain_pinecone import PineconeVectorStore
+from lib.embedding import EmbeddingBuilder
+from langchain_community.vectorstores import FAISS
+
+
+class FAISSBuilder:
+    def __init__(
+        self,
+        embedding_model: str = "intfloat/multilingual-e5-large",
+        local_path: str = "data/faiss_nk_28_05",
+    ):
+        self.embedding_model = embedding_model
+        self.local_path = local_path
+
+    def __enter__(self):
+        # Initialize resources
+        with EmbeddingBuilder(self.embedding_model) as rag_emb:
+            faiss_db = FAISS.load_local(
+                self.local_path, rag_emb, allow_dangerous_deserialization=True
+            )
+            return faiss_db.as_retriever()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # Cleanup resources if necessary
+        pass
+
+
+class PineConeBuilder:
+    def __init__(self, index_name, embedding_model):
+        self.index_name = index_name
+        self.embedding_model = embedding_model
+
+    def __enter__(self):
+        with EmbeddingBuilder(self.embedding_model) as embeddings:
+            pc_db = PineconeVectorStore.from_existing_index(self.index_name, embeddings)
+            return pc_db.as_retriever()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # Cleanup resources if necessary
+        pass