# 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:

   ```bash
   cd policy
   ```

2. **Create a virtual environment** (recommended):

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

3. **Install dependencies**:

   ```bash
   pip install -r requirements.txt
   ```

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

   ```env
   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:

```bash
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:

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

Response:

```json
{
  "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.