Update app.py

app.py

@@ -23,13 +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
-# Use the correct import path for the Python REPL tool
-from langchain.tools.utility import PythonREPL
+
+# Correct import for Python REPL tool
+from langchain.tools 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")
@@ -40,7 +40,6 @@ JWT_SECRET = os.getenv("JWT_SECRET", "changeme123")
 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")
@@ -53,18 +52,12 @@ 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:
-        # Load the specified model
         model_pipeline = pipeline("text-generation", model=MODEL_ID, use_auth_token=HF_TOKEN, **pipeline_kwargs)
     except Exception:
-        # 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()
@@ -76,31 +69,22 @@ embeddings = HuggingFaceEmbeddings()
 chroma_db = Chroma(persist_directory="./medini_memory", embedding_function=embeddings)
 retriever = chroma_db.as_retriever()
 
-# 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
 # ===========================================
@@ -108,26 +92,23 @@ search = SerpAPIWrapper(serpapi_api_key=SERPAPI_KEY)
 python_tool = PythonREPL()
 
 tools = [
-    # 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="Knowledge Recall", func=retrieve_and_answer, description="Retrieve info from Medini memory."),
     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)
-# 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
+    tools=tools,
+    llm=llm,
+    agent="conversational-react-description",
+    memory=memory,
+    verbose=True
 )
 
 # ===========================================
@@ -144,40 +125,23 @@ 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)
-    
-    # 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:
-        # 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
+    json_str = m.group(0) if m else raw
+    json_str = json_str.replace("```json", "").replace("```", "").strip()
+    plan = json.loads(json_str)
+    if 'steps' not in plan:
+        raise ValueError("Parsed JSON is missing the 'steps' array.")
+    return plan
 
 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
+    status = "error"
 
     try:
         tool_func = None
-        # FIX 4: Use string matching and TOOL_MAP lookup for robust execution
         if "recall" in hint:
             tool_func = TOOL_MAP.get("knowledgerecall")
         elif "search" in hint:
@@ -186,30 +150,22 @@ def execute_step(step: Dict[str, Any]) -> Dict[str, Any]:
             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 = f"Execution Error: {str(e)}"
         status = "error"
-        
-    # 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]:
-    """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}
 
@@ -220,13 +176,11 @@ 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
 
@@ -234,21 +188,16 @@ def goal_endpoint(goal: str, auth: bool = Depends(verify_jwt)):
 # GRADIO FRONTEND
 # ===========================================
 def gradio_chat(message, 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
+    return history, ""
 
 def gradio_execute_plan(goal):
-    """Gradio function to execute the full autonomous plan."""
     try:
-        report = execute_plan(goal)
-        return report
+        return execute_plan(goal)
     except Exception as e:
         return {"error": f"Failed to execute plan: {str(e)}"}
 
@@ -268,7 +217,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
         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")
+            goal_input = gr.Textbox(placeholder="Enter high-level goal.", label="Goal")
             run_goal_btn = gr.Button("Run Goal", variant="primary")
             gr.Markdown("---")
             gr.Markdown("### Execution Report")
@@ -281,12 +230,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 # ===========================================
 if __name__ == "__main__":
     def start_api():
-        """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()
-    
-    # Launch the Gradio interface
     demo.launch(share=False)