summarizer.py

#!/usr/bin/env python


from tg_reader import *
from llm_api import *
import yaml
from typing import Callable


class Summarizer:
    def __init__(self, config: dict):
        self.config = config

        ollama_config: dict = self.config["ollama"]
        self.__ollama = OllamaSession(
            model=ollama_config["model"],
            api_url=ollama_config["api_url"],
            timeout=ollama_config.get("timeout", 300),
            https_proxy=ollama_config.get("https_proxy")
        )

        telegram_config: dict = self.config["telegram_acc"]
        self.__telegram = TgReader(
            api_id=telegram_config["api_id"],
            api_hash=telegram_config["api_hash"]
        )

        self.__consumers = []

    def add_consumer(self, callback: Callable[[str], None]):
        self.__consumers.append(callback)

    def summarize(self):
        with self.__telegram:
            for chat_info in self.config.get("chats", []):
                messages = self.__telegram.get_messages(
                    chat_id=chat_info["id"],
                    topic_id=chat_info.get("topic"),
                    offset_date=chat_info.get("last_read", datetime(2026, 7, 9))
                )
                if len(messages) < 1:
                    continue
                chat_text: str = "\n".join(str(msg) for msg in messages)
                summary = self.__ollama.get_summary(
                    chat_history=chat_text,
                    chat_description=chat_info.get("descr")
                )
                chat_title = self.__telegram.get_chat_name(chat_info['id'])
                topic_id = chat_info.get("topic")
                if topic_id is not None:
                    topic_title = self.__telegram.get_topic_name(chat_info['id'], topic_id)
                    chat_title = f"{chat_title} / {topic_title}"
                summary = f"## {chat_title}\n\n{summary}\n\n---\n"
                assert isinstance(summary, str)
                for consumer in self.__consumers:
                    consumer(summary)
                chat_info["last_read"] = messages[-1].date
                self.__telegram.mark_as_read(
                    chat_info["id"],
                    messages[-1].id,
                    topic_id
                )


if __name__ == '__main__':
    import os
    from argparse import ArgumentParser
    from pathlib import Path
    arg_parser = ArgumentParser(
        description="Справка"
    )
    arg_parser.add_argument(
        "-c",
        "--config",
        default=Path(__file__).resolve().parent / "config.yaml",
        help="Путь к конфигурационному YAML-файлу"
    )
    args = arg_parser.parse_args()

    with open(args.config, "r", encoding="utf-8") as file:
        config: dict = yaml.safe_load(file.read())
    
    summarizer = Summarizer(config)
    summarizer.add_consumer(print)
    summarizer.summarize()

    with open(args.config, "w", encoding="utf-8") as file:
        yaml.safe_dump(
            config,
            file,
            default_flow_style=False
        )