fix lints, redirect mock_endpoint to agent-retriever

Signed-off-by: Jack Luar <[email protected]>

backend/src/api/models/response_model.py

@@ -9,7 +9,7 @@ class UserInput(BaseModel):
 
 
 class ContextSource(BaseModel):
-    sources: str = ""
+    source: str = ""
     context: str = ""
 
 

backend/src/api/routers/chains.py

@@ -114,30 +114,29 @@ async def get_hybrid_response(user_input: UserInput) -> ChatResponse:
     context_sources = []
     for i in result["context"]:
         if "url" in i.metadata:
-            context_sources.append(ContextSource(
-                context=i.page_content,
-                source=i.metadata["url"]
-            ))
+            context_sources.append(
+                ContextSource(context=i.page_content, source=i.metadata["url"])
+            )
         elif "source" in i.metadata:
-            context_sources.append(ContextSource(
-                context=i.page_content,
-                source=i.metadata["source"]
-            ))
+            context_sources.append(
+                ContextSource(context=i.page_content, source=i.metadata["source"])
+            )
 
     if user_input.list_sources and user_input.list_context:
-        response = {
-        "response": result["answer"],
-        "context_sources": context_sources
-        }
+        response = {"response": result["answer"], "context_sources": context_sources}
     elif user_input.list_sources:
         response = {
-        "response": result["answer"],
-        "context_sources": [ContextSource(context="", source=cs.source) for cs in context_sources]
+            "response": result["answer"],
+            "context_sources": [
+                ContextSource(context="", source=cs.source) for cs in context_sources
+            ],
         }
     elif user_input.list_context:
         response = {
             "response": result["answer"],
-            "context_sources": [ContextSource(context=cs.context, source="") for cs in context_sources]
+            "context_sources": [
+                ContextSource(context=cs.context, source="") for cs in context_sources
+            ],
         }
     else:
         response = {"response": result["answer"], "context_sources": []}
@@ -169,30 +168,29 @@ async def get_sim_response(user_input: UserInput) -> ChatResponse:
     context_sources = []
     for i in result["context"]:
         if "url" in i.metadata:
-            context_sources.append(ContextSource(
-                context=i.page_content,
-                source=i.metadata["url"]
-            ))
+            context_sources.append(
+                ContextSource(context=i.page_content, source=i.metadata["url"])
+            )
         elif "source" in i.metadata:
-            context_sources.append(ContextSource(
-                context=i.page_content,
-                source=i.metadata["source"]
-            ))
+            context_sources.append(
+                ContextSource(context=i.page_content, source=i.metadata["source"])
+            )
 
     if user_input.list_sources and user_input.list_context:
-        response = {
-            "response": result["answer"],
-            "context_sources": context_sources
-        }
+        response = {"response": result["answer"], "context_sources": context_sources}
     elif user_input.list_sources:
         response = {
-        "response": result["answer"],
-        "context_sources": [ContextSource(context="", source=cs.source) for cs in context_sources]
+            "response": result["answer"],
+            "context_sources": [
+                ContextSource(context="", source=cs.source) for cs in context_sources
+            ],
         }
     elif user_input.list_context:
         response = {
-        "response": result["answer"],
-        "context_sources": [ContextSource(context=cs.context, source="") for cs in context_sources]
+            "response": result["answer"],
+            "context_sources": [
+                ContextSource(context=cs.context, source="") for cs in context_sources
+            ],
         }
     else:
         response = {"response": result["answer"], "context_sources": []}

backend/src/api/routers/graphs.py

@@ -108,6 +108,7 @@ async def get_agent_response(user_input: UserInput) -> ChatResponse:
         raise ValueError("RetrieverGraph not initialized.")
     urls: list[str] = []
     context: list[str] = []
+    context_sources: list[ContextSource] = []
 
     if (
         isinstance(output, list)
@@ -121,14 +122,12 @@ async def get_agent_response(user_input: UserInput) -> ChatResponse:
 
         context_sources = []
         tool_index = 1
-        for tool in tools:
-            urls.extend(list(output[tool_index].values())[0]["urls"])
-            context.append(list(output[tool_index].values())[0]["context"])
-            tool_index += 1
-
-            for url, context in zip(urls, [context]):
-                context_sources.append(ContextSource(context=context, source=url))
-            tool_index += 1
+        for tool_index, tool in enumerate(tools):
+            urls = list(output[tool_index].values())[0]["urls"]
+            context = list(output[tool_index].values())[0]["context"]
+
+            for _url, _context in zip(urls, context):
+                context_sources.append(ContextSource(context=_context, source=_url))
     else:
         llm_response = "LLM response extraction failed"
         logging.error("LLM response extraction failed")
@@ -140,11 +139,27 @@ async def get_agent_response(user_input: UserInput) -> ChatResponse:
             "tool": tools,
         }
     elif user_input.list_sources:
-        response = {"response": llm_response, "context_sources": [ContextSource(context="", source=cs.source) for cs in context_sources], "tool": tools}
+        response = {
+            "response": llm_response,
+            "context_sources": [
+                ContextSource(context="", source=cs.source) for cs in context_sources
+            ],
+            "tool": tools,
+        }
     elif user_input.list_context:
-        response = {"response": llm_response, "context_sources": [ContextSource(context=cs.context, source="") for cs in context_sources], "tool": tools}
+        response = {
+            "response": llm_response,
+            "context_sources": [
+                ContextSource(context=cs.context, source="") for cs in context_sources
+            ],
+            "tool": tools,
+        }
     else:
-        response = {"response": llm_response, "context_sources": [ContextSource(context="", source="")], "tool": tools}
+        response = {
+            "response": llm_response,
+            "context_sources": [ContextSource(context="", source="")],
+            "tool": tools,
+        }
 
     return ChatResponse(**response)
 

frontend/streamlit_app.py

@@ -44,7 +44,6 @@ def translate_chat_history_to_api(chat_history, max_pairs=4):
 
 
 def display_sources_context(context_sources: list[dict[str, str]]):
-
     with st.expander("Sources and Context"):
         try:
             if context_sources:
@@ -91,7 +90,6 @@ def response_generator(user_input: str) -> tuple[str, str] | tuple[None, None]:
         context_sources = data.get("context_sources", [])
         st.session_state.metadata[user_input] = {
             "context_sources": context_sources,
-
         }
         return data.get("response", ""), context_sources
     except requests.exceptions.RequestException as e:
@@ -144,7 +142,9 @@ def main() -> None:
             user_message = st.session_state.chat_history[idx - 1]
             if user_message["role"] == "user":
                 user_input = user_message["content"]
-                context_sources = st.session_state.metadata.get(user_input, {}).get("context_sources", [])
+                context_sources = st.session_state.metadata.get(user_input, {}).get(
+                    "context_sources", []
+                )
                 display_sources_context(context_sources)
 
     user_input = st.chat_input("Enter your queries ...")
@@ -223,8 +223,12 @@ def update_state() -> None:
         # Handle thumbs up and thumbs down reactions
         if thumbs_up or thumbs_down:
             try:
-                selected_question = st.session_state.chat_history[-2]["content"]  # Last user question
-                gen_ans = st.session_state.chat_history[-1]["content"]  # Last AI response
+                selected_question = st.session_state.chat_history[-2][
+                    "content"
+                ]  # Last user question
+                gen_ans = st.session_state.chat_history[-1][
+                    "content"
+                ]  # Last AI response
                 metadata = st.session_state.metadata.get(selected_question, {})
                 context_sources = metadata.get("context_sources", [])
 

frontend/utils/feedback.py

@@ -8,8 +8,9 @@
 
 
 import sys
+
 sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
-from common.mongoClient import submit_feedback
+from common.mongoClient import submit_feedback  # type: ignore
 
 load_dotenv()
 
@@ -152,8 +153,8 @@ def submit_feedback_to_google_sheet(
         # Format sources and context as combined pairs
         formatted_pairs = []
         for cs in context_sources:
-            source = cs.get('source', '')
-            context = cs.get('context', '')
+            source = cs.get("source", "")
+            context = cs.get("context", "")
             if source or context:
                 formatted_pairs.append(f"Source: {source}\nContext: {context}")
 
@@ -169,7 +170,6 @@ def submit_feedback_to_google_sheet(
             reaction or "",
         ]
 
-        #
         headers = sheet.row_values(1)
         required_headers = [
             "Question",
@@ -193,15 +193,15 @@ def submit_feedback_to_google_sheet(
 
 def show_feedback_form(
     questions: dict[str, int],
-    metadata: dict[str, dict[str, str]],
+    metadata: dict[str, dict[str, list]],
     interactions: list[dict[str, str]],
 ) -> None:
     """
     Display feedback form in the sidebar.
 
     Args:
     - questions (dict[str, int]): Dictionary of questions and indices.
-    - metadata (dict[str, dict[str, str]]): Metadata contains sources and context for each question.
+    - metadata (dict[str, dict[str, list]]): Metadata contains sources and context for each question.
     - interactions (list[dict[str, str]]): List of chat interactions from st.session_state.chat_history
 
     Returns:

frontend/utils/mock_endpoint.py

@@ -15,7 +15,7 @@ def list_all_chains() -> Response:
     return jsonify(["/chains/mock"])
 
 
-@app.route("/chains/mock", methods=["POST"])
+@app.route("/graphs/agent-retriever", methods=["POST"])
 def chat_app() -> Response:
     """
     Endpoint to handle chat requests.
@@ -25,19 +25,26 @@ def chat_app() -> Response:
     """
     data: dict[str, Any] = request.get_json()
     user_query: str = data.get("query", "")
-    list_sources: bool = data.get("list_sources", False)
-    list_context: bool = data.get("list_context", False)
-
+    list_sources: bool = data.get("list_sources", True)
+    list_context: bool = data.get("list_context", True)
+
+    dummy_context_sources = [
+        {"source": "https://mocksource1.com", "context": "This is Mock Context 1"},
+        {"source": "https://mocksource2.com", "context": "This is Mock Context 2"},
+    ]
+    if not list_sources:
+        # drop the source keys
+        dummy_context_sources = [
+            {"source": "", "context": cs["context"]} for cs in dummy_context_sources
+        ]
+    if not list_context:
+        # drop the context keys
+        dummy_context_sources = [
+            {"source": cs["source"], "context": ""} for cs in dummy_context_sources
+        ]
     response = {
         "response": f"This is a mock response to your query: '{user_query}'",
-        "sources": [
-            "https://mocksource1.com",
-            "https://mocksource2.com",
-            "https://mocksource3.com",
-        ]
-        if list_sources
-        else [],
-        "context": ["This is Mock Context"] if list_context else [],
+        "context_sources": dummy_context_sources,
     }
 
     return jsonify(response)