api/routes/draft.py

"""
api/routes/draft.py — Draft generation and retrieval endpoints
"""

import json
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from loguru import logger

from api.dependencies import (
    vectorstore, bm25_index, reranker, drafter, pattern_store,
)
from retrieval.hybrid import retrieve
from feedback.preference_injector import build_preference_string

router = APIRouter()


class DraftRequest(BaseModel):
    query: str = "Provide a comprehensive analysis."
    draft_type: str = "Case Fact Summary"
    doc_names: str = ""


@router.post("/draft")
async def generate_draft(req: DraftRequest):
    """Generate a grounded legal draft from indexed documents."""
    # Retrieve evidence
    evidence = retrieve(
        query=req.query,
        vectorstore=vectorstore,
        bm25_index=bm25_index,
        reranker=reranker,
    )

    if not evidence:
        raise HTTPException(404, "No evidence found. Ingest documents first.")

    # Build preferences
    preferences = build_preference_string(pattern_store)

    # Generate draft
    draft, report = drafter.generate(
        evidence=evidence,
        draft_type=req.draft_type,
        doc_names=req.doc_names,
        user_query=req.query,
        injected_preferences=preferences,
    )

    # Save to history
    pattern_store.save_draft(
        draft_id=draft.draft_id,
        draft_type=draft.draft_type,
        doc_ids=", ".join(set(draft.doc_ids)),
        draft_text=draft.draft_text,
        validation_json=json.dumps(report.to_dict()),
    )

    return {
        "draft_id": draft.draft_id,
        "draft_text": draft.draft_text,
        "draft_type": draft.draft_type,
        "evidence_count": len(evidence),
        "citations_count": len(draft.citations),
        "validation": report.to_dict(),
    }


@router.get("/draft/{draft_id}")
async def get_draft(draft_id: str):
    """Retrieve a previously generated draft."""
    draft = pattern_store.get_draft(draft_id)
    if not draft:
        raise HTTPException(404, f"Draft {draft_id} not found.")
    return draft