pi-ai-python 1.3.0__py3-none-any.whl

pi_ai/utils/schema.py

@@ -0,0 +1,283 @@
+"""JSON Schema builders for tool parameters.
+
+:class:`Type` is a small TypeBox-style builder (the TypeScript package re-exports
+TypeBox's ``Type``). Every builder returns a plain ``dict`` subclass, so schemas
+serialize as ordinary JSON Schema and can be mixed with hand-written dicts::
+
+    from pi_ai import Type, StringEnum
+
+    params = Type.Object({
+        "location": Type.String(description="City name"),
+        "units": StringEnum(["celsius", "fahrenheit"], default="celsius"),
+        "days": Type.Optional(Type.Integer(minimum=1, maximum=14)),
+    })
+
+Tool parameters may also be Pydantic model classes; :func:`parameters_schema`
+converts them to JSON Schema once and caches the result.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Any
+
+__all__ = ["TSchema", "Type", "StringEnum", "parameters_schema", "is_pydantic_model", "clone_schema"]
+
+
+class TSchema(dict[str, Any]):
+    """A JSON Schema node produced by :class:`Type`.
+
+    ``optional`` marks a property as not required when used inside ``Type.Object``.
+    """
+
+    __slots__ = ("optional",)
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.optional = False
+
+
+def _node(base: Mapping[str, Any], options: Mapping[str, Any]) -> TSchema:
+    schema = TSchema(base)
+    for key, value in options.items():
+        if value is not None:
+            schema[_option_key(key)] = value
+    return schema
+
+
+_OPTION_KEYS = {
+    "min_length": "minLength",
+    "max_length": "maxLength",
+    "min_items": "minItems",
+    "max_items": "maxItems",
+    "unique_items": "uniqueItems",
+    "exclusive_minimum": "exclusiveMinimum",
+    "exclusive_maximum": "exclusiveMaximum",
+    "multiple_of": "multipleOf",
+    "additional_properties": "additionalProperties",
+    "min_properties": "minProperties",
+    "max_properties": "maxProperties",
+    "pattern_properties": "patternProperties",
+    "content_encoding": "contentEncoding",
+    "content_media_type": "contentMediaType",
+}
+
+
+def _option_key(key: str) -> str:
+    return _OPTION_KEYS.get(key, key)
+
+
+class Type:
+    """TypeBox-style JSON Schema builders. Keyword options map to JSON Schema keywords
+    (``min_length`` → ``minLength``; camelCase keywords are accepted as-is)."""
+
+    @staticmethod
+    def Object(properties: Mapping[str, Any] | None = None, **options: Any) -> TSchema:
+        props: dict[str, Any] = {}
+        required: list[str] = []
+        for name, value in (properties or {}).items():
+            props[name] = _strip_optional(value)
+            if not getattr(value, "optional", False):
+                required.append(name)
+        base: dict[str, Any] = {"type": "object", "properties": props}
+        if required:
+            base["required"] = required
+        return _node(base, options)
+
+    @staticmethod
+    def String(**options: Any) -> TSchema:
+        return _node({"type": "string"}, options)
+
+    @staticmethod
+    def Number(**options: Any) -> TSchema:
+        return _node({"type": "number"}, options)
+
+    @staticmethod
+    def Integer(**options: Any) -> TSchema:
+        return _node({"type": "integer"}, options)
+
+    @staticmethod
+    def Boolean(**options: Any) -> TSchema:
+        return _node({"type": "boolean"}, options)
+
+    @staticmethod
+    def Null(**options: Any) -> TSchema:
+        return _node({"type": "null"}, options)
+
+    @staticmethod
+    def Any(**options: Any) -> TSchema:
+        return _node({}, options)
+
+    @staticmethod
+    def Unknown(**options: Any) -> TSchema:
+        return _node({}, options)
+
+    @staticmethod
+    def Array(items: Mapping[str, Any], **options: Any) -> TSchema:
+        return _node({"type": "array", "items": _strip_optional(items)}, options)
+
+    @staticmethod
+    def Tuple(items: Sequence[Mapping[str, Any]], **options: Any) -> TSchema:
+        values = [_strip_optional(item) for item in items]
+        return _node(
+            {"type": "array", "prefixItems": values, "items": False, "minItems": len(values), "maxItems": len(values)},
+            options,
+        )
+
+    @staticmethod
+    def Literal(value: Any, **options: Any) -> TSchema:
+        base: dict[str, Any] = {"const": value}
+        json_type = _json_type_of(value)
+        if json_type:
+            base["type"] = json_type
+        return _node(base, options)
+
+    @staticmethod
+    def Union(schemas: Iterable[Mapping[str, Any]], **options: Any) -> TSchema:
+        return _node({"anyOf": [_strip_optional(schema) for schema in schemas]}, options)
+
+    @staticmethod
+    def Intersect(schemas: Iterable[Mapping[str, Any]], **options: Any) -> TSchema:
+        return _node({"allOf": [_strip_optional(schema) for schema in schemas]}, options)
+
+    @staticmethod
+    def Enum(values: Iterable[Any] | Mapping[str, Any], **options: Any) -> TSchema:
+        """An enum of literal values (``anyOf`` of ``const``). Prefer :func:`StringEnum` for Google APIs."""
+        items = list(values.values()) if isinstance(values, Mapping) else list(values)
+        return _node({"anyOf": [Type.Literal(item) for item in items]}, options)
+
+    @staticmethod
+    def Record(key: Mapping[str, Any], value: Mapping[str, Any], **options: Any) -> TSchema:
+        base: dict[str, Any] = {"type": "object", "additionalProperties": _strip_optional(value)}
+        pattern = key.get("pattern") if isinstance(key, Mapping) else None
+        if pattern:
+            base = {"type": "object", "patternProperties": {pattern: _strip_optional(value)}}
+        return _node(base, options)
+
+    @staticmethod
+    def Optional(schema: Mapping[str, Any]) -> TSchema:
+        """Mark a property of ``Type.Object`` as optional (not listed in ``required``)."""
+        node = TSchema(schema)
+        node.optional = True
+        return node
+
+    @staticmethod
+    def Nullable(schema: Mapping[str, Any], **options: Any) -> TSchema:
+        return _node({"anyOf": [_strip_optional(schema), {"type": "null"}]}, options)
+
+    @staticmethod
+    def Ref(ref: str, **options: Any) -> TSchema:
+        return _node({"$ref": ref}, options)
+
+    @staticmethod
+    def Unsafe(schema: Mapping[str, Any]) -> TSchema:
+        """Wrap an arbitrary JSON Schema dict."""
+        return TSchema(schema)
+
+
+def StringEnum(values: Sequence[str], *, description: str | None = None, default: str | None = None) -> TSchema:
+    """A string enum schema compatible with Google's API and other providers that do not
+    support ``anyOf``/``const`` patterns."""
+    schema = TSchema({"type": "string", "enum": list(values)})
+    if description:
+        schema["description"] = description
+    if default:
+        schema["default"] = default
+    return schema
+
+
+def _strip_optional(schema: Any) -> Any:
+    if isinstance(schema, TSchema) and schema.optional:
+        return TSchema(schema)
+    return schema
+
+
+def _json_type_of(value: Any) -> str | None:
+    if value is None:
+        return "null"
+    if isinstance(value, bool):
+        return "boolean"
+    if isinstance(value, int):
+        return "integer"
+    if isinstance(value, float):
+        return "number"
+    if isinstance(value, str):
+        return "string"
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Pydantic support
+# ---------------------------------------------------------------------------
+
+_pydantic_schema_cache: dict[type, dict[str, Any]] = {}
+
+
+def is_pydantic_model(value: Any) -> bool:
+    """True for Pydantic ``BaseModel`` subclasses (without importing pydantic)."""
+    return isinstance(value, type) and hasattr(value, "model_json_schema") and hasattr(value, "model_validate")
+
+
+def parameters_schema(parameters: Any) -> dict[str, Any]:
+    """The JSON Schema for tool ``parameters`` (a schema dict or a Pydantic model class)."""
+    if isinstance(parameters, dict):
+        return parameters
+    if is_pydantic_model(parameters):
+        cached = _pydantic_schema_cache.get(parameters)
+        if cached is None:
+            cached = _inline_pydantic_schema(parameters.model_json_schema())
+            _pydantic_schema_cache[parameters] = cached
+        return cached
+    if isinstance(parameters, Mapping):
+        return dict(parameters)
+    raise TypeError(f"Tool parameters must be a JSON Schema dict or a Pydantic model, got {type(parameters)!r}")
+
+
+def _inline_pydantic_schema(schema: dict[str, Any]) -> dict[str, Any]:
+    """Drop Pydantic's ``title`` noise and inline local ``$defs`` references.
+
+    Several providers (Google in particular) reject ``$ref``; inlining keeps the
+    schema portable. Recursive models keep their references.
+    """
+    defs = schema.get("$defs", {})
+    named_schema_maps = ("properties", "patternProperties", "dependentSchemas")
+    recursive = False
+
+    def resolve(node: Any, seen: frozenset[str]) -> Any:
+        if isinstance(node, list):
+            return [resolve(item, seen) for item in node]
+        if not isinstance(node, dict):
+            return node
+        ref = node.get("$ref")
+        if isinstance(ref, str) and ref.startswith("#/$defs/"):
+            name = ref[len("#/$defs/") :]
+            target = defs.get(name)
+            if target is not None and name not in seen:
+                resolved = resolve(target, seen | {name})
+                extra = {k: resolve(v, seen) for k, v in node.items() if k not in ("$ref", "title")}
+                return {**resolved, **extra}
+            nonlocal recursive
+            recursive = True
+            return node
+        result: dict[str, Any] = {}
+        for key, value in node.items():
+            if key in ("$defs", "title"):
+                continue
+            if key in named_schema_maps and isinstance(value, dict):
+                result[key] = {name: resolve(sub, seen) for name, sub in value.items()}
+            else:
+                result[key] = resolve(value, seen)
+        return result
+
+    result = resolve(schema, frozenset())
+    if recursive:
+        result = dict(schema)
+    if "type" not in result:
+        result["type"] = "object"
+    return result
+
+
+def clone_schema(schema: Any) -> Any:
+    """Deep copy a JSON-compatible schema (drops dict subclasses and non-JSON values)."""
+    return json.loads(json.dumps(parameters_schema(schema) if not isinstance(schema, (list, str)) else schema))

pi_ai/utils/text.py

@@ -0,0 +1,65 @@
+"""Text helpers for message content and system prompts."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+from ..types import SystemMessage
+
+__all__ = ["content_text", "get_system_message_text", "render_system_message_update", "append_text", "append_thinking"]
+
+
+# ``block.text += delta`` copies the whole string on every delta because the attribute holds a
+# second reference. Moving the string into a local first lets CPython resize it in place, which
+# keeps accumulating a long streamed response linear instead of quadratic.
+
+
+def append_text(block: Any, delta: str) -> None:
+    """Append a streamed delta to ``block.text`` in amortized O(len(delta))."""
+    text = block.text
+    block.text = ""
+    text += delta
+    block.text = text
+
+
+def append_thinking(block: Any, delta: str) -> None:
+    """Append a streamed delta to ``block.thinking`` in amortized O(len(delta))."""
+    text = block.thinking
+    block.thinking = ""
+    text += delta
+    block.thinking = text
+
+
+def content_text(content: str | Sequence[Any], separator: str = "\n") -> str:
+    """Extract and join the text blocks of message content."""
+    if isinstance(content, str):
+        return content
+    return separator.join(block.text for block in content if block.type == "text")
+
+
+def get_system_message_text(message: SystemMessage) -> str:
+    """Render a system message as a complete prompt: its content followed by its sections."""
+    parts = [content_text(message.content)]
+    if message.sections:
+        parts.extend(text for text in message.sections.values() if text is not None)
+    return "\n\n".join(part for part in parts if part)
+
+
+def render_system_message_update(message: SystemMessage) -> str:
+    """Render a later system message for APIs that accept system messages mid-conversation.
+
+    Section changes are framed by name so the model can relate them to the leading
+    prompt. This framing is request-time only and may change between versions.
+    """
+    parts: list[str] = []
+    text = content_text(message.content)
+    if text:
+        parts.append(text)
+    for name, value in (message.sections or {}).items():
+        parts.append(
+            f'Removed system prompt section "{name}".'
+            if value is None
+            else f'Updated system prompt section "{name}":\n\n{value}'
+        )
+    return "\n\n".join(parts)

pi_ai/utils/transcript.py

@@ -0,0 +1,271 @@
+"""Transcript helpers: system prompt and tool state carried by system messages."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from ..types import Context, Message, SystemMessage, Tool, ToolReference, TranscriptContext
+from .schema import parameters_schema
+from .text import content_text, get_system_message_text
+
+__all__ = [
+    "TranscriptContext",
+    "create_initial_system_message",
+    "normalize_context",
+    "get_initial_system_message",
+    "without_initial_system_message",
+    "get_current_tools",
+    "get_current_system_message",
+    "get_current_system_prompt",
+    "collapse_system_messages",
+    "resolve_transcript",
+    "to_tool_declaration",
+    "declarations_equal",
+    "ToolStateChanges",
+    "get_tool_state_changes",
+    "get_declared_tools",
+    "has_tool_redefinitions",
+    "has_non_additive_tool_changes",
+    "TranscriptTools",
+    "resolve_transcript_tools",
+]
+
+
+def create_initial_system_message(system_prompt: str | None, tools: Sequence[Tool] | None) -> SystemMessage | None:
+    """Build the leading system message for a prompt and tool set.
+
+    Returns ``None`` when both are empty, so an empty transcript stays empty.
+    """
+    has_prompt = bool(system_prompt)
+    has_tools = bool(tools)
+    if not has_prompt and not has_tools:
+        return None
+    return SystemMessage(system_prompt or "", tools_added=list(tools) if has_tools else None, timestamp=0)
+
+
+def normalize_context(context: Context | Mapping[str, Any] | TranscriptContext) -> TranscriptContext:
+    """Fold ``Context.system_prompt`` and ``Context.tools`` into a leading system message.
+
+    This is the entry point that produces a :class:`TranscriptContext`. Plain dicts
+    (``{"system_prompt": ..., "messages": [...], "tools": [...]}``, snake_case or pi's
+    camelCase keys) and dict messages are accepted and converted.
+    """
+    if isinstance(context, TranscriptContext):
+        return context
+    if isinstance(context, Context):
+        system_prompt, tools, messages = context.system_prompt, context.tools, context.messages
+    else:
+        system_prompt = context.get("system_prompt", context.get("systemPrompt"))
+        tools = context.get("tools")
+        messages = context.get("messages") or []
+    messages = _coerce_messages(messages)
+    if tools:
+        tools = [_coerce_tool(tool) for tool in tools]
+    initial = create_initial_system_message(system_prompt, tools)
+    return TranscriptContext([initial, *messages] if initial else list(messages))
+
+
+def _coerce_messages(messages: Sequence[Any]) -> list[Message]:
+    if all(not isinstance(message, Mapping) for message in messages):
+        return messages if isinstance(messages, list) else list(messages)
+    from ..serialization import message_from_dict
+
+    return [message_from_dict(message) if isinstance(message, Mapping) else message for message in messages]
+
+
+def _coerce_tool(tool: Any) -> Tool:
+    if isinstance(tool, Mapping):
+        from ..serialization import tool_from_dict
+
+        return tool_from_dict(tool)
+    return tool
+
+
+def _is_system(message: Any) -> bool:
+    return getattr(message, "role", None) == "system"
+
+
+def get_initial_system_message(messages: Sequence[Any]) -> SystemMessage | None:
+    """The leading system message, if the transcript starts with one."""
+    if messages and _is_system(messages[0]):
+        return messages[0]
+    return None
+
+
+def without_initial_system_message(messages: list[Message]) -> list[Message]:
+    """Drop the leading system message for APIs that carry the prompt outside the message list."""
+    return messages[1:] if get_initial_system_message(messages) else messages
+
+
+def get_current_tools(messages: Sequence[Any]) -> list[Tool]:
+    """Resolve the tools available after applying every transcript delta in order."""
+    tools: dict[str, Tool] = {}
+    for message in messages:
+        if not _is_system(message):
+            continue
+        for removed in message.tools_removed or ():
+            tools.pop(removed.name, None)
+        for added in message.tools_added or ():
+            tools[added.name] = added
+    return list(tools.values())
+
+
+def get_current_system_message(messages: Sequence[Any]) -> SystemMessage | None:
+    """Replay every system message into one leading system message.
+
+    Later ``content`` is appended to the base prompt, ``sections`` are patched by
+    name, and tools are resolved with :func:`get_current_tools`.
+    """
+    content: list[str] = []
+    sections: dict[str, str] = {}
+    timestamp: int | None = None
+    for message in messages:
+        if not _is_system(message):
+            continue
+        if timestamp is None:
+            timestamp = message.timestamp
+        text = content_text(message.content)
+        if text:
+            content.append(text)
+        for name, value in (message.sections or {}).items():
+            if value is None:
+                sections.pop(name, None)
+            else:
+                sections[name] = value
+    tools = get_current_tools(messages)
+    if timestamp is None and not tools:
+        return None
+    return SystemMessage(
+        "\n\n".join(content),
+        sections=dict(sections) if sections else None,
+        tools_added=tools if tools else None,
+        timestamp=timestamp if timestamp is not None else 0,
+    )
+
+
+def get_current_system_prompt(messages: Sequence[Any]) -> str:
+    """The current system prompt text after replaying every system message."""
+    message = get_current_system_message(messages)
+    return get_system_message_text(message) if message else ""
+
+
+def collapse_system_messages(context: TranscriptContext) -> TranscriptContext:
+    """Rebuild the transcript for APIs without mid-conversation system messages.
+
+    The replayed system message leads and every later system message is dropped.
+    """
+    head = get_current_system_message(context.messages)
+    messages = [message for message in context.messages if not _is_system(message)]
+    return TranscriptContext([head, *messages] if head else messages)
+
+
+def resolve_transcript(
+    context: TranscriptContext, supports_mid_convo_system_messages: bool | None
+) -> TranscriptContext:
+    """Keep later system messages in place when the model accepts them; otherwise collapse them."""
+    return context if supports_mid_convo_system_messages else collapse_system_messages(context)
+
+
+def to_tool_declaration(tool: Tool) -> Tool:
+    """Strip executable and display-only fields from a tool (e.g. an agent tool).
+
+    Parameters are deep-copied as plain JSON Schema (Pydantic models are converted).
+    """
+    parameters = json.loads(json.dumps(parameters_schema(tool.parameters)))
+    constrained = tool.constrained_sampling
+    return Tool(
+        tool.name,
+        tool.description,
+        parameters,
+        constrained_sampling=json.loads(json.dumps(constrained)) if isinstance(constrained, dict) else constrained,
+    )
+
+
+def _declaration_key(tool: Tool) -> str:
+    constrained = tool.constrained_sampling
+    return json.dumps(
+        [tool.name, tool.description, parameters_schema(tool.parameters), constrained if constrained else None],
+        sort_keys=True,
+        default=str,
+    )
+
+
+def declarations_equal(left: Tool, right: Tool) -> bool:
+    """Whether two tools declare the same interface to the model."""
+    if left is right:
+        return True
+    return _declaration_key(left) == _declaration_key(right)
+
+
+@dataclass(slots=True)
+class ToolStateChanges:
+    tools_added: list[Tool]
+    tools_removed: list[ToolReference]
+
+
+def get_tool_state_changes(previous: Sequence[Tool], current: Sequence[Tool]) -> ToolStateChanges:
+    """Compare two complete tool states. A changed definition is a removal followed by an addition."""
+    previous_keys = {tool.name: _declaration_key(tool) for tool in previous}
+    current_keys = {tool.name: _declaration_key(tool) for tool in current}
+    added = [to_tool_declaration(tool) for tool in current if previous_keys.get(tool.name) != current_keys[tool.name]]
+    removed = [ToolReference(tool.name) for tool in previous if current_keys.get(tool.name) != previous_keys[tool.name]]
+    return ToolStateChanges(added, removed)
+
+
+def get_declared_tools(messages: Sequence[Any]) -> list[Tool]:
+    """Every definition referenced by transcript tool state, in first-declaration order."""
+    definitions: dict[str, Tool] = {}
+    for message in messages:
+        if _is_system(message):
+            for tool in message.tools_added or ():
+                definitions[tool.name] = tool
+    return list(definitions.values())
+
+
+def has_tool_redefinitions(messages: Sequence[Any]) -> bool:
+    """Whether a tool name was declared twice with different definitions (deprecated in pi)."""
+    declared: dict[str, Tool] = {}
+    for message in messages:
+        if not _is_system(message):
+            continue
+        for tool in message.tools_added or ():
+            previous = declared.get(tool.name)
+            if previous is not None and not declarations_equal(previous, tool):
+                return True
+            declared[tool.name] = tool
+    return False
+
+
+def has_non_additive_tool_changes(messages: Sequence[Any]) -> bool:
+    """Whether tool history contains a removal or same-name redeclaration."""
+    declared: set[str] = set()
+    for message in messages:
+        if not _is_system(message):
+            continue
+        if message.tools_removed:
+            return True
+        for tool in message.tools_added or ():
+            if tool.name in declared:
+                return True
+            declared.add(tool.name)
+    return False
+
+
+@dataclass(slots=True)
+class TranscriptTools:
+    #: Tools sent in the top-level request field.
+    request_tools: list[Tool]
+    #: Whether later system messages carry their own ``tools_added`` as in-place additions.
+    anchors_additions: bool
+
+
+def resolve_transcript_tools(messages: Sequence[Any], supports_tool_additions: bool) -> TranscriptTools:
+    """Split tool declarations between the top-level request field and in-place additions."""
+    anchors = supports_tool_additions and not has_non_additive_tool_changes(messages)
+    if anchors:
+        initial = get_initial_system_message(messages)
+        return TranscriptTools(list(initial.tools_added or []) if initial else [], True)
+    return TranscriptTools(get_current_tools(messages), False)