Создание бота для суммаризации чатов и каналов в Telegram на LangChain и Foundation Models

# chat_bot/models/ai_config.py
from pydantic import BaseModel, Field, validator

class AIConfig(BaseModel):
    """Model for AI configuration."""

    api_key: str = Field(..., description="AI API key")
    model: str = Field("t-tech/T-pro-it-2.0", description="AI model name")
    base_url: Optional[str] = Field(None, description="Custom AI base URL")
    temperature: float = Field(0.3, ge=0.0, le=2.0, description="Generation temperature")
    max_tokens: int = Field(500, gt=0, description="Maximum tokens for generation")

    @classmethod
    @validator("api_key")
    def validate_api_key(cls, v: str) -> str:
        """Validate that API key is not empty."""
        if not v or not v.strip():
            raise ValueError("api_key cannot be empty")
        return v.strip()

# chat_bot/assistant.py
from langchain_openai import ChatOpenAI
from pydantic import SecretStr

def _init_llm(self) -> None:
    """Initialize the language model."""
    try:
        # Initialize with required parameters
        self.llm = ChatOpenAI(
            api_key=SecretStr(self.config.api_key),
            model=self.config.model,
            temperature=self.config.temperature,
            base_url=self.config.base_url,
        )
        logger.info(
            f"Initialized AI model: {self.config.model} "
            f"(temp: {self.config.temperature}, max_tokens: {self.config.max_tokens})"
        )
    except Exception as e:
        logger.error(f"Failed to initialize AI model: {e}")
        raise

def _load_prompts(self) -> None:
    """Load prompt templates from files."""
    try:
        prompts_dir = Path(__file__).parent / "prompts"
        summary_prompt_file = prompts_dir / "summary.txt"
        task_extraction_prompt_file = prompts_dir / "task_extraction.txt"

        # Load summary prompt
        if summary_prompt_file.exists():
            with open(summary_prompt_file, "r", encoding="utf-8") as f:
                summary_template = f.read()

            self.summary_prompt = ChatPromptTemplate.from_template(summary_template)
            logger.info("Loaded summary prompt from file")
        else:
            # Fallback to default prompt
            self.summary_prompt = ChatPromptTemplate.from_template(
                "You are an assistant for creating brief chat summaries. "
                "Please provide your response in Russian.\n\n{messages}\n\n"
                "Create a brief summary in Russian."
            )
    except Exception as e:
        logger.error(f"Failed to load prompts: {e}")
        # Fallback to default prompts

# chat_bot/models/summary_response.py
class SummaryOutput(BaseModel):
    """
    Structured output schema for summary generation from chat messages.

    This model is used with LangChain's structured output feature to ensure
    the AI model returns properly formatted summary data.
    """

    thoughts: str = Field(
        ...,
        description="The AI's reasoning process and thoughts about the messages before creating the summary. This should be in Russian.",
    )
    summary: str = Field(
        ...,
        description="The actual summary of the chat messages. This should be concise and in Russian.",
    )

# chat_bot/models/task_extraction_response.py
class TaskExtractionOutput(BaseModel):
    """
    Structured output schema for task extraction from chat messages.
    """

    tasks: List[Task] = Field(
        default_factory=list,
        description="List of tasks extracted from the chat messages. If no tasks are found, return an empty list.",
    )

class Task(BaseModel):
    """Represents a task extracted from chat messages."""

    assignee: str = Field(..., description="The person assigned to the task")
    title: str = Field(..., description="The title/description of the task")
    deadline: Optional[datetime] = Field(
        None, description="Optional deadline date/time for the task"
    )

async def summarize(self, messages_input: Union[str, Dict[str, Any], MessagesData]) -> SummaryResponse:
    """Summarize messages using LangChain's structured output."""
    import time
    start_time = time.time()

    try:
        # Handle different input types
        if isinstance(messages_input, str):
            data = json.loads(messages_input)
            messages_data = MessagesData(**data)
        elif isinstance(messages_input, dict):
            messages_data = MessagesData(**messages_input)
        elif isinstance(messages_input, MessagesData):
            messages_data = messages_input
        else:
            raise ValueError("Input must be either a JSON string, dictionary, or MessagesData object")

        # Format messages for summarization
        formatted_messages = MessageFormatter.format_messages_for_summary(messages_data)

        # Create the prompt
        prompt = self.summary_prompt.format(messages=formatted_messages)

        # Create model with structured output
        model_with_structure = self.llm.with_structured_output(SummaryOutput)

        # Generate summary response using structured output
        structured_output: SummaryOutput = await model_with_structure.ainvoke(prompt)

        processing_time = time.time() - start_time
        logger.info("Successfully generated summary")

        return SummaryResponse(
            summary=structured_output.summary,
            success=True,
            error_message=None,
            processing_time=processing_time,
        )
    except Exception as e:
        logger.error(f"Failed to generate summary: {e}")
        return SummaryResponse(
            summary="",
            success=False,
            error_message=f"Ошибка при создании сводки: {str(e)}",
            processing_time=time.time() - start_time,
        )