90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import logging
|
|||
|
|
|
|||
|
|
from yuxi.channel.extensions.workplace.outbound import WorkplaceOutbound
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class WorkplaceStreaming:
|
|||
|
|
streaming_mode: str = "block"
|
|||
|
|
block_streaming_enabled: bool = True
|
|||
|
|
block_streaming_break: str = "text_end"
|
|||
|
|
block_streaming_chunk_min_chars: int = 200
|
|||
|
|
block_streaming_chunk_max_chars: int = 2000
|
|||
|
|
block_streaming_chunk_break_preference: str = "paragraph"
|
|||
|
|
block_streaming_coalesce_defaults: dict = {"min_chars": 200, "max_chars": 2000, "idle_ms": 1000}
|
|||
|
|
|
|||
|
|
def __init__(self, outbound: WorkplaceOutbound | None = None):
|
|||
|
|
self._outbound = outbound
|
|||
|
|
|
|||
|
|
async def send_block_stream(
|
|||
|
|
self,
|
|||
|
|
target_id: str,
|
|||
|
|
full_text: str,
|
|||
|
|
*,
|
|||
|
|
thread_id: str | None = None,
|
|||
|
|
account_id: str | None = None,
|
|||
|
|
) -> list[str]:
|
|||
|
|
chunks = self._chunk_text(full_text)
|
|||
|
|
outbound = self._outbound or WorkplaceOutbound()
|
|||
|
|
message_ids: list[str] = []
|
|||
|
|
|
|||
|
|
for i, chunk in enumerate(chunks):
|
|||
|
|
if not chunk.strip():
|
|||
|
|
continue
|
|||
|
|
results = await outbound.send_text(target_id, chunk, thread_id=thread_id, account_id=account_id)
|
|||
|
|
for r in results:
|
|||
|
|
msg_id = r.get("message_id")
|
|||
|
|
if msg_id:
|
|||
|
|
message_ids.append(msg_id)
|
|||
|
|
|
|||
|
|
if i < len(chunks) - 1:
|
|||
|
|
await asyncio.sleep(0.3)
|
|||
|
|
|
|||
|
|
return message_ids
|
|||
|
|
|
|||
|
|
def create_draft_stream_session(self, **kwargs) -> object | None:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def _chunk_text(self, text: str) -> list[str]:
|
|||
|
|
if not text:
|
|||
|
|
return [""]
|
|||
|
|
if len(text) <= self.block_streaming_chunk_max_chars:
|
|||
|
|
return [text]
|
|||
|
|
|
|||
|
|
chunks: list[str] = []
|
|||
|
|
remaining = text
|
|||
|
|
|
|||
|
|
while remaining:
|
|||
|
|
if len(remaining) <= self.block_streaming_chunk_max_chars:
|
|||
|
|
chunks.append(remaining)
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
cutoff = self.block_streaming_chunk_max_chars
|
|||
|
|
chunk_text = remaining[:cutoff]
|
|||
|
|
|
|||
|
|
para_break = chunk_text.rfind("\n\n")
|
|||
|
|
if para_break > self.block_streaming_chunk_min_chars:
|
|||
|
|
cutoff = para_break + 2
|
|||
|
|
else:
|
|||
|
|
line_break = chunk_text.rfind("\n")
|
|||
|
|
if line_break > self.block_streaming_chunk_min_chars:
|
|||
|
|
cutoff = line_break + 1
|
|||
|
|
else:
|
|||
|
|
sentence_break = max(
|
|||
|
|
chunk_text.rfind("。"),
|
|||
|
|
chunk_text.rfind(". "),
|
|||
|
|
chunk_text.rfind("！"),
|
|||
|
|
chunk_text.rfind("？"),
|
|||
|
|
)
|
|||
|
|
if sentence_break > self.block_streaming_chunk_min_chars:
|
|||
|
|
cutoff = sentence_break + 1
|
|||
|
|
|
|||
|
|
chunks.append(remaining[:cutoff].strip())
|
|||
|
|
remaining = remaining[cutoff:].strip()
|
|||
|
|
|
|||
|
|
return chunks
|