src/cli_agent/stream_handler.py - master - Watashicuvu/agentic-tools

src/cli_agent/stream_handler.py
242 строки8 KB
Якуб
mts update
12 апр 2026, 21:19
12 апр 2026, 21:19bff6dcd
"""Обработка стриминга с возможностью пользовательского вмешательства."""

import asyncio
import json
import sys
from typing import (
    AsyncGenerator,
    Dict,
    Any,
    List,
    Optional,
    Callable,
    Awaitable
)
from dataclasses import dataclass, field
from pathlib import Path
from datetime import datetime

from .llm_client import LLMClient, StreamEvent


@dataclass
class SessionState:
    """Состояние сессии."""
    session_id: str
    messages: List[Dict[str, Any]] = field(default_factory=list)
    accumulated_text: str = ""
    tool_calls: List[Dict[str, Any]] = field(default_factory=list)
    is_paused: bool = False
    is_interrupted: bool = False
    start_time: Optional[datetime] = None
    end_time: Optional[datetime] = None


class StreamHandler:
    """Обработчик стриминга с интерактивным управлением."""
    
    def __init__(
        self,
        llm_client: LLMClient,
        save_path: Optional[str] = None,
        interrupt_callback: Optional[Callable[[], Awaitable[bool]]] = None
    ):
        """Инициализация обработчика.
        
        Args:
            llm_client: LLM клиент
            save_path: Путь для сохранения сессии (JSON)
            interrupt_callback: Коллбэк для проверки прерывания
        """
        self.llm_client = llm_client
        self.save_path = save_path
        self.interrupt_callback = interrupt_callback
        self.session: Optional[SessionState] = None
    
    async def run_stream(
        self,
        messages: List[Dict[str, Any]],
        tools: Optional[List[Dict[str, Any]]] = None,
        session_id: Optional[str] = None
    ) -> SessionState:
        """Запустить стриминг с сохранением состояния.
        
        Args:
            messages: Сообщения для LLM
            tools: Доступные инструменты
            session_id: ID сессии
            
        Returns:
            SessionState с результатами
        """
        from datetime import datetime
        
        # Инициализация сессии
        self.session = SessionState(
            session_id=session_id or f"session_{int(datetime.now().timestamp())}",
            messages=messages.copy(),
            start_time=datetime.now()
        )
        
        # Форматирование инструментов
        tools_format = None
        if tools:
            tools_format = self.llm_client.get_tools_format(tools)
        
        # Запуск стриминга
        async for event in self.llm_client.chat_stream(
            messages=messages,
            tools=tools_format
        ):
            # Проверка прерывания
            if self.interrupt_callback:
                should_interrupt = await self.interrupt_callback()
                if should_interrupt:
                    self.session.is_interrupted = True
                    self.session.end_time = datetime.now()
                    
                    # Сохранение промежуточного состояния
                    await self._save_session()
                    
                    return self.session
            
            # Обработка события
            await self._process_event(event)
        
        # Завершение
        self.session.end_time = datetime.now()
        
        # Финальное сохранение
        if self.save_path:
            await self._save_session()
        
        return self.session
    
    async def _process_event(self, event: StreamEvent):
        """Обработать событие из стрима.
        
        Args:
            event: StreamEvent
        """
        if event.type == "text":
            self.session.accumulated_text += event.content or ""
            
            # Вывод в реальном времени
            if event.content:
                print(event.content, end="", flush=True)
        
        elif event.type == "tool_call":
            if event.tool_calls:
                self.session.tool_calls.extend(event.tool_calls)
        
        elif event.type == "done":
            print()  # Новая строка после завершения
            self.session.messages.append({
                "role": "assistant",
                "content": event.content,
                "tool_calls": event.tool_calls
            })
        
        elif event.type == "error":
            print(f"\n[ERROR] {event.error}", file=sys.stderr)
            self.session.messages.append({
                "role": "system",
                "content": f"Error: {event.error}"
            })
    
    async def _save_session(self):
        """Сохранить состояние сессии в JSON."""
        if not self.save_path or not self.session:
            return
        
        try:
            path = Path(self.save_path).resolve()
            path.parent.mkdir(parents=True, exist_ok=True)
            
            session_data = {
                "session_id": self.session.session_id,
                "messages": self.session.messages,
                "accumulated_text": self.session.accumulated_text,
                "tool_calls": self.session.tool_calls,
                "is_interrupted": self.session.is_interrupted,
                "start_time": self.session.start_time.isoformat() if self.session.start_time else None,
                "end_time": self.session.end_time.isoformat() if self.session.end_time else None
            }
            
            path.write_text(
                json.dumps(session_data, indent=2, ensure_ascii=False),
                encoding='utf-8'
            )
            
            print(f"\n[SESSION SAVED] {path}", file=sys.stderr)
        except Exception as e:
            print(f"\n[SAVE ERROR] {str(e)}", file=sys.stderr)
    
    def get_accumulated_text(self) -> str:
        """Получить накопленный текст.
        
        Returns:
            Строка с теклом
        """
        return self.session.accumulated_text if self.session else ""
    
    def get_tool_calls(self) -> List[Dict[str, Any]]:
        """Получить вызовы инструментов.
        
        Returns:
            Список вызовов
        """
        return self.session.tool_calls if self.session else []


class InteractiveStreamHandler(StreamHandler):
    """Интерактивный обработчик с возможностью прерывания через Ctrl+C."""
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._interrupt_flag = False
        self._setup_signal_handler()
    
    def _setup_signal_handler(self):
        """Настроить обработчик сигналов."""
        import signal
        
        def handler(signum, frame):
            self._interrupt_flag = True
            print("\n\n[INTERRUPT] Ctrl+C detected! Saving session...", file=sys.stderr)
        
        try:
            signal.signal(signal.SIGINT, handler)
        except (ValueError, OSError):
            # Не работает в не-main потоке
            pass
    
    async def check_interrupt(self) -> bool:
        """Проверить флаг прерывания.
        
        Returns:
            True если нужно прервать
        """
        return self._interrupt_flag


def create_interactive_handler(
    llm_client: LLMClient,
    save_path: Optional[str] = None
) -> InteractiveStreamHandler:
    """Создать интерактивный обработчик.
    
    Args:
        llm_client: LLM клиент
        save_path: Путь для сохранения
        
    Returns:
        InteractiveStreamHandler
    """
    return InteractiveStreamHandler(
        llm_client=llm_client,
        save_path=save_path,
        interrupt_callback=lambda: asyncio.get_event_loop().run_in_executor(
            None, lambda: False  # Заменяется на реальную проверку
        )
    )