agentex-sdk 0.1.0a6__py3-none-any.whl

agentex/lib/utils/completions.py

@@ -0,0 +1,131 @@
+from copy import deepcopy
+from functools import reduce, singledispatch
+from typing import Any
+
+from agentex.lib.types.llm_messages import (
+    Choice,
+    Completion,
+    Delta,
+    ToolCall,
+    ToolCallRequest,
+)
+
+
+@singledispatch
+def _concat_chunks(a: None, b: Any):
+    return b
+
+
+@_concat_chunks.register
+def _(a: Completion, b: Completion) -> Completion:
+    a.choices = [_concat_chunks(*c) for c in zip(a.choices, b.choices, strict=False)]
+    return a
+
+
+@_concat_chunks.register
+def _(a: Choice, b: Choice) -> Choice:
+    if hasattr(a, "index") and hasattr(b, "index"):
+        assert a.index == b.index
+
+    if hasattr(a, "delta") and hasattr(b, "delta"):
+        a.delta = _concat_chunks(a.delta, b.delta)
+
+    a.finish_reason = a.finish_reason or b.finish_reason
+    return a
+
+
+@_concat_chunks.register
+def _(a: Delta, b: Delta) -> Delta:
+    a.content = a.content + b.content if a.content and b.content else a.content or b.content
+
+    if hasattr(a, "tool_calls") and hasattr(b, "tool_calls") and a.tool_calls and b.tool_calls:
+        # Group tool calls by index
+        grouped_tool_calls = {}
+        for tool_call in a.tool_calls + b.tool_calls:
+            if tool_call.index not in grouped_tool_calls:
+                grouped_tool_calls[tool_call.index] = tool_call
+            else:
+                grouped_tool_calls[tool_call.index] = _concat_chunks(
+                    grouped_tool_calls[tool_call.index], tool_call
+                )
+
+        a.tool_calls = list(grouped_tool_calls.values())
+    elif hasattr(b, "tool_calls") and b.tool_calls:
+        a.tool_calls = b.tool_calls
+
+    return a
+
+
+@_concat_chunks.register
+def _(a: ToolCallRequest, b: ToolCallRequest) -> ToolCallRequest:
+    # Preserve id from either a or b, with preference for a
+    id_val = a.id if a.id is not None else b.id
+
+    # Use index from either a or b, with preference for a's index
+    index_val = a.index if hasattr(a, "index") and a.index is not None else b.index
+
+    # Concatenate the function part
+    function_val = (
+        _concat_chunks(a.function, b.function)
+        if a.function and b.function
+        else a.function or b.function
+    )
+
+    # Set all properties
+    a.id = id_val
+    a.index = index_val
+    a.function = function_val
+
+    return a
+
+
+@_concat_chunks.register
+def _(a: ToolCall, b: ToolCall) -> ToolCall:
+    # Preserve name from either a or b, with preference for a
+    name_val = a.name or b.name
+
+    # Concatenate arguments string
+    args_val = ""
+    if a.arguments is not None and b.arguments is not None:
+        args_val = a.arguments + b.arguments
+    else:
+        args_val = a.arguments or b.arguments
+
+    # Set all properties
+    a.name = name_val
+    a.arguments = args_val
+
+    return a
+
+
+def concat_completion_chunks(chunks: list[Completion]) -> Completion:
+    """
+    Accumulates all chunks returned from a streaming completion call into a `Completion` message.
+    This is useful when you stream responses from an LLM and want to keep track of the context (i.e. previous messages + current message).
+
+    Args:
+        chunks: list of completion chunks returned from streamed completion
+    Returns:
+        Completion: same as type returned from non-streaming completion
+
+
+
+    To implement `concat_completion_chunks` we first implement a binary `_concat_chunks` function for each
+    type. Using `singledispatch` to dispatch the call to the appropriate function based on the type of the first argument.
+    Each nested type is then concatenated. We can then use reduce to accumulate the entire stream into a single a
+    single `CompletionChunk`. Finally we convert the type to the appropriate non-streaming type `Completion` and return it.
+    """
+    if not chunks:
+        return None
+
+    chunks_copy = chunks.copy()
+    chunks_copy[0] = deepcopy(chunks_copy[0])  # _concat_chunks mutates first argument
+    accumulated_chunks = reduce(_concat_chunks, chunks_copy)
+
+    data = accumulated_chunks.model_dump()
+    data["object"] = "chat.completion"
+    choices = data["choices"]
+    for choice in choices:
+        choice["message"] = choice.pop("delta")
+
+    return Completion.model_validate(data)

agentex/lib/utils/console.py

@@ -0,0 +1,14 @@
+from rich import box
+from rich.console import Console
+from rich.table import Table
+
+console = Console()
+
+
+def print_section(name: str, contents: list[str], subtitle: str | None = None):
+    console.print()
+    table = Table(box=box.SQUARE, caption=subtitle, show_header=False, expand=True)
+    table.title = name
+    table.add_column(name, style="dim", width=12)
+    table.add_row(*contents)
+    console.print(table)

agentex/lib/utils/io.py

@@ -0,0 +1,29 @@
+from typing import Any
+
+import yaml
+from yaml.scanner import ScannerError
+
+
+class InvalidYAMLError(ValueError):
+    """
+    Raised when trying to red a YAML file, but the file is not formatted correctly.
+    """
+
+
+def load_yaml_file(file_path: str) -> dict[str, Any]:
+    """
+    Loads a YAML file from the specified path.
+
+    :param file_path: The path of the YAML file to load.
+    :type file_path: str
+    :return: The contents of the YAML file.
+    :rtype: dict
+    """
+    try:
+        with open(file_path) as file:
+            yaml_dict = yaml.safe_load(file)
+        return yaml_dict
+    except ScannerError as error:
+        raise InvalidYAMLError(
+            f"The following file is not in valid YAML format: {file_path}"
+        ) from error

agentex/lib/utils/iterables.py

@@ -0,0 +1,14 @@
+from collections.abc import AsyncGenerator
+from typing import Any
+
+
+async def async_enumerate(
+    aiterable: AsyncGenerator, start: int = 0
+) -> AsyncGenerator[tuple[int, Any], None]:
+    """
+    Enumerate an async generator.
+    """
+    i = start
+    async for item in aiterable:
+        yield i, item
+        i += 1

agentex/lib/utils/json_schema.py

@@ -0,0 +1,23 @@
+from typing import Any
+
+import jsonref
+from jsonschema import validate as schema_validation
+
+
+def resolve_refs(schema: dict) -> dict:
+    """
+    Resolve JSON references in a schema.
+    """
+    resolved = jsonref.replace_refs(schema, proxies=False, lazy_load=False)
+    serializable = {
+        "type": resolved.get("type"),
+        "properties": resolved.get("properties"),
+        "required": list(resolved.get("required", [])),
+        "additionalProperties": resolved.get("additionalProperties", False),
+    }
+    return serializable
+
+
+def validate_payload(json_schema: dict[str, Any], payload: dict[str, Any]) -> None:
+    """Validate the payload against the JSON schema."""
+    schema_validation(instance=payload, schema=json_schema)

agentex/lib/utils/logging.py

@@ -0,0 +1,31 @@
+import logging
+
+from rich.console import Console
+from rich.logging import RichHandler
+
+
+def make_logger(name: str):
+    """
+    Creates a logger object with a RichHandler to print colored text.
+    :param name: The name of the module to create the logger for.
+    :return: A logger object.
+    """
+    # Create a console object to print colored text
+    console = Console()
+
+    # Create a logger object with the name of the current module
+    logger = logging.getLogger(name)
+
+    # Set the global log level to INFO
+    logger.setLevel(logging.INFO)
+
+    # Add the RichHandler to the logger to print colored text
+    handler = RichHandler(
+        console=console,
+        show_level=False,
+        show_path=False,
+        show_time=False,
+    )
+    logger.addHandler(handler)
+
+    return logger

agentex/lib/utils/mcp.py

@@ -0,0 +1,17 @@
+from typing import Any
+from mcp import StdioServerParameters
+
+
+def redact_mcp_server_params(
+    mcp_server_params: list[StdioServerParameters],
+) -> list[dict[str, Any]]:
+    """Redact MCP server params for logging."""
+    return [
+        {
+            **{k: v for k, v in server_param.model_dump().items() if k != "env"},
+            "env": dict.fromkeys(server_param.env, "********")
+            if server_param.env
+            else None,
+        }
+        for server_param in mcp_server_params
+    ]

agentex/lib/utils/model_utils.py

@@ -0,0 +1,46 @@
+from collections.abc import Iterable, Mapping
+from typing import Any, TypeVar
+
+from pydantic import BaseModel as PydanticBaseModel
+from pydantic import ConfigDict
+
+from agentex.lib.utils.io import load_yaml_file
+
+T = TypeVar("T", bound="BaseModel")
+
+
+class BaseModel(PydanticBaseModel):
+    model_config = ConfigDict(from_attributes=True, populate_by_name=True)
+
+    @classmethod
+    def from_yaml(cls: type[T], file_path: str) -> T:
+        """
+        Returns an instance of this class by deserializing from a YAML file.
+
+        :param file_path: The path to the YAML file.
+        :return: An instance of this class.
+        """
+        yaml_dict = load_yaml_file(file_path=file_path)
+        class_object = cls.model_validate(yaml_dict)
+        return class_object
+
+    def to_json(self, *args, **kwargs) -> str:
+        return self.model_dump_json(*args, **kwargs)
+
+    def to_dict(self, *args, **kwargs) -> dict[str, Any]:
+        return recursive_model_dump(self)
+
+
+def recursive_model_dump(obj: Any) -> Any:
+    if isinstance(obj, PydanticBaseModel):
+        # Serialize BaseModel to dict
+        return obj.model_dump(mode="json")
+    elif isinstance(obj, Mapping):
+        # Recursively serialize dictionary values
+        return {k: recursive_model_dump(v) for k, v in obj.items()}
+    elif isinstance(obj, Iterable) and not isinstance(obj, str | bytes):
+        # Recursively serialize items in lists, tuples, sets, etc.
+        return [recursive_model_dump(item) for item in obj]
+    else:
+        # Return primitive types as-is
+        return obj

agentex/lib/utils/parsing.py

@@ -0,0 +1,15 @@
+from urllib.parse import urlsplit, urlunsplit
+
+
+def remove_query_params(url):
+    split_url = urlsplit(url)
+    scheme, netloc, path, query, fragment = split_url
+
+    if query:
+        query = ''
+    else:
+        amp_index = path.find('&')
+        if amp_index != -1:
+            path = path[:amp_index]
+
+    return urlunsplit((scheme, netloc, path, query, fragment))

agentex/lib/utils/regex.py

@@ -0,0 +1,6 @@
+import re
+
+
+def camel_to_snake(camel_case_str: str) -> str:
+    # Substitute capital letters with an underscore followed by the lowercase letter
+    return re.sub(r'(?<!^)(?=[A-Z])', '_', camel_case_str).lower()

agentex/lib/utils/temporal.py

@@ -0,0 +1,13 @@
+from temporalio import activity, workflow
+
+
+def in_temporal_workflow():
+    try:
+        return workflow.in_workflow()
+    except RuntimeError:
+        return False
+
+
+def heartbeat_if_in_workflow(heartbeat_name: str):
+    if in_temporal_workflow():
+        activity.heartbeat(heartbeat_name)

agentex/resources/__init__.py

@@ -0,0 +1,103 @@
+# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+from .spans import (
+    SpansResource,
+    AsyncSpansResource,
+    SpansResourceWithRawResponse,
+    AsyncSpansResourceWithRawResponse,
+    SpansResourceWithStreamingResponse,
+    AsyncSpansResourceWithStreamingResponse,
+)
+from .tasks import (
+    TasksResource,
+    AsyncTasksResource,
+    TasksResourceWithRawResponse,
+    AsyncTasksResourceWithRawResponse,
+    TasksResourceWithStreamingResponse,
+    AsyncTasksResourceWithStreamingResponse,
+)
+from .agents import (
+    AgentsResource,
+    AsyncAgentsResource,
+    AgentsResourceWithRawResponse,
+    AsyncAgentsResourceWithRawResponse,
+    AgentsResourceWithStreamingResponse,
+    AsyncAgentsResourceWithStreamingResponse,
+)
+from .events import (
+    EventsResource,
+    AsyncEventsResource,
+    EventsResourceWithRawResponse,
+    AsyncEventsResourceWithRawResponse,
+    EventsResourceWithStreamingResponse,
+    AsyncEventsResourceWithStreamingResponse,
+)
+from .states import (
+    StatesResource,
+    AsyncStatesResource,
+    StatesResourceWithRawResponse,
+    AsyncStatesResourceWithRawResponse,
+    StatesResourceWithStreamingResponse,
+    AsyncStatesResourceWithStreamingResponse,
+)
+from .tracker import (
+    TrackerResource,
+    AsyncTrackerResource,
+    TrackerResourceWithRawResponse,
+    AsyncTrackerResourceWithRawResponse,
+    TrackerResourceWithStreamingResponse,
+    AsyncTrackerResourceWithStreamingResponse,
+)
+from .messages import (
+    MessagesResource,
+    AsyncMessagesResource,
+    MessagesResourceWithRawResponse,
+    AsyncMessagesResourceWithRawResponse,
+    MessagesResourceWithStreamingResponse,
+    AsyncMessagesResourceWithStreamingResponse,
+)
+
+__all__ = [
+    "AgentsResource",
+    "AsyncAgentsResource",
+    "AgentsResourceWithRawResponse",
+    "AsyncAgentsResourceWithRawResponse",
+    "AgentsResourceWithStreamingResponse",
+    "AsyncAgentsResourceWithStreamingResponse",
+    "TasksResource",
+    "AsyncTasksResource",
+    "TasksResourceWithRawResponse",
+    "AsyncTasksResourceWithRawResponse",
+    "TasksResourceWithStreamingResponse",
+    "AsyncTasksResourceWithStreamingResponse",
+    "MessagesResource",
+    "AsyncMessagesResource",
+    "MessagesResourceWithRawResponse",
+    "AsyncMessagesResourceWithRawResponse",
+    "MessagesResourceWithStreamingResponse",
+    "AsyncMessagesResourceWithStreamingResponse",
+    "SpansResource",
+    "AsyncSpansResource",
+    "SpansResourceWithRawResponse",
+    "AsyncSpansResourceWithRawResponse",
+    "SpansResourceWithStreamingResponse",
+    "AsyncSpansResourceWithStreamingResponse",
+    "StatesResource",
+    "AsyncStatesResource",
+    "StatesResourceWithRawResponse",
+    "AsyncStatesResourceWithRawResponse",
+    "StatesResourceWithStreamingResponse",
+    "AsyncStatesResourceWithStreamingResponse",
+    "EventsResource",
+    "AsyncEventsResource",
+    "EventsResourceWithRawResponse",
+    "AsyncEventsResourceWithRawResponse",
+    "EventsResourceWithStreamingResponse",
+    "AsyncEventsResourceWithStreamingResponse",
+    "TrackerResource",
+    "AsyncTrackerResource",
+    "TrackerResourceWithRawResponse",
+    "AsyncTrackerResourceWithRawResponse",
+    "TrackerResourceWithStreamingResponse",
+    "AsyncTrackerResourceWithStreamingResponse",
+]