chore: update codegen to use annotated / cleanup

Author: gersmann

README.md

@@ -101,23 +101,19 @@ class Summary(BaseModel):
     summary: str
 
 
-schema = {
-    "type": "object",
-    "properties": {"summary": {"type": "string"}},
-    "required": ["summary"],
-    "additionalProperties": False,
-}
-
 with AppServerClient.connect_stdio() as client:
     thread = client.start_thread()
     result = thread.run_model(
         "Summarize repository status",
         Summary,
-        outputSchema=schema,
     )
     print(result.summary)
 ```
 
+`run_model()` uses `Summary` both as the validation model and, by default, as the output schema sent
+to Codex. If you want JSON back without validation, you can also pass the model class directly to
+`outputSchema`, for example `thread.run_json(..., outputSchema=Summary)`.
+
 ## Streaming
 
 ### Exec stream

codex/app_server/_async_client.py

@@ -18,6 +18,7 @@
     extract_token_usage,
     extract_turn,
     extract_turn_id,
+    has_output_schema,
     merge_params,
     normalize_turn_input,
 )
@@ -341,6 +342,8 @@ async def run_model(
         **overrides: object,
     ) -> _ModelT:
         """Run a turn and validate the final assistant text with `model_type`."""
+        if "outputSchema" not in overrides and not has_output_schema(params):
+            overrides = {"outputSchema": model_type, **overrides}
         stream = await self.run(input, params, **overrides)
         await stream.wait()
         return stream.final_model(model_type)

codex/app_server/_helpers.py

@@ -9,7 +9,6 @@
 from codex.app_server.transports import JsonObject
 from codex.protocol import types as protocol
 
-JsonValue = str | int | float | bool | None | list["JsonValue"] | dict[str, "JsonValue"]
 InputItem = (
     str
     | Mapping[str, Any]
@@ -50,11 +49,13 @@ def request_id(message: BaseModel) -> str | int:
 def serialize_value(value: object) -> object:
     if isinstance(value, BaseModel):
         return value.model_dump(mode="json", by_alias=True, exclude_none=True)
+    if isinstance(value, type) and issubclass(value, BaseModel):
+        return value.model_json_schema()
     if isinstance(value, list):
         return [serialize_value(item) for item in value]
     if isinstance(value, tuple):
         return [serialize_value(item) for item in value]
-    if isinstance(value, dict):
+    if isinstance(value, Mapping):
         return {key: serialize_value(item) for key, item in value.items()}
     return value
 
@@ -94,6 +95,14 @@ def merge_params(
     return payload
 
 
+def has_output_schema(params: BaseModel | Mapping[str, Any] | None) -> bool:
+    if params is None:
+        return False
+    if isinstance(params, BaseModel):
+        return getattr(params, "outputSchema", None) is not None
+    return params.get("outputSchema") is not None or params.get("output_schema") is not None
+
+
 def parse_result(result: object, result_model: type[BaseModel]) -> BaseModel:
     if isinstance(result, result_model):
         return result

codex/events.py

@@ -5,6 +5,8 @@
 from dataclasses import dataclass
 from typing import Literal, Protocol, runtime_checkable
 
+from codex.output_schema import OutputSchemaInput
+
 ApprovalMode = Literal["never", "on-request", "on-failure", "untrusted"]
 SandboxMode = Literal["read-only", "workspace-write", "danger-full-access"]
 ModelReasoningEffort = Literal["minimal", "low", "medium", "high", "xhigh"]
@@ -61,5 +63,5 @@ class ThreadOptions:
 class TurnOptions:
     """Per-run options for the exec-based `Thread.run()` APIs."""
 
-    output_schema: dict[str, object] | None = None
+    output_schema: OutputSchemaInput | None = None
     signal: CancelSignal | None = None

codex/items.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import cast
+
+from pydantic import BaseModel
+
+type OutputSchemaInput = Mapping[str, object] | type[BaseModel]
+
+
+def normalize_output_schema(schema: OutputSchemaInput | None) -> dict[str, object] | None:
+    if schema is None:
+        return None
+    if isinstance(schema, type) and issubclass(schema, BaseModel):
+        return cast(dict[str, object], schema.model_json_schema())
+    if isinstance(schema, Mapping):
+        return dict(schema)
+    raise ValueError("output_schema must be a JSON object or a Pydantic model class")

codex/options.py

@@ -6,6 +6,8 @@
 from dataclasses import dataclass
 from pathlib import Path
 
+from codex.output_schema import OutputSchemaInput, normalize_output_schema
+
 
 @dataclass(slots=True, frozen=True)
 class OutputSchemaFile:
@@ -17,17 +19,15 @@ def cleanup(self) -> None:
             shutil.rmtree(self.schema_dir, ignore_errors=True)
 
 
-def create_output_schema_file(schema: object | None) -> OutputSchemaFile:
-    if schema is None:
+def create_output_schema_file(schema: OutputSchemaInput | None) -> OutputSchemaFile:
+    normalized_schema = normalize_output_schema(schema)
+    if normalized_schema is None:
         return OutputSchemaFile(schema_path=None, schema_dir=None)
 
-    if not isinstance(schema, dict):
-        raise ValueError("output_schema must be a plain JSON object")
-
     schema_dir = Path(tempfile.mkdtemp(prefix="codex-output-schema-"))
     schema_path = schema_dir / "schema.json"
     try:
-        schema_path.write_text(json.dumps(schema), encoding="utf-8")
+        schema_path.write_text(json.dumps(normalized_schema), encoding="utf-8")
     except Exception:
         shutil.rmtree(schema_dir, ignore_errors=True)
         raise