Update app.py

app.py

@@ -1,121 +1,77 @@
 import torch
 import re
-from typing import Annotated, TypedDict, Union
 from flask import Flask, request, Response, render_template
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from langgraph.graph import StateGraph, END
-from langchain_core.tools import tool
-from pydantic import BaseModel, Field
+from typing import TypedDict
 
 app = Flask(__name__)
 
-# 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):
-    """Useful for basic math calculations."""
-    try:
-        # Source 351: Tools provide deterministic results for agents.
-        return str(eval(expression, {"__builtins__": None}, {}))
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-@tool("get_time")
-def get_time():
-    """Returns the current system time."""
-    from datetime import datetime
-    return datetime.now().strftime("%H:%M:%S")
-
-tools = {"simple_calculator": simple_calculator, "get_time": get_time}
-
-# 2. LOAD REASONING ENGINE
+# 1. Loading the Cognitive Core [5, 6]
 model_id = "AshokGakr/model-tiny"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
 
-# 3. DEFINE THE AGENT STATE
+# 2. Defining State and Tools [7, 8]
 class AgentState(TypedDict):
-    # Source 828: StateGraph acts as the system's real-time workflow tracker.
     messages: list[dict]
-    next_step: str
+    next_action: str
 
-# 4. AGENT LOGIC NODES
+def get_time(query: str):
+    from datetime import datetime
+    return f"Observation: The current time is {datetime.now().strftime('%H:%M:%S')}."
+
+tools = {"get_time": get_time}
+
+# 3. The Reasoning Node [9, 10]
 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
-    )
+    # Context Engineering: Applying chat template for multi-turn coherence [11, 12]
+    inputs = tokenizer.apply_chat_template(state['messages'], add_generation_prompt=True, return_tensors="pt").to(model.device)
     
-    # 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]:]
+    # Generate and Slice: Correctly targeting the token dimension to avoid empty strings
+    output_ids = model.generate(**inputs, max_new_tokens=256, do_sample=True, temperature=0.7)
+    new_tokens = output_ids[0, inputs['input_ids'].shape[-1]:]
     response = tokenizer.decode(new_tokens, skip_special_tokens=True)
     
-    # Identify if a tool needs to be called [1, 6]
+    # Logic to identify if the agent needs a tool or has a Final Answer [13, 14]
     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"
+        "next_action": action_match.group(1) if action_match else "end"
     }
 
+# 4. The Action Node [14, 15]
 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 ""
-    
-    observation = tools[tool_name].run(arg)
-    return {"messages": state['messages'] + [{"role": "user", "content": f"Observation: {observation}"}]}
+    tool_name = state['next_action']
+    observation = tools[tool_name]("")
+    return {"messages": state['messages'] + [{"role": "user", "content": observation}]}
 
-# 5. CONSTRUCT THE GRAPH
-# Source 96: Nodes represent actions; edges define the control flow.
+# 5. Graph Construction [9, 16]
 workflow = StateGraph(AgentState)
 workflow.add_node("agent", call_model)
 workflow.add_node("tools", execute_tool)
-
 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]
-
-agent_app = workflow.compile()
-@app.route('/')
-def index():
-    return render_template('index.html')
+workflow.add_conditional_edges("agent", lambda x: "tools" if x["next_action"] in tools else "end", {"tools": "tools", "end": END})
+workflow.add_edge("tools", "agent")
+agent_executor = workflow.compile()
 
 @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}]}
+    user_msg = request.json.get("message")
+    # System Prompt: Establishing identity and rules [17, 18]
+    inputs = {"messages": [
+        {"role": "system", "content": "You are a ReAct agent. Use Thought:, Action:, and Final Answer: tags."},
+        {"role": "user", "content": user_msg}
+    ]}
     
     def run():
-        for output in agent_app.stream(inputs):
+        for output in agent_executor.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(run(), mimetype='text/plain')
 
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860)
+@app.route('/')
+def index(): return render_template('index.html')
+
+if __name__ == '__main__': app.run(host='0.0.0.0', port=7860)