agent.py

"""
Weather Activity Agent - LangGraph + Qwen (via Ollama)
"""

from langgraph.graph import StateGraph, END
from langchain_ollama import ChatOllama
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, ToolMessage
from typing import TypedDict, Annotated
import operator
import json

from tools.weather import get_weather
from tools.activities import get_activity_recommendations

# ── 1. State Definition ─────────────────────────────────────────────────────
# This is the "memory" that flows through every node in the graph.
class AgentState(TypedDict):
    messages: Annotated[list, operator.add]   # Full message history
    location: str                              # User-provided location
    weather_data: dict                         # Raw weather fetched from API
    final_answer: str                          # Agent's final response


# ── 2. LLM Setup ─────────────────────────────────────────────────────────────
# Point LangChain at your local Ollama instance running Qwen
llm = ChatOllama(
    model="qwen3.5:9b",     # Change to whatever Qwen model you've pulled
    base_url="http://localhost:11434",
    temperature=0.7,
)

# Bind our tools so the LLM knows it can call them
tools = [get_weather, get_activity_recommendations]
llm_with_tools = llm.bind_tools(tools)

SYSTEM_PROMPT = """You are a helpful weather and activity assistant.

When a user asks about activities or what to do, you should:
1. First use get_weather to fetch current conditions for their location
2. Then use get_activity_recommendations to suggest suitable activities
3. Finally, summarize both in a friendly, conversational response

Always be specific about WHY certain activities are recommended based on the weather.
"""


# ── 3. Node Functions ─────────────────────────────────────────────────────────
# Each node receives the full state, does something, and returns state updates.

def call_llm(state: AgentState) -> dict:
    """Main LLM node — decides what tool to call next, or generates final answer."""
    messages = [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}


def run_tools(state: AgentState) -> dict:
    """Tool executor node — runs whatever tool the LLM decided to call."""
    tool_map = {t.name: t for t in tools}
    last_message = state["messages"][-1]
    results = []

    for tool_call in last_message.tool_calls:
        tool_name = tool_call["name"]
        tool_args = tool_call["args"]

        print(f"  🔧 Running tool: {tool_name}({tool_args})")

        if tool_name in tool_map:
            result = tool_map[tool_name].invoke(tool_args)
        else:
            result = f"Error: tool '{tool_name}' not found."

        results.append(
            ToolMessage(
                content=str(result),
                tool_call_id=tool_call["id"],
                name=tool_name
            )
        )

    return {"messages": results}


def should_continue(state: AgentState) -> str:
    """
    Routing function — the 'brain' of the graph loop.
    Returns 'tools' to keep looping, or 'end' to finish.
    """
    last_message = state["messages"][-1]

    # If the last LLM message has tool calls → go run them
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        return "tools"

    # Otherwise the LLM is done — extract the final answer
    return "end"


def extract_final(state: AgentState) -> dict:
    """Pulls the final text response out of the last AI message."""
    last_message = state["messages"][-1]
    return {"final_answer": last_message.content}


# ── 4. Build the Graph ────────────────────────────────────────────────────────
def build_agent():
    graph = StateGraph(AgentState)

    # Add nodes
    graph.add_node("llm", call_llm)
    graph.add_node("tools", run_tools)
    graph.add_node("extract", extract_final)

    # Entry point
    graph.set_entry_point("llm")

    # Conditional routing after LLM responds
    graph.add_conditional_edges(
        "llm",
        should_continue,
        {
            "tools": "tools",   # LLM wants to call a tool → run it
            "end": "extract",   # LLM is done → extract answer
        }
    )

    # After tools run, always go back to the LLM to process results
    graph.add_edge("tools", "llm")
    graph.add_edge("extract", END)

    return graph.compile()


# ── 5. Run It ─────────────────────────────────────────────────────────────────
def run_agent(user_input: str, location: str):
    agent = build_agent()

    print(f"\n{'='*50}")
    print(f"📍 Location: {location}")
    print(f"💬 Query: {user_input}")
    print(f"{'='*50}\n")

    initial_state = {
        "messages": [HumanMessage(content=f"Location: {location}\n\nUser question: {user_input}")],
        "location": location,
        "weather_data": {},
        "final_answer": "",
    }

    result = agent.invoke(initial_state)

    print(f"\n{'='*50}")
    print("✅ FINAL ANSWER:")
    print(f"{'='*50}")
    print(result["final_answer"])
    return result["final_answer"]


if __name__ == "__main__":
    # Example run — change location to wherever you want
    run_agent(
        user_input="What should I do outside today?",
        location="Fremont, CA"
    )