Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 import json
 import threading
+import re
 from typing import Dict, Any
 
 import gradio as gr
@@ -8,6 +9,7 @@ from fastapi import FastAPI, Depends, HTTPException
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 import uvicorn
 
+# Note: The 'transformers' pipeline is required for the LLM loader
 from transformers import pipeline
 
 # ===== LangChain imports (fixed for 0.3.x + community modules) =====
@@ -21,11 +23,13 @@ from langchain.agents import initialize_agent, Tool
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
 from langchain.docstore.document import Document
-from langchain.tools import PythonREPL  # ✅ Corrected import
+# Use the correct import path for the Python REPL tool
+from langchain.tools.utility import PythonREPL
 
 # ===========================================
 # ENVIRONMENT VARIABLES
 # ===========================================
+# IMPORTANT: These environment variables must be set for the app to run outside of this environment
 HF_TOKEN = os.getenv("HF_TOKEN")
 SERPAPI_KEY = os.getenv("SERPAPI_API_KEY")
 JWT_SECRET = os.getenv("JWT_SECRET", "changeme123")
@@ -34,7 +38,9 @@ JWT_SECRET = os.getenv("JWT_SECRET", "changeme123")
 # AUTH
 # ===========================================
 security = HTTPBearer()
+
 def verify_jwt(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    """Verifies the custom JWT token (in this case, checking against JWT_SECRET)."""
     token = credentials.credentials
     if token != JWT_SECRET:
         raise HTTPException(status_code=403, detail="Invalid token")
@@ -47,10 +53,18 @@ MODEL_ID = "PuruAI/Medini_Intelligence"
 FALLBACK_MODEL = "gpt2"
 
 def load_llm():
+    """Loads the HuggingFace model pipeline, adding generation arguments for stability."""
+    # FIX 2: Added max_new_tokens for better agent reasoning and response length
+    pipeline_kwargs = {"max_new_tokens": 512, "temperature": 0.7}
     try:
-        model_pipeline = pipeline("text-generation", model=MODEL_ID, use_auth_token=HF_TOKEN)
+        # Load the specified model
+        model_pipeline = pipeline("text-generation", model=MODEL_ID, use_auth_token=HF_TOKEN, **pipeline_kwargs)
     except Exception:
-        model_pipeline = pipeline("text-generation", model=FALLBACK_MODEL)
+        # Fallback to a common model if the primary one fails
+        print(f"Warning: Failed to load {MODEL_ID}. Falling back to {FALLBACK_MODEL}.")
+        model_pipeline = pipeline("text-generation", model=FALLBACK_MODEL, **pipeline_kwargs)
+
+    # Wrap the pipeline in HuggingFacePipeline for LangChain integration
     return HuggingFacePipeline(pipeline=model_pipeline)
 
 llm = load_llm()
@@ -61,31 +75,60 @@ llm = load_llm()
 embeddings = HuggingFaceEmbeddings()
 chroma_db = Chroma(persist_directory="./medini_memory", embedding_function=embeddings)
 retriever = chroma_db.as_retriever()
-retrieval_qa = LLMChain(
-    llm=llm,
-    prompt=PromptTemplate(
-        input_variables=["question"],
-        template="{question}"
-    )
-)
+
+# FIX 1: Redefine the Retrieval QA components
+qa_prompt_template = """
+You are a question-answering system. Use the following context, which contains information retrieved from memory, to answer the user's question. 
+If the context is empty or does not contain the answer, state clearly that the information is not in memory.
+
+Context:
+{context}
+
+Question: {question}
+Answer:
+"""
+QA_PROMPT = PromptTemplate(template=qa_prompt_template, input_variables=["context", "question"])
+# This LLMChain is used specifically for answering questions based on retrieved context
+qa_chain = LLMChain(llm=llm, prompt=QA_PROMPT)
+
+def retrieve_and_answer(question: str) -> str:
+    """Retrieves context from Chroma DB and passes it to the QA Chain."""
+    # 1. Use the defined retriever to find relevant documents
+    docs = retriever.get_relevant_documents(question)
+    context = "\n---\n".join([d.page_content for d in docs])
+    
+    # 2. Run the QA chain with the retrieved context
+    return qa_chain.run(context=context, question=question)
+
 
 # ===========================================
 # TOOLS
 # ===========================================
 search = SerpAPIWrapper(serpapi_api_key=SERPAPI_KEY)
-python_tool = PythonREPL()  # ✅ Updated
+python_tool = PythonREPL()
 
 tools = [
-    Tool(name="Knowledge Recall", func=lambda q: retrieval_qa.run({"question": q}), description="Retrieve info from Medini memory."),
-    Tool(name="Web Search", func=search.run, description="Search the web for info."),
-    Tool(name="Python REPL", func=python_tool.run, description="Execute Python code."),
+    # FIX 1: Use the new function that correctly retrieves info from the vector store
+    Tool(name="Knowledge Recall", func=retrieve_and_answer, description="Retrieve info from Medini memory (Chroma DB). Use this when the answer might be in a previously executed step or private notes."),
+    Tool(name="Web Search", func=search.run, description="Search the web for up-to-date information."),
+    Tool(name="Python REPL", func=python_tool.run, description="Execute Python code, useful for math and data manipulation."),
 ]
 
+# FIX 4: Create a map for robust tool execution lookup
+TOOL_MAP = {tool.name.lower().replace(" ", ""): tool.func for tool in tools}
+
 # ===========================================
 # AGENT
 # ===========================================
 memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-agent = initialize_agent(tools=tools, llm=llm, agent="conversational-react-description", memory=memory, verbose=False)
+# The LLM must be passed to the agent
+agent = initialize_agent(
+    tools=tools, 
+    llm=llm, 
+    agent="conversational-react-description", 
+    memory=memory, 
+    verbose=True # Changed to True for better debugging/visibility
+)
 
 # ===========================================
 # PLANNER (Autonomous Goal)
@@ -93,7 +136,7 @@ agent = initialize_agent(tools=tools, llm=llm, agent="conversational-react-descr
 plan_prompt = PromptTemplate(
     input_variables=["goal"],
     template="""
-You are Medini Planner. Decompose the high-level goal into JSON steps (max 6) with: id, name, description, tool_hint.
+You are Medini Planner. Decompose the high-level goal into a JSON object containing a 'steps' array (max 6 steps). Each step must have: id (integer), name (short string), description (detailed instruction), and tool_hint (either 'recall', 'search', 'python', or 'agent').
 Return JSON only.
 Goal: {goal}
 """
@@ -101,34 +144,72 @@ Goal: {goal}
 planner_chain = LLMChain(llm=llm, prompt=plan_prompt)
 
 def create_plan(goal: str) -> Dict[str, Any]:
+    """Generates a structured plan using the planner chain."""
     raw = planner_chain.run(goal=goal)
-    import re
+    
+    # FIX 3: Robust JSON Parsing - Find the JSON block and clean up common LLM formatting
     m = re.search(r"\{.*\}", raw, flags=re.DOTALL)
     if not m:
-        raise ValueError("Planner did not return JSON")
-    return json.loads(m.group(0))
+        # Fallback to entire raw output if no braces are found
+        json_str = raw
+    else:
+        json_str = m.group(0)
+
+    # Clean up common markdown code fences
+    json_str = json_str.replace("```json", "").replace("```", "").strip() 
+
+    try:
+        plan = json.loads(json_str)
+        if 'steps' not in plan:
+             raise ValueError("Parsed JSON is missing the 'steps' array.")
+        return plan
+    except json.JSONDecodeError as e:
+        print(f"JSON Parsing Error: {e} in string: {json_str[:200]}...")
+        raise ValueError("Planner returned malformed JSON. Check the LLM's output format.") from e
 
 def execute_step(step: Dict[str, Any]) -> Dict[str, Any]:
+    """Executes a single step using the appropriate tool or the main agent."""
     hint = (step.get("tool_hint") or "").lower()
     input_text = step.get("description")
+    
+    output = "Execution skipped."
+    status = "error" # Default to error
+
     try:
+        tool_func = None
+        # FIX 4: Use string matching and TOOL_MAP lookup for robust execution
         if "recall" in hint:
-            output = tools[0].func(input_text)
+            tool_func = TOOL_MAP.get("knowledgerecall")
         elif "search" in hint:
-            output = tools[1].func(input_text)
+            tool_func = TOOL_MAP.get("websearch")
         elif "python" in hint:
-            output = tools[2].func(input_text)
+            tool_func = TOOL_MAP.get("pythonrepl")
+
+        if tool_func:
+            # Execute the specific tool
+            output = tool_func(input_text)
         else:
+            # Fallback to the main agent for generic reasoning/conversation
             output = agent.run(input_text)
+            
         status = "ok"
+
     except Exception as e:
-        output = str(e)
+        output = f"Execution Error: {str(e)}"
         status = "error"
-    chroma_db.add_documents([Document(page_content=f"Step {step['id']} - {step['name']}: {output}")])
+        
+    # Add the result of the step execution to the vector memory
+    chroma_db.add_documents([Document(page_content=f"Step {step['id']} - {step['name']} Result: {output}")])
+    
     return {"id": step['id'], "name": step['name'], "status": status, "output": output}
 
 def execute_plan(goal: str) -> Dict[str, Any]:
-    plan = create_plan(goal)
+    """Creates a plan and executes all steps sequentially."""
+    try:
+        plan = create_plan(goal)
+    except ValueError as e:
+        return {"goal": goal, "error": str(e)}
+        
     results = [execute_step(step) for step in plan.get("steps", [])]
     return {"goal": goal, "plan": plan, "results": results}
 
@@ -139,11 +220,13 @@ app = FastAPI(title="Medini Agent API")
 
 @app.post("/chat")
 def chat_endpoint(message: str, auth: bool = Depends(verify_jwt)):
+    """API endpoint for basic conversational chat."""
     response = agent.run(message)
     return {"response": response}
 
 @app.post("/goal")
 def goal_endpoint(goal: str, auth: bool = Depends(verify_jwt)):
+    """API endpoint for executing autonomous goals."""
     report = execute_plan(goal)
     return report
 
@@ -151,29 +234,59 @@ def goal_endpoint(goal: str, auth: bool = Depends(verify_jwt)):
 # GRADIO FRONTEND
 # ===========================================
 def gradio_chat(message, history):
-    response = agent.run(message)
-    history.append((message, response))
-    return history, history
+    """Gradio function for the chat interface."""
+    # The agent manages its own history/memory
+    try:
+        response = agent.run(message)
+        history.append((message, response))
+    except Exception as e:
+        history.append((message, f"An error occurred: {str(e)}"))
+        
+    return history, "" # Return history and clear the input box
+
+def gradio_execute_plan(goal):
+    """Gradio function to execute the full autonomous plan."""
+    try:
+        report = execute_plan(goal)
+        return report
+    except Exception as e:
+        return {"error": f"Failed to execute plan: {str(e)}"}
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🤖 Medini Autonomous Agent")
     gr.Markdown("Chat or submit high-level goals. Agentic AI handles reasoning, memory, and tool use.")
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(placeholder="Type your message...", label="Chat")
-    goal_input = gr.Textbox(placeholder="Enter high-level goal...", label="Goal")
-    run_goal_btn = gr.Button("Run Goal")
-    clear_btn = gr.Button("Clear Chat")
-    goal_output = gr.JSON()
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            gr.Markdown("## Conversational Chat")
+            chatbot = gr.Chatbot(height=400)
+            msg = gr.Textbox(placeholder="Type your message...", label="Chat Input")
+            clear_btn = gr.Button("Clear Chat")
+            
+            msg.submit(gradio_chat, [msg, chatbot], [chatbot, msg])
+            clear_btn.click(lambda: [], None, chatbot, queue=False)
 
-    msg.submit(gradio_chat, [msg, chatbot], [chatbot, chatbot])
-    clear_btn.click(lambda: None, None, chatbot, queue=False)
-    run_goal_btn.click(lambda g: execute_plan(g), [goal_input], goal_output)
+        with gr.Column(scale=1):
+            gr.Markdown("## Autonomous Goal Planner")
+            goal_input = gr.Textbox(placeholder="Enter high-level goal (e.g., 'Research the latest quarterly earnings of Tesla and save the key points').", label="Goal")
+            run_goal_btn = gr.Button("Run Goal", variant="primary")
+            gr.Markdown("---")
+            gr.Markdown("### Execution Report")
+            goal_output = gr.JSON(label="Plan and Results")
+            
+            run_goal_btn.click(gradio_execute_plan, [goal_input], goal_output)
 
 # ===========================================
 # LAUNCH
 # ===========================================
 if __name__ == "__main__":
     def start_api():
-        uvicorn.run(app, host="0.0.0.0", port=8000)
+        """Starts the FastAPI server in a separate thread."""
+        # Use log_level="critical" to reduce noisy logs from uvicorn in the console
+        uvicorn.run(app, host="0.0.0.0", port=8000, log_level="critical")
+        
+    # Start the API in the background
     threading.Thread(target=start_api, daemon=True).start()
-    demo.launch()
+    
+    # Launch the Gradio interface
+    demo.launch(share=False)