diff --git a/CHANGELOG.md b/CHANGELOG.md index 8319618..7386ba1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,23 @@ # Changelog +## Unreleased + +- Coze 集成补齐开发者文档全量能力：`AsyncCozeClient.upload_file` 走 `/v1/files/upload` + （multipart）；新增 `file_object`/`image_object`/`text_object` 与 `object_string` 消息编码， + 支持在 `additional_messages` 中携带文件/图片。 +- `CozeAgentNode`/`CozeChatModel` 流式输出 `reasoning_content`（思考信息，打 + `additional_kwargs={"reasoning": True}` 标记）并收集 `follow_up` 用户问题建议，写入最终 + `AIMessage.additional_kwargs`（`reasoning_content`/`follow_ups`）与 `response_metadata`。 + `CozeAgentNode(suggestions_key=...)` 可选把建议写入 state（默认 `None`，不破坏严格 schema）。 +- 修复：`conversation.message.delta`/`.completed` 流式路径此前未按 `data.type` 过滤， + `verbose`（多智能体 jump 信息）、`function_call`、`knowledge_recall` 等非正文消息会被 + 错误拼接进可见回答；新增 `_is_answer_delta` 只放行 `type in (None, "answer")`。 +- 新增会话/消息管理端点：`conversation_retrieve`、`conversation_message_create`、 + `conversation_message_list`（游标分页）、`conversation_message_retrieve`；新增 + `file_retrieve`（查询上传文件状态）与 `bot_retrieve`（bot 元信息）。 +- `conversation.chat.completed` 的 `usage`（`token_count`/`input_count`/`output_count`） + 现在会写入最终 `AIMessage.usage`，流式与轮询路径均覆盖。 + ## 2.0.0 - 完成 MVP P0/P1 硬化：强类型 state/output/工具参数校验、结构化输出修复、工具权限/审批/ diff --git a/docs/integrations-coze.md b/docs/integrations-coze.md index 96b3c18..c3ab70d 100644 --- a/docs/integrations-coze.md +++ b/docs/integrations-coze.md @@ -10,6 +10,46 @@ conversation_id 写入用户声明的 `coze_conversations` 状态键。SSE delta `astream(..., stream_mode="messages"|"events")` 已可消费，Chainlit 可直接逐 token 调用 `stream_token()`。requires_action 可交给本地工具或通过 `hitl=True` 触发耐久审批 interrupt。 +## 完整能力（对齐 Coze 开发者文档） + +- **流式返回**：SSE `conversation.message.delta` 的 `content` 逐 token 通过 messages 模式输出。 +- **思考信息流式输出**：delta 中的 `reasoning_content` 单独形成 `AIMessageChunk`，并在 + `additional_kwargs={"reasoning": True}` 上打标，Chainlit 侧据此渲染独立的“思考中”Step； + 完整思维链也汇总到最终 `AIMessage.additional_kwargs["reasoning_content"]`。 +- **用户问题建议（follow-up）**：`conversation.message.completed` 中 `type == "follow_up"` + 的消息被收集为建议问题，写入最终 `AIMessage.additional_kwargs["follow_ups"]` 与 + `response_metadata["follow_ups"]`；若 `CozeAgentNode(suggestions_key=...)` 指定了状态键 + （且图 state schema 已声明），也会写入该键。非流式轮询路径从 `chat/message/list` 的 + `type == "follow_up"` 条目提取，逻辑一致。 +- **文件上传**：`AsyncCozeClient.upload_file(content, filename=..., content_type=...)` 调用 + `/v1/files/upload`（multipart），返回含 `id` 的文件元数据。用 `file_object(file_id)` / + `image_object(file_id)` / `text_object(text)` 构造 `object_string` 项，放入 + `HumanMessage(additional_kwargs={"objects": [...]})`；`_message_to_coze` 会自动切换为 + `content_type="object_string"`，并在缺省时把消息正文补成首个 text 项。`file_retrieve(file_id)` + 可查询上传文件的处理状态/元数据。 +- **Token 用量**：`conversation.chat.completed`（流式）与轮询终态 `chat` 对象的 `usage` + 字段（`token_count`/`input_count`/`output_count`）会写入最终 `AIMessage.usage`。 + +`suggestions_key` 默认 `None`：只有显式指定且 state schema 声明该键时才写入，避免破坏严格 +schema 的图。 + +### 消息 type 过滤（重要） + +Coze `conversation.message.delta` / `.completed` 把 `answer`、`function_call`、 +`tool_output`、`verbose`（如多智能体 jump 信息）、`knowledge_recall`、`follow_up` 等不同 +语义的消息复用同一个事件名，仅靠 `event:` 无法区分。集成内部用 `_is_answer_delta(data)` +（即 `data["type"] in (None, "answer")`）过滤出真正要展示给用户的正文，避免 verbose/ +function_call/knowledge_recall 内容混入可见回答；`follow_up` 单独按 `_extract_follow_up` +收集。自定义消费 `chat_stream()`/`workflow_stream()` 原始事件时也应做同样的 `type` 判断。 + +### 会话与消息管理 + +除 `create_conversation` 外，还提供 `conversation_retrieve`、 +`conversation_message_create`（直接向会话追加消息，不触发 chat 运行）、 +`conversation_message_list`（`before_id`/`after_id`/`limit` 游标分页，返回 +`items`/`first_id`/`last_id`/`has_more`）、`conversation_message_retrieve`。 +`bot_retrieve(bot_id)` 对应 `/v1/bot/get_online_info`，用于读取 bot 元信息。 + `CozeWorkflowNode` 遇到工作流 interrupt 时返回 `coze_workflow_question`。恢复值必须回显 `event_id`、`interrupt_type` 和 `resume_data`，例如： diff --git a/src/lingxigraph/integrations/__init__.py b/src/lingxigraph/integrations/__init__.py index 4aad462..d4e704d 100644 --- a/src/lingxigraph/integrations/__init__.py +++ b/src/lingxigraph/integrations/__init__.py @@ -6,7 +6,15 @@ def __getattr__(name: str) -> Any: - if name in {"AsyncCozeClient", "CozeAgentNode", "CozeChatModel", "CozeWorkflowNode"}: + if name in { + "AsyncCozeClient", + "CozeAgentNode", + "CozeChatModel", + "CozeWorkflowNode", + "file_object", + "image_object", + "text_object", + }: from . import coze return getattr(coze, name) @@ -23,4 +31,7 @@ def __getattr__(name: str) -> Any: "CozeChatModel", "CozeWorkflowNode", "OpenAICompatChatModel", + "file_object", + "image_object", + "text_object", ] diff --git a/src/lingxigraph/integrations/coze.py b/src/lingxigraph/integrations/coze.py index 26e4b6f..6b58954 100644 --- a/src/lingxigraph/integrations/coze.py +++ b/src/lingxigraph/integrations/coze.py @@ -36,6 +36,13 @@ "submit_tool_outputs": "/v3/chat/submit_tool_outputs", "cancel_chat": "/v3/chat/cancel", "conversation": "/v1/conversation/create", + "conversation_retrieve": "/v1/conversation/retrieve", + "conversation_message_create": "/v1/conversation/message/create", + "conversation_message_list": "/v1/conversation/message/list", + "conversation_message_retrieve": "/v1/conversation/message/retrieve", + "files_upload": "/v1/files/upload", + "files_retrieve": "/v1/files/retrieve", + "bot_retrieve": "/v1/bot/get_online_info", "workflow": "/v1/workflow/run", "workflow_stream": "/v1/workflow/stream_run", "workflow_resume": "/v1/workflow/stream_resume", @@ -109,12 +116,15 @@ async def _headers( *, operation_key: str, last_event_id: str | None = None, + content_type: str | None = "application/json", ) -> dict[str, str]: token = self._token if token is None and self._token_provider is not None: provided = self._token_provider() token = await provided if inspect.isawaitable(provided) else provided - headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"} + headers = {"Authorization": f"Bearer {token}"} + if content_type is not None: + headers["Content-Type"] = content_type try: runtime = get_runtime() runtime.raise_if_cancelled() @@ -128,12 +138,16 @@ async def _headers( async def _request(self, method: str, endpoint: str, **kwargs: Any) -> dict[str, Any]: response: httpx.Response | None = None operation_key = str(uuid4()) + # httpx sets the multipart Content-Type (with boundary) itself, so drop ours. + content_type = None if "files" in kwargs else "application/json" for attempt in range(self.max_retries + 1): try: response = await self._client.request( method, _ENDPOINTS[endpoint], - headers=await self._headers(operation_key=operation_key), + headers=await self._headers( + operation_key=operation_key, content_type=content_type + ), **kwargs, ) if should_retry_status(response.status_code) and attempt < self.max_retries: @@ -266,6 +280,85 @@ async def cancel_chat(self, conversation_id: str, chat_id: str) -> dict[str, Any async def create_conversation(self, messages: Sequence[Mapping[str, Any]] = ()) -> dict[str, Any]: return await self._request("POST", "conversation", json={"messages": list(messages)}) + async def conversation_retrieve(self, conversation_id: str) -> dict[str, Any]: + return await self._request( + "GET", "conversation_retrieve", params={"conversation_id": conversation_id} + ) + + async def conversation_message_create( + self, conversation_id: str, *, role: str, content: str, content_type: str = "text" + ) -> dict[str, Any]: + """Append a message to a conversation without triggering a chat run.""" + return await self._request( + "POST", + "conversation_message_create", + params={"conversation_id": conversation_id}, + json={"role": role, "content": content, "content_type": content_type}, + ) + + async def conversation_message_list( + self, + conversation_id: str, + *, + order: str = "desc", + chat_id: str | None = None, + before_id: str | None = None, + after_id: str | None = None, + limit: int = 50, + ) -> dict[str, Any]: + """List messages in a conversation with cursor pagination. + + Returns a mapping with ``items``, ``first_id``, ``last_id`` and ``has_more``. + """ + body: dict[str, Any] = {"order": order, "limit": limit} + if chat_id: + body["chat_id"] = chat_id + if before_id: + body["before_id"] = before_id + if after_id: + body["after_id"] = after_id + return await self._request( + "POST", + "conversation_message_list", + params={"conversation_id": conversation_id}, + json=body, + ) + + async def conversation_message_retrieve( + self, conversation_id: str, message_id: str + ) -> dict[str, Any]: + return await self._request( + "GET", + "conversation_message_retrieve", + params={"conversation_id": conversation_id, "message_id": message_id}, + ) + + async def file_retrieve(self, file_id: str) -> dict[str, Any]: + return await self._request("GET", "files_retrieve", params={"file_id": file_id}) + + async def bot_retrieve(self, bot_id: str) -> dict[str, Any]: + return await self._request("GET", "bot_retrieve", params={"bot_id": bot_id}) + + async def upload_file( + self, + content: bytes, + *, + filename: str, + content_type: str = "application/octet-stream", + ) -> dict[str, Any]: + """Upload a file to Coze and return its metadata (including ``id``). + + The returned ``id`` is referenced from ``additional_messages`` via an + ``object_string`` content item (see :func:`file_object` / + :func:`image_object`). + """ + + return await self._request( + "POST", + "files_upload", + files={"file": (filename, content, content_type)}, + ) + async def workflow_run(self, workflow_id: str, parameters: Mapping[str, Any], **kwargs: Any) -> dict[str, Any]: return await self._request( "POST", "workflow", json={"workflow_id": workflow_id, "parameters": dict(parameters), **kwargs} @@ -298,9 +391,41 @@ async def aclose(self) -> None: await self._client.aclose() +def file_object(file_id: str) -> dict[str, str]: + """Build a Coze ``object_string`` item that references an uploaded file.""" + + return {"type": "file", "file_id": str(file_id)} + + +def image_object(file_id: str) -> dict[str, str]: + """Build a Coze ``object_string`` item that references an uploaded image.""" + + return {"type": "image", "file_id": str(file_id)} + + +def text_object(text: str) -> dict[str, str]: + """Build a Coze ``object_string`` text item.""" + + return {"type": "text", "text": str(text)} + + def _message_to_coze(message: AnyMessage) -> dict[str, Any]: role = {"human": "user", "ai": "assistant"}.get(message.type, message.type) - data: dict[str, Any] = {"role": role, "content": message.content, "content_type": "text"} + # Multimodal messages carry object_string items in additional_kwargs["objects"]; + # when present we emit a Coze object_string message (text + file/image refs). + objects = message.additional_kwargs.get("objects") + if objects: + items = list(objects) + text = str(message.content or "") + if text and not any(item.get("type") == "text" for item in items): + items = [text_object(text), *items] + data: dict[str, Any] = { + "role": role, + "content": json.dumps(items, ensure_ascii=False), + "content_type": "object_string", + } + else: + data = {"role": role, "content": message.content, "content_type": "text"} if isinstance(message, ToolMessage): data["tool_call_id"] = message.tool_call_id return data @@ -323,6 +448,45 @@ def _extract_tool_calls(data: Mapping[str, Any]) -> tuple[ToolCall, ...]: return tuple(calls) +def _is_answer_delta(data: Mapping[str, Any]) -> bool: + """True if a message.delta/completed event carries the visible answer text. + + Coze multiplexes ``function_call``, ``tool_output``, ``verbose`` (e.g. multi-agent + jump info) and ``knowledge_recall`` payloads onto the same event names as the + real answer; only ``type in (None, "answer")`` should be shown to the user. + """ + + return data.get("type") in (None, "answer") + + +def _extract_follow_up(data: Mapping[str, Any]) -> str | None: + """Return the follow-up question text from a completed message event, if any.""" + + if data.get("type") != "follow_up": + return None + content = str(data.get("content", "")).strip() + return content or None + + +def _split_answer_and_follow_ups( + items: Sequence[Mapping[str, Any]], +) -> tuple[str, list[str]]: + """Split polled chat messages into the assistant answer and follow-up prompts.""" + + answer_parts: list[str] = [] + follow_ups: list[str] = [] + for item in items: + if item.get("role") != "assistant": + continue + if item.get("type") == "follow_up": + text = str(item.get("content", "")).strip() + if text: + follow_ups.append(text) + elif item.get("type") in (None, "answer"): + answer_parts.append(str(item.get("content", ""))) + return "".join(answer_parts), follow_ups + + @dataclass(slots=True) class CozeAgentNode: bot_id: str @@ -330,6 +494,9 @@ class CozeAgentNode: user_id: str messages_key: str = "messages" conversation_key: str = "coze_conversations" + # State key for follow-up suggestions. Only written when the graph's state + # schema declares it (opt-in), so strict schemas are never broken. + suggestions_key: str | None = None stream: bool = True tools: Sequence[ToolSpec | Callable[..., Any]] | None = None hitl: bool = False @@ -342,8 +509,11 @@ async def __call__(self, state: Mapping[str, Any], runtime: Runtime[Any]) -> Map conversation_id = conversations.get(self.bot_id) messages = [_message_to_coze(item) for item in state.get(self.messages_key, ())] content = "" + reasoning = "" + suggestions: list[str] = [] chat_id: str | None = None tool_calls: tuple[ToolCall, ...] = () + usage: dict[str, Any] = {} if self.stream: try: async for event in self.client.chat_stream( @@ -360,11 +530,31 @@ async def __call__(self, state: Mapping[str, Any], runtime: Runtime[Any]) -> Map ) or None runtime.raise_if_cancelled() if kind == "conversation.message.delta": - delta = str(data.get("content", "")) - content += delta - runtime.emit_message( - AIMessageChunk(delta, id=chat_id), {"provider": "coze"} - ) + reasoning_delta = str(data.get("reasoning_content", "") or "") + if reasoning_delta: + reasoning += reasoning_delta + runtime.emit_message( + AIMessageChunk( + reasoning_delta, + id=chat_id, + additional_kwargs={"reasoning": True}, + ), + {"provider": "coze", "channel": "reasoning"}, + ) + delta = str(data.get("content", "")) if _is_answer_delta(data) else "" + if delta: + content += delta + runtime.emit_message( + AIMessageChunk(delta, id=chat_id), {"provider": "coze"} + ) + elif kind == "conversation.message.completed": + suggestion = _extract_follow_up(data) + if suggestion is not None: + suggestions.append(suggestion) + elif kind == "conversation.chat.completed": + raw_usage = data.get("usage") if isinstance(data, Mapping) else None + if isinstance(raw_usage, Mapping): + usage = dict(raw_usage) elif kind == "conversation.chat.requires_action": tool_calls = _extract_tool_calls(data) elif kind in {"conversation.chat.failed", "error"}: @@ -391,10 +581,10 @@ async def __call__(self, state: Mapping[str, Any], runtime: Runtime[Any]) -> Map if chat.get("status") == "requires_action": tool_calls = _extract_tool_calls(chat) else: + if isinstance(chat.get("usage"), Mapping): + usage = dict(chat["usage"]) items = await self.client.chat_messages(conversation_id, chat_id) - content = "".join( - str(item.get("content", "")) for item in items if item.get("role") == "assistant" - ) + content, suggestions = _split_answer_and_follow_ups(items) if runtime.cancelled and conversation_id and chat_id: await self.client.cancel_chat(conversation_id, chat_id) if tool_calls: @@ -442,12 +632,10 @@ async def __call__(self, state: Mapping[str, Any], runtime: Runtime[Any]) -> Map if chat.get("status") == "requires_action": tool_calls = _extract_tool_calls(chat) continue + if isinstance(chat.get("usage"), Mapping): + usage = dict(chat["usage"]) items = await self.client.chat_messages(conversation_id, chat_id) - content = "".join( - str(item.get("content", "")) - for item in items - if item.get("role") == "assistant" - ) + content, suggestions = _split_answer_and_follow_ups(items) tool_calls = () break else: @@ -456,16 +644,30 @@ async def __call__(self, state: Mapping[str, Any], runtime: Runtime[Any]) -> Map ) if conversation_id: conversations[self.bot_id] = conversation_id - return { + additional_kwargs: dict[str, Any] = {} + if reasoning: + additional_kwargs["reasoning_content"] = reasoning + if suggestions: + additional_kwargs["follow_ups"] = tuple(suggestions) + result: dict[str, Any] = { self.messages_key: [ AIMessage( content, tool_calls=tool_calls, - response_metadata={"conversation_id": conversation_id, "chat_id": chat_id}, + usage=usage, + additional_kwargs=additional_kwargs, + response_metadata={ + "conversation_id": conversation_id, + "chat_id": chat_id, + "follow_ups": tuple(suggestions), + }, ) ], self.conversation_key: conversations, } + if self.suggestions_key is not None: + result[self.suggestions_key] = tuple(suggestions) + return result @dataclass(slots=True) @@ -555,20 +757,25 @@ async def agenerate(self, messages: Sequence[AnyMessage], *, tools=None, **kwarg }, ) items = await self.client.chat_messages(str(conversation_id), str(chat_id)) + answer, tool_follow_ups = _split_answer_and_follow_ups(items) return AIMessage( - "".join( - str(item.get("content", "")) - for item in items - if item.get("role") == "assistant" + answer, + usage=dict(chat["usage"]) if isinstance(chat.get("usage"), Mapping) else {}, + additional_kwargs=( + {"follow_ups": tuple(tool_follow_ups)} if tool_follow_ups else {} ), response_metadata={ "conversation_id": conversation_id, "chat_id": chat_id, + "follow_ups": tuple(tool_follow_ups), }, ) content = "" + reasoning = "" + follow_ups: list[str] = [] calls: tuple[ToolCall, ...] = () metadata: dict[str, Any] = {} + usage: dict[str, Any] = {} async for event in self.client.chat_stream( self.bot_id, self.user_id, @@ -577,14 +784,36 @@ async def agenerate(self, messages: Sequence[AnyMessage], *, tools=None, **kwarg ): data = event.get("data", {}) if event["event"] == "conversation.message.delta": - content += str(data.get("content", "")) + if _is_answer_delta(data): + content += str(data.get("content", "")) + reasoning += str(data.get("reasoning_content", "") or "") + elif event["event"] == "conversation.message.completed": + suggestion = _extract_follow_up(data) + if suggestion is not None: + follow_ups.append(suggestion) + elif event["event"] == "conversation.chat.completed": + raw_usage = data.get("usage") if isinstance(data, Mapping) else None + if isinstance(raw_usage, Mapping): + usage = dict(raw_usage) elif event["event"] == "conversation.chat.requires_action": calls = _extract_tool_calls(data) if isinstance(data, Mapping): metadata.update( {key: data[key] for key in ("conversation_id", "chat_id") if key in data} ) - return AIMessage(content, tool_calls=calls, response_metadata=metadata) + additional_kwargs: dict[str, Any] = {} + if reasoning: + additional_kwargs["reasoning_content"] = reasoning + if follow_ups: + additional_kwargs["follow_ups"] = tuple(follow_ups) + metadata["follow_ups"] = tuple(follow_ups) + return AIMessage( + content, + tool_calls=calls, + usage=usage, + additional_kwargs=additional_kwargs, + response_metadata=metadata, + ) async def astream( self, messages: Sequence[AnyMessage], *, tools=None, **kwargs: Any @@ -598,7 +827,16 @@ async def astream( ): if event["event"] == "conversation.message.delta": data = event.get("data", {}) - yield AIMessageChunk(str(data.get("content", "")), id=data.get("id")) + reasoning_delta = str(data.get("reasoning_content", "") or "") + if reasoning_delta: + yield AIMessageChunk( + reasoning_delta, + id=data.get("id"), + additional_kwargs={"reasoning": True}, + ) + content_delta = str(data.get("content", "")) if _is_answer_delta(data) else "" + if content_delta: + yield AIMessageChunk(content_delta, id=data.get("id")) __all__ = [ @@ -606,4 +844,7 @@ async def astream( "CozeAgentNode", "CozeChatModel", "CozeWorkflowNode", + "file_object", + "image_object", + "text_object", ] diff --git a/tests/test_integrations_v2.py b/tests/test_integrations_v2.py index 919cfce..91671b3 100644 --- a/tests/test_integrations_v2.py +++ b/tests/test_integrations_v2.py @@ -23,6 +23,9 @@ CozeAgentNode, CozeChatModel, CozeWorkflowNode, + _message_to_coze, + file_object, + image_object, ) from lingxigraph.integrations.openai_compat import OpenAICompatChatModel @@ -262,5 +265,377 @@ async def run() -> None: asyncio.run(run()) +class CozeCompleteFeatureTests(unittest.TestCase): + def test_file_upload_multipart_and_object_string(self) -> None: + captured: dict[str, Any] = {} + + async def handler(request: httpx.Request) -> httpx.Response: + captured["path"] = request.url.path + captured["content_type"] = request.headers.get("content-type", "") + captured["body"] = request.content + return httpx.Response(200, json={"code": 0, "data": {"id": "file-123"}}) + + async def run() -> None: + client = AsyncCozeClient("token", transport=httpx.MockTransport(handler)) + result = await client.upload_file( + b"hello", filename="a.txt", content_type="text/plain" + ) + self.assertEqual(result["id"], "file-123") + await client.aclose() + + asyncio.run(run()) + self.assertEqual(captured["path"], "/v1/files/upload") + self.assertTrue(captured["content_type"].startswith("multipart/form-data")) + self.assertIn(b"hello", captured["body"]) + + # An object_string message carries text plus file/image references. + message = HumanMessage( + "look at these", + additional_kwargs={"objects": [file_object("f1"), image_object("i1")]}, + ) + encoded = _message_to_coze(message) + self.assertEqual(encoded["content_type"], "object_string") + items = json.loads(encoded["content"]) + self.assertEqual(items[0], {"type": "text", "text": "look at these"}) + self.assertEqual(items[1], {"type": "file", "file_id": "f1"}) + self.assertEqual(items[2], {"type": "image", "file_id": "i1"}) + + def test_reasoning_and_follow_up_streaming(self) -> None: + body = ( + "event: conversation.chat.created\n" + 'data: {"id":"chat-1","conversation_id":"conv-1"}\n\n' + "event: conversation.message.delta\n" + 'data: {"id":"m1","reasoning_content":"let me think"}\n\n' + "event: conversation.message.delta\n" + 'data: {"id":"m1","content":"answer"}\n\n' + "event: conversation.message.completed\n" + 'data: {"id":"m2","type":"follow_up","content":"want more?"}\n\n' + "event: done\ndata: [DONE]\n\n" + ) + + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, text=body, headers={"content-type": "text/event-stream"} + ) + + async def run() -> None: + client = AsyncCozeClient("token", transport=httpx.MockTransport(handler)) + model = CozeChatModel("bot", client=client, user_id="user") + # agenerate splits reasoning, content, and follow-ups. + result = await model.agenerate([HumanMessage("hi")]) + self.assertEqual(result.content, "answer") + self.assertEqual(result.additional_kwargs["reasoning_content"], "let me think") + self.assertEqual(result.additional_kwargs["follow_ups"], ("want more?",)) + self.assertEqual(result.response_metadata["follow_ups"], ("want more?",)) + # astream tags reasoning chunks so the UI can route them separately. + chunks = [chunk async for chunk in model.astream([HumanMessage("hi")])] + reasoning = [c.content for c in chunks if c.additional_kwargs.get("reasoning")] + answer = [c.content for c in chunks if not c.additional_kwargs.get("reasoning")] + self.assertEqual("".join(reasoning), "let me think") + self.assertEqual("".join(answer), "answer") + await client.aclose() + + asyncio.run(run()) + + def test_agent_node_surfaces_reasoning_and_follow_ups(self) -> None: + class FakeClient: + async def chat_stream(self, *args, **kwargs): + del args, kwargs + yield { + "event": "conversation.chat.created", + "data": {"id": "chat", "conversation_id": "conv"}, + } + yield { + "event": "conversation.message.delta", + "data": {"id": "chat", "reasoning_content": "thinking..."}, + } + yield { + "event": "conversation.message.delta", + "data": {"id": "chat", "content": "done"}, + } + yield { + "event": "conversation.message.completed", + "data": {"type": "follow_up", "content": "next?"}, + } + yield { + "event": "conversation.chat.completed", + "data": {"usage": {"token_count": 9, "input_count": 4, "output_count": 5}}, + } + + class AgentState(TypedDict): + messages: Annotated[list[Any], add_messages] + coze_conversations: dict[str, str] + coze_suggestions: tuple[str, ...] + + node = CozeAgentNode( + "bot", client=FakeClient(), user_id="user", suggestions_key="coze_suggestions" + ) + builder = StateGraph(AgentState) + builder.add_node("coze", node).add_edge(START, "coze").add_edge("coze", END) + result = builder.compile().invoke( + { + "messages": [HumanMessage("go")], + "coze_conversations": {}, + "coze_suggestions": (), + } + ) + final = result["messages"][-1] + self.assertEqual(final.content, "done") + self.assertEqual(final.additional_kwargs["reasoning_content"], "thinking...") + self.assertEqual(final.additional_kwargs["follow_ups"], ("next?",)) + self.assertEqual(result["coze_suggestions"], ("next?",)) + self.assertEqual(final.usage, {"token_count": 9, "input_count": 4, "output_count": 5}) + + def test_verbose_and_function_call_deltas_do_not_leak_into_answer(self) -> None: + # Coze multiplexes verbose (multi-agent jump), function_call, and + # knowledge_recall deltas onto the same event name as the real answer; + # only type in (None, "answer") should ever reach visible content. + body = ( + "event: conversation.chat.created\n" + 'data: {"id":"chat-1","conversation_id":"conv-1"}\n\n' + "event: conversation.message.delta\n" + 'data: {"id":"m0","type":"verbose","content":"{\\"msg_type\\":\\"jump_to\\"}"}\n\n' + "event: conversation.message.delta\n" + 'data: {"id":"m1","type":"answer","content":"real "}\n\n' + "event: conversation.message.delta\n" + 'data: {"id":"m2","type":"knowledge_recall","content":"[recalled chunk]"}\n\n' + "event: conversation.message.delta\n" + 'data: {"id":"m1","type":"answer","content":"answer"}\n\n' + "event: done\ndata: [DONE]\n\n" + ) + + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, text=body, headers={"content-type": "text/event-stream"} + ) + + async def run() -> None: + client = AsyncCozeClient("token", transport=httpx.MockTransport(handler)) + model = CozeChatModel("bot", client=client, user_id="user") + result = await model.agenerate([HumanMessage("hi")]) + self.assertEqual(result.content, "real answer") + chunks = [chunk async for chunk in model.astream([HumanMessage("hi")])] + self.assertEqual("".join(c.content for c in chunks), "real answer") + await client.aclose() + + asyncio.run(run()) + + class FakeClient: + async def chat_stream(self, *args, **kwargs): + del args, kwargs + yield { + "event": "conversation.chat.created", + "data": {"id": "chat", "conversation_id": "conv"}, + } + yield { + "event": "conversation.message.delta", + "data": {"id": "chat", "type": "verbose", "content": "jump info"}, + } + yield { + "event": "conversation.message.delta", + "data": {"id": "chat", "type": "answer", "content": "clean answer"}, + } + + class AgentState(TypedDict): + messages: Annotated[list[Any], add_messages] + coze_conversations: dict[str, str] + + node = CozeAgentNode("bot", client=FakeClient(), user_id="user") + builder = StateGraph(AgentState) + builder.add_node("coze", node).add_edge(START, "coze").add_edge("coze", END) + result = builder.compile().invoke( + {"messages": [HumanMessage("go")], "coze_conversations": {}} + ) + self.assertEqual(result["messages"][-1].content, "clean answer") + + def test_token_usage_extracted_from_chat_completed(self) -> None: + body = ( + "event: conversation.chat.created\n" + 'data: {"id":"chat-1","conversation_id":"conv-1"}\n\n' + "event: conversation.message.delta\n" + 'data: {"id":"m1","type":"answer","content":"hi"}\n\n' + "event: conversation.chat.completed\n" + 'data: {"id":"chat-1","conversation_id":"conv-1",' + '"usage":{"token_count":30,"input_count":10,"output_count":20}}\n\n' + "event: done\ndata: [DONE]\n\n" + ) + + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, text=body, headers={"content-type": "text/event-stream"} + ) + + async def run() -> None: + client = AsyncCozeClient("token", transport=httpx.MockTransport(handler)) + model = CozeChatModel("bot", client=client, user_id="user") + result = await model.agenerate([HumanMessage("hi")]) + self.assertEqual( + result.usage, + {"token_count": 30, "input_count": 10, "output_count": 20}, + ) + await client.aclose() + + asyncio.run(run()) + + class FakeClient: + async def chat_stream(self, *args, **kwargs): + del args, kwargs + yield { + "event": "conversation.chat.created", + "data": {"id": "chat", "conversation_id": "conv"}, + } + yield { + "event": "conversation.message.delta", + "data": {"id": "chat", "type": "answer", "content": "hi"}, + } + yield { + "event": "conversation.chat.completed", + "data": {"usage": {"token_count": 30, "input_count": 10, "output_count": 20}}, + } + + class AgentState(TypedDict): + messages: Annotated[list[Any], add_messages] + coze_conversations: dict[str, str] + + node = CozeAgentNode("bot", client=FakeClient(), user_id="user") + builder = StateGraph(AgentState) + builder.add_node("coze", node).add_edge(START, "coze").add_edge("coze", END) + result = builder.compile().invoke( + {"messages": [HumanMessage("go")], "coze_conversations": {}} + ) + self.assertEqual( + result["messages"][-1].usage, + {"token_count": 30, "input_count": 10, "output_count": 20}, + ) + + def test_conversation_message_and_file_and_bot_endpoints(self) -> None: + paths: list[str] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + paths.append(request.url.path) + if request.url.path == "/v1/conversation/message/list": + return httpx.Response( + 200, + json={ + "code": 0, + "data": { + "items": [{"id": "msg-1", "content": "hi"}], + "first_id": "msg-1", + "last_id": "msg-1", + "has_more": False, + }, + }, + ) + if request.url.path == "/v1/files/retrieve": + return httpx.Response( + 200, json={"code": 0, "data": {"id": "file-1", "status": "processed"}} + ) + if request.url.path == "/v1/bot/get_online_info": + return httpx.Response( + 200, json={"code": 0, "data": {"bot_id": "bot-1", "name": "Assistant"}} + ) + return httpx.Response(200, json={"code": 0, "data": {"id": "ok"}}) + + async def run() -> None: + client = AsyncCozeClient("token", transport=httpx.MockTransport(handler)) + await client.conversation_retrieve("conv") + await client.conversation_message_create("conv", role="user", content="hi") + listed = await client.conversation_message_list("conv", limit=10) + self.assertEqual(listed["items"][0]["id"], "msg-1") + self.assertFalse(listed["has_more"]) + await client.conversation_message_retrieve("conv", "msg-1") + retrieved = await client.file_retrieve("file-1") + self.assertEqual(retrieved["status"], "processed") + bot = await client.bot_retrieve("bot-1") + self.assertEqual(bot["name"], "Assistant") + await client.aclose() + + asyncio.run(run()) + self.assertEqual( + paths, + [ + "/v1/conversation/retrieve", + "/v1/conversation/message/create", + "/v1/conversation/message/list", + "/v1/conversation/message/retrieve", + "/v1/files/retrieve", + "/v1/bot/get_online_info", + ], + ) + + def test_agent_node_non_stream_path_polls_and_extracts_usage(self) -> None: + # CozeAgentNode(stream=False) exercises the chat()/chat_retrieve() polling + # loop instead of SSE; usage should be extracted once the chat completes. + calls = {"retrieve": 0} + + async def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/v3/chat" and request.method == "POST": + return httpx.Response( + 200, + json={ + "code": 0, + "data": { + "id": "chat-1", + "conversation_id": "conv-1", + "status": "in_progress", + }, + }, + ) + if request.url.path == "/v3/chat/retrieve": + calls["retrieve"] += 1 + return httpx.Response( + 200, + json={ + "code": 0, + "data": { + "id": "chat-1", + "conversation_id": "conv-1", + "status": "completed", + "usage": { + "token_count": 12, + "input_count": 5, + "output_count": 7, + }, + }, + }, + ) + if request.url.path == "/v3/chat/message/list": + return httpx.Response( + 200, + json={ + "code": 0, + "data": [ + {"role": "assistant", "type": "answer", "content": "polled answer"}, + {"role": "assistant", "type": "follow_up", "content": "and then?"}, + ], + }, + ) + raise AssertionError(f"unexpected request: {request.url.path}") + + class AgentState(TypedDict): + messages: Annotated[list[Any], add_messages] + coze_conversations: dict[str, str] + + async def run() -> None: + client = AsyncCozeClient("token", transport=httpx.MockTransport(handler)) + node = CozeAgentNode("bot", client=client, user_id="user", stream=False) + builder = StateGraph(AgentState) + builder.add_node("coze", node).add_edge(START, "coze").add_edge("coze", END) + result = await builder.compile().ainvoke( + {"messages": [HumanMessage("go")], "coze_conversations": {}} + ) + final = result["messages"][-1] + self.assertEqual(final.content, "polled answer") + self.assertEqual( + final.usage, {"token_count": 12, "input_count": 5, "output_count": 7} + ) + self.assertEqual(final.additional_kwargs["follow_ups"], ("and then?",)) + self.assertEqual(result["coze_conversations"], {"bot": "conv-1"}) + await client.aclose() + + asyncio.run(run()) + self.assertGreaterEqual(calls["retrieve"], 1) + + if __name__ == "__main__": unittest.main()