ai-tax-agent/apps/svc_rag_retriever/main.py

# FILE: apps/svc-rag-retriever/main.py
# mypy: disable-error-code=union-attr
# Hybrid search with KG fusion, reranking, and calibrated confidence

import os

# Import shared libraries
import sys
from datetime import datetime
from typing import Any

import structlog
from fastapi import Depends, HTTPException, Query, Request
from fastapi.responses import JSONResponse
from qdrant_client.models import SparseVector

sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))

from libs.app_factory import create_app
from libs.calibration import ConfidenceCalibrator
from libs.config import (
    BaseAppSettings,
    create_event_bus,
    create_neo4j_client,
    create_qdrant_client,
)
from libs.events import EventBus
from libs.neo import Neo4jClient
from libs.observability import get_metrics, get_tracer, setup_observability
from libs.rag import RAGRetriever
from libs.schemas import ErrorResponse, RAGSearchRequest, RAGSearchResponse
from libs.security import get_current_user, get_tenant_id

logger = structlog.get_logger()


class RAGRetrieverSettings(BaseAppSettings):
    """Settings for RAG retriever service"""

    service_name: str = "svc-rag-retriever"

    # Embedding configuration
    embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2"
    embedding_dimension: int = 384

    # Search configuration
    default_k: int = 10
    max_k: int = 100
    alpha: float = 0.5  # Dense/sparse balance
    beta: float = 0.3  # Vector/KG balance
    gamma: float = 0.2  # Reranking weight

    # Collections to search
    search_collections: list[str] = ["documents", "tax_rules", "guidance"]

    # Reranking
    reranker_model: str | None = None
    rerank_top_k: int = 50


# Create app and settings
app, settings = create_app(
    service_name="svc-rag-retriever",
    title="Tax Agent RAG Retriever Service",
    description="Hybrid search with KG fusion and reranking",
    settings_class=RAGRetrieverSettings,
)

# Global clients
qdrant_client = None
neo4j_client: Neo4jClient | None = None
rag_retriever: RAGRetriever | None = None
event_bus: EventBus | None = None
embedding_model = None
confidence_calibrator: ConfidenceCalibrator | None = None
tracer = get_tracer("svc-rag-retriever")
metrics = get_metrics()


@app.on_event("startup")
async def startup_event() -> None:
    """Initialize service dependencies"""
    global qdrant_client, neo4j_client, rag_retriever, event_bus, embedding_model, confidence_calibrator

    logger.info("Starting RAG retriever service")

    # Setup observability
    setup_observability(settings)

    # Initialize Qdrant client
    qdrant_client = create_qdrant_client(settings)

    # Initialize Neo4j client
    neo4j_driver = create_neo4j_client(settings)
    neo4j_client = Neo4jClient(neo4j_driver)

    # Initialize RAG retriever
    rag_retriever = RAGRetriever(
        qdrant_client=qdrant_client,
        neo4j_client=neo4j_client,
        reranker_model=settings.reranker_model,
    )

    # Initialize embedding model
    try:
        from sentence_transformers import SentenceTransformer

        embedding_model = SentenceTransformer(settings.embedding_model)
        logger.info("Embedding model loaded", model=settings.embedding_model)
    except ImportError:
        logger.warning("sentence-transformers not available, using mock embeddings")
        embedding_model = None

    # Initialize confidence calibrator
    confidence_calibrator = ConfidenceCalibrator(method="isotonic")

    # Initialize event bus
    event_bus = create_event_bus(settings)
    await event_bus.start() # fmt: skip # pyright: ignore[reportOptionalMemberAccess]

    logger.info("RAG retriever service started successfully")


@app.on_event("shutdown")
async def shutdown_event() -> None:
    """Cleanup service dependencies"""
    global neo4j_client, event_bus

    logger.info("Shutting down RAG retriever service")

    if neo4j_client:
        await neo4j_client.close()

    if event_bus:
        await event_bus.stop()

    logger.info("RAG retriever service shutdown complete")


@app.get("/health")
async def health_check() -> dict[str, Any]:
    """Health check endpoint"""
    return {
        "status": "healthy",
        "service": settings.service_name,
        "version": settings.service_version,
        "timestamp": datetime.utcnow().isoformat(),
        "search_collections": settings.search_collections,
    }


@app.post("/search", response_model=RAGSearchResponse)
async def search(
    request_data: RAGSearchRequest,
    current_user: dict[str, Any] = Depends(get_current_user),
    tenant_id: str = Depends(get_tenant_id),
) -> RAGSearchResponse:
    """Perform hybrid RAG search"""

    with tracer.start_as_current_span("rag_search") as span:
        span.set_attribute("query", request_data.query[:100])
        span.set_attribute("tenant_id", tenant_id)
        span.set_attribute("k", request_data.k)

        try:
            # Generate embeddings for query
            dense_vector = await _generate_embedding(request_data.query)
            sparse_vector = await _generate_sparse_vector(request_data.query)

            # Perform search
            search_results = await rag_retriever.search(  # fmt: skip # pyright: ignore[reportOptionalMemberAccess]
                query=request_data.query,
                collections=settings.search_collections,
                dense_vector=dense_vector,
                sparse_vector=sparse_vector,
                k=request_data.k,
                alpha=settings.alpha,
                beta=settings.beta,
                gamma=settings.gamma,
                tax_year=request_data.tax_year,
                jurisdiction=request_data.jurisdiction,
            )

            # Update metrics
            metrics.counter("searches_total").labels(tenant_id=tenant_id).inc()

            metrics.histogram("search_results_count").labels(
                tenant_id=tenant_id
            ).observe(len(search_results["chunks"]))

            metrics.histogram("search_confidence").labels(tenant_id=tenant_id).observe(
                search_results["calibrated_confidence"]
            )

            logger.info(
                "RAG search completed",
                query=request_data.query[:50],
                results=len(search_results["chunks"]),
                confidence=search_results["calibrated_confidence"],
            )

            return RAGSearchResponse(
                chunks=search_results["chunks"],
                citations=search_results["citations"],
                kg_hints=search_results["kg_hints"],
                calibrated_confidence=search_results["calibrated_confidence"],
            )

        except Exception as e:
            logger.error(
                "RAG search failed", query=request_data.query[:50], error=str(e)
            )

            # Update error metrics
            metrics.counter("search_errors_total").labels(
                tenant_id=tenant_id, error_type=type(e).__name__
            ).inc()

            raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}")


@app.get("/similar/{doc_id}")
async def find_similar_documents(
    doc_id: str,
    k: int = Query(default=10, le=settings.max_k),
    current_user: dict[str, Any] = Depends(get_current_user),
    tenant_id: str = Depends(get_tenant_id),
) -> dict[str, Any]:
    """Find documents similar to given document"""

    with tracer.start_as_current_span("find_similar") as span:
        span.set_attribute("doc_id", doc_id)
        span.set_attribute("tenant_id", tenant_id)
        span.set_attribute("k", k)

        try:
            # Get document content from vector database
            # This would search for the document by doc_id in metadata
            from qdrant_client.models import FieldCondition, Filter, MatchValue

            filter_conditions = Filter(
                must=[
                    FieldCondition(key="doc_id", match=MatchValue(value=doc_id)),
                    FieldCondition(key="tenant_id", match=MatchValue(value=tenant_id)),
                ]
            )

            # Search for the document
            doc_results = await rag_retriever.collection_manager.search_dense(  # fmt: skip # pyright: ignore[reportOptionalMemberAccess]
                collection_name="documents",
                query_vector=[0.0] * settings.embedding_dimension,  # Dummy vector
                limit=1,
                filter_conditions=filter_conditions,
            )

            if not doc_results:
                raise HTTPException(status_code=404, detail="Document not found")

            # Get the document's vector and use it for similarity search
            doc_vector = doc_results[0]["payload"].get("vector")
            if not doc_vector:
                raise HTTPException(status_code=400, detail="Document has no vector")

            # Find similar documents
            similar_results = await rag_retriever.collection_manager.search_dense(  # fmt: skip # pyright: ignore[reportOptionalMemberAccess]
                collection_name="documents",
                query_vector=doc_vector,
                limit=k + 1,  # +1 to exclude the original document
                filter_conditions=Filter(
                    must=[
                        FieldCondition(
                            key="tenant_id", match=MatchValue(value=tenant_id)
                        )
                    ],
                    must_not=[
                        FieldCondition(key="doc_id", match=MatchValue(value=doc_id))
                    ],
                ),
            )

            return {
                "doc_id": doc_id,
                "similar_documents": similar_results[:k],
                "count": len(similar_results[:k]),
            }

        except HTTPException:
            raise
        except Exception as e:
            logger.error("Similar document search failed", doc_id=doc_id, error=str(e))
            raise HTTPException(
                status_code=500, detail=f"Similar search failed: {str(e)}"
            )


@app.post("/explain")
async def explain_search(
    query: str,
    search_results: list[dict[str, Any]],
    current_user: dict[str, Any] = Depends(get_current_user),
    tenant_id: str = Depends(get_tenant_id),
) -> dict[str, Any]:
    """Explain search results and ranking"""

    with tracer.start_as_current_span("explain_search") as span:
        span.set_attribute("query", query[:100])
        span.set_attribute("tenant_id", tenant_id)
        span.set_attribute("results_count", len(search_results))

        try:
            explanations = []

            for i, result in enumerate(search_results):
                explanation = {
                    "rank": i + 1,
                    "chunk_id": result.get("id"),
                    "score": result.get("score", 0.0),
                    "dense_score": result.get("dense_score", 0.0),
                    "sparse_score": result.get("sparse_score", 0.0),
                    "collection": result.get("collection"),
                    "explanation": _generate_explanation(query, result),
                }
                explanations.append(explanation)

            return {
                "query": query,
                "explanations": explanations,
                "ranking_factors": {
                    "alpha": settings.alpha,
                    "beta": settings.beta,
                    "gamma": settings.gamma,
                },
            }

        except Exception as e:
            logger.error("Search explanation failed", error=str(e))
            raise HTTPException(status_code=500, detail=f"Explanation failed: {str(e)}")


async def _generate_embedding(text: str) -> list[float]:
    """Generate dense embedding for text"""
    if embedding_model:
        try:
            embedding = embedding_model.encode(text)
            return embedding.tolist()
        except Exception as e:
            logger.error("Failed to generate embedding", error=str(e))

    # Fallback: random embedding
    import random

    return [random.random() for _ in range(settings.embedding_dimension)]


async def _generate_sparse_vector(text: str) -> SparseVector:
    """Generate sparse vector for text (BM25-style)"""
    try:
        # This would use a proper sparse encoder like SPLADE
        # For now, create a simple sparse representation
        from qdrant_client.models import SparseVector

        # Simple word-based sparse vector
        words = text.lower().split()
        word_counts: dict[str, int] = {}
        for word in words:
            word_counts[word] = word_counts.get(word, 0) + 1

        # Convert to sparse vector format
        indices = []
        values = []

        for _i, (word, count) in enumerate(word_counts.items()):
            # Use hash of word as index
            word_hash = hash(word) % 10000  # Limit vocabulary size
            indices.append(word_hash)
            values.append(float(count))

        return SparseVector(indices=indices, values=values)

    except Exception as e:
        logger.error("Failed to generate sparse vector", error=str(e))
        # Return empty sparse vector
        from qdrant_client.models import SparseVector

        return SparseVector(indices=[], values=[])


def _generate_explanation(query: str, result: dict[str, Any]) -> str:
    """Generate human-readable explanation for search result"""

    explanations = []

    # Score explanation
    score = result.get("score", 0.0)
    dense_score = result.get("dense_score", 0.0)
    sparse_score = result.get("sparse_score", 0.0)

    explanations.append(f"Overall score: {score:.3f}")

    if dense_score > 0:
        explanations.append(f"Semantic similarity: {dense_score:.3f}")

    if sparse_score > 0:
        explanations.append(f"Keyword match: {sparse_score:.3f}")

    # Collection explanation
    collection = result.get("collection")
    if collection:
        explanations.append(f"Source: {collection}")

    # Metadata explanation
    payload = result.get("payload", {})
    doc_id = payload.get("doc_id")
    if doc_id:
        explanations.append(f"Document: {doc_id}")

    confidence = payload.get("confidence")
    if confidence:
        explanations.append(f"Extraction confidence: {confidence:.3f}")

    return "; ".join(explanations)


@app.get("/stats")
async def get_search_stats(
    current_user: dict[str, Any] = Depends(get_current_user),
    tenant_id: str = Depends(get_tenant_id),
) -> dict[str, Any]:
    """Get search statistics"""

    try:
        # This would aggregate metrics from Prometheus
        # For now, return mock stats
        stats = {
            "total_searches": 1000,
            "avg_results_per_search": 8.5,
            "avg_confidence": 0.75,
            "collections": {
                "documents": {"searches": 800, "avg_confidence": 0.78},
                "tax_rules": {"searches": 150, "avg_confidence": 0.85},
                "guidance": {"searches": 50, "avg_confidence": 0.70},
            },
            "top_queries": [
                {"query": "capital gains tax", "count": 45},
                {"query": "business expenses", "count": 38},
                {"query": "property income", "count": 32},
            ],
        }

        return stats

    except Exception as e:
        logger.error("Failed to get search stats", error=str(e))
        raise HTTPException(status_code=500, detail="Failed to get stats")


@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException) -> JSONResponse:
    """Handle HTTP exceptions with RFC7807 format"""
    return JSONResponse(
        status_code=exc.status_code,
        content=ErrorResponse(
            type=f"https://httpstatuses.com/{exc.status_code}",
            title=exc.detail,
            status=exc.status_code,
            detail=exc.detail,
            instance=str(request.url),
            trace_id=getattr(request.state, "trace_id", None),
        ).dict(),
    )


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("main:app", host="0.0.0.0", port=8007, reload=True, log_config=None)