agenthub-python 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

agenthub/abort_signal.py

@@ -0,0 +1,135 @@
+# Copyright 2025 Prism Shadow. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import threading
+from contextlib import suppress
+from typing import Any, Awaitable, TypeVar
+
+
+T = TypeVar("T")
+
+
+class AbortSignal:
+    """Abort signal that can also trigger its own aborted state."""
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._aborted = False
+        self._reason: Any = None
+        self._waiters: set[asyncio.Future[None]] = set()
+
+    @property
+    def aborted(self) -> bool:
+        with self._lock:
+            return self._aborted
+
+    @property
+    def reason(self) -> Any:
+        with self._lock:
+            return self._reason
+
+    def abort(self, reason: Any = None) -> None:
+        with self._lock:
+            if self._aborted:
+                return
+
+            self._aborted = True
+            self._reason = reason
+            waiters = tuple(self._waiters)
+            self._waiters.clear()
+
+        for waiter in waiters:
+            _notify_waiter(waiter)
+
+    async def wait(self) -> None:
+        loop = asyncio.get_running_loop()
+        waiter = loop.create_future()
+        with self._lock:
+            if self._aborted:
+                return
+
+            self._waiters.add(waiter)
+
+        try:
+            await waiter
+        finally:
+            with self._lock:
+                self._waiters.discard(waiter)
+
+    def throw_if_aborted(self) -> None:
+        with self._lock:
+            aborted = self._aborted
+            reason = self._reason
+
+        if aborted:
+            raise _cancelled_error(reason)
+
+
+async def run_with_abort(awaitable: Awaitable[T], signal: AbortSignal) -> T:
+    """Run an awaitable and cancel it when the signal is aborted."""
+
+    task = asyncio.ensure_future(awaitable)
+
+    if signal.aborted:
+        task.cancel(signal.reason)
+        with suppress(asyncio.CancelledError):
+            await task
+        raise _cancelled_error(signal.reason)
+
+    abort_task = asyncio.create_task(signal.wait())
+
+    try:
+        done, _ = await asyncio.wait((task, abort_task), return_when=asyncio.FIRST_COMPLETED)
+        if task in done:
+            return await task
+
+        task.cancel(signal.reason)
+        with suppress(asyncio.CancelledError):
+            await task
+        raise _cancelled_error(signal.reason)
+    except asyncio.CancelledError:
+        if not task.done():
+            task.cancel()
+            with suppress(asyncio.CancelledError):
+                await task
+        raise
+    finally:
+        if not abort_task.done():
+            abort_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await abort_task
+
+
+def _set_waiter_result(waiter: asyncio.Future[None]) -> None:
+    if not waiter.done():
+        waiter.set_result(None)
+
+
+def _notify_waiter(waiter: asyncio.Future[None]) -> None:
+    loop = waiter.get_loop()
+    if loop.is_closed():
+        return
+
+    if loop.is_running():
+        loop.call_soon_threadsafe(_set_waiter_result, waiter)
+    else:
+        _set_waiter_result(waiter)
+
+
+def _cancelled_error(reason: Any) -> asyncio.CancelledError:
+    if reason is None:
+        return asyncio.CancelledError()
+
+    return asyncio.CancelledError(reason)

agenthub/auto_client.py

@@ -15,6 +15,7 @@
 import os
 from typing import Any, AsyncIterator
 
+from .abort_signal import AbortSignal
 from .base_client import LLMClient
 from .types import UniConfig, UniEvent, UniMessage
 
@@ -45,35 +46,45 @@ class AutoLLMClient(LLMClient):
         self, model: str, api_key: str | None = None, base_url: str | None = None, client_type: str | None = None
     ) -> LLMClient:
         """Create the appropriate client for the given model."""
-        client_type = client_type or os.getenv("CLIENT_TYPE", model.lower())
-        if "gemini-3-" in client_type or "gemini-3.1-" in client_type:  # e.g., gemini-3-flash-preview
+        client_type = (client_type or os.getenv("CLIENT_TYPE", model)).lower()
+        if any(
+            prefix in client_type for prefix in ("gemini-3", "gemini-embedding")
+        ):  # e.g., gemini-3-flash-preview, gemini-embedding-2
             from .gemini3 import Gemini3Client
 
             return Gemini3Client(model=model, api_key=api_key, base_url=base_url)
+        elif "claude" in client_type and ("4-7" in client_type or "4-8" in client_type):  # e.g., claude-opus-4-7
+            from .claude4_8 import Claude4_8Client
+
+            return Claude4_8Client(model=model, api_key=api_key, base_url=base_url)
         elif "claude" in client_type and "4-6" in client_type:  # e.g., claude-sonnet-4-6
             from .claude4_6 import Claude4_6Client
 
             return Claude4_6Client(model=model, api_key=api_key, base_url=base_url)
-        elif "gpt-5.4" in client_type:  # e.g., gpt-5.4
-            from .gpt5_4 import GPT5_4Client
+        elif "gpt-5.4" in client_type or "gpt-5.5" in client_type:  # e.g., gpt-5.5
+            from .gpt5_5 import GPT5_5Client
+
+            return GPT5_5Client(model=model, api_key=api_key, base_url=base_url)
+        elif "glm-5" in client_type or "glm-5.1" in client_type:
+            from .glm5_1 import GLM5_1Client
 
-            return GPT5_4Client(model=model, api_key=api_key, base_url=base_url)
-        elif "glm-5" in client_type:
-            from .glm5 import GLM5Client
+            return GLM5_1Client(model=model, api_key=api_key, base_url=base_url)
+        elif "kimi-k2.5" in client_type or "kimi-k2.6" in client_type:
+            from .kimi_k2_6 import KimiK2_6Client
 
-            return GLM5Client(model=model, api_key=api_key, base_url=base_url)
-        elif "kimi-k2.5" in client_type:
-            from .kimi_k2_5 import KimiK2_5Client
+            return KimiK2_6Client(model=model, api_key=api_key, base_url=base_url)
+        elif "deepseek-v4" in client_type:
+            from .deepseek_v4 import DeepSeekV4Client
 
-            return KimiK2_5Client(model=model, api_key=api_key, base_url=base_url)
-        elif "qwen3" in client_type:
-            from .qwen3 import Qwen3Client
+            return DeepSeekV4Client(model=model, api_key=api_key, base_url=base_url)
+        elif "openai" in client_type:
+            from .openai import OpenaiClient
 
-            return Qwen3Client(model=model, api_key=api_key, base_url=base_url)
+            return OpenaiClient(model=model, api_key=api_key, base_url=base_url)
         else:
             raise ValueError(
                 f"{client_type} is not supported. "
-                "Supported client types: gemini-3, claude-4-6, gpt-5.4, glm-5, kimi-k2.5, qwen3."
+                "Supported client types: gemini-3, claude-4-8, claude-4-7, claude-4-6, gpt-5.5, gpt-5.4, glm-5.1, kimi-k2.6, kimi-k2.5, deepseek-v4, openai."
             )
 
     def transform_uni_config_to_model_config(self, config: UniConfig) -> Any:
@@ -99,11 +110,13 @@ class AutoLLMClient(LLMClient):
         self,
         messages: list[UniMessage],
         config: UniConfig,
+        signal: AbortSignal | None = None,
     ) -> AsyncIterator[UniEvent]:
         """Route to underlying client's streaming_response."""
         async for event in self._client.streaming_response(
             messages=messages,
             config=config,
+            signal=signal,
         ):
             yield event
 
@@ -111,11 +124,13 @@ class AutoLLMClient(LLMClient):
         self,
         message: UniMessage,
         config: UniConfig,
+        signal: AbortSignal | None = None,
     ) -> AsyncIterator[UniEvent]:
         """Route to underlying client's streaming_response_stateful."""
         async for event in self._client.streaming_response_stateful(
             message=message,
             config=config,
+            signal=signal,
         ):
             yield event
 
@@ -126,3 +141,7 @@ class AutoLLMClient(LLMClient):
     def get_history(self) -> list[UniMessage]:
         """Get history from the underlying client."""
         return self._client.get_history()
+
+    def set_history(self, history: list[UniMessage]) -> None:
+        """Set history in the underlying client."""
+        self._client.set_history(history)

agenthub/base_client.py

@@ -12,10 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
+import time
 from abc import ABC, abstractmethod
+from contextlib import suppress
 from typing import Any, AsyncIterator
 
-from .types import ContentItem, FinishReason, UniConfig, UniEvent, UniMessage, UsageMetadata
+from .abort_signal import AbortSignal
+from .types import (
+    ContentItem,
+    FinishReason,
+    UniConfig,
+    UniEvent,
+    UniMessage,
+    UsageMetadata,
+)
 
 
 class LLMClient(ABC):
@@ -84,6 +95,7 @@ class LLMClient(ABC):
         content_items: list[ContentItem] = []
         usage_metadata: UsageMetadata | None = None
         finish_reason: FinishReason | None = None
+        created_at: int | None = None
 
         for event in events:
             # Merge content_items from all events
@@ -119,12 +131,14 @@ class LLMClient(ABC):
 
             usage_metadata = event.get("usage_metadata")  # usage_metadata is taken from the last event
             finish_reason = event.get("finish_reason")  # finish_reason is taken from the last event
+            created_at = event.get("created_at")  # created_at is taken from the last event
 
         return {
             "role": "assistant",
             "content_items": content_items,
             "usage_metadata": usage_metadata,
             "finish_reason": finish_reason,
+            "created_at": created_at,
         }
 
     @abstractmethod
@@ -152,6 +166,7 @@ class LLMClient(ABC):
         self,
         messages: list[UniMessage],
         config: UniConfig,
+        signal: AbortSignal | None = None,
     ) -> AsyncIterator[UniEvent]:
         """
         Generate content in streaming mode (stateless).
@@ -163,21 +178,86 @@ class LLMClient(ABC):
         Args:
             messages: List of universal message dictionaries containing conversation history
             config: Universal configuration dict
+            signal: Optional abort signal used to cancel the active request
 
         Yields:
             Universal events from the streaming response
         """
+        # Stamp any messages that don't yet have a created_at timestamp
+        for msg in messages:
+            if "created_at" not in msg:
+                msg["created_at"] = int(time.time() * 1000)
+
         last_event: UniEvent | None = None
-        async for event in self._streaming_response_internal(messages, config):
-            last_event = event
-            yield event
+        events = []
+        if signal is not None:
+            signal.throw_if_aborted()
+
+        stream = self._streaming_response_internal(messages, config)
+        abort_task: asyncio.Task[None] | None = None
+        waiting_for_stream = False
+        if signal is not None:
+            streaming_task = asyncio.current_task()
+            abort_task = asyncio.create_task(signal.wait())
+
+            def cancel_streaming_task(task: asyncio.Task[None]) -> None:
+                if (
+                    task.cancelled()
+                    or not signal.aborted
+                    or not waiting_for_stream
+                    or streaming_task is None
+                    or streaming_task.done()
+                ):
+                    return
+
+                streaming_task.cancel(signal.reason)
+
+            abort_task.add_done_callback(cancel_streaming_task)
+
+        try:
+            while True:
+                try:
+                    if signal is not None:
+                        signal.throw_if_aborted()
+                        waiting_for_stream = True
+                        signal.throw_if_aborted()
+
+                    event = await anext(stream)
+                except StopAsyncIteration:
+                    break
+                except asyncio.CancelledError:
+                    if signal is not None and signal.aborted:
+                        signal.throw_if_aborted()
+                    raise
+                finally:
+                    waiting_for_stream = False
+
+                event["created_at"] = int(time.time() * 1000)
+                last_event = event
+                events.append(event)
+                yield event
+        finally:
+            if abort_task is not None and not abort_task.done():
+                abort_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await abort_task
+            await stream.aclose()
 
         self._validate_last_event(last_event)
 
+        # Save history to file if trace_id is specified
+        if config.get("trace_id") and events:
+            from .integration.tracer import Tracer
+
+            assistant_message = self.concat_uni_events_to_uni_message(events)
+            tracer = Tracer()
+            tracer.save_history(self._model, messages + [assistant_message], config["trace_id"], config)
+
     async def streaming_response_stateful(
         self,
         message: UniMessage,
         config: UniConfig,
+        signal: AbortSignal | None = None,
     ) -> AsyncIterator[UniEvent]:
         """
         Generate content in streaming mode (stateful).
@@ -189,6 +269,7 @@ class LLMClient(ABC):
         Args:
             message: Latest universal message dictionary to add to conversation
             config: Universal configuration dict
+            signal: Optional abort signal used to cancel the active request
 
         Yields:
             Universal events from the streaming response
@@ -198,23 +279,17 @@ class LLMClient(ABC):
 
         # Collect all events for history
         events = []
-        async for event in self.streaming_response(messages=temp_messages, config=config):
+        async for event in self.streaming_response(messages=temp_messages, config=config, signal=signal):
             events.append(event)
             yield event
 
         # Only update history after successful inference
+        # temp_messages[-1] is the user message, now stamped with created_at by streaming_response
         if events:
             assistant_message = self.concat_uni_events_to_uni_message(events)
-            self._history.append(message)
+            self._history.append(temp_messages[-1])
             self._history.append(assistant_message)
 
-        # Save history to file if trace_id is specified
-        if config.get("trace_id"):
-            from .integration.tracer import Tracer
-
-            tracer = Tracer()
-            tracer.save_history(self._model, self._history, config["trace_id"], config)
-
     @staticmethod
     def _validate_last_event(last_event: UniEvent | None) -> None:
         """Validate that the last event has usage_metadata and finish_reason.
@@ -244,3 +319,11 @@ class LLMClient(ABC):
     def get_history(self) -> list[UniMessage]:
         """Get the current message history."""
         return self._history.copy()
+
+    def set_history(self, history: list[UniMessage]) -> None:
+        """Replace the message history with a copy of the provided history.
+
+        Args:
+            history: List of universal message dictionaries to set as the new history
+        """
+        self._history = list(history)

agenthub/claude4_6/client.py

@@ -110,6 +110,7 @@ class Claude4_6Client(LLMClient):
             ThinkingLevel.LOW: {"thinking": {"type": "adaptive"}, "output_config": {"effort": "low"}},
             ThinkingLevel.MEDIUM: {"thinking": {"type": "adaptive"}, "output_config": {"effort": "medium"}},
             ThinkingLevel.HIGH: {"thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}},
+            ThinkingLevel.XHIGH: {"thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}},
         }
         return mapping.get(thinking_level)
 
@@ -145,7 +146,7 @@ class Claude4_6Client(LLMClient):
         if config.get("max_tokens") is not None:
             claude_config["max_tokens"] = config["max_tokens"]
         else:
-            claude_config["max_tokens"] = 32768  # Claude requires max_tokens to be specified
+            claude_config["max_tokens"] = 64000  # Claude requires max_tokens to be specified
 
         if config.get("temperature") is not None:
             claude_config["temperature"] = config["temperature"]
@@ -171,6 +172,15 @@ class Claude4_6Client(LLMClient):
         if config.get("tool_choice") is not None:
             claude_config["tool_choice"] = self._convert_tool_choice(config["tool_choice"])
 
+        # Add cache_control if prompt caching is enabled
+        # TODO: wait for bedrock to support cache_control in config
+        if not self._use_bedrock:
+            prompt_caching = config.get("prompt_caching", PromptCaching.ENABLE)
+            if prompt_caching == PromptCaching.ENABLE:
+                claude_config["cache_control"] = {"type": "ephemeral"}
+            elif prompt_caching == PromptCaching.ENHANCE:
+                claude_config["cache_control"] = {"type": "ephemeral", "ttl": "1h"}
+
         return claude_config
 
     async def transform_uni_message_to_model_input(self, messages: list[UniMessage]) -> list[BetaMessageParam]:
@@ -334,18 +344,20 @@ class Claude4_6Client(LLMClient):
         # Use unified message conversion
         claude_messages = await self.transform_uni_message_to_model_input(messages)
 
-        # Add cache_control to last user message's last item if enabled
-        prompt_caching = config.get("prompt_caching", PromptCaching.ENABLE)
-        if prompt_caching != PromptCaching.DISABLE and claude_messages:
-            try:
-                last_user_message = next(filter(lambda x: x["role"] == "user", claude_messages[::-1]))
-                last_content_item = last_user_message["content"][-1]
-                last_content_item["cache_control"] = {
-                    "type": "ephemeral",
-                    "ttl": "1h" if prompt_caching == PromptCaching.ENHANCE else "5m",
-                }
-            except StopIteration:
-                pass
+        # Add cache_control to last user message's last item if using bedrock and enabled prompt caching
+        # TODO: remove after bedrock supports cache_control in config
+        if self._use_bedrock:
+            prompt_caching = config.get("prompt_caching", PromptCaching.ENABLE)
+            if prompt_caching != PromptCaching.DISABLE and claude_messages:
+                try:
+                    last_user_message = next(filter(lambda x: x["role"] == "user", claude_messages[::-1]))
+                    last_content_item = last_user_message["content"][-1]
+                    last_content_item["cache_control"] = {
+                        "type": "ephemeral",
+                        "ttl": "1h" if prompt_caching == PromptCaching.ENHANCE else "5m",
+                    }
+                except StopIteration:
+                    pass
 
         # Stream generate
         partial_tool_call = {}
@@ -362,7 +374,9 @@ class Claude4_6Client(LLMClient):
                             "arguments": "",
                             "tool_call_id": item["tool_call_id"],
                         }
-                        yield event
+
+                if event["content_items"]:
+                    yield event
 
                 if event["usage_metadata"] is not None:
                     # initialize partial_usage

agenthub/claude4_8/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2025 Prism Shadow. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .client import Claude4_8Client
+
+
+__all__ = ["Claude4_8Client"]