parser.py

import asyncio
from datetime import datetime, timezone

from telethon import TelegramClient
from telethon.tl.functions.messages import GetHistoryRequest

from database import add_message, save_event, get_stats, check_sent
from rabbitmq import connect_rabbit, setup_exchange, publish_event, EXCHANGE_EVENTS, ROUTING_EVENT
from models import load_model, classify
from config import (
    TG_API_ID, TG_API_HASH,
    RABBITMQ_HOST, RABBITMQ_PORT, RABBITMQ_USER, RABBITMQ_PASS
)

async def connect_tg():
    client = TelegramClient("sessions/extractor_session", TG_API_ID, TG_API_HASH)
    await client.connect()
    try:
        authorized = await client.is_user_authorized()
    except TypeError:
        authorized = client.is_user_authorized()
    if not authorized:
        await client.disconnect()
        raise RuntimeError("Telegram session missing; run auth.py to authorize")
    return client

def init_rabbit():
    connect_rabbit(RABBITMQ_HOST, RABBITMQ_PORT, RABBITMQ_USER, RABBITMQ_PASS)
    setup_exchange()

async def parse_messages(client, model_name, channel_name, start_date, end_date, check_duplicates=True):
    entity = await client.get_entity(channel_name)
    
    parsed_count = 0
    offset_id = 0
    offset_date = end_date
    
    while True:
        history = await client(GetHistoryRequest(
            peer=entity,
            offset_id=offset_id,
            offset_date=offset_date,
            add_offset=0,
            limit=100,
            max_id=0,
            min_id=0,
            hash=0
        ))
        
        if not history.messages:
            break
        
        stop_parsing = False
        
        for message in history.messages:
            if not message.message:
                continue
            
            msg_date = message.date.replace(tzinfo=timezone.utc)
            
            if msg_date < start_date:
                stop_parsing = True
                continue
            
            if msg_date > end_date:
                continue
            
            msg_link = "https://t.me/" + channel_name + "/" + str(message.id)
            
            add_message(channel_name, message.id, message.message, msg_date, msg_link)

            is_event, confidence = classify(message.message, model_name)
            
            save_event(channel_name, message.id, is_event, confidence, "мероприятие" if is_event else None)
            
            if is_event:
                send = True
                if check_duplicates:
                    send = not check_sent(channel_name, message.id)
                
                if send:
                    event = {
                        'id': message.id,
                        'channel': channel_name,
                        'text': message.message,
                        'date': msg_date.isoformat(),
                        'link': msg_link,
                        'confidence': confidence,
                        'event_type': "мероприятие"
                    }
                    publish_event(event)
            
            parsed_count += 1
        
        if stop_parsing or len(history.messages) < 100:
            break
        
        offset_id = history.messages[-1].id
        offset_date = history.messages[-1].date
        
        await asyncio.sleep(0.2)
    
    return parsed_count

async def parse_channel(channel, start_date, end_date, model_name=None, check_duplicates=True):
    start_date = start_date.replace(tzinfo=timezone.utc)
    end_date = end_date.replace(tzinfo=timezone.utc)
    
    model_name = model_name or 'rubert'
    load_model(model_name)
    
    client = await connect_tg()
    init_rabbit()
    
    extracted = await parse_messages(
        client, model_name,
        channel, start_date, end_date, check_duplicates
    )
    
    await client.disconnect()
    
    stats = get_stats(channel, start_date, end_date)
    
    return extracted, stats['events'], stats['sent']