xah 0.0.1__py3-none-any.whl

xah/__init__.py

@@ -0,0 +1,4 @@
+from xah.harness.base import Harness
+from xah.toolsets.base.toolset import Toolset
+from xah.toolsets.base.tool import Tool
+from xah.class_collections import classes

xah/caps/retry.py

@@ -0,0 +1,99 @@
+from dataclasses import dataclass
+from pydantic_ai import RunContext
+
+from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.messages import SystemPromptPart
+from pydantic_deep import DeepAgentDeps
+
+from corio import logs, strings, yml
+from xah.constants import HARNESS_FEEDBACK
+from corio.jsn import from_json, to_json
+from xah.obs import logger
+
+@dataclass
+class ToolCallArgsDeserialzer(AbstractCapability[DeepAgentDeps]):
+    """Log schema-validation failures before Pydantic AI retries the model."""
+
+    label: str
+
+    def _enqueue(self, ctx, *messages: str):
+        content = strings.trim(HARNESS_FEEDBACK, *messages)
+        ctx.enqueue(SystemPromptPart(content=content))
+
+    def _log(self, ctx: RunContext[DeepAgentDeps], *, kind: str, error, tool_name: str | None = None):
+        payload = {
+            "agent": self.label,
+            "part_kind": "model-retry",
+            "retry_kind": kind,
+            "retry": ctx.retry,
+            "max_retries": ctx.max_retries,
+            "error": str(error),
+        }
+        if tool_name is not None:
+            payload["tool_name"] = tool_name
+        logger.info(logs.sanitize(yml.to_yaml(payload)))
+
+    async def wrap_tool_validate(self, ctx, *, call, tool_def, args, handler):
+        try:
+            return await handler(args)
+        except Exception as validation_error:
+            if isinstance(args, dict):
+                properties = tool_def.parameters_json_schema.get("properties", {})
+                recovered_args = dict(args)
+                recovered = False
+
+                for name, value in args.items():
+                    expected_type = properties.get(name, {}).get("type")
+                    if not isinstance(value, str) or expected_type not in {"array", "object"}:
+                        continue
+
+                    try:
+                        decoded_value = from_json(value)
+                    except (TypeError, ValueError):
+                        continue
+
+                    expected_python_type = list if expected_type == "array" else dict
+                    if isinstance(decoded_value, expected_python_type):
+                        recovered_args[name] = decoded_value
+                        recovered = True
+
+                if not recovered:
+                    raise validation_error
+
+                self._enqueue(
+                    ctx,
+                    "In future, pass nested array and object arguments as JSON values, not escaped JSON strings.",
+                )
+                return await handler(recovered_args)
+
+            if not isinstance(args, str):
+                raise
+
+            try:
+                decoded_args = from_json(args)
+            except (TypeError, ValueError):
+                raise validation_error
+
+            # Regular JSON objects have already had their fair validation attempt.
+            # Only retry when the outer JSON value was itself an encoded JSON string.
+            if not isinstance(decoded_args, str):
+                raise validation_error
+
+            self._enqueue(
+                ctx,
+                "In future, pass tool arguments as regular JSON, not as a nested/escaped JSON string.",
+            )
+            return await handler(decoded_args)
+
+    async def on_tool_validate_error(self, ctx, *, call, tool_def, args, error):
+        self._log(ctx, kind="tool-schema", tool_name=call.tool_name, error=error)
+        self._enqueue(
+            ctx,
+            "Remember the tool definition and make sure the arguments match its schema:",
+            to_json(tool_def.parameters_json_schema),
+        )
+        raise error
+
+    async def on_output_validate_error(self, ctx, *, output_context, output, error):
+        self._log(ctx, kind="output-schema", error=error)
+        raise error

xah/caps/traffic.py

@@ -0,0 +1,398 @@
+import json
+import re
+from dataclasses import dataclass, field
+from typing import Any, ClassVar
+
+from pydantic_ai import ModelMessagesTypeAdapter, RunContext
+from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.messages import (
+    FunctionToolResultEvent,
+    InstructionPart,
+    ModelMessage,
+    ModelRequest,
+    ModelRequestPart,
+    ModelResponse,
+)
+from pydantic_ai.models import ModelRequestContext
+from pydantic_deep import DeepAgentDeps
+
+from corio import logs, yml
+from corio.jsn import to_json
+from xah.constants import TODO_STATE_END, TODO_STATE_START
+from xah.obs import logger
+
+
+TODO_STATE_PATTERN = re.compile(
+    f"{re.escape(TODO_STATE_START)}(.*?){re.escape(TODO_STATE_END)}",
+    re.DOTALL,
+)
+UNREADABLE_CHARACTER_PATTERN = re.compile(r"[\x00-\x08\x0b-\x1f\x7f-\x9f\ufeff]")
+
+
+@dataclass(frozen=True)
+class StructuredDataDeserializer:
+    value: Any
+
+    @property
+    def data(self) -> Any:
+        if self.is_binary:
+            return self.binary_data
+        if isinstance(self.value, dict):
+            return {key: type(self)(value).data for key, value in self.value.items()}
+        if isinstance(self.value, list | tuple):
+            return [type(self)(value).data for value in self.value]
+        if isinstance(self.value, str):
+            return self.decoded_string
+        return self.value
+
+    @property
+    def is_binary(self) -> bool:
+        return isinstance(self.value, dict) and self.value.get("kind") == "binary"
+
+    @property
+    def binary_data(self) -> dict[str, Any]:
+        encoded = self.value["data"]
+        byte_count = (len(encoded) * 3) // 4 - encoded.count("=")
+        return {
+            "kind": "binary",
+            "media_type": self.value["media_type"],
+            "bytes": byte_count,
+            "identifier": self.value["identifier"],
+        }
+
+    @property
+    def decoded_string(self) -> Any:
+        value = self.value
+        while isinstance(value, str):
+            try:
+                value = json.loads(value)
+            except json.JSONDecodeError:
+                return value
+
+        if isinstance(value, dict | list):
+            return type(self)(value).data
+        return self.value
+
+
+@dataclass(frozen=True)
+class PartDeserializer:
+    source: dict[str, Any]
+
+    fields: ClassVar[tuple[str, ...]] = ("content",)
+    ignored_fields: ClassVar[frozenset[str]] = frozenset(
+        {
+            "id",
+            "provider_details",
+            "provider_name",
+            "signature",
+            "timestamp",
+        }
+    )
+
+    @classmethod
+    def from_source(cls, source: dict[str, Any]) -> "PartDeserializer":
+        deserializer = cls.deserializers.get(source["part_kind"], GenericPartDeserializer)
+        return deserializer(source)
+
+    deserializers: ClassVar[dict[str, type["PartDeserializer"]]]
+
+    @property
+    def kind(self) -> str:
+        return self.source["part_kind"]
+
+    @property
+    def data(self) -> dict[str, Any]:
+        data = {"kind": self.kind}
+        data.update(self.selected_data)
+        return data
+
+    @property
+    def selected_data(self) -> dict[str, Any]:
+        return {
+            name: StructuredDataDeserializer(self.source[name]).data
+            for name in self.fields
+            if name in self.source
+        }
+
+
+@dataclass(frozen=True)
+class ToolPartDeserializer(PartDeserializer):
+    fields: ClassVar[tuple[str, ...]] = ("tool_name", "tool_call_id")
+
+
+@dataclass(frozen=True)
+class ToolCallDeserializer(ToolPartDeserializer):
+    fields: ClassVar[tuple[str, ...]] = (*ToolPartDeserializer.fields, "args")
+
+
+@dataclass(frozen=True)
+class ToolResultDeserializer(ToolPartDeserializer):
+    fields: ClassVar[tuple[str, ...]] = (*ToolPartDeserializer.fields, "outcome", "content")
+
+
+@dataclass(frozen=True)
+class GenericPartDeserializer(PartDeserializer):
+    @property
+    def selected_data(self) -> dict[str, Any]:
+        return {
+            name: StructuredDataDeserializer(value).data
+            for name, value in self.source.items()
+            if name != "part_kind" and name not in self.ignored_fields
+        }
+
+
+PartDeserializer.deserializers = {
+    "builtin-tool-call": ToolCallDeserializer,
+    "builtin-tool-return": ToolResultDeserializer,
+    "retry-prompt": ToolResultDeserializer,
+    "tool-call": ToolCallDeserializer,
+    "tool-return": ToolResultDeserializer,
+}
+
+
+@dataclass(frozen=True)
+class MessageDeserializer:
+    messages: list[ModelMessage]
+
+    @property
+    def serialized_messages(self) -> list[dict[str, Any]]:
+        return ModelMessagesTypeAdapter.dump_python(self.messages, mode="json", exclude_none=True)
+
+    @property
+    def current_request_messages(self) -> list[dict[str, Any]]:
+        current: list[dict[str, Any]] = []
+        for message in reversed(self.serialized_messages):
+            if message["kind"] == "response":
+                break
+            current.append(message)
+        return list(reversed(current))
+
+    @property
+    def current_request_parts(self) -> list[dict[str, Any]]:
+        return self.parts(self.current_request_messages)
+
+    @classmethod
+    def parts(cls, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        return [
+            PartDeserializer.from_source(part).data
+            for message in messages
+            for part in message["parts"]
+        ]
+
+
+@dataclass(frozen=True)
+class ResponseDeserializer(MessageDeserializer):
+    @property
+    def response_parts(self) -> list[dict[str, Any]]:
+        return self.parts(self.serialized_messages)
+
+
+@dataclass(frozen=True)
+class EventPartDeserializer:
+    part: ModelRequestPart
+
+    @property
+    def data(self) -> dict[str, Any]:
+        message = ModelRequest(parts=[self.part])
+        source = ModelMessagesTypeAdapter.dump_python([message], mode="json", exclude_none=True)[0]["parts"][0]
+        return PartDeserializer.from_source(source).data
+
+
+@dataclass(frozen=True)
+class SetupSnapshot:
+    instructions: str
+    tools: tuple[str, ...]
+    todos: str
+
+    @classmethod
+    def from_context(cls, request_context: ModelRequestContext) -> "SetupSnapshot":
+        parameters = request_context.model_request_parameters
+        parts = parameters.instruction_parts or []
+        instructions = InstructionDeserializer(parts)
+        return cls(
+            instructions=instructions.instructions,
+            tools=tuple(sorted(parameters.tool_defs)),
+            todos=instructions.todos,
+        )
+
+    def changes(self, previous: "SetupSnapshot | None") -> dict[str, Any]:
+        if previous is None:
+            return self.initial_data
+        changes: dict[str, Any] = {}
+        if self.instructions != previous.instructions:
+            changes["instructions"] = self.instructions
+        if self.tools != previous.tools:
+            changes["tools"] = list(self.tools)
+        if self.todos != previous.todos:
+            changes["todos"] = self.todos
+        return changes
+
+    @property
+    def initial_data(self) -> dict[str, Any]:
+        data: dict[str, Any] = {}
+        if self.instructions:
+            data["instructions"] = self.instructions
+        if self.tools:
+            data["tools"] = list(self.tools)
+        if self.todos:
+            data["todos"] = self.todos
+        return data
+
+
+@dataclass(frozen=True)
+class InstructionDeserializer:
+    parts: list[InstructionPart]
+
+    @property
+    def instruction_texts(self) -> tuple[str, ...]:
+        texts = [self.instruction_content(part) for part in self.parts]
+        return tuple(text for text in texts if text)
+
+    @property
+    def instructions(self) -> str:
+        texts = tuple(dict.fromkeys(self.instruction_texts))
+        texts = tuple(
+            text
+            for text in texts
+            if not any(text != other and text in other for other in texts)
+        )
+        return "\n\n".join(texts)
+
+    @property
+    def todos(self) -> str:
+        blocks = [
+            match.group(1).strip()
+            for part in self.parts
+            for match in TODO_STATE_PATTERN.finditer(part.content)
+            if match.group(1).strip()
+        ]
+        return "\n\n".join(blocks)
+
+    @staticmethod
+    def instruction_content(part: InstructionPart) -> str:
+        content = TODO_STATE_PATTERN.sub("", part.content)
+        return re.sub(r"\n{3,}", "\n\n", content).strip()
+
+@dataclass
+class SetupTracker:
+    snapshot: SetupSnapshot | None = None
+
+
+@dataclass
+class ModelTrafficLogger(AbstractCapability[DeepAgentDeps]):
+    label: str
+    setup_tracker: SetupTracker = field(default_factory=SetupTracker)
+    request_signature: str | None = None
+    pending_event_parts: set[str] = field(default_factory=set)
+
+    async def for_run(self, ctx: RunContext[DeepAgentDeps]) -> "ModelTrafficLogger":
+        return type(self)(label=self.label, setup_tracker=self.setup_tracker)
+
+    @staticmethod
+    def readable_payload(value: Any) -> Any:
+        if isinstance(value, dict):
+            return {key: ModelTrafficLogger.readable_payload(item) for key, item in value.items()}
+        if isinstance(value, list):
+            return [ModelTrafficLogger.readable_payload(item) for item in value]
+        if not isinstance(value, str):
+            return value
+
+        if "\n" not in value and "\\n" in value:
+            value = value.replace("\\r\\n", "\n").replace("\\n", "\n")
+
+        value = value.replace("\r\n", "\n").replace("\r", "\n").replace("\t", "    ")
+        value = "\n".join(line.rstrip() for line in value.split("\n"))
+        return UNREADABLE_CHARACTER_PATTERN.sub("�", value)
+
+    def log_payload(self, payload: dict[str, Any]) -> None:
+        payload = self.readable_payload(payload)
+        message = yml.to_yaml({"agent": self.label, **payload})
+        logger.info(logs.sanitize(message))
+
+    def log_setup(self, request_context: ModelRequestContext) -> None:
+        snapshot = SetupSnapshot.from_context(request_context)
+        changes = snapshot.changes(self.setup_tracker.snapshot)
+        self.setup_tracker.snapshot = snapshot
+        if changes:
+            self.log_payload({"event": "setup", **changes})
+
+    def log_request(self, request_context: ModelRequestContext) -> None:
+        message = MessageDeserializer(request_context.messages)
+        parts = self.unlogged_request_parts(message)
+        signature = to_json(parts)
+        if self.request_signature == signature:
+            return
+        self.request_signature = signature
+        if parts:
+            self.log_payload(
+                {
+                    "event": "model-request",
+                    "model": request_context.model.model_name,
+                    "messages": parts,
+                }
+            )
+
+    def unlogged_request_parts(self, message: MessageDeserializer) -> list[dict[str, Any]]:
+        parts: list[dict[str, Any]] = []
+        for part in message.current_request_parts:
+            signature = to_json(part)
+            if signature in self.pending_event_parts:
+                self.pending_event_parts.remove(signature)
+                continue
+            parts.append(part)
+        return parts
+
+    def log_response(self, response: ModelResponse) -> None:
+        message = ResponseDeserializer([response])
+        self.log_payload(
+            {
+                "event": "model-response",
+                "model": response.model_name,
+                "messages": message.response_parts,
+            }
+        )
+
+    def log_compaction(
+        self,
+        messages: list[ModelMessage],
+        summarized: bool,
+        summary: str | None,
+    ) -> None:
+        if not summarized:
+            return
+
+        self.log_payload(
+            {
+                "event": "context-compaction",
+                "summary": summary,
+                "messages": MessageDeserializer(messages).serialized_messages,
+            }
+        )
+
+    async def wrap_model_request(
+        self,
+        ctx: RunContext[DeepAgentDeps],
+        *,
+        request_context: ModelRequestContext,
+        handler,
+    ) -> ModelResponse:
+        self.log_setup(request_context)
+        self.log_request(request_context)
+        response = await handler(request_context)
+        self.log_response(response)
+        self.request_signature = None
+        return response
+
+    async def wrap_run_event_stream(self, ctx: RunContext[DeepAgentDeps], *, stream):
+        async for event in stream:
+            if isinstance(event, FunctionToolResultEvent):
+                self.log_event_part(event)
+            yield event
+
+    def log_event_part(self, event: FunctionToolResultEvent) -> None:
+        part = EventPartDeserializer(event.part).data
+        signature = to_json(part)
+        if signature in self.pending_event_parts:
+            return
+        self.pending_event_parts.add(signature)
+        self.log_payload({"event": "tool-result", **part})

xah/class_collections.py

@@ -0,0 +1,12 @@
+from collections.abc import Callable
+
+from corio.function import ccp
+from corio.iterator import ilist
+
+
+def classes(function: Callable) -> ccp:
+    @ccp
+    def declaration(cls):
+        return ilist(function(cls))
+
+    return declaration

xah/constants.py

@@ -0,0 +1,10 @@
+PROJECTION_READ_TOOL_NAMES = "read_tool_names"
+PROJECTION_WRITE_TOOL_NAMES = "write_tool_names"
+PROJECTION_EDIT_TOOL_NAMES = "edit_tool_names"
+PROJECTION_SEARCH_TOOL_NAMES = "search_tool_names"
+PROJECTION_BASH_TOOL_NAMES = "bash_tool_names"
+
+HARNESS_FEEDBACK = "[HARNESS FEEDBACK]"
+
+TODO_STATE_START = "<todo-state>"
+TODO_STATE_END = "</todo-state>"

xah/entrypoint.py

@@ -0,0 +1,15 @@
+
+
+from xah.settings import Settings
+
+
+def main():
+
+    Settings().run()
+    from xah.harness.base import Harness
+    Harness.launch()
+
+
+
+if __name__ == "__main__":
+    main()

xah/examples/m01/e01_hello.py

@@ -0,0 +1,10 @@
+from xah import Harness
+
+
+class HelloWorld(Harness):
+    # (1)!
+    instructions = 'You are an agent who only says "Hello World"!'
+
+
+if __name__ == "__main__":
+    HelloWorld.launch()

xah/examples/m02/e01_tool.py

@@ -0,0 +1,21 @@
+from random import randint
+
+from pydantic import Field
+
+from xah import Tool
+
+
+class AgeTool(Tool):
+    """Use this tool to provide a randomly generated age for a named user."""
+    # (1)!
+
+    # (2)!
+    user: str = Field(
+        ...,
+        description="The name of the user whose age you want to know.",  # (3)!
+    )
+
+    async def run(self):
+        # (4)!
+        age = randint(1, 100)
+        return f"{self.user} is {age} years old."

xah/examples/m02/e02_toolset.py

@@ -0,0 +1,11 @@
+from xah import Toolset, classes
+from xah.examples.m02.e01_tool import AgeTool
+
+
+class UserFactsToolset(Toolset):
+    """Tools for fetching facts about users."""
+
+    @classes
+    def Tools(cls):
+        # (1)!
+        return [AgeTool]

xah/examples/m02/e03_harness.py

@@ -0,0 +1,15 @@
+from xah import Harness, classes
+from xah.examples.m02.e02_toolset import UserFactsToolset
+
+
+class DemoHarness(Harness):
+    instructions = "Use the tools to tell the user facts about themself."
+
+    @classes
+    def Toolsets(cls):
+        # (1)!
+        return [UserFactsToolset]
+
+
+if __name__ == "__main__":
+    DemoHarness.launch()

xah/examples/m02/e04_subclassing.py

@@ -0,0 +1,13 @@
+from random import choice
+
+from xah.examples.m02.e01_tool import AgeTool
+
+
+# (1)!
+class PetTool(AgeTool):
+    """Return a randomly chosen pet for a user."""
+
+    # (2)!
+    async def run(self):
+        pet = choice(["cat", "otter", "axolotl"])
+        return f"{self.user} has a {pet} as a pet."

xah/examples/m02/e05_adding.py

@@ -0,0 +1,15 @@
+from xah.examples.m02.e01_tool import AgeTool
+from xah.examples.m02.e02_toolset import UserFactsToolset
+from xah.examples.m02.e03_harness import DemoHarness
+from xah.examples.m02.e04_subclassing import PetTool
+
+
+# (1)!
+UserFactsToolset.Tools.append(PetTool)
+
+# (2)!
+UserFactsToolset.Tools.remove(AgeTool)
+
+
+if __name__ == "__main__":
+    DemoHarness.launch()