Integration of MongoDB for feedbacks - Revised

.gitignore

@@ -31,3 +31,6 @@ creds.json
 temp_test_run_data.json
 **/llm_tests_output.txt
 **/error_log.txt
+
+# backend
+faiss_db

backend/.env.example

@@ -35,4 +35,7 @@ TOKENIZERS_PARALLELISM=false
 LOGLEVEL="INFO"
 
 BACKEND_WORKERS=4
-BACKEND_URL="0.0.0.0"
+BACKEND_URL="0.0.0.0"
+
+# Set FAST_MODE=true for fast prototyping
+FAST_MODE=false

backend/src/agents/retriever_graph.py

@@ -1,6 +1,6 @@
 import os
 import logging
-from typing import TypedDict, Annotated, Union, Optional
+from typing import Any, TypedDict, Annotated, Union, Optional
 
 from langchain_core.messages import AnyMessage
 from langchain_core.output_parsers import JsonOutputParser
@@ -29,6 +29,7 @@
 class AgentState(TypedDict):
     messages: Annotated[list[AnyMessage], add_messages]
     context: Annotated[list[AnyMessage], add_messages]
+    context_list: Annotated[list[str], add_messages]
     tools: list[str]
     sources: Annotated[list[str], add_messages]
     urls: Annotated[list[str], add_messages]
@@ -39,12 +40,12 @@ class ToolNode:
     def __init__(self, tool_fn: BaseTool) -> None:
         self.tool_fn = tool_fn
 
-    def get_node(self, state: AgentState) -> dict[str, list[str]]:
+    def get_node(self, state: AgentState) -> dict[str, Any]:
         query = state["messages"][-1].content
         if query is None:
             raise ValueError("Query is None")
 
-        response, sources, urls = self.tool_fn.invoke(query)  # type: ignore
+        response, sources, urls, context_list = self.tool_fn.invoke(query)  # type: ignore
 
         if response != []:
             response = (
@@ -64,7 +65,12 @@ def get_node(self, state: AgentState) -> dict[str, list[str]]:
                 if isinstance(urls[0], list)
                 else urls
             )
-        return {"context": response, "sources": sources, "urls": urls}
+        return {
+            "context": response,
+            "sources": sources,
+            "urls": urls,
+            "context_list": context_list,
+        }
 
 
 class RetrieverGraph:
@@ -75,13 +81,15 @@ def __init__(
         reranking_model_name: str,
         inbuilt_tool_calling: bool,
         use_cuda: bool = False,
+        fast_mode: bool = False,
     ):
         self.llm = llm_model
         self.retriever_tools: RetrieverTools = RetrieverTools()
         self.retriever_tools.initialize(
             embeddings_config=embeddings_config,
             reranking_model_name=reranking_model_name,
             use_cuda=use_cuda,
+            fast_mode=fast_mode,
         )
 
         self.tools = [
@@ -178,9 +186,10 @@ def agent(self, state: AgentState) -> dict[str, list[str]]:
 
             return {"tools": tool_calls}
 
-    def generate(self, state: AgentState) -> dict[str, list[AnyMessage]]:
+    def generate(self, state: AgentState) -> dict[str, Any]:
         query = state["messages"][-1].content
         context = state["context"][-1].content
+        print("state keys", state.keys())
 
         ans = self.llm_chain.invoke({"context": context, "question": query})
 

backend/src/agents/retriever_tools.py

@@ -44,21 +44,58 @@ def initialize(
         embeddings_config: dict[str, str],
         reranking_model_name: str,
         use_cuda: bool = False,
+        fast_mode: bool = False,
     ) -> None:
-        general_retriever_chain = HybridRetrieverChain(
-            embeddings_config=embeddings_config,
-            reranking_model_name=reranking_model_name,
-            use_cuda=use_cuda,
-            html_docs_path=["./data/html/or_website/"],
-            markdown_docs_path=[
+        markdown_docs_map = {
+            "general": [
                 "./data/markdown/OR_docs",
                 "./data/markdown/ORFS_docs",
                 "./data/markdown/gh_discussions",
                 "./data/markdown/manpages/man1",
                 "./data/markdown/manpages/man2",
                 "./data/markdown/OpenSTA_docs",
             ],
-            other_docs_path=["./data/pdf"],
+            "install": [
+                "./data/markdown/ORFS_docs/installation",
+                "./data/markdown/OR_docs/installation",
+                "./data/markdown/gh_discussions/Build",
+                "./data/markdown/gh_discussions/Installation",
+                "./data/markdown/OpenSTA_docs",
+            ],
+            "commands": [
+                "./data/markdown/OR_docs/tools",
+                "./data/markdown/ORFS_docs/general",
+                "./data/markdown/gh_discussions/Query",
+                "./data/markdown/gh_discussions/Runtime",
+                "./data/markdown/gh_discussions/Documentation",
+                "./data/markdown/manpages/man1",
+                "./data/markdown/manpages/man2",
+                "./data/markdown/OpenSTA_docs",
+            ],
+            "errinfo": [
+                "./data/markdown/manpages/man3",
+                "./data/markdown/gh_discussions/Bug",
+            ],
+        }
+        fastmode_docs_map = {
+            "general": [markdown_docs_map["general"][0]],
+            "install": [markdown_docs_map["install"][0]],
+            "commands": [markdown_docs_map["commands"][0]],
+            "errinfo": [markdown_docs_map["errinfo"][1]],
+            "yosys": [
+                "./data/html/yosys_docs/yosyshq.readthedocs.io/projects/yosys/en/latest/getting_started"
+            ],
+            "klayout": ["./data/html/klayout_docs/www.klayout.de/examples"],
+        }
+        general_retriever_chain = HybridRetrieverChain(
+            embeddings_config=embeddings_config,
+            reranking_model_name=reranking_model_name,
+            use_cuda=use_cuda,
+            html_docs_path=[] if fast_mode else ["./data/html/or_website/"],
+            markdown_docs_path=fastmode_docs_map["general"]
+            if fast_mode
+            else markdown_docs_map["general"],
+            other_docs_path=[] if fast_mode else ["./data/pdf"],
             weights=[0.6, 0.2, 0.2],
             contextual_rerank=True,
             search_k=search_k,
@@ -71,13 +108,9 @@ def initialize(
             embeddings_config=embeddings_config,
             reranking_model_name=reranking_model_name,
             use_cuda=use_cuda,
-            markdown_docs_path=[
-                "./data/markdown/ORFS_docs/installation",
-                "./data/markdown/OR_docs/installation",
-                "./data/markdown/gh_discussions/Build",
-                "./data/markdown/gh_discussions/Installation",
-                "./data/markdown/OpenSTA_docs/",
-            ],
+            markdown_docs_path=fastmode_docs_map["install"]
+            if fast_mode
+            else markdown_docs_map["install"],
             weights=[0.6, 0.2, 0.2],
             contextual_rerank=True,
             search_k=search_k,
@@ -90,17 +123,10 @@ def initialize(
             embeddings_config=embeddings_config,
             reranking_model_name=reranking_model_name,
             use_cuda=use_cuda,
-            markdown_docs_path=[
-                "./data/markdown/OR_docs/tools",
-                "./data/markdown/ORFS_docs/general",
-                "./data/markdown/gh_discussions/Query",
-                "./data/markdown/gh_discussions/Runtime",
-                "./data/markdown/gh_discussions/Documentation",
-                "./data/markdown/manpages/man1",
-                "./data/markdown/manpages/man2",
-                "./data/markdown/OpenSTA_docs",
-            ],
-            other_docs_path=["./data/pdf"],
+            markdown_docs_path=fastmode_docs_map["commands"]
+            if fast_mode
+            else markdown_docs_map["commands"],
+            other_docs_path=[] if fast_mode else ["./data/pdf"],
             weights=[0.6, 0.2, 0.2],
             contextual_rerank=True,
             search_k=search_k,
@@ -113,7 +139,9 @@ def initialize(
             embeddings_config=embeddings_config,
             reranking_model_name=reranking_model_name,
             use_cuda=use_cuda,
-            html_docs_path=["./data/html/yosys_docs"],
+            html_docs_path=fastmode_docs_map["yosys"]
+            if fast_mode
+            else ["./data/html/yosys_docs"],
             weights=[0.6, 0.2, 0.2],
             contextual_rerank=True,
             search_k=search_k,
@@ -126,7 +154,9 @@ def initialize(
             embeddings_config=embeddings_config,
             reranking_model_name=reranking_model_name,
             use_cuda=use_cuda,
-            html_docs_path=["./data/html/klayout_docs"],
+            html_docs_path=fastmode_docs_map["klayout"]
+            if fast_mode
+            else ["./data/html/klayout_docs"],
             weights=[0.6, 0.2, 0.2],
             contextual_rerank=True,
             search_k=search_k,
@@ -139,10 +169,9 @@ def initialize(
             embeddings_config=embeddings_config,
             reranking_model_name=reranking_model_name,
             use_cuda=use_cuda,
-            markdown_docs_path=[
-                "./data/markdown/manpages/man3",
-                "./data/markdown/gh_discussions/Bug",
-            ],
+            markdown_docs_path=fastmode_docs_map["errinfo"]
+            if fast_mode
+            else markdown_docs_map["errinfo"],
             weights=[0.6, 0.2, 0.2],
             contextual_rerank=True,
             search_k=search_k,
@@ -153,7 +182,7 @@ def initialize(
 
     @staticmethod
     @tool
-    def retrieve_general(query: str) -> Tuple[str, list[str], list[str]]:
+    def retrieve_general(query: str) -> Tuple[str, list[str], list[str], list[str]]:
         """
         Retrieve comprehensive and detailed information pertaining to the OpenROAD project, OpenROAD-Flow-Scripts and OpenSTA.\
         This includes, but is not limited to, general information, specific functionalities, usage guidelines,\
@@ -168,7 +197,7 @@ def retrieve_general(query: str) -> Tuple[str, list[str], list[str]]:
 
     @staticmethod
     @tool
-    def retrieve_cmds(query: str) -> Tuple[str, list[str], list[str]]:
+    def retrieve_cmds(query: str) -> Tuple[str, list[str], list[str], list[str]]:
         """
         Retrieve information on the commands available in OpenROAD, OpenROAD-Flow-Scripts and OpenSTA.\
         This includes usage guidelines, command syntax, examples, and best practices about commands that cover various \
@@ -197,7 +226,7 @@ def retrieve_cmds(query: str) -> Tuple[str, list[str], list[str]]:
 
     @staticmethod
     @tool
-    def retrieve_install(query: str) -> Tuple[str, list[str], list[str]]:
+    def retrieve_install(query: str) -> Tuple[str, list[str], list[str], list[str]]:
         """
         Retrieve comprehensive and detailed information pertaining to the installaion of OpenROAD, OpenROAD-Flow-Scripts and OpenSTA.\
         This includes, but is not limited to, various dependencies, system requirements, installation methods such as,\
@@ -213,7 +242,7 @@ def retrieve_install(query: str) -> Tuple[str, list[str], list[str]]:
 
     @staticmethod
     @tool
-    def retrieve_errinfo(query: str) -> Tuple[str, list[str], list[str]]:
+    def retrieve_errinfo(query: str) -> Tuple[str, list[str], list[str], list[str]]:
         """
         Retrieve descriptions and details regarding the various warning/error messages encountered while using the OpenROAD.\
         An error code usually is identified by the tool, followed by a number.\
@@ -228,7 +257,9 @@ def retrieve_errinfo(query: str) -> Tuple[str, list[str], list[str]]:
 
     @staticmethod
     @tool
-    def retrieve_yosys_rtdocs(query: str) -> Tuple[str, list[str], list[str]]:
+    def retrieve_yosys_rtdocs(
+        query: str,
+    ) -> Tuple[str, list[str], list[str], list[str]]:
         """
         Retrieve detailed information regarding the Yosys application.\
         This tool provides information pertaining to the installation, usage, and troubleshooting of Yosys.\
@@ -248,7 +279,9 @@ def retrieve_yosys_rtdocs(query: str) -> Tuple[str, list[str], list[str]]:
 
     @staticmethod
     @tool
-    def retrieve_klayout_docs(query: str) -> Tuple[str, list[str], list[str]]:
+    def retrieve_klayout_docs(
+        query: str,
+    ) -> Tuple[str, list[str], list[str], list[str]]:
         """
         Retrieve detailed information regarding the KLayout application.\
         This tool provides information pertaining to the installation, usage, and troubleshooting of KLayout.\

backend/src/api/models/response_model.py

@@ -8,10 +8,14 @@ class UserInput(BaseModel):
     list_context: bool = False
 
 
+class ContextSource(BaseModel):
+    source: str = ""
+    context: str = ""
+
+
 class ChatResponse(BaseModel):
     response: str
-    sources: list[str] = []
-    context: list[str] = []
+    context_sources: list[ContextSource] = []
     tools: list[str] = []