Source code for effectful.handlers.llm.encoding

import ast
import base64
import dataclasses
import functools
import inspect
import io
import json
import linecache
import textwrap
import types
import typing
import uuid
from collections.abc import (
    Callable,
    Mapping,
    MutableMapping,
)
from typing import Any

import litellm
import pydantic
from litellm import (
    ChatCompletionImageObject,
    ChatCompletionMessageToolCall,
    ChatCompletionTextObject,
    ChatCompletionToolParam,
    OpenAIMessageContentListBlock,
)
from openai.lib._pydantic import _ensure_strict_json_schema
from openai.types.chat import (
    ChatCompletionMessageToolCall as OpenAIChatCompletionMessageToolCall,
)
from PIL import Image

import effectful.handlers.llm.evaluation as evaluation
from effectful.handlers.llm.template import Tool
from effectful.internals.unification import GenericAlias, TypeEvaluator, nested_type
from effectful.ops.types import Operation, Term

type ToolCallID = str

# Reserved key under which the type-check anchor (the enclosing Template's
# underlying function) rides in the Pydantic decoding context, alongside the
# lexical environment. `decode` reads it to type-check a synthesized function
# against the Template's source; absent (tool-argument decoding) means skip.
# Deliberately not a valid identifier so `LexicalReaders` skips it (no tool leak)
# and it can never collide with a lexical name.
TYPE_CHECK_ANCHOR_KEY = "<type_check_anchor>"

# Type-check anchor for REPL `exec_code` snippets, separate from the Callable/result
# synthesis anchor (TYPE_CHECK_ANCHOR_KEY): the two decoders check against different
# contracts -- a REPL snippet against the Template body, a synthesized Callable tool
# argument against its own parameter type.
REPL_ANCHOR_KEY = "<repl_anchor>"

CONTENT_BLOCK_TYPES: frozenset[str] = frozenset(
    literal
    for member in typing.get_args(OpenAIMessageContentListBlock)
    for literal in typing.get_args(typing.get_type_hints(member).get("type", str))
    if isinstance(literal, str)
)


[docs] @pydantic.validate_call(validate_return=True) def to_content_blocks(value: typing.Any) -> list[OpenAIMessageContentListBlock]: """Convert an encoded JSON-compatible value into a flat list of content blocks. Walks the value tree, extracting content-block-shaped dicts (identified by their ``type`` discriminator) and emitting JSON syntax as text around them. Top-level strings are emitted bare (for natural template rendering). Inside JSON structures, separators match ``json.dumps`` defaults so that the linearization law holds for non-string encoded values: ``linearize(to_content_blocks(v)) == json.dumps(v)``. """ if isinstance(value, str): return [ChatCompletionTextObject(type="text", text=value)] buf: list[str] = [] blocks: list[OpenAIMessageContentListBlock] = [] def flush() -> None: if buf: blocks.append(ChatCompletionTextObject(type="text", text="".join(buf))) buf.clear() def walk(v: typing.Any) -> None: if isinstance(v, dict) and v.get("type") in CONTENT_BLOCK_TYPES: flush() blocks.append(typing.cast(OpenAIMessageContentListBlock, v)) elif isinstance(v, dict): buf.append("{") for i, (k, val) in enumerate(v.items()): if i: buf.append(", ") buf.append(json.dumps(k) + ": ") walk(val) buf.append("}") elif isinstance(v, list): buf.append("[") for i, item in enumerate(v): if i: buf.append(", ") walk(item) buf.append("]") else: buf.append(json.dumps(v)) walk(value) flush() return blocks
[docs] @dataclasses.dataclass(frozen=True, eq=True) class DecodedToolCall[T]: """ Structured representation of a tool call decoded from an LLM response. """ tool: Tool[..., T] bound_args: inspect.BoundArguments id: ToolCallID name: str
if typing.TYPE_CHECKING: type Encodable[T] = typing.Annotated[T, "encoded"] else:
[docs] class Encodable: def __class_getitem__(cls, item): return TypeToPydanticType().evaluate(item)
[docs] class TypeToPydanticType(TypeEvaluator): """Substitute custom types with their Pydantic Annotated equivalents. Recursively walks a type annotation tree, replacing leaf types that have registered Pydantic annotations (e.g., Image.Image -> PydanticImage) and reconstructing the full generic type. The result can be passed to pydantic.TypeAdapter() for automatic validation and serialization of nested structures. """ @staticmethod @functools.singledispatch def _registry(ty: type): raise RuntimeError("should not be here!")
[docs] @classmethod def register(cls, *args, **kwargs): return cls._registry.register(*args, **kwargs)
[docs] def evaluate(self, ty): app = super().evaluate(ty) origin = typing.get_origin(app) # Only dispatch on regular types. Special forms (Literal, Annotated, # Union) have non-type origins that singledispatch can't resolve; pass # them through for Pydantic to handle natively. if isinstance(app, type | GenericAlias) and ( origin is None or isinstance(origin, type) ): return self._registry.dispatch(origin or app)(app) else: return app
@TypeToPydanticType.register(str) def _pydantic_type_str[T](ty: type[T]) -> type[T]: return ty @TypeToPydanticType.register(object) def _pydantic_type_base(ty: type) -> Any: return ty class _ComplexModel(typing.TypedDict): real: float imag: float @pydantic.validate_call(validate_return=True) def _validate_complex(value: _ComplexModel) -> complex: return complex(value["real"], value["imag"]) @pydantic.validate_call(validate_return=True) def _serialize_complex(value: complex) -> _ComplexModel: return {"real": value.real, "imag": value.imag} @TypeToPydanticType.register(complex) def _pydantic_type_complex(ty): """Encode ``complex`` as ``{"real": float, "imag": float}``.""" adapted_schema = pydantic.TypeAdapter(_ComplexModel).json_schema() return typing.Annotated[ ty, pydantic.PlainValidator(_validate_complex), pydantic.PlainSerializer(_serialize_complex), pydantic.WithJsonSchema({**adapted_schema, "additionalProperties": False}), ] _CODE_FILENAME_PREFIX = "<exec_code-" @TypeToPydanticType.register(types.CodeType) def _pydantic_type_code(ty): """Encode a `types.CodeType` as a JSON string of Python source. This is the internal `Encodable` implementation for code objects -- the public type is `types.CodeType`, with no separate model (analogous to `_ComplexModel`). Decoding compiles the source through the `parse`/`compile` effect operations under a unique per-snippet filename, so invalid source is rejected here rather than at run time and the snippet's source lands in `linecache` (keeping each snippet's tracebacks resolvable). A decoded value is therefore a ready-to-run code object; re-encoding recovers its source from `linecache`, which carries everything the source string did. """ def validate(value: object, info: pydantic.ValidationInfo) -> types.CodeType: if isinstance(value, types.CodeType): return value if not isinstance(value, str): raise ValueError( f"expected Python source as a string, got {type(value).__name__}" ) filename = f"{_CODE_FILENAME_PREFIX}{uuid.uuid4()}>" try: module = evaluation.parse(value, filename) # Reject `__future__`/star imports: both are `SyntaxError` once nested in a # function body, so such a snippet can't be spliced into the Template for # type checking. evaluation.scan_non_nestable(module) except (SyntaxError, ValueError) as exc: raise ValueError(f"source is not valid REPL code: {exc}") from exc # Type-check the snippet in its execution context, exactly as a synthesized # `Callable` is (see `_pydantic_callable`): when the enclosing Template is the # type-check anchor in the decode context, splice the accumulated REPL session (the # `_repl_session` op is in scope during the response decode) plus this snippet into # the Template body and check it. A type error raises here -> the tool-call decode # fails -> `RetryLLMHandler` retries, so ill-typed code never reaches `runcode`. ctx = info.context or {} anchor = ctx.get(REPL_ANCHOR_KEY) if anchor is not None: # Pass an empty env (not `ctx`): the managed session ignores it, and a fresh # fallback session must not be seeded from the decode context (which holds tool # names and the anchor key). The decoder only reads `prior_snippets`. prior = evaluation._repl_session({}).prior_snippets checked = evaluation._splice_repl(prior, value, anchor) if checked is not None: evaluation.type_check(*checked, lenient=True) try: return evaluation.compile(module, filename) except (SyntaxError, ValueError) as exc: raise ValueError(f"source does not compile: {exc}") from exc return typing.Annotated[ ty, pydantic.PlainValidator(validate), pydantic.PlainSerializer( lambda value: "".join(linecache.getlines(value.co_filename)) ), pydantic.WithJsonSchema({"type": "string"}), ] def _inline_refs(schema: dict) -> dict: """Inline ``$ref`` pointers so ``WithJsonSchema`` never emits orphan refs. Workaround for https://github.com/pydantic/pydantic/issues/12145 — Pydantic's ``GenerateJsonSchema`` does not merge user-provided ``$defs`` into its internal ref map, so any ``$ref`` in a ``WithJsonSchema`` value causes a ``KeyError`` when the annotated type is composed into a model. """ defs = schema.get("$defs", {}) def _resolve(obj): if isinstance(obj, dict): if "$ref" in obj: ref_name = obj["$ref"].split("/")[-1] if ref_name in defs: return _resolve(defs[ref_name]) return {k: _resolve(v) for k, v in obj.items() if k != "$defs"} if isinstance(obj, list): return [_resolve(item) for item in obj] return obj return _resolve(schema) @TypeToPydanticType.register(tuple) def _pydantic_type_tuple(ty): """Convert finitary tuples to object-based schemas (``properties/required``). OpenAI's strict mode rejects the ``prefixItems`` array schema that Pydantic emits for fixed-length tuples. We convert them to a Pydantic model with positional ``item_0``, ``item_1``, … fields instead. NamedTuples are handled similarly using their field names. Bare ``tuple`` and variadic ``tuple[T, ...]`` are passed through unchanged. """ # NamedTuple subclasses dispatch here via MRO; use field names. if isinstance(ty, type) and hasattr(ty, "_fields"): hints = typing.get_type_hints(ty) nt_fields: list[str] = list(ty._fields) nt_types = [hints.get(f, typing.Any) for f in nt_fields] nt_adapters = [pydantic.TypeAdapter(t) for t in nt_types] nt_model = pydantic.create_model( ty.__name__, __config__={"extra": "forbid"}, **{f: (t, ...) for f, t in zip(nt_fields, nt_types)}, ) def _nt_validate(value, info: pydantic.ValidationInfo): if isinstance(value, tuple | list): value = dict(zip(nt_fields, value)) return ty( **{ f: nt_adapters[i].validate_python(value[f], context=info.context) for i, f in enumerate(nt_fields) } ) def _nt_serialize(value, info: pydantic.SerializationInfo): return { f: nt_adapters[i].dump_python( getattr(value, f), mode="json", context=info.context ) for i, f in enumerate(nt_fields) } return typing.Annotated[ ty, pydantic.PlainValidator(_nt_validate), pydantic.PlainSerializer(_nt_serialize), pydantic.WithJsonSchema(_inline_refs(nt_model.model_json_schema())), ] args = typing.get_args(ty) # Bare tuple or tuple[T, ...] — Pydantic's native handling is fine. # Note: tuple[()] also has get_args() == (), but has origin=tuple. if (not args and typing.get_origin(ty) is None) or ( len(args) == 2 and args[1] is Ellipsis ): return ty # tuple[()] (empty args with origin) maps to zero fields; otherwise use args. effective: list[typing.Any] = list(args) adapters = [pydantic.TypeAdapter(a) for a in effective] model = pydantic.create_model( "TupleItems", __config__={"extra": "forbid"}, **{f"item_{i}": (a, ...) for i, a in enumerate(effective)}, ) def _validate(value, info: pydantic.ValidationInfo): if isinstance(value, tuple | list): value = {f"item_{i}": v for i, v in enumerate(value)} return tuple( adapters[i].validate_python(value[f"item_{i}"], context=info.context) for i in range(len(effective)) ) def _serialize(value, info: pydantic.SerializationInfo): return { f"item_{i}": adapters[i].dump_python(v, mode="json", context=info.context) for i, v in enumerate(value) } return typing.Annotated[ ty, pydantic.PlainValidator(_validate), pydantic.PlainSerializer(_serialize), pydantic.WithJsonSchema(_inline_refs(model.model_json_schema())), ] @TypeToPydanticType.register(Term) def _pydantic_type_term(ty: type[Term]): raise TypeError("Terms cannot be converted to Pydantic types.") @TypeToPydanticType.register(Operation) def _pydantic_type_operation(ty: type[Operation]): raise TypeError("Operations cannot be converted to Pydantic types.") @pydantic.validate_call(validate_return=False) def _validate_image(value: ChatCompletionImageObject) -> Image.Image: value = pydantic.TypeAdapter(ChatCompletionImageObject).validate_python(value) image_url: litellm.ChatCompletionImageUrlObject | str = value["image_url"] url: str = image_url["url"] if isinstance(image_url, dict) else image_url prefix, data = url.split(",") if not prefix.startswith("data:image/"): raise ValueError(f"expected base64 encoded image as data uri, received {url}") return Image.open(fp=io.BytesIO(base64.b64decode(data))) def _serialize_image(value: Image.Image) -> ChatCompletionImageObject: buf = io.BytesIO() value.save(buf, format="PNG") url = f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}" return pydantic.TypeAdapter(ChatCompletionImageObject).validate_python( {"type": "image_url", "image_url": {"detail": "auto", "url": url}} ) @TypeToPydanticType.register(Image.Image) def _pydantic_type_image(ty: type[Image.Image]): adapter = pydantic.TypeAdapter(ChatCompletionImageObject) return typing.Annotated[ ty, pydantic.PlainValidator(_validate_image), pydantic.PlainSerializer(_serialize_image), pydantic.WithJsonSchema(_inline_refs(adapter.json_schema())), ]
[docs] class SynthesizedFunction(pydantic.BaseModel): """Structured output for function synthesis. Pydantic model representing synthesized code with function name and module code. """ module_code: str = pydantic.Field( ..., description="Complete Python module code (no imports needed)", )
def _create_typed_synthesized_function( callable_type: type[Callable], ) -> type[SynthesizedFunction]: """Create a SynthesizedFunction subclass with type signature in the model description. Uses pydantic.create_model to ensure the description is included in the JSON schema sent to the LLM, informing it of the expected function signature. """ if not typing.get_args(callable_type): type_signature = "Callable" # Callable[[arg1, arg2, ...], return_type] elif len(typing.get_args(callable_type)) >= 2: param_types = typing.get_args(callable_type)[0] return_type = typing.get_args(callable_type)[-1] if param_types is ...: params_str = "..." elif isinstance(param_types, list | tuple): params_str = ", ".join(getattr(t, "__name__", str(t)) for t in param_types) else: params_str = str(param_types) return_str = getattr(return_type, "__name__", str(return_type)) type_signature = f"Callable[[{params_str}], {return_str}]" else: type_signature = str(callable_type) description = f"""Given the specification above, generate a Python function satisfying the following specification and type signature. <signature>{type_signature}</signature> <instructions> 1. Produce one block of Python code. 2. The function MUST have type annotations for all parameters and the return type. 3. The function definition must be the LAST statement - do not add any code after it. 4. Do not include usage examples or function calls. </instructions> """ # Use pydantic.create_model to create a proper model with the description # The __doc__ becomes the model's description in the JSON schema model = pydantic.create_model( "TypedSynthesizedFunction", __base__=SynthesizedFunction, __doc__=description, ) return model def _validate_signature_ast( func_ast: ast.FunctionDef | ast.AsyncFunctionDef, expected_params: list[type] | None, ) -> None: """Validate the function signature from AST before execution.""" if expected_params is not None: ast_params = func_ast.args.args + func_ast.args.posonlyargs if len(ast_params) != len(expected_params): params_str = ", ".join( getattr(t, "__name__", str(t)) for t in expected_params ) raise ValueError( f"synthesized function must take exactly {len(expected_params)} " f"parameter(s) ({params_str}), but got {len(ast_params)}" ) def _validate_signature_callable( func: Callable, expected_params: list[type] | None, expected_return: type, ) -> None: """Validate the function signature from runtime callable after execution. The synthesized function must have type annotations for parameters and return type. """ sig = inspect.signature(func) if expected_params is not None: actual_params = list(sig.parameters.values()) if len(actual_params) != len(expected_params): params_str = ", ".join( getattr(t, "__name__", str(t)) for t in expected_params ) return_str = getattr(expected_return, "__name__", str(expected_return)) raise ValueError( f"synthesized function must match Callable[[{params_str}], {return_str}] " f"-- exactly {len(expected_params)} parameter(s) -- " f"but got {len(actual_params)}" ) actual_return = sig.return_annotation if actual_return is inspect.Parameter.empty: raise ValueError( "decode() requires synthesized function to have a return type annotation" ) @TypeToPydanticType.register(Callable) def _pydantic_callable(callable_type: Any) -> Any: """Create a Pydantic-compatible Annotated type for a parameterized Callable. Usage: PydanticCallable(Callable[[int, str], bool]) """ type_args = typing.get_args(callable_type) if not type_args: typed_enc = _create_typed_synthesized_function(Callable[..., typing.Any]) # type: ignore[arg-type] expected_params = None expected_return = None else: if len(type_args) < 2: raise TypeError( f"Callable type signature incomplete: {callable_type}. " "Expected Callable[[ParamTypes...], ReturnType] or Callable[..., ReturnType]." ) param_types, expected_return = type_args[0], type_args[-1] typed_enc = _create_typed_synthesized_function(callable_type) if param_types is not ... and isinstance(param_types, list | tuple): expected_params = list(param_types) else: expected_params = None def _validate(value: Any, info: pydantic.ValidationInfo) -> Callable: if callable(value) and not isinstance(value, dict): return value if isinstance(value, SynthesizedFunction): encoded = value elif isinstance(value, dict): encoded = typed_enc.model_validate(value) elif isinstance(value, str): encoded = typed_enc.model_validate_json(value) else: raise ValueError( f"Expected callable, SynthesizedFunction dict, or JSON string, " f"got {type(value)}" ) if expected_return is None: raise TypeError( "Cannot decode/synthesize callable without a concrete type signature. " "Use Callable[[ParamTypes...], ReturnType] or Callable[..., ReturnType] " "with a concrete return type (not Any)." ) ctx = info.context or {} filename = f"<synthesis:{id(encoded)}>" module: ast.AST = evaluation.parse(encoded.module_code, filename) if not isinstance(module, ast.Module) or not module.body: raise ValueError( "decode() requires module code with at least one statement." ) last_stmt = module.body[-1] if not isinstance(last_stmt, ast.FunctionDef): raise ValueError( f"decode() requires the last statement to be a function definition, " f"got {type(last_stmt).__name__}" ) _validate_signature_ast(last_stmt, expected_params) # The anchor (Template's underlying function) rides in the decoding context # under TYPE_CHECK_ANCHOR_KEY; absent for tool-argument decoding, whose # synthesized Callables are contracted by the tool param's type, not the # Template's return type, so the Template anchor doesn't apply. When # present, the code is spliced into the Template body, so first reject # constructs illegal once nested (star / `__future__` imports), then check. anchor = ctx.get(TYPE_CHECK_ANCHOR_KEY) if anchor is not None: evaluation.scan_non_nestable(module) spliced = evaluation.splice_into_source(module, anchor) if spliced is not None: evaluation.type_check(*spliced) g: MutableMapping[str, Any] = {} g.update( { k: v for k, v in ctx.items() if k.isidentifier() and k != TYPE_CHECK_ANCHOR_KEY } ) bytecode: types.CodeType = evaluation.compile(module, filename) evaluation.exec(bytecode, g) func_name = last_stmt.name if func_name not in g: raise ValueError( f"decode() expected function '{func_name}' to be defined in globals" ) result = g[func_name] if not callable(result): raise ValueError( f"decode() expected '{func_name}' to be callable, got {type(result)}" ) _validate_signature_callable(result, expected_params, expected_return) return result def _serialize(value: Callable) -> dict: if not callable(value): raise TypeError(f"Expected callable, got {type(value)}") try: source = inspect.getsource(value) except (OSError, TypeError): source = None if source: return typed_enc(module_code=textwrap.dedent(source)).model_dump() name = getattr(value, "__name__", None) docstring = inspect.getdoc(value) if name is None or docstring is None: raise ValueError( f"Cannot encode callable {value}: no source code and no __name__ or docstring" ) try: sig = inspect.signature(value) sig_str = str(sig) except (ValueError, TypeError): sig_str = "(...)" stub_code = f'''def {name}{sig_str}: """{docstring}""" ... ''' return typed_enc(module_code=stub_code).model_dump() return typing.Annotated[ callable_type, pydantic.PlainValidator(_validate), pydantic.PlainSerializer(_serialize), pydantic.WithJsonSchema( _inline_refs(pydantic.TypeAdapter(typed_enc).json_schema()) ), ] def _validate_tool( value: ChatCompletionToolParam, info: pydantic.ValidationInfo ) -> Tool: assert isinstance(info.context, Mapping), "Tool decoding requires context" value = pydantic.TypeAdapter(ChatCompletionToolParam).validate_python(value) try: return info.context[value["function"]["name"]] except KeyError as e: raise NotImplementedError(f"Unknown tool: {value['function']['name']}") from e def _serialize_tool( value: Tool, info: pydantic.SerializationInfo ) -> ChatCompletionToolParam: fields: dict[str, Any] = { name: TypeToPydanticType().evaluate(param.annotation) for name, param in inspect.signature(value).parameters.items() } sig_model = pydantic.create_model( "Params", __config__={"extra": "forbid"}, **fields, ) response_format = litellm.utils.type_to_response_format_param(sig_model) assert response_format is not None assert value.__default__.__doc__ is not None # Advertise under the context key, since decode (`_validate_tool`) resolves the call by that name. tool_name = value.__name__ context = info.context if isinstance(context, Mapping): for key, tool in context.items(): if tool is value: tool_name = key break return pydantic.TypeAdapter(ChatCompletionToolParam).validate_python( { "type": "function", "function": { "name": tool_name, "description": textwrap.dedent(value.__default__.__doc__), "parameters": response_format["json_schema"]["schema"], "strict": True, }, } ) @TypeToPydanticType.register(Tool) def _pydantic_type_tool(ty: type[Tool]): schema = _inline_refs(pydantic.TypeAdapter(ChatCompletionToolParam).json_schema()) schema = _ensure_strict_json_schema(schema, path=(), root={}) return typing.Annotated[ ty, pydantic.PlainValidator(_validate_tool), pydantic.PlainSerializer(_serialize_tool), pydantic.WithJsonSchema(schema), ] def _validate_tool_call( value: ChatCompletionMessageToolCall, info: pydantic.ValidationInfo, ) -> DecodedToolCall: if isinstance(value, dict): value = OpenAIChatCompletionMessageToolCall.model_validate(value) ctx = info.context or {} assert value.function.name is not None tool = ctx[value.function.name] assert isinstance(tool, Tool) sig = inspect.signature(tool) decoded_args = {} for name, raw_arg in json.loads(value.function.arguments).items(): assert name in sig.parameters, ( f"Unexpected argument {name} for tool {tool.__name__}" ) param = sig.parameters[name] arg_enc: pydantic.TypeAdapter[Any] = pydantic.TypeAdapter( Encodable[param.annotation] # type: ignore[name-defined] ) decoded_args[name] = arg_enc.validate_python(raw_arg, context=ctx) return DecodedToolCall( tool=tool, bound_args=sig.bind(**decoded_args), id=value.id, name=value.function.name, ) def _serialize_tool_call( value: DecodedToolCall, info: pydantic.SerializationInfo ) -> dict: ctx = info.context or {} encoded_args = {} for k, v in value.bound_args.arguments.items(): v_enc: pydantic.TypeAdapter[Any] = pydantic.TypeAdapter( Encodable[nested_type(v).value] # type: ignore[misc] ) encoded_args[k] = v_enc.dump_python(v, mode="json", context=ctx) return OpenAIChatCompletionMessageToolCall.model_validate( { "type": "function", "id": value.id, "function": { "name": value.tool.__name__, "arguments": json.dumps(encoded_args), }, } ).model_dump(mode="json") @TypeToPydanticType.register(DecodedToolCall) def _pydantic_type_tool_call(ty: type[DecodedToolCall]): # Use OpenAI's ChatCompletionMessageToolCall (has actual fields: id, function, # type) rather than litellm's (empty dict with extra="allow"). schema = _inline_refs(OpenAIChatCompletionMessageToolCall.model_json_schema()) schema = _ensure_strict_json_schema(schema, path=(), root={}) return typing.Annotated[ ty, pydantic.PlainValidator(_validate_tool_call), pydantic.PlainSerializer(_serialize_tool_call), pydantic.WithJsonSchema(schema), ]