ForcePilot/backend/package/yuxi/channel/extensions/msteams/streaming.py

from __future__ import annotations

import asyncio
import logging
import time

from .sdk import BotFrameworkAdapter, build_message_activity, build_typing_activity
from .types import StoredConversationReference

logger = logging.getLogger(__name__)

TEAMS_HTTP_STREAM_THROTTLE_MS = 1500
TEAMS_HTTP_STREAM_TIMEOUT_S = 45


class MSTeamsStreaming:
    streaming_mode = "block"

    block_streaming_enabled = True
    block_streaming_chunk_min_chars = 1500
    block_streaming_chunk_max_chars = 4000
    block_streaming_chunk_break_preference = "paragraph"

    block_streaming_coalesce_defaults = {
        "min_chars": 1500,
        "max_chars": 4000,
        "idle_ms": 1000,
    }

    teams_http_stream_enabled = False
    teams_http_stream_throttle_ms = TEAMS_HTTP_STREAM_THROTTLE_MS
    teams_http_stream_timeout_s = TEAMS_HTTP_STREAM_TIMEOUT_S

    def create_block_chunker(self) -> dict:
        return {
            "mode": "length",
            "min_chars": self.block_streaming_chunk_min_chars,
            "max_chars": self.block_streaming_chunk_max_chars,
        }

    def create_stream_session(self, target_id: str) -> dict:
        return {
            "target_id": target_id,
            "mode": "block",
            "chunks_sent": 0,
        }


class TeamsHttpStreamSession:
    def __init__(
        self,
        adapter: BotFrameworkAdapter,
        ref: StoredConversationReference,
        *,
        throttle_ms: int = TEAMS_HTTP_STREAM_THROTTLE_MS,
        timeout_s: int = TEAMS_HTTP_STREAM_TIMEOUT_S,
    ):
        self._adapter = adapter
        self._ref = ref
        self._throttle_ms = throttle_ms
        self._timeout_s = timeout_s
        self._message_id: str | None = None
        self._last_update: float = 0.0
        self._last_text: str = ""
        self._start_time: float = 0.0
        self._finished: bool = False
        self._cancelled: bool = False

    @property
    def message_id(self) -> str | None:
        return self._message_id

    @property
    def finished(self) -> bool:
        return self._finished

    async def start(self, initial_text: str = "") -> str:
        activity = {
            "type": "message",
            "text": initial_text or "...",
            "textFormat": "markdown",
        }
        if self._ref.tenant_id:
            activity.setdefault("channelData", {})
            activity["channelData"]["tenant"] = {"id": self._ref.tenant_id}

        result = await self._adapter.send_activity(self._ref, activity)
        self._message_id = result.get("id", "")
        self._last_text = initial_text
        self._start_time = time.monotonic()
        return self._message_id

    async def update(self, text: str) -> bool:
        if self._cancelled or self._finished:
            return False

        if time.monotonic() - self._start_time > self._timeout_s:
            logger.warning("TeamsHttpStream timeout (%ds), falling back to block streaming", self._timeout_s)
            self._cancelled = True
            return False

        if len(text) <= len(self._last_text):
            return True

        if not text.startswith(self._last_text):
            logger.debug("TeamsHttpStream prefix constraint violated, falling back")
            self._cancelled = True
            return False

        elapsed = (time.monotonic() - self._last_update) * 1000
        if elapsed < self._throttle_ms:
            await asyncio.sleep((self._throttle_ms - elapsed) / 1000.0)

        if not self._message_id:
            return False

        activity = build_message_activity(text, tenant_id=self._ref.tenant_id, ai_generated=True)
        try:
            await self._adapter.update_activity(
                self._ref.service_url,
                self._ref.conversation_id,
                self._message_id,
                activity,
            )
            self._last_text = text
            self._last_update = time.monotonic()
            return True
        except Exception as e:
            logger.warning("TeamsHttpStream update failed: %s", e)
            self._cancelled = True
            return False

    async def finish(self, final_text: str) -> bool:
        if self._cancelled:
            return False

        activity = build_message_activity(final_text, tenant_id=self._ref.tenant_id, ai_generated=True)
        try:
            await self._adapter.update_activity(
                self._ref.service_url,
                self._ref.conversation_id,
                self._message_id,
                activity,
            )
            self._finished = True
            return True
        except Exception as e:
            logger.warning("TeamsHttpStream finish failed: %s", e)
            return False

    async def cancel(self) -> None:
        self._cancelled = True


async def send_streaming_typing(
    adapter: BotFrameworkAdapter,
    ref: StoredConversationReference,
) -> None:
    try:
        activity = build_typing_activity()
        await adapter.send_activity(ref, activity)
    except Exception:
        pass