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

pi_agent/proxy.py

@@ -0,0 +1,269 @@
+"""Proxy stream function for apps that route LLM calls through a server.
+
+The server manages auth and proxies requests to providers. It strips the
+``partial`` field from delta events to reduce bandwidth; the partial message is
+reconstructed client-side.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+from typing import Any, TypedDict
+
+import httpx
+
+from pi_ai import (
+    AbortSignal,
+    AssistantMessage,
+    AssistantMessageEventStream,
+    DoneEvent,
+    ErrorEvent,
+    Model,
+    StartEvent,
+    StreamingJsonParser,
+    TextContent,
+    TextDeltaEvent,
+    TextEndEvent,
+    TextStartEvent,
+    ThinkingContent,
+    ThinkingDeltaEvent,
+    ThinkingEndEvent,
+    ThinkingStartEvent,
+    ToolCall,
+    ToolCallDeltaEvent,
+    ToolCallEndEvent,
+    ToolCallStartEvent,
+    TranscriptContext,
+    Usage,
+    abort_scope,
+)
+from pi_ai.serialization import content_from_dict, from_dict, to_json_value
+from pi_ai.utils.http import get_http_client
+from pi_ai.utils.text import append_text, append_thinking
+
+__all__ = ["ProxyStreamOptions", "stream_proxy", "process_proxy_event"]
+
+#: Simple stream options forwarded to the proxy server (snake_case keys are sent as pi's camelCase).
+_SERIALIZABLE_OPTIONS = (
+    "temperature",
+    "sampling_params",
+    "max_tokens",
+    "reasoning",
+    "cache_retention",
+    "session_id",
+    "headers",
+    "metadata",
+    "transport",
+    "thinking_budgets",
+    "max_retry_delay_ms",
+)
+
+
+class ProxyStreamOptions(TypedDict, total=False):
+    #: Local abort signal for the proxy request.
+    signal: AbortSignal
+    #: Auth token for the proxy server.
+    auth_token: str
+    #: Proxy server URL, e.g. ``"https://genai.example.com"``.
+    proxy_url: str
+    http_client: httpx.AsyncClient
+    temperature: float
+    sampling_params: dict[str, Any]
+    max_tokens: int
+    reasoning: str
+    cache_retention: str
+    session_id: str
+    headers: dict[str, str | None]
+    metadata: dict[str, Any]
+    transport: str
+    thinking_budgets: dict[str, int]
+    max_retry_delay_ms: float
+
+
+def _request_options(options: Mapping[str, Any]) -> dict[str, Any]:
+    result: dict[str, Any] = {}
+    for key in _SERIALIZABLE_OPTIONS:
+        value = options.get(key)
+        if value is None:
+            continue
+        camel = "".join(part if index == 0 else part.capitalize() for index, part in enumerate(key.split("_")))
+        result[camel] = value
+    return result
+
+
+def stream_proxy(model: Model, context: TranscriptContext, options: Mapping[str, Any]) -> AssistantMessageEventStream:
+    """Stream through a pi proxy server (``POST {proxy_url}/api/stream``) instead of a provider.
+
+    Use it as the agent's ``stream_fn``::
+
+        Agent(stream_fn=lambda model, context, options: stream_proxy(
+            model, context, {**(options or {}), "auth_token": token, "proxy_url": "https://genai.example.com"}))
+    """
+    options = dict(options)
+    stream = AssistantMessageEventStream()
+
+    async def run() -> None:
+        partial = AssistantMessage(
+            [], api=model.api, provider=model.provider, model=model.id, usage=Usage(), stop_reason="pending"
+        )
+        parsers: dict[int, StreamingJsonParser] = {}
+        signal: AbortSignal | None = options.get("signal")
+        try:
+            async with abort_scope(signal):
+                proxy_url = str(options.get("proxy_url", "")).rstrip("/")
+                client = get_http_client(proxy_url, options)
+                body = json.dumps(
+                    {
+                        "model": to_json_value(model),
+                        "context": to_json_value(context),
+                        "options": _request_options(options),
+                    }
+                )
+                request = client.build_request(
+                    "POST",
+                    f"{proxy_url}/api/stream",
+                    headers={
+                        "Authorization": f"Bearer {options.get('auth_token', '')}",
+                        "Content-Type": "application/json",
+                    },
+                    content=body,
+                )
+                response = await client.send(request, stream=True)
+                try:
+                    if response.status_code >= 400:
+                        text = (await response.aread()).decode("utf-8", "replace")
+                        message = f"Proxy error: {response.status_code} {response.reason_phrase}"
+                        try:
+                            data = json.loads(text)
+                            if isinstance(data, dict) and data.get("error"):
+                                message = f"Proxy error: {data['error']}"
+                        except ValueError:
+                            pass
+                        raise RuntimeError(message)
+                    saw_terminal = False
+                    async for line in response.aiter_lines():
+                        if not line.startswith("data: "):
+                            continue
+                        data = line[6:].strip()
+                        if not data:
+                            continue
+                        event = process_proxy_event(json.loads(data), partial, parsers)
+                        if event is not None:
+                            if event.type in ("done", "error"):
+                                saw_terminal = True
+                            stream.push(event)
+                finally:
+                    await response.aclose()
+            if signal is not None and signal.aborted:
+                raise RuntimeError("Request aborted by user")
+            if not saw_terminal:
+                # A clean EOF without done/error means the server dropped the response mid-stream.
+                partial.stop_reason = "error"
+                partial.error_message = "Connection closed by proxy server before the response completed"
+                stream.push(ErrorEvent("error", partial))
+            stream.end()
+        except Exception as error:
+            reason = "aborted" if signal is not None and signal.aborted else "error"
+            partial.stop_reason = reason
+            partial.error_message = str(error) or type(error).__name__
+            stream.push(ErrorEvent(reason, partial))
+            stream.end()
+
+    stream.start(run())
+    return stream
+
+
+def _usage(data: Any) -> Usage:
+    return from_dict(Usage, data) if isinstance(data, Mapping) else Usage()
+
+
+def process_proxy_event(
+    event: Mapping[str, Any], partial: AssistantMessage, parsers: dict[int, StreamingJsonParser] | None = None
+) -> Any:
+    """Apply one proxy event to ``partial`` and return the matching assistant message event."""
+    parsers = parsers if parsers is not None else {}
+    kind = event.get("type")
+    index = event.get("contentIndex", 0)
+    content = partial.content
+
+    def block() -> Any:
+        return content[index] if 0 <= index < len(content) else None
+
+    def place(value: Any) -> None:
+        while len(content) <= index:
+            content.append(TextContent(""))
+        content[index] = value
+
+    if kind == "start":
+        return StartEvent(partial)
+    if kind == "text_start":
+        place(TextContent(""))
+        return TextStartEvent(index, partial)
+    if kind == "text_delta":
+        current = block()
+        if isinstance(current, TextContent):
+            append_text(current, event["delta"])
+            return TextDeltaEvent(index, event["delta"], partial)
+        raise ValueError("Received text_delta for non-text content")
+    if kind == "text_end":
+        current = block()
+        if isinstance(current, TextContent):
+            current.text_signature = event.get("contentSignature")
+            return TextEndEvent(index, current.text, partial)
+        raise ValueError("Received text_end for non-text content")
+    if kind == "thinking_start":
+        place(ThinkingContent(""))
+        return ThinkingStartEvent(index, partial)
+    if kind == "thinking_delta":
+        current = block()
+        if isinstance(current, ThinkingContent):
+            append_thinking(current, event["delta"])
+            return ThinkingDeltaEvent(index, event["delta"], partial)
+        raise ValueError("Received thinking_delta for non-thinking content")
+    if kind == "thinking_end":
+        current = block()
+        if isinstance(current, ThinkingContent):
+            current.thinking_signature = event.get("contentSignature")
+            return ThinkingEndEvent(index, current.thinking, partial)
+        raise ValueError("Received thinking_end for non-thinking content")
+    if kind == "toolcall_start":
+        place(ToolCall(event.get("id", ""), event.get("toolName", ""), {}))
+        parsers[index] = StreamingJsonParser()
+        return ToolCallStartEvent(index, partial)
+    if kind == "toolcall_delta":
+        current = block()
+        if isinstance(current, ToolCall):
+            parser = parsers.setdefault(index, StreamingJsonParser())
+            parser.feed(event["delta"])
+            current.arguments = parser.value()
+            return ToolCallDeltaEvent(index, event["delta"], partial)
+        raise ValueError("Received toolcall_delta for non-toolCall content")
+    if kind == "toolcall_end":
+        current = block()
+        if isinstance(current, ToolCall):
+            final = content_from_dict(event.get("toolCall") or {})
+            if isinstance(final, ToolCall):
+                current.id = final.id or current.id
+                current.name = final.name or current.name
+                current.arguments = final.arguments
+                current.thought_signature = final.thought_signature
+                current.namespace = final.namespace
+                current.arguments_error = final.arguments_error
+            parsers.pop(index, None)
+            return ToolCallEndEvent(index, current, partial)
+        return None
+    if kind == "done":
+        partial.stop_reason = event["reason"]
+        partial.usage = _usage(event.get("usage"))
+        if event.get("providerThinkingLevel") is not None:
+            partial.provider_thinking_level = event["providerThinkingLevel"]
+        return DoneEvent(event["reason"], partial)
+    if kind == "error":
+        partial.stop_reason = event["reason"]
+        partial.error_message = event.get("errorMessage")
+        partial.usage = _usage(event.get("usage"))
+        if event.get("providerThinkingLevel") is not None:
+            partial.provider_thinking_level = event["providerThinkingLevel"]
+        return ErrorEvent(event["reason"], partial)
+    return None

pi_agent/stream_fn.py

@@ -0,0 +1,27 @@
+"""Process-wide fallback stream function."""
+
+from __future__ import annotations
+
+from .types import StreamFn
+
+__all__ = ["set_default_stream_fn", "get_default_stream_fn"]
+
+_default_stream_fn: StreamFn | None = None
+
+
+def set_default_stream_fn(stream_fn: StreamFn | None) -> None:
+    """Configure the fallback used by ``Agent`` and the low-level loops when ``stream_fn`` is omitted.
+
+    Hosts that provide a default model runtime install its stream function here, e.g.
+    ``set_default_stream_fn(models.stream_simple)``.
+    """
+    global _default_stream_fn
+    _default_stream_fn = stream_fn
+
+
+def get_default_stream_fn() -> StreamFn:
+    if _default_stream_fn is None:
+        raise RuntimeError(
+            "No default stream function configured. Pass stream_fn explicitly or call set_default_stream_fn()."
+        )
+    return _default_stream_fn

pi_agent/tools.py

@@ -0,0 +1,327 @@
+"""Build agent tools from plain Python functions.
+
+::
+
+    @agent_tool
+    async def get_weather(city: str, units: Literal["celsius", "fahrenheit"] = "celsius") -> str:
+        \"\"\"Get the current weather.
+
+        Args:
+            city: City name.
+            units: Temperature units.
+        \"\"\"
+        return f"Sunny in {city}"
+
+The JSON Schema comes from the signature (``str``, ``int``, ``float``, ``bool``,
+``list[...]``, ``dict``, ``Literal``, ``Enum``, ``X | None``, Pydantic models and
+dataclasses), descriptions from a Google-style ``Args:`` docstring section.
+Parameters named ``tool_call_id``, ``signal`` and ``on_update`` are injected
+instead of exposed to the model. Return values are converted: ``str`` becomes a
+text block, :class:`AgentToolResult` passes through, other values are sent as JSON.
+Synchronous functions run in a worker thread.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import collections.abc
+import dataclasses
+import enum
+import inspect
+import json
+import re
+import types
+import typing
+from collections.abc import Callable
+from typing import Any, Literal, get_args, get_origin
+
+from pi_ai import ImageContent, TextContent
+from pi_ai.utils.schema import is_pydantic_model, parameters_schema
+
+from .types import AgentTool, AgentToolResult, ToolExecutionMode
+
+__all__ = ["agent_tool", "tool_from_function", "schema_from_type"]
+
+_INJECTED = frozenset({"tool_call_id", "signal", "on_update"})
+
+
+def schema_from_type(annotation: Any) -> dict[str, Any]:
+    """JSON Schema for a Python type annotation."""
+    if annotation is inspect.Parameter.empty or annotation is Any:
+        return {}
+    origin = get_origin(annotation)
+    args = get_args(annotation)
+    if origin is typing.Annotated:
+        schema = schema_from_type(args[0])
+        for extra in args[1:]:
+            if isinstance(extra, str):
+                schema["description"] = extra
+            elif isinstance(extra, dict):
+                schema.update(extra)
+        return schema
+    if origin in (typing.Union, types.UnionType):
+        members = [arg for arg in args if arg is not type(None)]
+        nullable = len(members) != len(args)
+        if len(members) == 1:
+            schema = schema_from_type(members[0])
+        else:
+            schema = {"anyOf": [schema_from_type(member) for member in members]}
+        if nullable:
+            return {"anyOf": [schema, {"type": "null"}]} if schema else schema
+        return schema
+    if origin is Literal:
+        values = list(args)
+        if all(isinstance(value, str) for value in values):
+            return {"type": "string", "enum": values}
+        return {"enum": values}
+    # ``get_origin`` reports ``collections.abc`` classes for the ``typing`` aliases (``Sequence[int]``).
+    if origin in (list, tuple, set, frozenset, collections.abc.Sequence) or annotation in (list, tuple, set):
+        item = args[0] if args else Any
+        if origin is tuple and len(args) > 1 and args[1] is not Ellipsis:
+            return {"type": "array", "prefixItems": [schema_from_type(arg) for arg in args], "items": False}
+        return {"type": "array", "items": schema_from_type(item)}
+    if origin in (dict, collections.abc.Mapping) or annotation is dict:
+        value = args[1] if len(args) == 2 else Any
+        schema: dict[str, Any] = {"type": "object"}
+        value_schema = schema_from_type(value)
+        if value_schema:
+            schema["additionalProperties"] = value_schema
+        return schema
+    if annotation is str:
+        return {"type": "string"}
+    if annotation is bool:
+        return {"type": "boolean"}
+    if annotation is int:
+        return {"type": "integer"}
+    if annotation is float:
+        return {"type": "number"}
+    if annotation is type(None):
+        return {"type": "null"}
+    if isinstance(annotation, type) and issubclass(annotation, enum.Enum):
+        values = [member.value for member in annotation]
+        if all(isinstance(value, str) for value in values):
+            return {"type": "string", "enum": values}
+        return {"enum": values}
+    if is_pydantic_model(annotation):
+        return dict(parameters_schema(annotation))
+    if dataclasses.is_dataclass(annotation) and isinstance(annotation, type):
+        hints = typing.get_type_hints(annotation)
+        properties: dict[str, Any] = {}
+        required: list[str] = []
+        for item in dataclasses.fields(annotation):
+            properties[item.name] = schema_from_type(hints.get(item.name, Any))
+            if item.default is dataclasses.MISSING and item.default_factory is dataclasses.MISSING:
+                required.append(item.name)
+        result: dict[str, Any] = {"type": "object", "properties": properties}
+        if required:
+            result["required"] = required
+        return result
+    return {}
+
+
+_ARGS_HEADER = re.compile(r"^\s*(Args|Arguments|Parameters):\s*$")
+_SECTION_HEADER = re.compile(r"^\s*[A-Z][A-Za-z ]+:\s*$")
+_ARG_LINE = re.compile(r"^\s*(\*{0,2}\w+)\s*(?:\(([^)]*)\))?\s*:\s*(.*)$")
+
+
+def _parse_docstring(doc: str | None) -> tuple[str, dict[str, str]]:
+    if not doc:
+        return "", {}
+    lines = inspect.cleandoc(doc).splitlines()
+    summary: list[str] = []
+    params: dict[str, str] = {}
+    index = 0
+    while index < len(lines) and not _ARGS_HEADER.match(lines[index]):
+        if _SECTION_HEADER.match(lines[index]):
+            break
+        summary.append(lines[index])
+        index += 1
+    current: str | None = None
+    if index < len(lines) and _ARGS_HEADER.match(lines[index]):
+        index += 1
+        while index < len(lines):
+            line = lines[index]
+            if line.strip() and not line.startswith((" ", "\t")) and _SECTION_HEADER.match(line):
+                break
+            match = _ARG_LINE.match(line)
+            if match and (line.startswith((" ", "\t")) or current is None):
+                name = match.group(1).lstrip("*")
+                params[name] = match.group(3).strip()
+                current = name
+            elif current is not None and line.strip():
+                params[current] = f"{params[current]} {line.strip()}".strip()
+            index += 1
+    return "\n".join(summary).strip(), params
+
+
+def _convert_argument(annotation: Any, value: Any) -> Any:
+    """Convert a validated JSON value to its annotated type (enums, Pydantic models, dataclasses).
+
+    The schema advertises these types, so the function receives instances, not raw JSON.
+    """
+    if value is None or annotation is inspect.Parameter.empty or annotation is Any:
+        return value
+    origin = get_origin(annotation)
+    args = get_args(annotation)
+    if origin is typing.Annotated:
+        return _convert_argument(args[0], value)
+    if origin in (typing.Union, types.UnionType):
+        members = [arg for arg in args if arg is not type(None)]
+        return _convert_argument(members[0], value) if len(members) == 1 else value
+    if origin in (list, tuple, set, frozenset, collections.abc.Sequence) and isinstance(value, list):
+        if origin is tuple and len(args) > 1 and args[1] is not Ellipsis:
+            return tuple(_convert_argument(arg, item) for arg, item in zip(args, value, strict=False))
+        items = [_convert_argument(args[0], item) for item in value] if args else list(value)
+        return origin(items) if origin in (tuple, set, frozenset) else items
+    if origin in (dict, collections.abc.Mapping) and isinstance(value, dict) and len(args) == 2:
+        return {key: _convert_argument(args[1], item) for key, item in value.items()}
+    if not isinstance(annotation, type):
+        return value
+    if issubclass(annotation, enum.Enum):
+        return value if isinstance(value, annotation) else annotation(value)
+    if is_pydantic_model(annotation) and isinstance(value, dict):
+        return annotation.model_validate(value)
+    if dataclasses.is_dataclass(annotation) and isinstance(value, dict):
+        hints = typing.get_type_hints(annotation)
+        return annotation(**{key: _convert_argument(hints.get(key, Any), item) for key, item in value.items()})
+    return value
+
+
+def _to_result(value: Any) -> AgentToolResult:
+    if isinstance(value, AgentToolResult):
+        return value
+    if value is None:
+        return AgentToolResult([], None)
+    if isinstance(value, str):
+        return AgentToolResult([TextContent(value)], None)
+    if isinstance(value, (TextContent, ImageContent)):
+        return AgentToolResult([value], None)
+    if isinstance(value, list) and value and all(isinstance(item, (TextContent, ImageContent)) for item in value):
+        return AgentToolResult(list(value), None)
+    from pi_ai.serialization import to_json_value
+
+    data = to_json_value(value)
+    return AgentToolResult([TextContent(json.dumps(data, ensure_ascii=False))], data)
+
+
+def tool_from_function(
+    fn: Callable[..., Any],
+    *,
+    name: str | None = None,
+    description: str | None = None,
+    label: str | None = None,
+    parameters: Any = None,
+    execution_mode: ToolExecutionMode | None = None,
+) -> AgentTool:
+    """Create an :class:`AgentTool` from a function (see the module docstring)."""
+    signature = inspect.signature(fn)
+    try:
+        hints = typing.get_type_hints(fn, include_extras=True)
+    except Exception:
+        hints = {}
+    summary, param_docs = _parse_docstring(fn.__doc__)
+    exposed = [param for param in signature.parameters.values() if param.name not in _INJECTED]
+    injected = {param.name for param in signature.parameters.values() if param.name in _INJECTED}
+
+    single_model = (
+        parameters is None
+        and len(exposed) == 1
+        and is_pydantic_model(hints.get(exposed[0].name))
+        and exposed[0].kind is not inspect.Parameter.VAR_KEYWORD
+    )
+    if parameters is None and not single_model:
+        properties: dict[str, Any] = {}
+        required: list[str] = []
+        for param in exposed:
+            if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
+                continue
+            schema = schema_from_type(hints.get(param.name, param.annotation))
+            if param.name in param_docs and "description" not in schema:
+                schema["description"] = param_docs[param.name]
+            if param.default is inspect.Parameter.empty:
+                required.append(param.name)
+            elif param.default is not None and isinstance(param.default, (str, int, float, bool)):
+                schema.setdefault(
+                    "default", param.default.value if isinstance(param.default, enum.Enum) else param.default
+                )
+            properties[param.name] = schema
+        parameters = {"type": "object", "properties": properties}
+        if required:
+            parameters["required"] = required
+    elif single_model:
+        parameters = hints[exposed[0].name]
+
+    is_async = inspect.iscoroutinefunction(fn)
+    model_param = exposed[0].name if single_model else None
+
+    async def execute(tool_call_id: str, params: Any, signal: Any = None, on_update: Any = None) -> AgentToolResult:
+        if model_param is not None:
+            kwargs: dict[str, Any] = {model_param: params}
+        elif isinstance(params, dict):
+            kwargs = dict(params)
+            for param in exposed:
+                if param.name in kwargs:
+                    kwargs[param.name] = _convert_argument(hints.get(param.name, Any), kwargs[param.name])
+        else:
+            kwargs = {"params": params}
+        if "tool_call_id" in injected:
+            kwargs["tool_call_id"] = tool_call_id
+        if "signal" in injected:
+            kwargs["signal"] = signal
+        if "on_update" in injected:
+            kwargs["on_update"] = on_update
+        if is_async:
+            value = await fn(**kwargs)
+        else:
+            value = await asyncio.to_thread(fn, **kwargs)
+        return _to_result(value)
+
+    tool_name = name or fn.__name__
+    return AgentTool(
+        tool_name,
+        description if description is not None else summary,
+        parameters,
+        execute=execute,
+        label=label or tool_name,
+        execution_mode=execution_mode,
+    )
+
+
+@typing.overload
+def agent_tool(fn: Callable[..., Any], /) -> AgentTool: ...
+
+
+@typing.overload
+def agent_tool(
+    *,
+    name: str | None = None,
+    description: str | None = None,
+    label: str | None = None,
+    parameters: Any = None,
+    execution_mode: ToolExecutionMode | None = None,
+) -> Callable[[Callable[..., Any]], AgentTool]: ...
+
+
+def agent_tool(
+    fn: Callable[..., Any] | None = None,
+    /,
+    *,
+    name: str | None = None,
+    description: str | None = None,
+    label: str | None = None,
+    parameters: Any = None,
+    execution_mode: ToolExecutionMode | None = None,
+) -> Any:
+    """Decorator form of :func:`tool_from_function` (usable with or without arguments)."""
+
+    def wrap(function: Callable[..., Any]) -> AgentTool:
+        return tool_from_function(
+            function,
+            name=name,
+            description=description,
+            label=label,
+            parameters=parameters,
+            execution_mode=execution_mode,
+        )
+
+    return wrap(fn) if fn is not None else wrap