src/README.md

Leave Policy Assistant

An AI-powered HR Leave Policy Assistant that answers leave policy questions and checks leave eligibility. Built with Google ADK (Agent Development Kit), LiteLLM (Groq), FastAPI, and Snowflake.

Features

• Leave policy Q&A — Answers questions about PTO, Sick Leave, Casual Leave, allowances, carryover, notice periods, and eligibility.
• Leave eligibility checks — Uses employee and policy data (Snowflake or in-memory) to determine if a leave request is valid.
• Conversational agent — Maintains session context and asks for missing info (e.g. employee ID, leave type).
• REST API — /chat and /health endpoints for integration with frontends or other services.
• ADK callbacks — Lifecycle hooks (before/after agent, model, tool) for logging and optional audit.
• Optional audit — Persist callback events to Snowflake via audit_db.py when ENABLE_AUDIT_SINK=1.

Project Structure

policy/
├── main.py         # FastAPI app: /chat, /health; optional audit sink at startup
├── agent.py        # ADK agent (LiteLLM/Groq), tools + callbacks
├── callback.py     # ADK callbacks (before/after agent, model, tool); audit sink wiring
├── prompt.py       # Agent name, description, system instruction
├── sf_tools.py     # Snowflake tools: get_employee, get_leave_policy, check_leave_eligibility (default)
├── policy_tools.py # In-memory tools using data.py (optional; switch in agent.py)
├── data.py         # Static LEAVE_POLICIES and EMPLOYEE_DATA (used by policy_tools)
├── audit_db.py     # Optional Snowflake audit sink; table agent_audit_events
├── .env            # Secrets (not committed)
├── requirements.txt
└── README.md

Prerequisites

• Python 3.10+
• Groq API key — for the LLM (e.g. llama-3.1-8b-instant)
• Snowflake — account with:
  • Tables: employees, leave_balances, leave_policies
  • Optional: agent_audit_events if using audit

Setup

1. Clone / open the project and go to the policy folder:

   cd policy
   

2. Create a virtual environment (recommended):

   python -m venv venv
   venv\Scripts\activate   # Windows
   # source venv/bin/activate   # Linux/macOS
   

3. Install dependencies:

   pip install -r requirements.txt
   

4. Configure environment — Create a .env file in policy/:

   MODEL=llama-3.1-8b-instant
   GROQ_API_KEY=your_groq_api_key
   
   SNOWFLAKE_USER=your_user
   SNOWFLAKE_PASSWORD=your_password
   SNOWFLAKE_ACCOUNT=your_account
   SNOWFLAKE_ROLE=your_role
   SNOWFLAKE_WAREHOUSE=your_warehouse
   SNOWFLAKE_DATABASE=your_database
   SNOWFLAKE_SCHEMA=your_schema
   
   # Optional: persist ADK callback events to Snowflake (table agent_audit_events)
   ENABLE_AUDIT_SINK=1
   

   Replace placeholders with your Groq and Snowflake credentials. Do not commit .env.

Running the App

Start the FastAPI server:

uvicorn main:app --reload --host 0.0.0.0 --port 8000

• API docs: http://localhost:8000/docs
• Health: http://localhost:8000/health

API Usage

Chat

POST /chat

Request body:

{
  "user_id": "user_123",
  "message": "What is my PTO balance? My employee ID is E001."
}

Response:

{
  "response": "Your PTO balance is 15 days..."
}

Sessions are keyed by user_id (one session per user). The agent uses tools to fetch employee data, leave policies, and check eligibility as needed.

Health

GET /health

Returns {"status": "ok"}.

Agent Tools

The agent is wired to Snowflake by default (sf_tools.py). You can switch to in-memory data by using policy_tools and data.py in agent.py.

Tool	Description
get_employee	Fetches employee details and leave balances (Snowflake or data.py).
get_leave_policy	Fetches leave policy for a country and leave type.
check_leave_eligibility	Checks if an employee can take a given number of days (balance, policy rules).

The agent is instructed to always use these tools for employee/policy data and not to guess or fabricate information.

• sf_tools.py — Requires Snowflake and .env credentials; reads from employees, leave_balances, leave_policies.
• policy_tools.py — Uses in-memory LEAVE_POLICIES and EMPLOYEE_DATA from data.py; no Snowflake needed for tools (Snowflake still needed for optional audit).

ADK Callbacks

The agent uses ADK callbacks (see callback.py) to observe and optionally persist events:

• before_agent / after_agent — before and after the agent runs for a request
• before_model / after_model — before and after each LLM call
• before_tool / after_tool — before and after each tool execution
• on_model_error / on_tool_error — when the model or a tool fails

Callbacks log at DEBUG and, if an audit sink is set, call sink.store(event) for each event. This allows logging, monitoring, and optional persistence without changing agent logic.

Optional: Audit to Snowflake

To persist callback events to Snowflake:

1. Add ENABLE_AUDIT_SINK=1 to your .env.
2. On startup, main.py will call callback.set_audit_sink(SnowflakeAuditSink()) and ensure_table(), creating the agent_audit_events table if needed. Events (before_agent, after_tool, etc.) are then stored in Snowflake.

Troubleshooting

• 503 — Agent did not return a response; check that the LLM (Groq) and, if used, Snowflake are reachable and credentials in .env are correct.
• 500 — Unexpected server error; check the uvicorn terminal logs for the full traceback (logger.exception).
• 400 — Invalid request (e.g. missing session or invalid runner params); the detail field in the response has the message.

License

Use as needed for your organization.