botrun-flow-lang 6.2.61__tar.gz → 6.2.62__tar.gz

{botrun_flow_lang-6.2.61 → botrun_flow_lang-6.2.62}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: botrun-flow-lang
-Version: 6.2.61
+Version: 6.2.62
 Summary: A flow language for botrun
 Author-email: sebastian-hsu <sebastian.hsu@gmail.com>
 License: MIT

{botrun_flow_lang-6.2.61 → botrun_flow_lang-6.2.62}/botrun_flow_lang/langgraph_agents/agents/util/custom_vertex_claude.py

@@ -8,7 +8,7 @@ Supports tool calling for LangGraph react agent compatibility.
 """
 
 import json
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 import httpx
 
@@ -18,12 +18,13 @@ from langchain_core.callbacks import CallbackManagerForLLMRun
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import (
     AIMessage,
+    AIMessageChunk,
     BaseMessage,
     HumanMessage,
     SystemMessage,
     ToolMessage,
 )
-from langchain_core.outputs import ChatGeneration, ChatResult
+from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
 from pydantic import ConfigDict
 
 from botrun_flow_lang.utils.botrun_logger import get_default_botrun_logger
@@ -193,23 +194,20 @@ class ChatVertexAIClaude(BaseChatModel):
 
         return system, merged
 
-    def _generate(
+    def _build_payload(
         self,
         messages: List[BaseMessage],
         stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs,
-    ) -> ChatResult:
-        """Call Vertex AI Claude via rawPredict (non-streaming)."""
+    ) -> Tuple[str, Dict[str, Any], int]:
+        """Build API payload shared by _generate and _stream.
+
+        Returns:
+            (access_token, payload, tools_count)
+        """
         system, api_messages = self._convert_messages(messages)
         access_token = self._get_access_token()
 
-        url = (
-            f"https://{self.location}-aiplatform.googleapis.com/v1/"
-            f"projects/{self.project_id}/locations/{self.location}/"
-            f"publishers/anthropic/models/{self.model}:rawPredict"
-        )
-
         payload: Dict[str, Any] = {
             "anthropic_version": "vertex-2023-10-16",
             "messages": api_messages,
@@ -221,12 +219,10 @@ class ChatVertexAIClaude(BaseChatModel):
         if stop:
             payload["stop_sequences"] = stop
 
-        # Tools from bind_tools()
         tools = kwargs.get("tools", [])
         if tools:
             payload["tools"] = tools
 
-        # Tool choice
         tool_choice = kwargs.get("tool_choice")
         if tool_choice:
             if isinstance(tool_choice, str):
@@ -244,13 +240,32 @@ class ChatVertexAIClaude(BaseChatModel):
             elif isinstance(tool_choice, dict):
                 payload["tool_choice"] = tool_choice
 
+        return access_token, payload, len(tools)
+
+    def _generate(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs,
+    ) -> ChatResult:
+        """Call Vertex AI Claude via rawPredict (non-streaming)."""
+        access_token, payload, tools_count = self._build_payload(
+            messages, stop, **kwargs
+        )
+
+        url = (
+            f"https://{self.location}-aiplatform.googleapis.com/v1/"
+            f"projects/{self.project_id}/locations/{self.location}/"
+            f"publishers/anthropic/models/{self.model}:rawPredict"
+        )
+
         logger.info(
             f"[ChatVertexAIClaude] rawPredict: model={self.model}, "
-            f"location={self.location}, messages={len(api_messages)}, "
-            f"tools={len(tools)}"
+            f"location={self.location}, messages={len(payload['messages'])}, "
+            f"tools={tools_count}"
         )
 
-        # Make API call via httpx
         data = _http_post_json(url, payload, access_token)
 
         # Parse response
@@ -305,6 +320,192 @@ class ChatVertexAIClaude(BaseChatModel):
             },
         )
 
+    def _stream(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs,
+    ) -> Iterator[ChatGenerationChunk]:
+        """Call Vertex AI Claude via streamRawPredict (streaming).
+
+        Yields ChatGenerationChunk with AIMessageChunk for each SSE event.
+        Handles both text and tool_use content blocks.
+        """
+        access_token, payload, tools_count = self._build_payload(
+            messages, stop, **kwargs
+        )
+        payload["stream"] = True
+
+        url = (
+            f"https://{self.location}-aiplatform.googleapis.com/v1/"
+            f"projects/{self.project_id}/locations/{self.location}/"
+            f"publishers/anthropic/models/{self.model}:streamRawPredict"
+        )
+
+        logger.info(
+            f"[ChatVertexAIClaude] streamRawPredict: model={self.model}, "
+            f"location={self.location}, messages={len(payload['messages'])}, "
+            f"tools={tools_count}"
+        )
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {access_token}",
+        }
+
+        # Track state across SSE events
+        usage_metadata: Dict[str, int] = {}
+        # Map block index -> tool info for tool_use streaming
+        tool_use_blocks: Dict[int, Dict[str, Any]] = {}
+
+        try:
+            with httpx.Client(timeout=300.0) as client:
+                with client.stream(
+                    "POST", url, headers=headers, json=payload
+                ) as response:
+                    if response.status_code != 200:
+                        error_body = response.read().decode("utf-8", errors="ignore")
+                        error_msg = f"Vertex AI API error: {response.status_code} - {error_body}"
+                        logger.error(f"[ChatVertexAIClaude] {error_msg}")
+                        raise Exception(error_msg)
+
+                    for line in response.iter_lines():
+                        if isinstance(line, bytes):
+                            line = line.decode("utf-8")
+                        line = line.strip()
+
+                        if not line or line.startswith("event:"):
+                            continue
+
+                        if line.startswith("data:"):
+                            data_str = line[5:].strip()
+                        else:
+                            data_str = line
+
+                        try:
+                            data = json.loads(data_str)
+                        except json.JSONDecodeError:
+                            continue
+
+                        if not data:
+                            continue
+
+                        event_type = data.get("type", "")
+
+                        if event_type == "message_start":
+                            msg = data.get("message", {})
+                            if "usage" in msg:
+                                usage_metadata["input_tokens"] = msg["usage"].get(
+                                    "input_tokens", 0
+                                )
+
+                        elif event_type == "content_block_start":
+                            block = data.get("content_block", {})
+                            index = data.get("index", 0)
+                            if block.get("type") == "tool_use":
+                                # Start of a tool_use block
+                                tool_use_blocks[index] = {
+                                    "id": block.get("id", ""),
+                                    "name": block.get("name", ""),
+                                    "args_json": "",
+                                }
+                                chunk = AIMessageChunk(
+                                    content="",
+                                    tool_call_chunks=[
+                                        {
+                                            "name": block.get("name", ""),
+                                            "args": "",
+                                            "id": block.get("id", ""),
+                                            "index": index,
+                                        }
+                                    ],
+                                )
+                                yield ChatGenerationChunk(message=chunk)
+                                if run_manager:
+                                    run_manager.on_llm_new_token(
+                                        "", chunk=chunk
+                                    )
+
+                        elif event_type == "content_block_delta":
+                            delta = data.get("delta", {})
+                            index = data.get("index", 0)
+                            delta_type = delta.get("type", "")
+
+                            if delta_type == "text_delta":
+                                text = delta.get("text", "")
+                                if text:
+                                    chunk = AIMessageChunk(content=text)
+                                    yield ChatGenerationChunk(message=chunk)
+                                    if run_manager:
+                                        run_manager.on_llm_new_token(
+                                            text, chunk=chunk
+                                        )
+
+                            elif delta_type == "input_json_delta":
+                                partial_json = delta.get("partial_json", "")
+                                if index in tool_use_blocks:
+                                    tool_use_blocks[index][
+                                        "args_json"
+                                    ] += partial_json
+                                chunk = AIMessageChunk(
+                                    content="",
+                                    tool_call_chunks=[
+                                        {
+                                            "name": None,
+                                            "args": partial_json,
+                                            "id": None,
+                                            "index": index,
+                                        }
+                                    ],
+                                )
+                                yield ChatGenerationChunk(message=chunk)
+                                if run_manager:
+                                    run_manager.on_llm_new_token(
+                                        "", chunk=chunk
+                                    )
+
+                        elif event_type == "message_delta":
+                            delta = data.get("delta", {})
+                            if "usage" in data:
+                                usage_metadata["output_tokens"] = data[
+                                    "usage"
+                                ].get("output_tokens", 0)
+
+                        elif event_type == "message_stop":
+                            # Yield final chunk with usage metadata
+                            input_tokens = usage_metadata.get("input_tokens", 0)
+                            output_tokens = usage_metadata.get(
+                                "output_tokens", 0
+                            )
+                            chunk = AIMessageChunk(
+                                content="",
+                                usage_metadata={
+                                    "input_tokens": input_tokens,
+                                    "output_tokens": output_tokens,
+                                    "total_tokens": input_tokens
+                                    + output_tokens,
+                                },
+                                response_metadata={
+                                    "model": self.model,
+                                },
+                            )
+                            yield ChatGenerationChunk(message=chunk)
+
+                            logger.info(
+                                f"[ChatVertexAIClaude] Stream complete: "
+                                f"tokens=({input_tokens}+{output_tokens}"
+                                f"={input_tokens + output_tokens})"
+                            )
+
+        except httpx.HTTPStatusError as e:
+            error_body = e.response.text if e.response else ""
+            error_msg = (
+                f"Vertex AI API error: {e.response.status_code} - {error_body}"
+            )
+            logger.error(f"[ChatVertexAIClaude] {error_msg}")
+            raise Exception(error_msg) from e
+
 
 def _http_post_json(
     url: str, payload: Dict[str, Any], access_token: str

{botrun_flow_lang-6.2.61 → botrun_flow_lang-6.2.62}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "botrun-flow-lang"
-version = "6.2.61"
+version = "6.2.62"
 description = "A flow language for botrun"
 authors = [
     { name = "sebastian-hsu", email = "sebastian.hsu@gmail.com" }