backend/services/summarizer.py

"""Financial document summarization service."""

from typing import Optional

from db.faiss_client import FaissDB
from models.llm import MiniCPMLLM


from utils.response_cleaner import clean_model_response


class SummarizerService:
    def __init__(self, llm: MiniCPMLLM, db: FaissDB):
        self.llm = llm
        self.db = db

    def summarize(
        self,
        document_id: Optional[str] = None,
        text: Optional[str] = None,
        summary_type: str = "financial",
    ) -> dict:
        if text:
            doc_text = text
        elif document_id:
            chunks = self.db.fetch_chunks_by_document_id(document_id)
            if not chunks:
                return {"error": f"No chunks found for document {document_id}"}
            doc_text = "\n\n".join(c["text"] for c in chunks)
        else:
            return {"error": "Provide either document_id or text"}

        summary = clean_model_response(self.llm.generate_summary(doc_text, summary_type))
        return {"summary": summary, "type": summary_type}