backend/routers/entities.py

"""Financial entity extraction endpoint."""

import asyncio
from typing import Optional

from fastapi import APIRouter, Depends
from pydantic import BaseModel

router = APIRouter(prefix="/api")


class EntityRequest(BaseModel):
    document_id: Optional[str] = None
    text: Optional[str] = None


def get_llm():
    from core.bootstrap import get_llm as _get
    return _get()


def get_db():
    from core.bootstrap import get_db as _get
    return _get()


@router.post("/extract-entities")
async def extract_entities(
    request: EntityRequest,
    llm=Depends(get_llm),
    db=Depends(get_db),
):
    if request.text:
        doc_text = request.text
    elif request.document_id:
        chunks = db.fetch_chunks_by_document_id(request.document_id)
        if not chunks:
            return {"error": f"No chunks found for document {request.document_id}"}
        doc_text = "\n\n".join(c["text"] for c in chunks)
    else:
        return {"error": "Provide either document_id or text"}

    entities = await asyncio.to_thread(llm.extract_entities, doc_text)
    return {"entities": entities}