Update app.py

app.py

@@ -1,100 +1,118 @@
 import torch
-import json
 import re
-import datetime
+from typing import Annotated, TypedDict, Union
 from flask import Flask, request, Response, render_template
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from langgraph.graph import StateGraph, END
+from langchain_core.tools import tool
+from pydantic import BaseModel, Field
 
 app = Flask(__name__)
 
-# 1. TOOL DEFINITIONS
-def get_current_datetime(query: str = ""):
-    """Returns the current date and time."""
-    return f"Observation: The current date and time is {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}."
+# 1. DEFINE STRUCTURED TOOLS WITH PYDANTIC
+class CalcInput(BaseModel):
+    expression: str = Field(description="The math expression to evaluate, e.g., '2 + 2'")
 
+@tool("simple_calculator", args_schema=CalcInput)
 def simple_calculator(expression: str):
-    """An easy-to-construct tool for basic math (add, sub, mult, div)."""
+    """Useful for basic math calculations."""
     try:
-        # Source 351: Calculators are essential tools for deterministic results.
-        # Note: In production, use a safer math parser instead of eval.
-        result = eval(expression, {"__builtins__": None}, {})
-        return f"Observation: The calculation result is {result}."
+        # Source 351: Tools provide deterministic results for agents.
+        return str(eval(expression, {"__builtins__": None}, {}))
     except Exception as e:
-        return f"Observation: Error in calculation: {str(e)}."
+        return f"Error: {str(e)}"
 
-# Tool Registry
-tools = {
-    "get_current_datetime": get_current_datetime,
-    "simple_calculator": simple_calculator
-}
+@tool("get_time")
+def get_time():
+    """Returns the current system time."""
+    from datetime import datetime
+    return datetime.now().strftime("%H:%M:%S")
 
-# Load Model
+tools = {"simple_calculator": simple_calculator, "get_time": get_time}
+
+# 2. LOAD REASONING ENGINE
 model_id = "AshokGakr/model-tiny"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
 
-SYSTEM_PROMPT = """
-ROLE: You are a ReAct Agent. You solve tasks using this loop:
-Thought: (Reasoning about what to do)
-Action: (Tool name: 'get_current_datetime' or 'simple_calculator')
-Action Input: (Parameter for the tool)
-Observation: (Result from the tool - provided to you)
-... (Repeat Thought/Action/Observation if needed)
-Final Answer: (The final response to the user)
-
-AVAILABLE TOOLS:
-- get_current_datetime: Use this for any questions about the current date or time. No input needed.
-- simple_calculator: Use this for any math calculations. Input should be a math expression (e.g., '10 + 5').
-"""
+# 3. DEFINE THE AGENT STATE
+class AgentState(TypedDict):
+    # Source 828: StateGraph acts as the system's real-time workflow tracker.
+    messages: list[dict]
+    next_step: str
 
-@app.route('/')
-def index():
-    return render_template('index.html')
+# 4. AGENT LOGIC NODES
+def call_model(state: AgentState):
+    # 1. Apply Chat Template to format the history
+    # This prepares the context for the reasoning engine [5].
+    inputs = tokenizer.apply_chat_template(
+        state['messages'], 
+        add_generation_prompt=True, 
+        return_tensors="pt"
+    ).to(model.device)
+    
+    # 2. FIX: Unpack the inputs using ** to pass tensors correctly
+    # This prevents the KeyError: 'shape' by giving generate the specific tensors it needs.
+    output_ids = model.generate(
+        **inputs,  # <--- CRITICAL FIX: Unpack the dictionary
+        max_new_tokens=256,
+        do_sample=True,
+        temperature=0.7
+    )
+    
+    # 3. Decode only the newly generated tokens (skipping the original prompt)
+    # inputs['input_ids'].shape[-1] provides the length of the input tokens.
+    new_tokens = output_ids[inputs['input_ids'].shape[-1]:]
+    response = tokenizer.decode(new_tokens, skip_special_tokens=True)
+    
+    # Identify if a tool needs to be called [1, 6]
+    action_match = re.search(r"Action:\s*(\w+)", response)
+    
+    return {
+        "messages": state['messages'] + [{"role": "assistant", "content": response}],
+        "next_step": action_match.group(1) if action_match else "end"
+    }
 
-@app.route('/chat', methods=['POST'])
-def chat():
-    data = request.json
-    user_query = data.get("message", "")
+def execute_tool(state: AgentState):
+    tool_name = state['next_step']
+    last_message = state['messages'][-1]['content']
+    
+    # Parse input (simplified for this model-tiny example)
+    input_match = re.search(r"Action Input:\s*(.*)", last_message)
+    arg = input_match.group(1).strip() if input_match else ""
     
-    def generate_agent_response():
-        # Source 13: Episodic memory maintains the conversation trajectory.
-        history = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_query}]
-        
-        for i in range(5): # Limit iterations to prevent infinite loops [5]
-            input_ids = tokenizer.apply_chat_template(history, add_generation_prompt=True, return_tensors="pt").to(model.device)
-            streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-            
-            thread = Thread(target=model.generate, kwargs={"input_ids": input_ids, "streamer": streamer, "max_new_tokens": 256})
-            thread.start()
+    observation = tools[tool_name].run(arg)
+    return {"messages": state['messages'] + [{"role": "user", "content": f"Observation: {observation}"}]}
+
+# 5. CONSTRUCT THE GRAPH
+# Source 96: Nodes represent actions; edges define the control flow.
+workflow = StateGraph(AgentState)
+workflow.add_node("agent", call_model)
+workflow.add_node("tools", execute_tool)
 
-            full_turn_output = ""
-            for new_text in streamer:
-                full_turn_output += new_text
-                yield new_text # Stream thoughts to the UI [6]
+workflow.set_entry_point("agent")
+workflow.add_conditional_edges(
+    "agent",
+    lambda x: "tools" if x["next_step"] in tools else "end",
+    {"tools": "tools", "end": END}
+)
+workflow.add_edge("tools", "agent") # Create the ReAct Loop cycle [10]
 
-            # Check for Action [7]
-            action_match = re.search(r"Action:\s*(\w+)", full_turn_output)
-            input_match = re.search(r"Action Input:\s*(.*)", full_turn_output)
+agent_app = workflow.compile()
 
-            if action_match and input_match:
-                tool_name = action_match.group(1).strip()
-                tool_input = input_match.group(1).strip()
-                
-                if tool_name in tools:
-                    obs = tools[tool_name](tool_input)
-                    yield f"\n{obs}\n"
-                    # Feed observation back into history [8, 9]
-                    history.append({"role": "assistant", "content": full_turn_output})
-                    history.append({"role": "user", "content": obs})
-                else:
-                    break
-            elif "Final Answer:" in full_turn_output:
-                break
-            else:
-                break
+@app.route('/chat', methods=['POST'])
+def chat():
+    user_input = request.json.get("message")
+    # Execute the graph [5, 11]
+    inputs = {"messages": [{"role": "user", "content": user_input}]}
+    
+    def run():
+        for output in agent_app.stream(inputs):
+            for key, value in output.items():
+                # Stream the latest message content to the UI [12]
+                yield value['messages'][-1]['content'] + "\n"
 
-    return Response(generate_agent_response(), mimetype='text/plain')
+    return Response(run(), mimetype='text/plain')
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860)