wisruntime 0.1.0__py3-none-any.whl

wisruntime/server/dify_handler.py

@@ -0,0 +1,416 @@
+"""DifyHandler — Dify SSE 协议处理器。
+
+实现 Dify 兼容的对话 API:
+- POST /dify/v1/chat-messages (SSE Streaming + Blocking)
+- 心跳 ping 每 10 秒
+- agent_thought 的 id/position 状态维护
+
+CanonicalEvent → Dify SSE 映射:
+  TextDelta      → agent_message
+  ReasoningDelta → reasoning_chunk
+  ToolCallStart  → agent_thought (tool + tool_input)
+  ToolCallEnd    → agent_thought (observation)
+  SubAgentStart  → subagent_start
+  SubAgentEnd    → subagent_end
+  FileOutput     → message_file
+  Usage          → 累积到 message_end.metadata.usage
+  Error          → error
+  Custom         → {event: name, ...payload}
+  Done           → message_end
+"""
+
+import asyncio
+import json
+import logging
+import uuid
+from collections.abc import AsyncGenerator, AsyncIterator
+
+from fastapi import APIRouter, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+from sse_starlette.sse import EventSourceResponse
+from starlette.responses import Response
+
+from wisruntime.model import (
+    AgentRequest,
+    Message,
+    UsageEvent,
+)
+from wisruntime.server.invoker import AgentInvoker
+from wisruntime.server.protocol import ProtocolHandler
+
+logger = logging.getLogger(__name__)
+
+# SSE 心跳间隔 (秒)
+PING_INTERVAL = 10
+
+
+class DifyHandler(ProtocolHandler):
+    """Dify SSE 协议处理器。
+
+    将 CanonicalEvent 流序列化为 Dify 兼容的 SSE 格式。
+    内部维护 agent_thought 的 id/position 状态 (协议特有)。
+
+    Demo:
+        >>> handler = DifyHandler()
+        >>> handler.install(app, invoker)
+        >>> # → 注册 POST /dify/v1/chat-messages
+    """
+
+    def get_prefix(self) -> str:
+        return "/dify/v1"
+
+    async def parse_request(self, request: Request) -> AgentRequest:
+        """解析 Dify 请求 JSON 为归一化 AgentRequest。
+
+        字段映射:
+        - query → messages[0] (user message)
+        - response_mode → stream (streaming=True, blocking=False)
+        - conversation_id → conversation_id
+        - user → user_id
+        - inputs, files → protocol_extensions
+
+        Args:
+            request: Starlette Request。
+
+        Returns:
+            AgentRequest 归一化请求。
+        """
+        body = await request.json()
+
+        query = body.get("query", "")
+        response_mode = body.get("response_mode", "streaming")
+
+        # 构建归一化消息列表 — Dify 请求中是单轮 query
+        messages = [Message(role="user", content=query)]
+
+        # 收集 Dify 特有字段到 protocol_extensions
+        protocol_extensions = {}
+        for key in ("inputs", "files"):
+            if key in body:
+                protocol_extensions[key] = body[key]
+
+        return AgentRequest(
+            messages=messages,
+            stream=response_mode == "streaming",
+            conversation_id=body.get("conversation_id") or None,
+            user_id=body.get("user") or None,
+            headers=dict(request.headers),
+            raw_request=request,
+            protocol_extensions=protocol_extensions,
+        )
+
+    # ========================================================================
+    # SSE 流式序列化
+    # ========================================================================
+
+    async def _format_stream(
+        self,
+        events: AsyncGenerator,
+        task_id: str,
+    ) -> AsyncGenerator[str, None]:
+        """将 CanonicalEvent 流格式化为 Dify SSE 字符串。
+
+        维护内部状态:
+        - _position_counter: agent_thought 的位置序号 (从 1 开始递增)
+        - _accumulated_usage: 累积的 token 用量 (用于 message_end)
+
+        Args:
+            events: CanonicalEvent 异步流。
+            task_id: 本次请求的唯一标识。
+
+        Yields:
+            SSE 数据行字符串。
+        """
+        position = 0
+        accumulated_usage: dict = {}
+
+        async for event in events:
+            match event.type:
+                case "text_delta":
+                    yield self._sse_line(
+                        event="agent_message",
+                        task_id=task_id,
+                        answer=event.delta,
+                    )
+
+                case "reasoning_delta":
+                    yield self._sse_line(
+                        event="reasoning_chunk",
+                        task_id=task_id,
+                        reasoning=event.delta,
+                    )
+
+                case "tool_call_start":
+                    position += 1
+                    thought_id = str(uuid.uuid4())
+                    yield self._sse_line(
+                        event="agent_thought",
+                        task_id=task_id,
+                        id=thought_id,
+                        position=position,
+                        tool=event.name,
+                        tool_input=event.args_delta,
+                    )
+
+                case "tool_call_end":
+                    yield self._sse_line(
+                        event="agent_thought",
+                        task_id=task_id,
+                        id=str(uuid.uuid4()),
+                        position=position,
+                        observation=str(event.result) if event.result else event.error or "",
+                    )
+
+                case "subagent_start":
+                    yield self._sse_line(
+                        event="subagent_start",
+                        task_id=task_id,
+                        agent_name=event.agent_name,
+                        run_id=event.run_id,
+                        task_description=event.task_description,
+                    )
+
+                case "subagent_end":
+                    yield self._sse_line(
+                        event="subagent_end",
+                        task_id=task_id,
+                        agent_name=event.agent_name,
+                        run_id=event.run_id,
+                        status=event.status,
+                    )
+
+                case "file_output":
+                    yield self._sse_line(
+                        event="message_file",
+                        task_id=task_id,
+                        id=event.file_id,
+                        type=event.file_type,
+                        url=event.url,
+                        belongs_to=event.belongs_to,
+                    )
+
+                case "usage":
+                    accumulated_usage = {
+                        "prompt_tokens": event.prompt_tokens,
+                        "completion_tokens": event.completion_tokens,
+                        "total_tokens": event.total_tokens,
+                    }
+
+                case "error":
+                    yield self._sse_line(
+                        event="error",
+                        task_id=task_id,
+                        message=event.message,
+                        code=event.code or "UNKNOWN",
+                        status=500,
+                    )
+                    return  # 错误后停止流，不发送 message_end
+
+                case "custom":
+                    yield self._sse_line(
+                        event=event.name,
+                        task_id=task_id,
+                        **event.payload,
+                    )
+
+                case "done":
+                    yield self._sse_line(
+                        event="message_end",
+                        task_id=task_id,
+                        metadata={
+                            "usage": accumulated_usage,
+                        },
+                    )
+                    return  # 流正常结束
+
+                case _:
+                    # 未知事件类型 → 跳过 (向前兼容)
+                    logger.debug("跳过未知 CanonicalEvent 类型: %s", event.type)
+
+    # ========================================================================
+    # 心跳 ping
+    # ========================================================================
+
+    @staticmethod
+    async def _ping_generator(event_generator: AsyncIterator[str]) -> AsyncGenerator[str, None]:
+        """包装事件流，每 PING_INTERVAL 秒插入心跳 ping。
+
+        当内部 event_generator 结束时，ping task 自动取消。
+        使用 asyncio.create_task + Event 来实现"有数据时重置计时器"模式。
+
+        Args:
+            event_generator: 格式化后的事件 SSE 字符串流。
+
+        Yields:
+            包含心跳 ping 的完整 SSE 流。
+        """
+        ping_event = DifyHandler._sse_line_static(event="ping")
+
+        async def event_producer(queue: asyncio.Queue):
+            """从 event_generator 读取事件放入队列。"""
+            try:
+                async for event_str in event_generator:
+                    await queue.put(("data", event_str))
+                await queue.put(("done", None))
+            except Exception as exc:
+                await queue.put(("error", exc))
+
+        queue: asyncio.Queue = asyncio.Queue()
+        producer_task = asyncio.create_task(event_producer(queue))
+
+        try:
+            while True:
+                try:
+                    # 等待事件或超时
+                    msg_type, payload = await asyncio.wait_for(queue.get(), timeout=PING_INTERVAL)
+
+                    if msg_type == "data":
+                        yield payload
+                    elif msg_type == "done":
+                        break
+                    elif msg_type == "error":
+                        raise payload  # type: ignore[misc]
+                except asyncio.TimeoutError:
+                    # 超时 → 发送心跳
+                    yield ping_event
+
+        finally:
+            # 确保 producer task 被取消
+            if not producer_task.done():
+                producer_task.cancel()
+                try:
+                    await producer_task
+                except asyncio.CancelledError:
+                    pass
+
+    # ========================================================================
+    # Blocking 模式
+    # ========================================================================
+
+    async def _format_blocking(
+        self,
+        events: AsyncGenerator,
+        task_id: str,
+    ) -> dict:
+        """收集全部 CanonicalEvent，组装为 Blocking 模式响应。
+
+        Args:
+            events: CanonicalEvent 异步流。
+            task_id: 请求标识。
+
+        Returns:
+            Blocking 模式的 JSON 响应字典。
+        """
+        answer_parts: list[str] = []
+        metadata: dict = {}
+        tool_calls: list[dict] = []
+
+        async for event in events:
+            match event.type:
+                case "text_delta":
+                    answer_parts.append(event.delta)
+                case "usage":
+                    metadata["usage"] = {
+                        "prompt_tokens": event.prompt_tokens,
+                        "completion_tokens": event.completion_tokens,
+                        "total_tokens": event.total_tokens,
+                    }
+                case "tool_call_start":
+                    tool_calls.append({"name": event.name, "args": event.args_delta})
+                case "tool_call_end":
+                    if tool_calls:
+                        tool_calls[-1]["result"] = event.result
+                case "error":
+                    return {
+                        "task_id": task_id,
+                        "error": {"message": event.message, "code": event.code},
+                    }
+                case "done":
+                    break
+
+        return {
+            "task_id": task_id,
+            "answer": "".join(answer_parts),
+            "metadata": metadata,
+            "tool_calls": tool_calls,
+        }
+
+    # ========================================================================
+    # 路由注册
+    # ========================================================================
+
+    def _build_router(self) -> APIRouter:
+        """构建 Dify 路由。
+
+        注册 POST /chat-messages 端点，支持 streaming 和 blocking 两种模式。
+        """
+        router = APIRouter()
+        handler = self  # 闭包引用
+
+        @router.post("/chat-messages")
+        async def chat_messages(request: Request):
+            """Dify 对话接口 — SSE Streaming / Blocking。
+
+            根据请求中的 response_mode 自动选择:
+            - streaming: SSE 流式响应
+            - blocking: 阻塞等待完整结果后返回 JSON
+            """
+            # 解析请求
+            agent_request = await handler.parse_request(request)
+            task_id = str(uuid.uuid4())
+
+            if agent_request.stream:
+                # SSE Streaming 模式
+                async def generate_sse():
+                    """生成 SSE 事件流 (含心跳)。"""
+                    events = handler._invoker.invoke(handler._fn, agent_request)
+                    formatted = handler._format_stream(events, task_id)
+                    async for sse_line in handler._ping_generator(formatted):
+                        yield sse_line
+
+                return StreamingResponse(
+                    generate_sse(),
+                    media_type="text/event-stream",
+                    headers={
+                        "Cache-Control": "no-cache",
+                        "Connection": "keep-alive",
+                        "X-Accel-Buffering": "no",
+                    },
+                )
+            else:
+                # Blocking 模式
+                events = handler._invoker.invoke(handler._fn, agent_request)
+                result = await handler._format_blocking(events, task_id)
+                return JSONResponse(content=result)
+
+        return router
+
+    # ========================================================================
+    # SSE 格式化工具
+    # ========================================================================
+
+    @staticmethod
+    def _sse_line(**kwargs) -> str:
+        """将 dict 序列化为 Dify SSE 行: `data: {json}\\n\\n`。
+
+        Args:
+            **kwargs: 要序列化的字段。
+
+        Returns:
+            SSE 格式化字符串。
+
+        Demo:
+            >>> DifyHandler._sse_line(event="ping")
+            'data: {"event": "ping"}\\n\\n'
+        """
+        return f"data: {json.dumps(kwargs, ensure_ascii=False)}\n\n"
+
+    @staticmethod
+    def _sse_line_static(**kwargs) -> str:
+        """与 _sse_line 相同，但作为静态方法独立调用 (用于心跳)。
+
+        Demo:
+            >>> DifyHandler._sse_line_static(event="ping")
+            'data: {"event": "ping"}\\n\\n'
+        """
+        return f"data: {json.dumps(kwargs, ensure_ascii=False)}\n\n"

wisruntime/server/invoker.py

@@ -0,0 +1,161 @@
+"""AgentInvoker — 用户函数归一化引擎。
+
+负责调用用户编写的 invoke_agent 函数，将其产出的各种类型
+(str, CanonicalEvent, 异常) 统一转换为 AsyncGenerator[CanonicalEvent]。
+
+核心职责:
+1. 自动检测函数类型 (async generator / async func / sync func)
+2. 归一化输出: str → TextDeltaEvent, CanonicalEvent → 透传, None → 跳过
+3. 异常处理: Exception → ErrorEvent, 不 emit Done
+4. 自动插入 Done 事件 (正常结束时)
+
+参考: AgentRun SDK 的 AgentInvoker 设计。
+"""
+
+import asyncio
+import contextvars
+import inspect
+import logging
+from collections.abc import AsyncGenerator, Callable
+from typing import Any
+
+from wisruntime.exceptions import InvokerError
+from wisruntime.model import (
+    AgentRequest,
+    CanonicalEvent,
+    DoneEvent,
+    ErrorEvent,
+    TextDeltaEvent,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class AgentInvoker:
+    """用户函数归一化引擎。
+
+    接收用户编写的 invoke_agent 函数，调用它并将输出归一化为 CanonicalEvent 流。
+    用户函数可以是:
+    - 异步生成器 (async def + yield) — 最常见
+    - 异步函数 (async def + return) — 单次返回值
+    - 同步函数 (def) — 在线程池中执行
+
+    Demo:
+        >>> invoker = AgentInvoker()
+        >>> async def my_agent(request: AgentRequest):
+        ...     yield "你好"
+        ...     yield "世界"
+        >>> async for event in invoker.invoke(my_agent, request):
+        ...     print(event.type)
+        text_delta
+        text_delta
+        done
+    """
+
+    def __init__(self):
+        self._fn: Callable | None = None
+        self._is_async_gen: bool = False
+        self._is_async_func: bool = False
+
+    async def invoke(
+        self,
+        fn: Callable,
+        request: AgentRequest,
+    ) -> AsyncGenerator[CanonicalEvent, None]:
+        """调用用户函数并归一化输出。
+
+        Args:
+            fn: 用户编写的 invoke_agent 函数。
+            request: 归一化的 AgentRequest。
+
+        Yields:
+            CanonicalEvent 流，以 DoneEvent 结束 (正常) 或 ErrorEvent (异常)。
+        """
+        self._detect_function_type(fn)
+
+        try:
+            if self._is_async_gen:
+                # 异步生成器 — 最常见的场景，直接迭代
+                async for item in fn(request):
+                    event = self._normalize_item(item)
+                    if event is not None:
+                        yield event
+            elif self._is_async_func:
+                # 异步函数 (非生成器) — await 后处理返回值
+                result = await fn(request)
+                event = self._normalize_item(result)
+                if event is not None:
+                    yield event
+            else:
+                # 同步函数 — 在线程池中执行，保持 contextvars
+                ctx = contextvars.copy_context()
+                items = await asyncio.get_event_loop().run_in_executor(
+                    None,
+                    lambda: ctx.run(lambda: list(fn(request))),
+                )
+                for item in items:
+                    event = self._normalize_item(item)
+                    if event is not None:
+                        yield event
+
+            # 正常结束 — 自动插入 Done 事件
+            yield DoneEvent()
+
+        except Exception as exc:
+            logger.exception("invoke_agent 执行异常")
+            yield ErrorEvent(
+                message=str(exc),
+                code=type(exc).__name__,
+                recoverable=False,
+            )
+            # 异常后不 emit Done — 流以 Error 终止
+
+    def _detect_function_type(self, fn: Callable) -> None:
+        """检测函数类型: async generator / async function / sync function。
+
+        Args:
+            fn: 用户的 invoke_agent 函数。
+
+        Raises:
+            InvokerError: 如果是同步生成器 (不支持)，抛出此异常。
+        """
+        if inspect.isasyncgenfunction(fn):
+            self._is_async_gen = True
+            self._is_async_func = False
+        elif inspect.iscoroutinefunction(fn):
+            self._is_async_gen = False
+            self._is_async_func = True
+        elif inspect.isgeneratorfunction(fn):
+            # 同步生成器不支持 — 会阻塞事件循环
+            raise InvokerError(
+                "invoke_agent 不支持同步生成器 (def + yield)。"
+                "请使用异步生成器 (async def + yield) 或同步函数返回列表。"
+            )
+        else:
+            # 普通同步函数
+            self._is_async_gen = False
+            self._is_async_func = False
+
+    def _normalize_item(self, item: Any) -> CanonicalEvent | None:
+        """归一化单个输出项。
+
+        转换规则:
+        - str → TextDeltaEvent
+        - CanonicalEvent 子类 → 透传
+        - None → 跳过 (返回 None)
+
+        Args:
+            item: 用户函数 yield/return 的单个值。
+
+        Returns:
+            CanonicalEvent 或 None (表示跳过此项)。
+        """
+        if item is None:
+            return None
+        if isinstance(item, str):
+            return TextDeltaEvent(delta=item)
+        # 已经是 CanonicalEvent 的实例 (如 ToolCallStartEvent 等)，直接透传
+        if isinstance(item, type(TextDeltaEvent)) or hasattr(item, "type"):
+            return item
+        # 未知类型 — 转为 str 后包装为 TextDeltaEvent
+        return TextDeltaEvent(delta=str(item))

wisruntime/server/protocol.py

@@ -0,0 +1,104 @@
+"""协议处理器抽象基类和注册中心。
+
+ProtocolHandler 定义协议处理器的统一接口:
+- parse_request: HTTP 请求 → AgentRequest
+- _format_stream: CanonicalEvent 流 → 协议特定的 SSE/JSON 格式
+- install: 在 FastAPI app 上注册路由
+
+ProtocolRegistry 管理多个 handler 的注册和查找。
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator, Callable
+
+from fastapi import APIRouter
+
+from wisruntime.model import AgentRequest, CanonicalEvent
+from wisruntime.server.invoker import AgentInvoker
+
+
+class ProtocolHandler(ABC):
+    """协议处理器抽象基类。
+
+    每个协议 (Dify SSE, AG-UI, ...) 实现自己的 Handler，
+    负责请求解析和响应序列化。
+
+    install 时由 Server 注入 AgentInvoker 和用户函数 fn，
+    路由处理时可以调用 invoker.invoke(fn, request)。
+    """
+
+    @abstractmethod
+    def get_prefix(self) -> str:
+        """返回此协议的路由前缀，如 '/dify/v1'。"""
+        ...
+
+    @abstractmethod
+    async def parse_request(self, request) -> AgentRequest:
+        """解析原始 HTTP 请求为归一化 AgentRequest。"""
+        ...
+
+    @abstractmethod
+    async def _format_stream(
+        self,
+        events: AsyncGenerator[CanonicalEvent, None],
+        task_id: str,
+    ) -> AsyncGenerator[str, None]:
+        """将 CanonicalEvent 流格式化为协议特定的 SSE 字符串。"""
+        ...
+
+    def install(self, app, invoker: AgentInvoker, fn: Callable) -> None:
+        """在 FastAPI app 上注册此协议的路由。
+
+        子类可以覆写此方法以自定义路由注册逻辑。
+        默认实现调用 _build_router()。
+
+        Args:
+            app: FastAPI 应用实例。
+            invoker: AgentInvoker 实例，由 Server 注入。
+            fn: 用户的 invoke_agent 函数。
+        """
+        self._invoker = invoker
+        self._fn = fn
+        router = self._build_router()
+        prefix = self.get_prefix()
+        app.include_router(router, prefix=prefix)
+
+    @abstractmethod
+    def _build_router(self) -> APIRouter:
+        """构建包含此协议路由的 APIRouter。"""
+        ...
+
+
+class ProtocolRegistry:
+    """协议处理器注册中心。
+
+    管理多个 ProtocolHandler 的注册和查找。
+
+    Demo:
+        >>> registry = ProtocolRegistry()
+        >>> registry.register(DifyHandler())
+        >>> len(registry.handlers)
+        1
+    """
+
+    def __init__(self):
+        self._handlers: list[ProtocolHandler] = []
+
+    @property
+    def handlers(self) -> list[ProtocolHandler]:
+        return self._handlers
+
+    def register(self, handler: ProtocolHandler) -> None:
+        """注册一个协议处理器。"""
+        self._handlers.append(handler)
+
+    def install_all(self, app, invoker: AgentInvoker, fn: Callable) -> None:
+        """将所有已注册的 handler 安装到 FastAPI app 上。
+
+        Args:
+            app: FastAPI 应用实例。
+            invoker: AgentInvoker 实例。
+            fn: 用户的 invoke_agent 函数。
+        """
+        for handler in self._handlers:
+            handler.install(app, invoker, fn)