agent.py

from typing import List, Dict, Optional
from langgraph.graph import StateGraph, END
from pydantic import BaseModel
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
from os import getenv
from utils import AgentState
from agents_utils import make_generate_sql_node, run_sql_node, make_generate_summary_node, make_router_node, chat_node

load_dotenv()


llm = ChatOpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=getenv("OPENROUTER_API_KEY"),
    model="openai/gpt-oss-20b:free"
) 

router_node = make_router_node(llm)
generate_sql_node = make_generate_sql_node(llm)
generate_summary_node = make_generate_summary_node(llm)

graph = StateGraph(AgentState)

graph.add_node("route", router_node)
graph.add_node("chat", chat_node)

graph.add_node("generate_sql", generate_sql_node)
graph.add_node("run_sql", run_sql_node)
graph.add_node("summary", generate_summary_node)

graph.set_entry_point("route")

graph.add_conditional_edges(
    "route",
    lambda state: state.route,
    {
        "chat": "chat",
        "sql": "generate_sql",
    },
)

graph.add_edge("chat", END)

graph.add_edge("generate_sql", "run_sql")
graph.add_edge("run_sql", "summary")
graph.add_edge("summary", END)

agent = graph.compile()