Update app.py

app.py

@@ -15,10 +15,10 @@ nltk.download('punkt_tab')
 
 
 @st.cache_resource
-def init_agent(namespace1: str, namespace2: str):
+def init_agent(namespace1: str, top_k1: int, namespace2: str, top_k2: int):
     """
     Initialize the LangGraph agent with two Pinecone retriever tools,
-    each configured with a different namespace.
+    each configured with a specified namespace and top_k value.
     """
     # Retrieve API keys from environment variables
     OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
@@ -37,39 +37,39 @@ def init_agent(namespace1: str, namespace2: str):
     index_name = 'autogen'
     pc = Pinecone(api_key=PINE_API_KEY)
     index = pc.Index(index_name)
-    # Allow the index to connect
+    # Allow a moment for the index to connect
     time.sleep(1)
     index.describe_index_stats()
 
     # --- BM25 Sparse Encoder ---
     bm25_encoder = BM25Encoder().default()
 
-    # --- Create first Pinecone Hybrid Retriever Tool using namespace1 ---
+    # --- Create first Pinecone Hybrid Retriever Tool using namespace1 and top_k1 ---
     retriever1 = PineconeHybridSearchRetriever(
         embeddings=embed,
         sparse_encoder=bm25_encoder,
         index=index,
         namespace=namespace1,
-        top_k=4
+        top_k=top_k1
     )
     retriever_tool1 = create_retriever_tool(
         retriever1,
         "retrieve_context_1",
-        f"Search and return information from Autogen's codebase and documentation using namespace '{namespace1}'.",
+        f"Search and return information from Autogen's codebase and documentation using namespace '{namespace1}' with top_k = {top_k1}.",
     )
 
-    # --- Create second (duplicate) Pinecone Hybrid Retriever Tool using namespace2 ---
+    # --- Create second Pinecone Hybrid Retriever Tool using namespace2 and top_k2 ---
     retriever2 = PineconeHybridSearchRetriever(
         embeddings=embed,
         sparse_encoder=bm25_encoder,
         index=index,
         namespace=namespace2,
-        top_k=4
+        top_k=top_k2
     )
     retriever_tool2 = create_retriever_tool(
         retriever2,
         "retrieve_context_2",
-        f"Search and return information from Autogen's codebase and documentation using namespace '{namespace2}'.",
+        f"Search and return information from Autogen's codebase and documentation using namespace '{namespace2}' with top_k = {top_k2}.",
     )
 
     # Both retriever tools are added to the list of available tools.
@@ -127,14 +127,20 @@ Now, please help the user with their coding query for LangGraph:
     return graph
 
 
-# ----------------- Sidebar: Namespace Selection ----------------- #
-st.sidebar.header("Namespace Selection for Retriever Tools")
+# ----------------- Sidebar: Namespace & Top_K Selection ----------------- #
+st.sidebar.header("Retriever Tool Settings")
+
+# Dropdown and slider for Retriever Tool 1
 namespace_options = ["langgraph-main", "autogen"]
 namespace1 = st.sidebar.selectbox("Select namespace for Retriever Tool 1:", namespace_options, index=0)
+top_k1 = st.sidebar.slider("Select top_k for Retriever Tool 1:", min_value=1, max_value=4, value=1, step=1)
+
+# Dropdown and slider for Retriever Tool 2
 namespace2 = st.sidebar.selectbox("Select namespace for Retriever Tool 2:", namespace_options, index=0)
+top_k2 = st.sidebar.slider("Select top_k for Retriever Tool 2:", min_value=1, max_value=4, value=1, step=1)
 
-# Initialize the agent with the selected namespaces.
-graph = init_agent(namespace1, namespace2)
+# Initialize the agent with the selected namespaces and top_k values.
+graph = init_agent(namespace1, top_k1, namespace2, top_k2)
 
 
 # ----------------- Main Chat App UI ----------------- #
@@ -145,7 +151,7 @@ if "chat_history" not in st.session_state:
     st.session_state.chat_history = []  # Each entry is a tuple: (role, message)
 
 def display_conversation():
-    """Display the chat history."""
+    """Display the chat conversation."""
     for role, message in st.session_state.chat_history:
         if role == "user":
             st.markdown(f"**You:** {message}")
@@ -167,7 +173,7 @@ with st.form("chat_form", clear_on_submit=True):
 if st.session_state.chat_history and st.session_state.chat_history[-1][0] == "user":
     inputs = {"messages": st.session_state.chat_history}
 
-    # Placeholder to update the response stream in real time.
+    # Placeholder for real-time streaming of the response.
     response_placeholder = st.empty()
     assistant_message = ""
 
@@ -182,6 +188,6 @@ if st.session_state.chat_history and st.session_state.chat_history[-1][0] == "us
         assistant_message += text
         response_placeholder.markdown(f"**Assistant:** {assistant_message}")
 
-    # Append the full assistant response to the chat history.
+    # Append the complete assistant response to the chat history.
     st.session_state.chat_history.append(("assistant", assistant_message))
     st.experimental_rerun()