Update agent.py

agent.py

@@ -3,27 +3,106 @@ import os
 import re
 import ast
 import requests
+import time
 import tempfile
+
 from pathlib import Path
 from enum import Enum
-from typing import Optional, Dict, List, Any, TypedDict
+from typing import TypedDict, Dict, List, Any
 
 from dotenv import load_dotenv
 from langgraph.graph import StateGraph, END
-from langchain.tools import Tool as LangTool
+from langchain.tools import Tool as LangTool, StructuredTool
 from langchain_core.runnables import RunnableLambda
-from langchain.tools import StructuredTool
 
-# --- Integrazione Zephyr LLM ---
+# ——— Zephyr LLM setup ———
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain_community.llms import HuggingFacePipeline
 import torch
 
-# 
-
-
-
 
+# Function to get the desired model Zephyr
+def get_zephyr_llm():
+    model_id = 'HuggingFaceH4/zephyr-7b-alpha'
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        torch_dtype = torch.float16,
+        device_map = 'auto'
+    )
+    gen = pipeline(
+        'text-generation',
+        model = model,
+        tokenizer = tokenizer,
+        max_new_tokens = 256, 
+        temperature = 0.7, 
+        top_p = 0.9
+    )
+    return HuggingFacePipeline(pipeline=gen)
+
+# Declare the llm 
+llm = get_zephyr_llm()
+
+
+# HF APIs
+load_dotenv()
+DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+QUESTIONS_URL = f"{DEFAULT_API_URL}/questions"
+SUBMIT_URL   = f"{DEFAULT_API_URL}/submit"
+FILE_PATH    = f"{DEFAULT_API_URL}/files/"
+
+# --- TOOLS --- #
+
+
+
+# Agent Steps Class
+class AgentStep(Enum):
+    ANALYZE         = "analyze"
+    SELECT_TOOLS    = "select_tools"
+    EXECUTE_TOOLS   = "execute_tools"
+    SYNTHESIZE      = "synthesize_answer"
+    ERROR_RECOVERY  = "error_recovery"
+    COMPLETE        = "complete"
+
+
+
+# Agent State Class
+class AgentState(TypedDict):
+    question: str
+    original_question: str
+    selected_tools: List[str]
+    tool_results: Dict[str, Any]
+    final_answer: str
+    current_step: str
+    error_count: int
+    max_errors: int
+
+
+# Initialize state 
+def initialize_state(question: str) -> str:
+    return {
+        'question': question,
+        'original_question': question,
+        'selected_tools': [],
+        'tool_result': {},
+        'final_answer': '',
+        'current_step': AgentStep.ANALYZE.value,
+        "error_count": 0,
+        "max_errors": 3
+    }
+
+
+def analyze_question(state: AgentState) -> AgentState:
+    state["current_step"] = AgentStep.SELECT_TOOLS.value
+    return state
+
+# Agent Class
+class GaiaAgent:
+    def __init__(self):
+        self.graph = graph 
+
+    def __call__(self, task_id: str, question: str) -> str: 
+