rp/rp/core/streaming.py

import json
import logging
import sys
import time
from dataclasses import dataclass
from typing import Any, Callable, Dict, Generator, Optional

import requests

from rp.config import STREAMING_ENABLED, TOKEN_THROUGHPUT_TARGET
from rp.ui import Colors

logger = logging.getLogger("rp")


@dataclass
class StreamingChunk:
    content: str
    delta: str
    finish_reason: Optional[str]
    tool_calls: Optional[list]
    usage: Optional[Dict[str, int]]
    timestamp: float


@dataclass
class StreamingMetrics:
    start_time: float
    tokens_received: int
    chunks_received: int
    first_token_time: Optional[float]
    last_token_time: Optional[float]

    @property
    def time_to_first_token(self) -> Optional[float]:
        if self.first_token_time:
            return self.first_token_time - self.start_time
        return None

    @property
    def tokens_per_second(self) -> float:
        if self.last_token_time and self.first_token_time and self.tokens_received > 0:
            duration = self.last_token_time - self.first_token_time
            if duration > 0:
                return self.tokens_received / duration
        return 0.0

    @property
    def total_duration(self) -> float:
        if self.last_token_time:
            return self.last_token_time - self.start_time
        return time.time() - self.start_time


class StreamingResponseHandler:
    def __init__(
        self,
        on_token: Optional[Callable[[str], None]] = None,
        on_tool_call: Optional[Callable[[dict], None]] = None,
        on_complete: Optional[Callable[[str, StreamingMetrics], None]] = None,
        syntax_highlighting: bool = True,
        visible: bool = True
    ):
        self.on_token = on_token
        self.on_tool_call = on_tool_call
        self.on_complete = on_complete
        self.syntax_highlighting = syntax_highlighting
        self.visible = visible
        self.content_buffer = []
        self.tool_calls_buffer = []
        self.reasoning_buffer = []
        self.in_reasoning_block = False
        self.metrics: Optional[StreamingMetrics] = None

    def process_stream(self, response_stream: Generator) -> Dict[str, Any]:
        self.metrics = StreamingMetrics(
            start_time=time.time(),
            tokens_received=0,
            chunks_received=0,
            first_token_time=None,
            last_token_time=None
        )
        full_content = ""
        tool_calls = []
        finish_reason = None
        usage = None

        try:
            for chunk in response_stream:
                self.metrics.chunks_received += 1
                self.metrics.last_token_time = time.time()

                if chunk.delta:
                    if self.metrics.first_token_time is None:
                        self.metrics.first_token_time = time.time()

                    full_content += chunk.delta
                    self.metrics.tokens_received += len(chunk.delta.split())
                    self._process_delta(chunk.delta)

                if chunk.tool_calls:
                    tool_calls = chunk.tool_calls

                if chunk.finish_reason:
                    finish_reason = chunk.finish_reason

                if chunk.usage:
                    usage = chunk.usage

        except KeyboardInterrupt:
            logger.info("Streaming interrupted by user")
            if self.visible:
                sys.stdout.write(f"\n{Colors.YELLOW}[Interrupted]{Colors.RESET}\n")
                sys.stdout.flush()

        if self.visible:
            sys.stdout.write("\n")
            sys.stdout.flush()

        if self.on_complete:
            self.on_complete(full_content, self.metrics)

        return {
            'content': full_content,
            'tool_calls': tool_calls,
            'finish_reason': finish_reason,
            'usage': usage,
            'metrics': {
                'tokens_received': self.metrics.tokens_received,
                'time_to_first_token': self.metrics.time_to_first_token,
                'tokens_per_second': self.metrics.tokens_per_second,
                'total_duration': self.metrics.total_duration
            }
        }

    def _process_delta(self, delta: str):
        if '[THINKING]' in delta:
            self.in_reasoning_block = True
            if self.visible:
                sys.stdout.write(f"\n{Colors.BLUE}[THINKING]{Colors.RESET}\n")
                sys.stdout.flush()
            delta = delta.replace('[THINKING]', '')

        if '[/THINKING]' in delta:
            self.in_reasoning_block = False
            if self.visible:
                sys.stdout.write(f"\n{Colors.BLUE}[/THINKING]{Colors.RESET}\n")
                sys.stdout.flush()
            delta = delta.replace('[/THINKING]', '')

        if delta:
            if self.in_reasoning_block:
                self.reasoning_buffer.append(delta)
                if self.visible:
                    sys.stdout.write(f"{Colors.CYAN}{delta}{Colors.RESET}")
            else:
                self.content_buffer.append(delta)
                if self.visible:
                    sys.stdout.write(delta)

            if self.visible:
                sys.stdout.flush()

            if self.on_token:
                self.on_token(delta)


class StreamingHTTPClient:
    def __init__(self, timeout: float = 600.0):
        self.timeout = timeout
        self.session = requests.Session()

    def stream_request(
        self,
        url: str,
        data: Dict[str, Any],
        headers: Dict[str, str]
    ) -> Generator[StreamingChunk, None, None]:
        data['stream'] = True

        try:
            response = self.session.post(
                url,
                json=data,
                headers=headers,
                stream=True,
                timeout=self.timeout
            )
            response.raise_for_status()

            for line in response.iter_lines():
                if line:
                    line_str = line.decode('utf-8')
                    if line_str.startswith('data: '):
                        json_str = line_str[6:]
                        if json_str.strip() == '[DONE]':
                            break

                        try:
                            chunk_data = json.loads(json_str)
                            yield self._parse_chunk(chunk_data)
                        except json.JSONDecodeError as e:
                            logger.warning(f"Failed to parse streaming chunk: {e}")
                            continue

        except requests.exceptions.RequestException as e:
            logger.error(f"Streaming request failed: {e}")
            raise

    def _parse_chunk(self, chunk_data: Dict[str, Any]) -> StreamingChunk:
        choices = chunk_data.get('choices', [])
        delta = ""
        finish_reason = None
        tool_calls = None

        if choices:
            choice = choices[0]
            delta_data = choice.get('delta', {})
            delta = delta_data.get('content', '')
            finish_reason = choice.get('finish_reason')
            tool_calls = delta_data.get('tool_calls')

        return StreamingChunk(
            content=delta,
            delta=delta,
            finish_reason=finish_reason,
            tool_calls=tool_calls,
            usage=chunk_data.get('usage'),
            timestamp=time.time()
        )


def create_streaming_handler(
    syntax_highlighting: bool = True,
    visible: bool = True
) -> StreamingResponseHandler:
    return StreamingResponseHandler(
        syntax_highlighting=syntax_highlighting,
        visible=visible
    )


def stream_api_response(
    url: str,
    data: Dict[str, Any],
    headers: Dict[str, str],
    handler: Optional[StreamingResponseHandler] = None
) -> Dict[str, Any]:
    if not STREAMING_ENABLED:
        return None

    if handler is None:
        handler = create_streaming_handler()

    client = StreamingHTTPClient()
    stream = client.stream_request(url, data, headers)
    return handler.process_stream(stream)