DB-GPT/dbgpt/rag/transformer/llm_summarizer.py

"""LLMSummarizer class."""
import logging
from abc import ABC

from dbgpt.core import HumanPromptTemplate, LLMClient, ModelMessage, ModelRequest
from dbgpt.rag.transformer.base import SummarizerBase

logger = logging.getLogger(__name__)


class LLMSummarizer(SummarizerBase, ABC):
    """LLMSummarizer class."""

    def __init__(self, llm_client: LLMClient, model_name: str, prompt_template: str):
        """Initialize the LLMSummarizer."""
        self._llm_client = llm_client
        self._model_name = model_name
        self._prompt_template = prompt_template

    async def summarize(self, **args) -> str:
        """Summarize by LLM."""
        template = HumanPromptTemplate.from_template(self._prompt_template)
        messages = template.format_messages(**args)

        # use default model if needed
        if not self._model_name:
            models = await self._llm_client.models()
            if not models:
                raise Exception("No models available")
            self._model_name = models[0].model
            logger.info(f"Using model {self._model_name} to extract")

        model_messages = ModelMessage.from_base_messages(messages)
        request = ModelRequest(model=self._model_name, messages=model_messages)
        response = await self._llm_client.generate(request=request)

        if not response.success:
            code = str(response.error_code)
            reason = response.text
            logger.error(f"request llm failed ({code}) {reason}")

        return response.text

    def truncate(self):
        """Do nothing by default."""

    def drop(self):
        """Do nothing by default."""