agentengine-sdk-python 0.6.9__py3-none-any.whl → 0.7.0__py3-none-any.whl

ksadk/runners/langchain_runner.py

@@ -110,9 +110,66 @@ class LangChainRunner(BaseRunner):
 
         accumulated_text = ""
         last_chunk: Any = None
+        final_output_text = ""
+        emitted_non_text_event = False
+        message_snapshots: dict[str, str] = {}
 
         try:
-            if hasattr(self._agent, "astream"):
+            if self._should_stream_events(payload):
+                kwargs = self._build_optional_call_kwargs(
+                    self._agent.astream_events,
+                    config=config,
+                    context=native_context,
+                )
+                kwargs["version"] = "v2"
+                async for event in self._agent.astream_events(payload, **kwargs):
+                    if not isinstance(event, dict):
+                        continue
+                    event_kind = event.get("event", "")
+                    data = event.get("data") or {}
+
+                    if event_kind == "on_chat_model_stream":
+                        chunk = data.get("chunk") if isinstance(data, dict) else None
+                        if chunk is None:
+                            continue
+                        last_chunk = chunk
+                        delta, chunk_type = self._extract_chunk(chunk)
+                        if delta:
+                            if chunk_type == "text":
+                                accumulated_text += delta
+                            else:
+                                emitted_non_text_event = True
+                            yield {"delta": delta, "type": chunk_type}
+                    elif event_kind == "on_tool_start":
+                        emitted_non_text_event = True
+                        yield {
+                            "type": "tool_call",
+                            "tool_name": event.get("name", "unknown"),
+                            "tool_args": data.get("input", {}) if isinstance(data, dict) else {},
+                            "run_id": event.get("run_id"),
+                        }
+                    elif event_kind == "on_tool_end":
+                        emitted_non_text_event = True
+                        tool_output = data.get("output", "") if isinstance(data, dict) else ""
+                        yield {
+                            "type": "tool_result",
+                            "tool_name": event.get("name", "unknown"),
+                            "tool_args": data.get("input", {}) if isinstance(data, dict) else {},
+                            "tool_output": (
+                                tool_output
+                                if isinstance(tool_output, dict)
+                                else (str(tool_output) if tool_output else "")
+                            ),
+                            "run_id": event.get("run_id"),
+                        }
+                    elif event_kind == "on_chain_end":
+                        output = data.get("output") if isinstance(data, dict) else None
+                        extracted_output = self._extract_recognized_output(output)
+                        if extracted_output:
+                            final_output_text = extracted_output
+                        if self._extract_usage(output):
+                            last_chunk = output
+            elif hasattr(self._agent, "astream"):
                 kwargs = self._build_optional_call_kwargs(
                     self._agent.astream,
                     config=config,
@@ -120,7 +177,15 @@ class LangChainRunner(BaseRunner):
                 )
                 async for chunk in self._agent.astream(payload, **kwargs):
                     last_chunk = chunk
-                    delta, chunk_type = self._extract_chunk(chunk)
+                    message_state = self._extract_message_state(chunk)
+                    if message_state:
+                        content, message_key = message_state
+                        previous = message_snapshots.get(message_key, "")
+                        delta = self._snapshot_delta(content, previous)
+                        message_snapshots[message_key] = content
+                        chunk_type = "text"
+                    else:
+                        delta, chunk_type = self._extract_chunk(chunk)
                     if delta:
                         accumulated_text += delta
                         yield {"delta": delta, "type": chunk_type}
@@ -132,14 +197,32 @@ class LangChainRunner(BaseRunner):
                 )
                 for chunk in self._agent.stream(payload, **kwargs):
                     last_chunk = chunk
-                    delta, chunk_type = self._extract_chunk(chunk)
+                    message_state = self._extract_message_state(chunk)
+                    if message_state:
+                        content, message_key = message_state
+                        previous = message_snapshots.get(message_key, "")
+                        delta = self._snapshot_delta(content, previous)
+                        message_snapshots[message_key] = content
+                        chunk_type = "text"
+                    else:
+                        delta, chunk_type = self._extract_chunk(chunk)
                     if delta:
                         accumulated_text += delta
                         yield {"delta": delta, "type": chunk_type}
         except Exception as exc:
-            print(f"\n⚠️ 流式调用失败: {exc}，回退到同步模式")
+            logger.warning("LangChain stream failed: %s", exc)
 
         if not accumulated_text:
+            if final_output_text or emitted_non_text_event:
+                final_chunk = {"output": final_output_text, "type": "final"}
+                usage = self._extract_usage(last_chunk)
+                if usage:
+                    final_chunk["usage"] = usage
+                last_usage = self._extract_last_usage(last_chunk)
+                if last_usage:
+                    final_chunk.setdefault("metadata", {})["last_usage"] = last_usage
+                yield final_chunk
+                return
             result = await self.invoke(input_data)
             final_chunk = {"output": result.get("output", ""), "type": "final"}
             usage = self._extract_usage(result)
@@ -160,6 +243,12 @@ class LangChainRunner(BaseRunner):
             final_chunk.setdefault("metadata", {})["last_usage"] = last_usage
         yield final_chunk
 
+    def _should_stream_events(self, payload: dict[str, Any]) -> bool:
+        """Use LangGraph events for LangChain create_agent message-state agents."""
+        return isinstance(payload.get("messages"), list) and hasattr(
+            self._agent, "astream_events"
+        )
+
     def _resolve_request_path(self) -> str:
         module = getattr(self, "_module", None)
         if callable(getattr(module, "ksadk_prepare_input", None)):
@@ -168,7 +257,9 @@ class LangChainRunner(BaseRunner):
             return "runnable_with_message_history"
         return "replay"
 
-    def _prepare_with_standard_hook(self, input_data: Dict[str, Any], session_id: str) -> dict[str, Any]:
+    def _prepare_with_standard_hook(
+        self, input_data: Dict[str, Any], session_id: str
+    ) -> dict[str, Any]:
         module = getattr(self, "_module", None)
         prepare_input = getattr(module, "ksadk_prepare_input", None)
         if not callable(prepare_input):
@@ -195,7 +286,11 @@ class LangChainRunner(BaseRunner):
     @staticmethod
     def _ksadk_builtin_tool_context() -> dict[str, Any]:
         try:
-            from ksadk.toolsets import builtin_tool_descriptors_for_runtime, builtin_tools_mode, builtin_tools_profile
+            from ksadk.toolsets import (
+                builtin_tool_descriptors_for_runtime,
+                builtin_tools_mode,
+                builtin_tools_profile,
+            )
 
             mode = builtin_tools_mode(default="off")
             descriptors = builtin_tool_descriptors_for_runtime(mode=mode)
@@ -231,7 +326,11 @@ class LangChainRunner(BaseRunner):
     def _ambient_context_text(input_data: Dict[str, Any]) -> str:
         sections: list[str] = []
         kb_context = input_data.get("kb_context") or {}
-        kb_text = str(kb_context.get("formatted_text") or "").strip() if isinstance(kb_context, dict) else ""
+        kb_text = (
+            str(kb_context.get("formatted_text") or "").strip()
+            if isinstance(kb_context, dict)
+            else ""
+        )
         if kb_text:
             sections.append(f"Knowledge base context:\n{kb_text}")
 
@@ -518,6 +617,119 @@ class LangChainRunner(BaseRunner):
             return str(content)
         return str(result)
 
+    @classmethod
+    def _extract_recognized_output(cls, result: Any) -> str:
+        if result is None:
+            return ""
+        if isinstance(result, str):
+            return result
+
+        if isinstance(result, dict):
+            for key in ("output", "text"):
+                if key not in result:
+                    continue
+                value = result[key]
+                if isinstance(value, str):
+                    return value
+                extracted = cls._extract_recognized_output(value)
+                if extracted:
+                    return extracted
+                if value is not None and not isinstance(value, (dict, list, tuple)):
+                    return str(value)
+
+            message_state = cls._extract_message_state(result)
+            return message_state[0] if message_state else ""
+
+        if isinstance(result, (list, tuple)):
+            for item in reversed(result):
+                extracted = cls._extract_recognized_output(item)
+                if extracted:
+                    return extracted
+            return ""
+
+        command_update = getattr(result, "update", None)
+        if isinstance(command_update, dict):
+            return cls._extract_recognized_output(command_update)
+
+        content = cls._ai_message_content(result)
+        return content if content is not None else ""
+
+    @classmethod
+    def _extract_message_state(cls, chunk: Any) -> tuple[str, str] | None:
+        """Extract the newest AI message from LangGraph values/updates snapshots."""
+
+        def visit(value: Any, path: tuple[str, ...]) -> tuple[str, str] | None:
+            if not isinstance(value, dict):
+                return None
+
+            messages = value.get("messages")
+            if isinstance(messages, list):
+                for index in range(len(messages) - 1, -1, -1):
+                    message = messages[index]
+                    content = cls._ai_message_content(message)
+                    if content is None:
+                        continue
+                    message_id = (
+                        message.get("id")
+                        if isinstance(message, dict)
+                        else getattr(message, "id", None)
+                    )
+                    fallback_key = "/".join((*path, "messages", str(index)))
+                    return content, str(message_id or fallback_key)
+
+            for key, nested in value.items():
+                if key == "messages" or not isinstance(nested, dict):
+                    continue
+                result = visit(nested, (*path, str(key)))
+                if result:
+                    return result
+            return None
+
+        return visit(chunk, ())
+
+    @staticmethod
+    def _ai_message_content(message: Any) -> str | None:
+        if isinstance(message, dict):
+            role = str(message.get("role") or message.get("type") or "").lower()
+            if role not in {"ai", "assistant", "model", "aimessage", "aimessagechunk"}:
+                return None
+            content = message.get("content")
+        else:
+            role = str(getattr(message, "type", "") or "").lower()
+            class_name = type(message).__name__.lower()
+            if role not in {
+                "ai",
+                "assistant",
+                "model",
+                "aimessage",
+                "aimessagechunk",
+            } and not class_name.startswith("aimessage"):
+                return None
+            content = getattr(message, "content", None)
+
+        if isinstance(content, str):
+            return content
+        if not isinstance(content, list):
+            return None
+
+        parts: list[str] = []
+        for part in content:
+            if isinstance(part, str):
+                parts.append(part)
+            elif isinstance(part, dict):
+                text = part.get("text")
+                if isinstance(text, str):
+                    parts.append(text)
+        return "".join(parts)
+
+    @staticmethod
+    def _snapshot_delta(content: str, previous: str) -> str:
+        if content.startswith(previous):
+            return content[len(previous) :]
+        if previous.startswith(content):
+            return ""
+        return content
+
     def _extract_chunk(self, chunk: Any) -> tuple[Optional[str], Optional[str]]:
         if isinstance(chunk, dict):
             if "output" in chunk:

ksadk/runners/langgraph_runner.py

@@ -12,6 +12,7 @@ import base64
 from pathlib import Path
 
 from ksadk.runners.base_runner import BaseRunner
+from ksadk.runners.usage_accumulator import accumulate_usage
 from ksadk.sessions.continuity import LangGraphSessionAdapter
 from ksadk.runners.utils import get_langfuse_callbacks, get_langfuse_metadata, load_agent_module
 from langgraph.types import Command
@@ -259,6 +260,15 @@ class LangGraphRunner(BaseRunner):
             return self._extract_usage(values)
         return {}
 
+    @staticmethod
+    def _latest_checkpoint_config(config: dict[str, Any]) -> dict[str, Any]:
+        """Select the newest state in a resumed thread, not the source checkpoint."""
+        latest_config = dict(config)
+        configurable = dict(latest_config.get("configurable") or {})
+        configurable.pop("checkpoint_id", None)
+        latest_config["configurable"] = configurable
+        return latest_config
+
     @staticmethod
     def _ambient_context_text(payload: Dict[str, Any]) -> str:
         sections: list[str] = []
@@ -479,6 +489,51 @@ class LangGraphRunner(BaseRunner):
         )
         return self._agent.invoke(payload, **kwargs)
 
+    async def _invoke_from_stream_events(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+        chunks: list[dict[str, Any]] = []
+        accumulated_text = ""
+        final_chunk: dict[str, Any] | None = None
+        metadata: dict[str, Any] = {}
+
+        async for chunk in self.stream(input_data):
+            chunks.append(dict(chunk))
+            chunk_type = chunk.get("type")
+            if chunk_type == "text":
+                accumulated_text += str(chunk.get("delta") or "")
+            elif chunk_type == "final":
+                final_chunk = dict(chunk)
+            elif chunk_type == "checkpoint":
+                chunk_metadata = chunk.get("metadata")
+                if isinstance(chunk_metadata, Mapping):
+                    metadata.update(dict(chunk_metadata))
+            elif chunk_type == "interrupt":
+                return {
+                    "type": "interrupt",
+                    "interrupt_info": chunk.get("interrupt_info"),
+                    "session_id": chunk.get("session_id") or input_data.get("session_id"),
+                    "output": (
+                        chunk.get("interrupt_info", {}).get("message", "需要用户确认")
+                        if isinstance(chunk.get("interrupt_info"), Mapping)
+                        else "需要用户确认"
+                    ),
+                    "raw": {"chunks": chunks},
+                }
+
+        if final_chunk:
+            output_text = str(final_chunk.get("output") or accumulated_text)
+            final_metadata = final_chunk.get("metadata")
+            if isinstance(final_metadata, Mapping):
+                metadata = {**dict(final_metadata), **metadata}
+            result: dict[str, Any] = {"output": output_text, "raw": {"chunks": chunks}}
+            usage = final_chunk.get("usage")
+            if isinstance(usage, Mapping) and usage:
+                result["usage"] = dict(usage)
+            if metadata:
+                result["metadata"] = metadata
+            return result
+
+        return {"output": accumulated_text, "raw": {"chunks": chunks}}
+
     async def invoke(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
         """调用 LangGraph 图
         
@@ -487,6 +542,10 @@ class LangGraphRunner(BaseRunner):
         2. 原生格式: {"messages": [...]} 或自定义 State - 直接透传
         """
         payload = dict(input_data)
+        force_graph_invoke = bool(payload.pop("_ksadk_force_graph_invoke", False))
+        if not force_graph_invoke and hasattr(self._agent, "astream_events"):
+            return await self._invoke_from_stream_events(payload)
+
         session_id = payload.pop("session_id", None) or str(uuid.uuid4())[:8]
         is_resume = payload.pop("resume", False)
         is_checkpoint_resume = bool(payload.pop("checkpoint_resume", False))
@@ -618,14 +677,15 @@ class LangGraphRunner(BaseRunner):
             else:
                 yield {"type": "graph_update", "output": update}
 
-        state_usage = await self._latest_state_usage(config)
+        latest_config = self._latest_checkpoint_config(config)
+        state_usage = await self._latest_state_usage(latest_config)
         state_output = ""
         try:
             state = None
             if callable(getattr(self._agent, "aget_state", None)):
-                state = await self._agent.aget_state(config)
+                state = await self._agent.aget_state(latest_config)
             elif callable(getattr(self._agent, "get_state", None)):
-                state = self._agent.get_state(config)
+                state = self._agent.get_state(latest_config)
             values = getattr(state, "values", None)
             if values is not None:
                 state_output = self._extract_output(values)
@@ -639,13 +699,14 @@ class LangGraphRunner(BaseRunner):
             **({"resume_noop": True} if not emitted_update else {}),
         }
 
-        metadata = await self._latest_checkpoint_metadata(config)
+        metadata = await self._latest_checkpoint_metadata(latest_config)
         if metadata:
             yield {"type": "checkpoint", "metadata": metadata}
 
     async def stream(self, input_data: Dict[str, Any]) -> AsyncIterator[Dict[str, Any]]:
         """流式调用 LangGraph 图"""
         payload = dict(input_data)
+        payload.pop("_ksadk_force_graph_invoke", None)
         session_id = payload.pop("session_id", None) or str(uuid.uuid4())[:8]
         history = payload.pop("history", [])
         is_resume = payload.pop("resume", False)
@@ -689,6 +750,51 @@ class LangGraphRunner(BaseRunner):
         emitted_non_text_event = False
         final_output_text = ""
         final_output_usage: dict[str, Any] = {}
+        final_output_last_usage: dict[str, Any] = {}
+        model_run_usages: dict[str, dict[str, Any]] = {}
+        model_run_order: list[str] = []
+        stream_usage_run_keys: set[str] = set()
+        latest_stream_usage: dict[str, Any] = {}
+
+        def model_run_key(
+            event: Mapping[str, Any],
+            *,
+            fallback_key: str | None = None,
+        ) -> str:
+            raw_run_id = event.get("run_id")
+            return (
+                str(raw_run_id)
+                if raw_run_id
+                else fallback_key or f"model-event-{len(model_run_order)}"
+            )
+
+        def record_model_usage(
+            event: Mapping[str, Any],
+            usage: dict[str, Any],
+            *,
+            fallback_key: str | None = None,
+        ) -> None:
+            if not usage:
+                return
+            run_key = model_run_key(event, fallback_key=fallback_key)
+            if run_key not in model_run_usages:
+                model_run_order.append(run_key)
+            model_run_usages[run_key] = dict(usage)
+
+        def accumulated_model_usage() -> dict[str, Any]:
+            if len(model_run_order) == 1:
+                return dict(model_run_usages.get(model_run_order[0]) or {})
+            usage: dict[str, Any] = {}
+            for run_key in model_run_order:
+                usage = accumulate_usage(usage, model_run_usages.get(run_key) or {})
+            return usage
+
+        def latest_model_usage() -> dict[str, Any]:
+            for run_key in reversed(model_run_order):
+                usage = model_run_usages.get(run_key)
+                if usage:
+                    return dict(usage)
+            return {}
 
         if is_checkpoint_resume and callable(getattr(self._agent, "astream", None)):
             try:
@@ -731,6 +837,19 @@ class LangGraphRunner(BaseRunner):
                     chunk = event.get("data", {}).get("chunk")
                     if not chunk:
                         continue
+                    chunk_usage = self._extract_usage(chunk)
+                    if chunk_usage:
+                        # Some LangChain providers attach cumulative usage to
+                        # every stream chunk, and LangChain may then sum those
+                        # cumulative snapshots into an inflated
+                        # on_chat_model_end usage. For a concrete model run,
+                        # keep the latest stream snapshot and ignore the later
+                        # end usage for that same run_id.
+                        latest_stream_usage = dict(chunk_usage)
+                        if event.get("run_id"):
+                            run_key = model_run_key(event)
+                            stream_usage_run_keys.add(run_key)
+                            record_model_usage(event, latest_stream_usage)
 
                     # 推理内容
                     reasoning = getattr(chunk, "reasoning_content", None)
@@ -760,6 +879,15 @@ class LangGraphRunner(BaseRunner):
                                     accumulated_text += part.text
                                     yield {"delta": part.text, "type": "text"}
 
+                elif event_kind == "on_chat_model_end":
+                    data = event.get("data") or {}
+                    output = data.get("output") if isinstance(data, Mapping) else None
+                    usage = self._extract_usage(output) or self._extract_usage(data)
+                    last_usage = self._extract_last_usage(output) or self._extract_last_usage(data)
+                    run_key = model_run_key(event)
+                    if run_key not in stream_usage_run_keys:
+                        record_model_usage(event, last_usage or usage)
+
                 elif event_kind == "on_tool_start":
                     emitted_non_text_event = True
                     yield {
@@ -811,13 +939,17 @@ class LangGraphRunner(BaseRunner):
         if not accumulated_text:
             if final_output_text:
                 final_chunk = {"output": final_output_text, "type": "final"}
-                if final_output_usage:
-                    final_chunk["usage"] = final_output_usage
-                if final_output_last_usage:
-                    final_chunk.setdefault("metadata", {})["last_usage"] = final_output_last_usage
+                usage = accumulated_model_usage() or final_output_usage or latest_stream_usage
+                last_usage = latest_model_usage() or final_output_last_usage or latest_stream_usage or usage
+                if usage:
+                    final_chunk["usage"] = usage
+                if last_usage:
+                    final_chunk.setdefault("metadata", {})["last_usage"] = last_usage
                 yield final_chunk
             elif not emitted_non_text_event:
-                result = await self.invoke(invoke_payload)
+                result = await self.invoke(
+                    {**invoke_payload, "_ksadk_force_graph_invoke": True}
+                )
                 final_chunk = {"output": result.get("output", ""), "type": "final"}
                 usage = self._extract_usage(result)
                 if usage:
@@ -832,11 +964,12 @@ class LangGraphRunner(BaseRunner):
                     return
         else:
             final_chunk = {"output": accumulated_text, "type": "final"}
-            usage = await self._latest_state_usage(config)
+            state_usage = await self._latest_state_usage(config)
+            usage = accumulated_model_usage() or state_usage or final_output_usage or latest_stream_usage
             if usage:
                 final_chunk["usage"] = usage
-                # _latest_state_usage 返回末个 message usage,单次调用场景即 last_usage
-                final_chunk.setdefault("metadata", {})["last_usage"] = usage
+                last_usage = latest_model_usage() or state_usage or final_output_last_usage or latest_stream_usage or usage
+                final_chunk.setdefault("metadata", {})["last_usage"] = last_usage
             yield final_chunk
 
         metadata = await self._latest_checkpoint_metadata(config)