Upload 5 files

src/__init__.py

@@ -0,0 +1 @@
+from . import agent

src/agent.py

@@ -0,0 +1,57 @@
+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 policy_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
+
+# 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],
+# )
+
+import sf_tools as tls # replace with actual Tool class in your agent
+
+# Now create the agent
+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],
+)
+
+
+
+
+
+
+
+
+
+

src/main.py

@@ -0,0 +1,99 @@
+import asyncio
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+from google.adk.sessions import InMemorySessionService
+from google.adk.runners import Runner
+from google.genai import types
+
+from agent import root_agent  # your agent file
+
+# =========================
+# FASTAPI APP
+# =========================
+app = FastAPI(
+    title="Leave Policy Assistant",
+    description="AI agent to answer leave policy questions",
+    version="1.0.0"
+)
+
+# =========================
+# ADK SESSION SETUP
+# =========================
+APP_NAME = "leave_policy_app"
+session_service = InMemorySessionService()
+
+runner = Runner(
+    agent=root_agent,
+    app_name=APP_NAME,
+    session_service=session_service,
+)
+
+# =========================
+# REQUEST / RESPONSE MODELS
+# =========================
+class ChatRequest(BaseModel):
+    user_id: str
+    message: str
+
+
+class ChatResponse(BaseModel):
+    response: str
+
+
+# =========================
+# CHAT ENDPOINT
+# =========================
+@app.post("/chat", response_model=ChatResponse)
+async def chat(req: ChatRequest):
+    try:
+        # Generate session_id automatically per user
+        session_id = f"{req.user_id}_session"
+
+        # Check if session exists, create if not
+        session = await session_service.get_session(
+            app_name=APP_NAME,
+            user_id=req.user_id,
+            session_id=session_id,
+        )
+        if not session:
+            await session_service.create_session(
+                app_name=APP_NAME,
+                user_id=req.user_id,
+                session_id=session_id,
+            )
+
+        # Prepare the user message
+        user_content = types.Content(
+            role="user",
+            parts=[types.Part(text=req.message)],
+        )
+
+        final_response = None
+
+        # Run the agent asynchronously
+        async for event in runner.run_async(
+            user_id=req.user_id,
+            session_id=session_id,
+            new_message=user_content,
+        ):
+            if event.is_final_response():
+                if event.content and event.content.parts:
+                    final_response = event.content.parts[0].text
+                break
+
+        if not final_response:
+            raise HTTPException(status_code=500, detail="No response from agent")
+
+        return ChatResponse(response=final_response)
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+# =========================
+# HEALTH CHECK
+# =========================
+@app.get("/health")
+def health():
+    return {"status": "ok"}

src/prompt.py

@@ -0,0 +1,33 @@
+AGENT_CONFIG = {
+    "name": "vishal",
+    "description": "A Leave Policy Assistant Agent built using LiteLLM",
+    "instruction": """
+You are Vishal, an intelligent and reliable HR Leave Policy Assistant.
+
+Your primary role is to help employees understand company leave policies and
+check their leave eligibility accurately.
+
+What you should do:
+- Answer questions related to leave policies (PTO, Sick Leave, Casual Leave, etc.).
+- Check leave eligibility using the available tools.
+- Explain rules such as allowances, carryover limits, notice periods, and eligibility criteria.
+- Provide clear, concise, and employee-friendly responses.
+
+How to behave:
+- Always use tools when employee data or leave policy data is required.
+- Never guess or fabricate employee information.
+- If the employee ID is missing, politely ask for it.
+- If the leave type is unknown, suggest valid leave types for the employee’s country.
+- If the employee is inactive, clearly state that leave is not applicable.
+- Handle invalid inputs (wrong dates, negative days, unknown leave types) gracefully.
+
+Conversation handling:
+- Maintain context across multiple messages.
+- Ask follow-up questions when required information is missing.
+- Be professional, calm, and helpful at all times.
+
+Important:
+- Accuracy is more important than speed.
+- If something cannot be determined, explain why instead of guessing.
+"""
+}

src/sf_tools.py

@@ -0,0 +1,158 @@
+from snowflake.connector import connect  # type: ignore[import-untyped]
+from datetime import datetime
+import os
+import json
+from dotenv import load_dotenv
+
+# 🔑 Load .env file
+load_dotenv()
+# -----------------------------
+# Snowflake connection config
+# -----------------------------
+
+def get_snowflake_connection():
+    conn = connect(
+        user=os.getenv("SNOWFLAKE_USER"),
+        password=os.getenv("SNOWFLAKE_PASSWORD"),
+        account=os.getenv("SNOWFLAKE_ACCOUNT"),
+        role=os.getenv("SNOWFLAKE_ROLE"),
+        warehouse=os.getenv("SNOWFLAKE_WAREHOUSE"),
+        database=os.getenv("SNOWFLAKE_DATABASE"),
+        schema=os.getenv("SNOWFLAKE_SCHEMA")
+    )
+    return conn
+
+# -----------------------------
+# Fetch Employee Details
+# -----------------------------
+def get_employee(employee_id: str):
+    """Fetch employee details from Snowflake."""
+    print("get_employee called")
+    conn = get_snowflake_connection()
+    cur = conn.cursor()
+    try:
+        cur.execute("""
+            SELECT employee_id, country, employment_status
+            FROM employees
+            WHERE employee_id = %s
+        """, (employee_id,))
+        row = cur.fetchone()
+        if not row:
+            return {"error": f"Employee '{employee_id}' not found"}
+        
+        employee = {
+            "employee_id": row[0],
+            "country": row[1],
+            "employment_status": row[2]
+        }
+
+        if employee["employment_status"] != "active":
+            return {"error": f"Employee '{employee_id}' is not active"}
+
+        # Fetch leave balances
+        cur.execute("""
+            SELECT leave_type, balance
+            FROM leave_balances
+            WHERE employee_id = %s
+        """, (employee_id,))
+        balances = {r[0]: r[1] for r in cur.fetchall()}
+        employee["leave_balances"] = balances
+
+        return employee
+    finally:
+        cur.close()
+        conn.close()
+
+
+# -----------------------------
+# Fetch Leave Policy
+# -----------------------------
+def get_leave_policy(country: str, leave_type: str):
+    """Fetch leave policy from Snowflake."""
+    print("get_leave_policy called")
+    conn = get_snowflake_connection()
+    cur = conn.cursor()
+    try:
+        cur.execute("""
+            SELECT policy
+            FROM leave_policies
+            WHERE country = %s AND leave_type = %s
+        """, (country, leave_type))
+        row = cur.fetchone()
+        if not row:
+            # fetch available leave types
+            cur.execute("""
+                SELECT leave_type
+                FROM leave_policies
+                WHERE country = %s
+            """, (country,))
+            leave_types = [r[0] for r in cur.fetchall()]
+            return {
+                "error": f"Leave type '{leave_type}' not found",
+                "available_leave_types": leave_types
+            }
+        # policy is VARIANT JSON in Snowflake (may be dict or JSON string)
+        raw = row[0]
+        if isinstance(raw, dict):
+            return raw
+        if isinstance(raw, str):
+            return json.loads(raw)
+        return dict(raw)
+    finally:
+        cur.close()
+        conn.close()
+
+
+# -----------------------------
+# Check Leave Eligibility
+# -----------------------------
+def check_leave_eligibility(employee_id: str, leave_type: str, requested_days: int):
+    """Check if employee is eligible to take leave."""
+    print("check_leave_eligibility called")
+    if requested_days <= 0:
+        return {"eligible": False, "reason": "Requested leave days must be > 0"}
+
+    employee = get_employee(employee_id)
+    if "error" in employee:
+        return employee
+
+    country = employee["country"]
+    policy = get_leave_policy(country, leave_type)
+    if "error" in policy:
+        return policy
+
+    balance = employee["leave_balances"].get(leave_type, 0)
+
+    # Max consecutive days
+    if "max_consecutive_days" in policy:
+        if requested_days > policy["max_consecutive_days"]:
+            return {
+                "eligible": False,
+                "reason": f"Maximum consecutive days allowed for {leave_type} is {policy['max_consecutive_days']}"
+            }
+
+    # Annual allowance
+    if "annual_allowance" in policy:
+        if requested_days > policy["annual_allowance"]:
+            return {
+                "eligible": False,
+                "reason": "Requested days exceed annual allowance"
+            }
+
+    # Available balance
+    if requested_days > balance:
+        return {
+            "eligible": False,
+            "reason": "Insufficient leave balance",
+            "available_balance": balance
+        }
+
+    return {
+        "eligible": True,
+        "employee_id": employee_id,
+        "leave_type": leave_type,
+        "approved_days": requested_days,
+        "remaining_balance": balance - requested_days
+    }
+
+