src/components/summaryEngine.py

from ..utils.functions import getConfig, getYaml
from ..utils.logger import logger
import litellm
import os

class SummaryEngine:
    def __init__(self):
        logger.info("INITIALIZING SUMMARY ENGINE")
        self.config = getConfig(os.path.join(os.getcwd(), "config.ini"))
        self.prompts = getYaml(os.path.join(os.getcwd(), "prompts.yaml"))

    def summarize(self, texts: list[str]) -> str:
        """
        Summarize a text
        Args:
            texts: list of texts to summarize
        Returns:
            summary: summary of the texts
        """
        try:
            logger.info("Summarizing the details extracted from the images")
            allSummaries = str()
            for i in range(len(texts)):
                batchSize = self.config.getint("DETAIL EXTRACTOR", "BATCHSIZE")
                startPage = i * batchSize + 1
                endPage = startPage + batchSize - 1
                allSummaries += f"# SUMMARY FOR PAGES {startPage}-{endPage}:\n" + texts[i]
            completion = litellm.completion(
                model = self.config.get("SUMMARIZER", "LLM"),
                api_key = os.environ["GROQ_API_KEY"],
                api_base = self.config.get("GROQ CONFIG", "BASEURL"),
                messages = [
                    {"role": "system", "content": self.prompts["summaryEnginePrompt"]},
                    {"role": "user", "content": f"AGGEREGATED SUMMARIES: {allSummaries}"}
                ]
            )
            response = completion["choices"][0]["message"]["content"]
            logger.info("Summary generated successfully")
            return response
        except Exception as e:
            logger.exception(f"Error summarizing the text: {e}")
            return None