src/agent.py

"""
ADK agent for the Handbook Generator.
Uses Google ADK + LiteLLM (OpenAI) + RAG tools.
"""
import os
from dotenv import load_dotenv
from google.adk.agents.llm_agent import Agent
from google.adk.models.lite_llm import LiteLlm

import prompt as prmpt
import callback as cb
import rag_tools as tls

# 🔑 Load .env file
load_dotenv()

openai_api_key = os.getenv("OPENAI_API_KEY")
model = os.getenv("MODEL", "openai/gpt-4o")

if not openai_api_key:
    raise ValueError("OPENAI_API_KEY is not set in .env")

os.environ["OPENAI_API_KEY"] = openai_api_key

# Create the agent (same pattern as Policy)
root_agent = Agent(
    model=LiteLlm(
        model=model,
    ),
    name=prmpt.AGENT_CONFIG["name"],
    description=prmpt.AGENT_CONFIG["description"],
    instruction=prmpt.AGENT_CONFIG["instruction"],
    tools=[tls.query_uploaded_documents],
    before_agent_callback=cb.before_agent_callback,
    after_agent_callback=cb.after_agent_callback,
    before_model_callback=cb.before_model_callback,
    after_model_callback=cb.after_model_callback,
    before_tool_callback=cb.before_tool_callback,
    after_tool_callback=cb.after_tool_callback,
    on_model_error_callback=cb.on_model_error_callback,
    on_tool_error_callback=cb.on_tool_error_callback,
)