Update src/streamlit_app.py

src/streamlit_app.py

@@ -387,7 +387,7 @@ def initialize_agent():
         base_retriever = vector_store.as_retriever(
             search_type="similarity",
             search_kwargs={
-                "k": 5,
+                "k": 3,
                 
             }
         )
@@ -415,86 +415,127 @@ def initialize_agent():
             except Exception as e:
                 return f"Error retrieving information: {str(e)}"
         
-        # Document search tool
+        # Retriever tool (core RAG function)
+        retriever_tool = create_retriever_tool(
+            retriever=base_retriever,
+            name="retriever",
+            description=(
+                "Use this tool to answer ANY question that might be related to or found in the uploaded or provided documents. "
+                "Always call this tool FIRST whenever the question could possibly require information from those documents. "
+                "If the question asks about facts, data, summaries, policies, reports, or anything that may come from the user's documents, "
+                "use this tool to retrieve the relevant content before answering."
+            ),
+        )
+
         Retriver_tool = Tool(
-            name="document_search",
-            func=qa_with_sources,
-            description="Search and answer questions based on uploaded documents. Use this for ANY question about companies, acquisitions, financial data, or specific information that might be in the documents.",
+            name="retriever",
+            func=retriever_tool,
+            description=(
+                "Retrieves relevant context from the user's uploaded or stored documents. "
+                "Use this tool for any question that might involve the content of the documents, "
+                "such as document summaries, factual answers, or topic-specific details."
+            ),
         )
-        
-        # General QA tool
-        def general_qa(query):
-            """General question answering"""
-            try:
-                return llm.invoke(query).content
-            except Exception as e:
-                return f"Error: {str(e)}"
-        
+
+        # QA tool
         qa_tool = Tool(
-            name="general_question",
-            func=general_qa,
-            description="Answer general knowledge questions NOT related to the uploaded documents.",
+            name="Question Answering",
+            func=llm.invoke,
+            description=(
+                "A general-purpose question answering tool. "
+                "Use this ONLY for casual or open-ended questions that are NOT related to the provided documents. "
+                "Examples: greetings, opinions, or general world knowledge questions (e.g., 'How are you?', 'What is AI?'). "
+                "Do NOT use this if the question might depend on the document contents."
+            ),
         )
-        
+
         # Summary tool
-        def summarize_text(text):
-            """Summarize text"""
-            try:
-                prompt = f"Summarize the following concisely:\n\n{text}"
-                return llm.invoke(prompt).content
-            except Exception as e:
-                return f"Error: {str(e)}"
-        
         summary_tool = Tool(
-            name="summarize",
-            func=summarize_text,
-            description="Summarize text or information.",
+            name="Summary",
+            func=llm.invoke,
+            description="Summarizes long text passages into concise summaries using a structured summarization prompt.",
+            prompt=PromptTemplate(
+                input_variables=["input"],
+                template="""
+                You are a summarization assistant. Follow these steps to summarize the text:
+                1. Read the text carefully.
+                2. Identify the main points and key details.
+                3. Write a concise summary that captures the essence of the text.
+
+                Text: {input}
+
+                Summary:
+                """,
+            ),
         )
-        
+
         # Explanation tool
-        def explain_concept(concept):
-            """Explain concepts"""
-            try:
-                prompt = f"Explain clearly:\n\n{concept}"
-                return llm.invoke(prompt).content
-            except Exception as e:
-                return f"Error: {str(e)}"
-        
         explanation_tool = Tool(
-            name="explain",
-            func=explain_concept,
-            description="Explain concepts or ideas in detail.",
+            name="Explanation",
+            func=llm.invoke,
+            description="Explains complex concepts in simple, clear terms using examples or analogies when appropriate.",
+            prompt=PromptTemplate(
+                input_variables=["input"],
+                template="""
+                You are an explanation assistant. Follow these steps to explain the concept:
+                1. Understand the concept thoroughly.
+                2. Break down the concept into simpler parts.
+                3. Provide a clear and detailed explanation with examples.
+
+                Concept: {input}
+
+                Explanation:
+                """,
+            ),
         )
-        
-        tools = [Retriver_tool, qa_tool, summary_tool, explanation_tool]
+
+        # Tool list (retriever first for prioritization)
+        tools = [Retriver_tool, summary_tool, explanation_tool, qa_tool]
         tool_names = ", ".join([tool.name for tool in tools])
-        
+
+        example = """
+            Example:
+            Thought: I should use the retriever tool to find relevant info.
+            Action: retriever
+            Action Input: current head of the American Red Cross
+            Observation: The documents do not mention the head of the American Red Cross.
+            Thought: The information is not in the documents.
+            Final Answer: I'm sorry, but I couldn’t find information about that in the provided documents.
+            """
         # Custom ReAct prompt
         react_prompt = PromptTemplate.from_template(
-            """Answer the following question as best you can. You have access to the following tools:
+            example + """
+            You are a retrieval-augmented assistant that answers questions ONLY using the information
+            found in the user's provided documents.
+
+            You have access to the following tools:
+            {tools}
 
-{tools}
+            Follow this reasoning format:
 
-Use this format STRICTLY:
+            Thought: Think about what the question is asking and whether you can find the answer in the user's documents.
+            Action: The action to take, must be one of [{tool_names}]
+            Action Input: The input to the action (be specific)
+            Observation: The result of the action
+            ... (You may repeat this Thought/Action/Observation cycle as needed)
+            Final Answer: Your final grounded answer to the user's question.
 
-Thought: Think about what needs to be done
-Action: The exact tool name from [{tool_names}]
-Action Input: The specific input for the tool
-Observation: The result of the action
-... (repeat Thought/Action/Action Input/Observation as needed)
-Thought: I now know the final answer
-Final Answer: Provide a clear, complete answer
+            ### Important Grounding Rules:
+            - You MUST first use the 'retriever' tool to search for relevant information in the user's documents.
+            - Only use the information retrieved from the documents to answer the question.
+            - If the retrieved information does not contain a clear or relevant answer, respond with:
+            "I'm sorry, but I couldn’t find information about that in the provided documents."
+            - Do NOT use your own general knowledge or external world knowledge.
+            - Use the 'Question Answering' tool only for generic greetings (like 'hi', 'how are you') or clarification.
+            - You may use multiple tools in sequence before providing the final answer.
 
-IMPORTANT:
-1. For questions about documents, companies, or data, ALWAYS use "document_search" FIRST
-2. Action Input should be the question/text only - no quotes or special formatting
-3. Always provide a Final Answer
+            Previous conversation:
+            {chat_history}
 
-Previous conversation:
-{chat_history}
+            Question: {input}
 
-Question: {input}
-{agent_scratchpad}"""
+            {agent_scratchpad}
+            """
         ).partial(
             tools="\n".join([f"{tool.name}: {tool.description}" for tool in tools]),
             tool_names=tool_names