Emil Shanaty
0cfb4277b5
- protocol.py: pydantic models for all AICC message types (envelope, session_init, tool_call, tool_result, event, error, heartbeat) - bridge.py: Bridge with @bridge.tool registration, capability checks, session management, serve_forever - client.py: AICCClient with single background reader (safe on concurrent transports like WebSocket), call_tool, events, manifest - tool.py: @tool decorator with schema generation from type hints - schema.py: JSON Schema generation (str/int/float/bool, list, dict, Optional, pydantic models) - transport: Transport protocol, InProcessTransport, WebSocket client+server - tests: 12 passing (integration, schema, websocket roundtrip) - examples: bridge_minimal.py + agent_minimal.py (verified end-to-end)
114 lines
3.9 KiB
Python
114 lines
3.9 KiB
Python
"""JSON Schema generation utilities."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import inspect
|
|
from typing import Any, Callable, get_args, get_origin, get_type_hints
|
|
|
|
from pydantic import BaseModel
|
|
|
|
|
|
def function_schema(fn: Callable[..., Any]) -> dict[str, Any]:
|
|
"""Generate a JSON Schema for a function's input from its signature.
|
|
|
|
Supports:
|
|
- Plain Python type hints (str, int, float, bool, list, dict, etc.)
|
|
- Pydantic BaseModel subclasses as root input
|
|
- Optional[X] / Union[X, None]
|
|
- Default values from the signature
|
|
|
|
Returns a JSON Schema dict with type=object, properties, and required.
|
|
"""
|
|
sig = inspect.signature(fn)
|
|
hints = get_type_hints(fn)
|
|
|
|
# If the return / first param is a Pydantic model, use its schema as root.
|
|
for param_name, param in sig.parameters.items():
|
|
ann = hints.get(param_name, param.annotation)
|
|
if ann is inspect.Parameter.empty:
|
|
continue
|
|
if isinstance(ann, type) and issubclass(ann, BaseModel):
|
|
schema = ann.model_json_schema()
|
|
schema.pop("title", None)
|
|
return _strip_unsupported(schema)
|
|
|
|
properties: dict[str, Any] = {}
|
|
required: list[str] = []
|
|
|
|
for param_name, param in sig.parameters.items():
|
|
if param_name == "self":
|
|
continue
|
|
ann = hints.get(param_name, param.annotation)
|
|
if ann is inspect.Parameter.empty:
|
|
continue
|
|
properties[param_name] = _annotation_to_schema(ann)
|
|
if param.default is inspect.Parameter.empty:
|
|
required.append(param_name)
|
|
|
|
return {"type": "object", "properties": properties, "required": required, "additionalProperties": False}
|
|
|
|
|
|
def return_schema(fn: Callable[..., Any]) -> dict[str, Any]:
|
|
"""Generate JSON Schema for a function's return type."""
|
|
hints = get_type_hints(fn)
|
|
ret = hints.get("return", inspect.signature(fn).return_annotation)
|
|
if ret is inspect.Signature.empty:
|
|
return {"type": "object", "additionalProperties": True}
|
|
if isinstance(ret, type) and issubclass(ret, BaseModel):
|
|
schema = ret.model_json_schema()
|
|
schema.pop("title", None)
|
|
return _strip_unsupported(schema)
|
|
return _annotation_to_schema(ret)
|
|
|
|
|
|
def _annotation_to_schema(ann: Any) -> dict[str, Any]:
|
|
origin = get_origin(ann)
|
|
args = get_args(ann)
|
|
|
|
# Plain types
|
|
if ann is str:
|
|
return {"type": "string"}
|
|
if ann is int:
|
|
return {"type": "integer"}
|
|
if ann is float:
|
|
return {"type": "number"}
|
|
if ann is bool:
|
|
return {"type": "boolean"}
|
|
|
|
# Containers — check BEFORE generic Union handling, since
|
|
# dict[str, X] / list[X] carry args too.
|
|
if origin is dict or ann is dict:
|
|
if args and len(args) == 2:
|
|
return {
|
|
"type": "object",
|
|
"additionalProperties": _annotation_to_schema(args[1]),
|
|
}
|
|
return {"type": "object", "additionalProperties": True}
|
|
if origin is list or ann is list:
|
|
if args and len(args) == 1:
|
|
return {"type": "array", "items": _annotation_to_schema(args[0])}
|
|
return {"type": "array"}
|
|
|
|
# Optional[X] / Union[X, None] / Union[X, Y]
|
|
if origin is not None and args:
|
|
non_none = [a for a in args if a is not type(None)]
|
|
if len(non_none) == 1 and len(args) > len(non_none):
|
|
sub = _annotation_to_schema(non_none[0])
|
|
sub["nullable"] = True
|
|
return sub
|
|
if len(non_none) > 1:
|
|
return {"anyOf": [_annotation_to_schema(a) for a in non_none]}
|
|
|
|
if isinstance(ann, type) and issubclass(ann, BaseModel):
|
|
schema = ann.model_json_schema()
|
|
schema.pop("title", None)
|
|
return _strip_unsupported(schema)
|
|
|
|
return {"type": "object", "additionalProperties": True}
|
|
|
|
|
|
def _strip_unsupported(schema: dict[str, Any]) -> dict[str, Any]:
|
|
"""Remove JSON Schema keys that some LLM tool-use pipelines reject."""
|
|
schema.pop("title", None)
|
|
return schema
|