agent/tools.py

from __future__ import annotations

from typing import Any

from rag.reranker import rerank
from rag.vector_store import retrieve

MOCK_CLAIMS = {
    "CLM-1001": {"status": "approved", "updated_at": "2026-04-08", "amount": 1240.0},
    "CLM-1002": {"status": "under_review", "updated_at": "2026-04-09", "amount": None},
    "CLM-1003": {"status": "denied", "updated_at": "2026-04-07", "amount": 0.0},
    "CLM-1004": {"status": "paid", "updated_at": "2026-04-10", "amount": 5600.0},
}


def search_policy(
    query: str,
    use_reranker: bool = True,
    top_k: int = 4,
) -> dict[str, Any]:
    chunks = retrieve(query, k=top_k)
    if use_reranker and len(chunks) > 1:
        chunks = rerank(query, chunks, top_k=min(3, len(chunks)))

    return {
        "query": query,
        "count": len(chunks),
        "results": [
            {
                "source": c["source"],
                "section": c.get("section", "Unknown section"),
                "score": c.get("score"),
                "rerank_score": c.get("rerank_score"),
                "content": c["content"],
            }
            for c in chunks
        ],
    }


def calculate_premium(coverage: float, risk_score: float) -> dict[str, Any]:
    coverage_val = float(coverage)
    risk_val = float(risk_score)
    monthly = round(coverage_val * 0.002 * risk_val, 2)
    return {
        "formula": "coverage * 0.002 * risk_score",
        "coverage": coverage_val,
        "risk_score": risk_val,
        "monthly_premium": monthly,
    }


def check_claim_status(claim_id: str) -> dict[str, Any]:
    claim = MOCK_CLAIMS.get(claim_id.strip().upper())
    if claim is None:
        return {
            "claim_id": claim_id,
            "status": "not_found",
            "message": "Claim ID not found in mock system.",
        }
    return {"claim_id": claim_id.strip().upper(), **claim}


TOOL_SCHEMAS = {
    "search_policy": {
        "description": "Query insurance policy knowledge via RAG retrieval.",
        "arguments": {"query": "string"},
    },
    "calculate_premium": {
        "description": "Calculate mock monthly premium.",
        "arguments": {"coverage": "number", "risk_score": "number"},
    },
    "check_claim_status": {
        "description": "Lookup mock claim status from hardcoded map.",
        "arguments": {"claim_id": "string"},
    },
}