Update agents.py

agents.py

@@ -14,9 +14,11 @@ import sys
 import tempfile
 from typing import Annotated, Literal, TypedDict
 
-from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from langchain_core.messages import AnyMessage, AIMessage, HumanMessage, SystemMessage
 from langchain_core.tools import tool
-from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
 from langchain_tavily import TavilySearch
 from langgraph.graph import StateGraph, END
 from langgraph.graph.message import add_messages
@@ -28,16 +30,44 @@ MAX_TURNS = 6
 
 
 # ---------------------------------------------------------------------------
-# Model
+# Local model, run on a ZeroGPU-allocated GPU
 # ---------------------------------------------------------------------------
-def get_llm():
-    endpoint = HuggingFaceEndpoint(
-        repo_id="Qwen/Qwen2.5-72B-Instruct",
-        huggingfacehub_api_token=os.environ.get("HF_TOKEN"),
-        max_new_tokens=1024,
-        temperature=0.2,
-    )
-    return ChatHuggingFace(llm=endpoint)
+# Pick a model that fits comfortably in a ZeroGPU allocation. 3B is a safe
+# starting point; move up to 7B/8B if quality is the bottleneck, but watch
+# the `duration` below — bigger models need more time per call.
+MODEL_ID = "Qwen/Qwen2.5-3B-Instruct"
+
+# Loaded once at import time, kept on CPU. ZeroGPU only attaches a GPU for
+# the duration of a function decorated with @spaces.GPU — the model is
+# moved onto that GPU inside `generate()` below and this is the standard
+# pattern for ZeroGPU Spaces (see huggingface.co/docs/hub/spaces-zerogpu).
+_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+_model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
+
+
+@spaces.GPU(duration=60)
+def generate(messages: list[dict], max_new_tokens: int = 512) -> str:
+    """Run one generation on a ZeroGPU-allocated GPU. `messages` is a plain
+    list of {"role": ..., "content": ...} dicts (chat-template format)."""
+    _model.to("cuda")
+    input_ids = _tokenizer.apply_chat_template(
+        messages, add_generation_prompt=True, return_tensors="pt"
+    ).to("cuda")
+    with torch.no_grad():
+        output_ids = _model.generate(
+            input_ids,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+        )
+    new_tokens = output_ids[0][input_ids.shape[-1]:]
+    return _tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
+
+
+def _lc_to_dicts(messages: list[AnyMessage]) -> list[dict]:
+    """Convert langchain message objects into the plain role/content dicts
+    apply_chat_template expects."""
+    role_map = {"human": "user", "ai": "assistant", "system": "system"}
+    return [{"role": role_map.get(m.type, "user"), "content": m.content} for m in messages]
 
 
 # ---------------------------------------------------------------------------
@@ -58,14 +88,32 @@ class AgentState(TypedDict):
 # Requires TAVILY_API_KEY as a Space secret. Tavily is built for LLM agents
 # (returns clean, citeable snippets rather than raw HTML), which is why it's
 # preferred here over scraping a generic search engine yourself.
-_tavily = TavilySearch(max_results=5)
+#
+# Lazily constructed on first use rather than at import time: TavilySearch
+# validates its API key eagerly in __init__, so building it at module load
+# means a missing/bad key crashes the entire app before Gradio even starts.
+# Deferring it means only the search tool fails, with a clear message,
+# while the rest of the app (code_agent, data_agent) still works.
+_tavily_instance = None
+
+
+def _get_tavily():
+    global _tavily_instance
+    if _tavily_instance is None:
+        _tavily_instance = TavilySearch(max_results=5)
+    return _tavily_instance
 
 
 @tool
 def web_search(query: str) -> str:
     """Search the web for a query and return a short summary of results."""
     try:
-        results = _tavily.invoke({"query": query})
+        tavily = _get_tavily()
+    except Exception as e:
+        return f"Search unavailable: TAVILY_API_KEY missing or invalid ({e})"
+
+    try:
+        results = tavily.invoke({"query": query})
     except Exception as e:
         return f"Search failed: {e}"
 
@@ -177,8 +225,8 @@ Choose "finish" once you have enough information to answer the user directly.
 def manager_node(state: AgentState):
     turn_count = state.get("turn_count", 0) + 1
 
-    # Hard stop: don't even ask the LLM once we've hit the cap — force finish
-    # so a manager stuck in a routing loop can't burn turns indefinitely.
+    # Hard stop: don't even call the model once we've hit the cap — force
+    # finish so a manager stuck in a routing loop can't burn turns indefinitely.
     if turn_count > MAX_TURNS:
         return {
             "next": "finish",
@@ -186,23 +234,22 @@ def manager_node(state: AgentState):
             "messages": [SystemMessage(content=f"Max turns ({MAX_TURNS}) reached, finalizing with available results.")],
         }
 
-    llm = get_llm()
     messages = [SystemMessage(content=MANAGER_SYSTEM_PROMPT)] + state["messages"]
-    response = llm.invoke(messages)
+    raw_response = generate(_lc_to_dicts(messages), max_new_tokens=200)
 
     try:
-        raw = response.content.strip().strip("`")
+        raw = raw_response.strip().strip("`")
         if raw.startswith("json"):
             raw = raw[4:]
         decision = json.loads(raw)
         next_step = decision.get("next", "finish")
-    except (json.JSONDecodeError, AttributeError):
+    except json.JSONDecodeError:
         next_step = "finish"
 
     if next_step not in {"search_agent", "code_agent", "data_agent", "finish"}:
         next_step = "finish"
 
-    return {"next": next_step, "turn_count": turn_count, "messages": [response]}
+    return {"next": next_step, "turn_count": turn_count, "messages": [AIMessage(content=raw_response)]}
 
 
 def route_from_manager(state: AgentState) -> Literal["search_agent", "code_agent", "data_agent", "finalize"]:
@@ -214,19 +261,57 @@ def route_from_manager(state: AgentState) -> Literal["search_agent", "code_agent
 # ---------------------------------------------------------------------------
 # Specialist nodes — each does its work, appends a message, returns to manager
 # ---------------------------------------------------------------------------
-def make_specialist_node(name: str, tools, system_prompt: str):
+def _tool_spec_text(tools) -> str:
+    lines = []
+    for t in tools:
+        lines.append(f'- {t.name}: {t.description} Args schema: {t.args}')
+    return "\n".join(lines)
+
+
+def make_specialist_node(name: str, tools, role_description: str):
+    tool_spec = _tool_spec_text(tools)
+    system_prompt = (
+        f"{role_description}\n\n"
+        f"Available tools:\n{tool_spec}\n\n"
+        'Respond with ONLY a JSON object. Either call a tool: '
+        '{"tool": "<tool_name>", "args": {...}} — or, if no tool is needed, '
+        'give the answer directly: {"final": "<your answer>"}'
+    )
+
     def node(state: AgentState):
-        llm = get_llm().bind_tools(tools)
-        messages = [SystemMessage(content=system_prompt), HumanMessage(content=state["task"])]
-        ai_msg = llm.invoke(messages)
+        messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": state["task"]}]
+        raw_response = generate(messages, max_new_tokens=300)
 
-        tool_outputs = []
-        for call in getattr(ai_msg, "tool_calls", []) or []:
-            tool_fn = next(t for t in tools if t.name == call["name"])
-            output = tool_fn.invoke(call["args"])
-            tool_outputs.append(f"{call['name']}({call['args']}) -> {output}")
+        summary = None
+        try:
+            raw = raw_response.strip().strip("`")
+            if raw.startswith("json"):
+                raw = raw[4:]
+            decision = json.loads(raw)
+        except json.JSONDecodeError:
+            decision = None
+
+        if decision and "tool" in decision:
+            tool_fn = next((t for t in tools if t.name == decision["tool"]), None)
+            if tool_fn is not None:
+                try:
+                    tool_output = tool_fn.invoke(decision.get("args", {}))
+                except Exception as e:
+                    tool_output = f"Tool error: {e}"
+                # One follow-up call to turn the raw tool output into a
+                # concise summary, rather than dumping it unprocessed.
+                followup = messages + [
+                    {"role": "assistant", "content": raw_response},
+                    {"role": "user", "content": f"Tool result: {tool_output}\n\nSummarize this for the manager."},
+                ]
+                summary = generate(followup, max_new_tokens=250)
+            else:
+                summary = f"Requested unknown tool '{decision['tool']}'."
+        elif decision and "final" in decision:
+            summary = decision["final"]
+        else:
+            summary = raw_response  # model didn't follow the JSON protocol — use raw text as a fallback
 
-        summary = ai_msg.content or "\n".join(tool_outputs) or "(no output)"
         state["results"][name] = summary
         return {
             "messages": [HumanMessage(content=f"[{name}] {summary}")],
@@ -250,14 +335,13 @@ data_node = make_specialist_node(
 # Finalize node — synthesizes a final answer from gathered specialist results
 # ---------------------------------------------------------------------------
 def finalize_node(state: AgentState):
-    llm = get_llm()
     context = "\n".join(f"{k}: {v}" for k, v in state["results"].items())
-    prompt = [
-        SystemMessage(content="Write a final answer for the user based on the specialist findings below."),
-        HumanMessage(content=f"Task: {state['task']}\n\nFindings:\n{context}"),
+    messages = [
+        {"role": "system", "content": "Write a final answer for the user based on the specialist findings below."},
+        {"role": "user", "content": f"Task: {state['task']}\n\nFindings:\n{context}"},
     ]
-    response = llm.invoke(prompt)
-    return {"messages": [response]}
+    response = generate(messages, max_new_tokens=500)
+    return {"messages": [AIMessage(content=response)]}
 
 
 # ---------------------------------------------------------------------------
@@ -309,4 +393,4 @@ def run_task(app, task: str, session_state: dict) -> str:
     answer = final_state["messages"][-1].content
 
     session_state.setdefault("history", []).append({"task": task, "answer": answer})
-    return answer
+    return answer