refactor: code refactoring

app.py

@@ -1,9 +1,9 @@
 import chainlit as cl
-from langchain.callbacks.base import BaseCallbackHandler
-from langchain.chains.query_constructor.schema import AttributeInfo
-from langchain.retrievers.self_query.base import SelfQueryRetriever
+from langchain.retrievers import ParentDocumentRetriever
 from langchain.schema import StrOutputParser
 from langchain.schema.runnable import Runnable, RunnableConfig, RunnablePassthrough
+from langchain.storage import InMemoryStore
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores.chroma import Chroma
 from langchain_google_genai import (
     GoogleGenerativeAI,
@@ -11,12 +11,10 @@ from langchain_google_genai import (
     HarmBlockThreshold,
     HarmCategory,
 )
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.retrievers import ParentDocumentRetriever
-from langchain.storage import InMemoryStore
+
 import config
 from prompts import prompt
-import tiktoken
+from utils import PostMessageHandler, format_docs
 
 model = GoogleGenerativeAI(
     model=config.GOOGLE_CHAT_MODEL,
@@ -34,34 +32,20 @@ embedding = embeddings_model = GoogleGenerativeAIEmbeddings(
 vectordb = Chroma(persist_directory=config.STORAGE_PATH, embedding_function=embedding)
 
 ## retriever
-
 text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, separators=["\n"])
 
-
 # The storage layer for the parent documents
 store = InMemoryStore()
 retriever = ParentDocumentRetriever(
     vectorstore=vectordb,
     docstore=store,
     child_splitter=text_splitter,
-) 
-
+)
 
 
 @cl.on_chat_start
 async def on_chat_start():
 
-    def format_docs(documents, max_context_size= 100000, separator= "\n\n"):
-        context = ""
-        encoder = tiktoken.get_encoding("cl100k_base")
-        i=0
-        for doc in documents:
-            i+=1
-            if len(encoder.encode(context)) < max_context_size:
-                source = doc.metadata['link']
-                context += f"Article{i}:\n"+doc.page_content + f"\nSource: {source}" + separator
-        return context
-
     rag_chain = (
         {
             "context": retriever | format_docs,
@@ -71,7 +55,6 @@ async def on_chat_start():
         | model
         | StrOutputParser()
     )
-    
 
     cl.user_session.set("rag_chain", rag_chain)
 
@@ -86,33 +69,6 @@ async def on_message(message: cl.Message):
     runnable = cl.user_session.get("rag_chain")  # type: Runnable # type: ignore
     msg = cl.Message(content="")
 
-    class PostMessageHandler(BaseCallbackHandler):
-        """
-        Callback handler for handling the retriever and LLM processes.
-        Used to post the sources of the retrieved documents as a Chainlit element.
-        """
-
-        def __init__(self, msg: cl.Message):
-            BaseCallbackHandler.__init__(self)
-            self.msg = msg
-            self.sources = []
-
-        def on_retriever_end(self, documents, *, run_id, parent_run_id, **kwargs):
-            for d in documents:
-                source_doc = d.page_content + "\nSource: " + d.metadata["link"]
-                self.sources.append(source_doc)
-
-        def on_llm_end(self, response, *, run_id, parent_run_id, **kwargs):
-            if len(self.sources):
-                # Display the reference docs with a Text widget
-                sources_element = [
-                    cl.Text(name=f"source_{idx+1}", content=content)
-                    for idx, content in enumerate(self.sources)
-                ]
-                source_names = [el.name for el in sources_element]
-                self.msg.elements += sources_element
-                self.msg.content += f"\nSources: {', '.join(source_names)}"
-
     async with cl.Step(type="run", name="QA Assistant"):
         async for chunk in runnable.astream(
             message.content,

scrape_data.py

@@ -120,7 +120,7 @@ def process_docs(
         documents=splits,
         embedding=embeddings_model,
         persist_directory=persist_directory,
-    ) 
+    )
 
     return doc_search
 

utils.py

@@ -0,0 +1,45 @@
+import chainlit as cl
+import tiktoken
+from langchain.callbacks.base import BaseCallbackHandler
+
+
+def format_docs(documents, max_context_size=100000, separator="\n\n"):
+    context = ""
+    encoder = tiktoken.get_encoding("cl100k_base")
+    i = 0
+    for doc in documents:
+        i += 1
+        if len(encoder.encode(context)) < max_context_size:
+            source = doc.metadata["link"]
+            context += (
+                f"Article{i}:\n" + doc.page_content + f"\nSource: {source}" + separator
+            )
+    return context
+
+
+class PostMessageHandler(BaseCallbackHandler):
+    """
+    Callback handler for handling the retriever and LLM processes.
+    Used to post the sources of the retrieved documents as a Chainlit element.
+    """
+
+    def __init__(self, msg: cl.Message):
+        BaseCallbackHandler.__init__(self)
+        self.msg = msg
+        self.sources = []
+
+    def on_retriever_end(self, documents, *, run_id, parent_run_id, **kwargs):
+        for d in documents:
+            source_doc = d.page_content + "\nSource: " + d.metadata["link"]
+            self.sources.append(source_doc)
+
+    def on_llm_end(self, response, *, run_id, parent_run_id, **kwargs):
+        if len(self.sources):
+            # Display the reference docs with a Text widget
+            sources_element = [
+                cl.Text(name=f"source_{idx+1}", content=content)
+                for idx, content in enumerate(self.sources)
+            ]
+            source_names = [el.name for el in sources_element]
+            self.msg.elements += sources_element
+            self.msg.content += f"\nSources: {', '.join(source_names)}"