telegram_parser_async.py

import json
import asyncio
import aiohttp
import pandas as pd
import re
from typing import List, Dict, Any
from config import *

class AsyncTelegramChatParser:
    def __init__(self):
        self.api_key = OPENROUTER_API_KEY
        self.base_url = OPENROUTER_BASE_URL
        self.model = OPENROUTER_MODEL
        
        if not self.api_key:
            raise ValueError("OpenRouter API key not found. Please set OPENROUTER_API_KEY in your environment or .env file")
    
    def load_telegram_json(self, file_path: str) -> Dict[str, Any]:
        """Load Telegram export JSON file"""
        print(f"Loading Telegram export from {file_path}...")
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        print(f"Loaded {len(data.get('messages', []))} messages")
        return data
    
    def extract_text_from_message(self, message: Dict[str, Any]) -> str:
        """Extract plain text from Telegram message structure"""
        text_content = message.get('text', '')
        
        if isinstance(text_content, list):
            # Handle structured text with entities
            text_parts = []
            for item in text_content:
                if isinstance(item, dict):
                    text_parts.append(item.get('text', ''))
                else:
                    text_parts.append(str(item))
            return ''.join(text_parts)
        else:
            return str(text_content)
    
    def prepare_messages(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Prepare messages for processing (no filtering, just format and limit)"""
        prepared_messages = []
        
        for message in messages:
            # Skip service messages
            if message.get('type') != 'message':
                continue
                
            text = self.extract_text_from_message(message)
            
            # Skip very short messages (less than 20 characters)
            if len(text.strip()) < 20:
                continue
                
            prepared_messages.append({
                'id': message.get('id'),
                'date': message.get('date'),
                'text': text
            })
        
        # Apply limit if specified
        if MAX_MESSAGES_TO_PROCESS:
            prepared_messages = prepared_messages[:MAX_MESSAGES_TO_PROCESS]
            print(f"Limited to first {len(prepared_messages)} messages")
        
        return prepared_messages
    
    async def call_openrouter_api_async(self, session: aiohttp.ClientSession, prompt: str) -> str:
        """Make async API call to OpenRouter with structured outputs"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": MAX_TOKENS,
            "temperature": TEMPERATURE,
            "response_format": {
                "type": "json_schema",
                "json_schema": RESPONSE_SCHEMA
            }
        }
        
        try:
            async with session.post(f"{self.base_url}/chat/completions", 
                                  headers=headers, json=data, timeout=60) as response:
                
                if response.status != 200:
                    error_text = await response.text()
                    print(f"API Error {response.status}: {error_text}")
                    return ""
                
                result = await response.json()
                return result['choices'][0]['message']['content']
        
        except asyncio.TimeoutError:
            print("API request timeout")
            return ""
        except Exception as e:
            print(f"API request error: {e}")
            return ""
    
    def create_analysis_prompt(self, message_text: str, message_date: str) -> str:
        """Create prompt for analyzing a message"""
        questions_text = "\n".join([f"{i+1}. {q}" for i, q in enumerate(QUESTIONS)])
        
        prompt = SYSTEM_PROMPT.format(
            questions=questions_text,
            num_questions=len(QUESTIONS)
        )
        
        prompt += f"\n\nДата сообщения: {message_date}\n\nСообщение:\n{message_text}"
        
        return prompt
    
    def parse_api_response(self, response: str) -> List[Dict[str, Any]]:
        """Parse OpenRouter API response - now guaranteed to be valid JSON"""
        try:
            # With structured outputs, we get guaranteed valid JSON
            result = json.loads(response)
            
            if 'stories' in result:
                return result['stories']
            else:
                print(f"No 'stories' key in response: {response[:200]}...")
                return []
            
        except json.JSONDecodeError as e:
            print(f"JSON decode error (this shouldn't happen with structured outputs): {e}")
            print(f"Response: {response[:200]}...")
            return []
    
    async def analyze_message_async(self, session: aiohttp.ClientSession, message: Dict[str, Any], 
                                   index: int, total: int) -> List[Dict[str, Any]]:
        """Analyze a single message asynchronously"""
        print(f"Analyzing message {index+1}/{total} (ID: {message['id']})")
        
        prompt = self.create_analysis_prompt(message['text'], message['date'])
        response = await self.call_openrouter_api_async(session, prompt)
        
        stories = []
        if response:
            stories = self.parse_api_response(response)
            for story in stories:
                story['message_id'] = message['id']
                story['message_date'] = message['date']
        
        return stories
    
    async def analyze_messages_parallel(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Analyze messages using parallel async API calls"""
        all_stories = []
        
        # Create semaphore to limit concurrent requests
        semaphore = asyncio.Semaphore(CONCURRENT_REQUESTS)
        
        async def analyze_with_semaphore(session, message, index, total):
            async with semaphore:
                return await self.analyze_message_async(session, message, index, total)
        
        # Process messages in batches to avoid overwhelming the API
        async with aiohttp.ClientSession() as session:
            print(f"Processing {len(messages)} messages with {CONCURRENT_REQUESTS} concurrent requests...")
            
            # Create tasks for all messages
            tasks = [
                analyze_with_semaphore(session, message, i, len(messages))
                for i, message in enumerate(messages)
            ]
            
            # Execute all tasks and collect results
            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            # Process results
            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    print(f"Error processing message {i+1}: {result}")
                elif result:
                    all_stories.extend(result)
                
                # Add small delay between batches to be nice to the API
                if (i + 1) % CONCURRENT_REQUESTS == 0:
                    await asyncio.sleep(REQUEST_DELAY)
        
        print(f"Extracted {len(all_stories)} stories total")
        return all_stories
    
    def save_to_excel(self, stories: List[Dict[str, Any]], output_file: str):
        """Save analyzed stories to Excel file"""
        data = []
        
        for story in stories:
            row = {
                'Message ID': story.get('message_id', ''),
                'Message Date': story.get('message_date', ''),
                'Story Date': story.get('date', ''),
                'Story': story.get('story_text', ''),
            }
            
            # Add answers as separate columns
            answers = story.get('answers', [0] * len(QUESTIONS))
            for i, question in enumerate(QUESTIONS):
                row[question] = answers[i] if i < len(answers) else 0
            
            data.append(row)
        
        df = pd.DataFrame(data)
        df.to_excel(output_file, index=False)
        print(f"Results saved to {output_file}")
    
    async def run_async(self, input_file: str, output_file: str = "telegram_analysis_async.xlsx"):
        """Main async method to run the entire analysis pipeline"""
        try:
            # Load and parse Telegram data
            telegram_data = self.load_telegram_json(input_file)
            all_messages = telegram_data.get('messages', [])
            
            print(f"Total messages in file: {len(all_messages)}")
            
            # Prepare messages for processing
            prepared_messages = self.prepare_messages(all_messages)
            
            print(f"Messages to process: {len(prepared_messages)}")
            
            if not prepared_messages:
                print("No messages to process")
                return
            
            # Analyze messages with parallel AI processing
            import time
            start_time = time.time()
            
            stories = await self.analyze_messages_parallel(prepared_messages)
            
            end_time = time.time()
            print(f"Processing completed in {end_time - start_time:.2f} seconds")
            
            if not stories:
                print("No stories extracted")
                return
            
            # Save to Excel
            self.save_to_excel(stories, output_file)
            
        except Exception as e:
            print(f"Error during processing: {e}")
            import traceback
            traceback.print_exc()
    
    def run(self, input_file: str, output_file: str = "telegram_analysis_async.xlsx"):
        """Synchronous wrapper for the async run method"""
        asyncio.run(self.run_async(input_file, output_file))

if __name__ == "__main__":
    parser = AsyncTelegramChatParser()
    parser.run("пограничный_контроль_эксопрт.json")