Update app.py

app.py

@@ -250,6 +250,38 @@ def should_continue(state: AgentState) -> str:
     return "reason"
 
 
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
+
+# ====== DOCUMENT PROCESSING SETUP ======
+def create_vector_store():
+    """Create vector store with predefined documents"""
+    # Define the documents
+    documents = [
+        Document(page_content="The capital of France is Paris.", metadata={"source": "geography"}),
+        Document(page_content="Python is a popular programming language created by Guido van Rossum.", metadata={"source": "tech"}),
+        Document(page_content="The Eiffel Tower is located in Paris, France.", metadata={"source": "landmarks"}),
+    ]
+    
+    # Initialize embedding model
+    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+    
+    # Split documents into chunks
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=500,  # Smaller chunks for better precision
+        chunk_overlap=100
+    )
+    chunks = text_splitter.split_documents(documents)
+    
+    # Create in-memory vector store
+    return Chroma.from_documents(
+        documents=chunks,
+        embedding=embeddings
+    )
+
+
 def reasoning_node(state: AgentState) -> AgentState:
     """
     Node for the agent to analyze the question, determine next steps,
@@ -304,7 +336,42 @@ def reasoning_node(state: AgentState) -> AgentState:
         f"- **{t.name}**: {t.description}" for t in state.get("tools", [])
     ])
 
-    # Enhanced system prompt for better JSON adherence
+    # ====== RAG RETRIEVAL ======
+    # Initialize vector store if not present
+    if "vector_store" not in state["context"]:
+        state["context"]["vector_store"] = create_vector_store()
+    
+    vector_store = state["context"]["vector_store"]
+    
+    # Perform retrieval
+    relevant_docs = vector_store.similarity_search(
+        state["question"], 
+        k=3  # Retrieve top 3 most relevant chunks
+    )
+    
+    # Format context for LLM
+    rag_context = "\n\n[Relevant Knowledge]\n"
+    rag_context += "\n---\n".join([doc.page_content for doc in relevant_docs])
+    
+    # ====== RAG RETRIEVAL ======
+    # Initialize vector store if not present
+    if "vector_store" not in state["context"]:
+        state["context"]["vector_store"] = create_vector_store()
+    
+    vector_store = state["context"]["vector_store"]
+    
+    # Perform retrieval
+    relevant_docs = vector_store.similarity_search(
+        state["question"], 
+        k=3  # Retrieve top 3 most relevant chunks
+    )
+    
+    # Format context for LLM
+    rag_context = "\n\n[Relevant Knowledge]\n"
+    rag_context += "\n---\n".join([doc.page_content for doc in relevant_docs])
+    
+    # ====== MODIFIED PROMPT ======
+    # Add RAG context to system prompt
     system_prompt = (
         "You are an expert problem solver, designed to provide concise and accurate answers. "
         "Your process involves analyzing the question, intelligently selecting and using tools, "
@@ -318,21 +385,12 @@ def reasoning_node(state: AgentState) -> AgentState:
         "- Use **document_qa** when the question explicitly refers to a specific document or when you have content to query. Input format: 'document_text||question'.\n"
         "- Use **python_execution** for complex calculations, data manipulation, or logical operations that cannot be done with simple reasoning. Always provide the full Python code, ensuring it's valid and executable, and assign the final result to a variable named '_result_value' (e.g., '_result_value = 1 + 1').\n"
         "- Use **transcript_video** for any question involving video or audio content (e.g., YouTube). Provide the full YouTube URL or video ID.\n\n"
+        "**Retrieved Context:**\n{rag_context}\n\n"  # ADDED RAG CONTEXT
         "**Current Context:**\n{context}\n\n"
         "**Previous Reasoning Steps:**\n{reasoning}\n\n"
         "**Current Task:** {current_task}\n"
         "**Current Thoughts:** {current_thoughts}\n\n"
-        "**Your Response MUST be a valid JSON object with the following keys:**\n"
-        "```json\n"
-        "{\n"
-        "  \"Reasoning\": \"Your detailed analysis of the question and why you chose a specific action. Focus on the logical steps.\",\n"
-        "  \"Action\": \"[Tool name OR 'Final Answer']\",\n"
-        "  \"Action Input\": \"[Input for the selected tool OR the complete final answer]\"\n"
-        "}\n"
-        "```\n"
-        "**CRITICAL RULE: 'Action' and 'Action Input' MUST NOT be empty strings, unless 'Action' is 'Final Answer' and 'Action Input' is the conclusive response.**\n"
-        "If you cannot determine a suitable tool or a conclusive final answer after exhausting options, return Action: 'Final Answer' with a message like 'I cannot answer this question with the available tools.' or 'More information is needed.'\n"
-        "Ensure 'Action Input' is always the complete, valid input for the chosen 'Action'. If 'Action' is 'Final Answer', provide the complete, concise answer."
+        # ... [rest of prompt remains same] ...
     )
 
     prompt = ChatPromptTemplate.from_messages([
@@ -341,6 +399,7 @@ def reasoning_node(state: AgentState) -> AgentState:
     ])
 
     formatted_messages = prompt.format_messages(
+        rag_context=rag_context,  # ADD THIS ARGUMENT
         context=state["context"],
         reasoning=state["reasoning"],
         question=state["question"],
@@ -498,14 +557,19 @@ class BasicAgent:
             python_execution,
             VideoTranscriptionTool() 
         ]
+        
+        # Pre-initialize RAG vector store
+        self.vector_store = create_vector_store()
         self.workflow = create_agent_workflow(self.tools)
-
+    
     def __call__(self, question: str) -> str:
         print(f"\n--- Agent received question: {question[:80]}{'...' if len(question) > 80 else ''} ---")
 
         state = {
             "question": question,
-            "context": {},
+            "context": {
+                "vector_store": self.vector_store  # Include vector store in context
+            },
             "reasoning": "",
             "iterations": 0,
             "history": [HumanMessage(content=question)],