src/agent.py

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 sf_tools as tls

# 🔑 Load .env file
load_dotenv()

groq_api_key = os.getenv("GROQ_API_KEY")
groq_model = os.getenv("MODEL")

if not groq_api_key:
    raise ValueError("GROQ_API_KEY is not set in .env")

if not groq_model:
    raise ValueError("MODEL is not set in .env")

# Optional (only needed if LiteLLM expects env var)
os.environ["GROQ_API_KEY"] = groq_api_key

# Create the agent with ADK callbacks (logging + optional audit sink)
root_agent = Agent(
    model=LiteLlm(
        model=groq_model,
        custom_llm_provider="groq"
    ),
    name=prmpt.AGENT_CONFIG["name"],
    description=prmpt.AGENT_CONFIG["description"],
    instruction=prmpt.AGENT_CONFIG["instruction"],
    tools=[tls.check_leave_eligibility, tls.get_employee, tls.get_leave_policy],
    # ADK callbacks: observe agent/model/tool lifecycle and optionally persist to audit sink
    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,
)