|
| 1 | +"""Trace-helper codegen — one ``record_<entity>.py`` per concrete ``object.entity`` |
| 2 | +that (a) transitively ``extends`` ``metaobjects::ai::LlmCallBase`` AND (b) nests a |
| 3 | +``template.prompt`` carrying ``@responseRef`` and/or ``@payloadRef``. |
| 4 | +
|
| 5 | +AI LLM-call trace persistence — Unit 2 Slice 2 (Python port). Cross-port parity |
| 6 | +with the TypeScript ``trace-helper-file.ts`` (``record<Entity>`` + ``call<Entity>``) |
| 7 | +and the Java ``LlmTraceHelperGenerator`` (``record<Entity>`` only). This Python |
| 8 | +port emits the ``record_<entity>`` half ONLY — the render→call→record loop needs an |
| 9 | +``LlmClient`` seam that is BYO / vendor-neutral (ADR-0024), so ``call<Entity>`` is |
| 10 | +intentionally NOT emitted (matching Java). |
| 11 | +
|
| 12 | +The emitted helper exposes a single ``record_<snake>(recorder, input, redact=None)`` |
| 13 | +that: |
| 14 | +
|
| 15 | +1. runs the tolerant extract (``extract`` from ``metaobjects.render.extract``, |
| 16 | + which NEVER raises) of ``input.llm_response_text`` against a baked |
| 17 | + ``_RESPONSE_SCHEMA`` (the response VO's field shape — emitted via the SAME |
| 18 | + ``extract_schema_emitter`` path the output-parser generator uses for its |
| 19 | + tolerant ``extract_lenient_*`` twin); |
| 20 | +2. derives ``status``/``error_detail`` from the extract report's lost-required gate |
| 21 | + (a lost ``@required`` field → ``status="error"`` + a ``"lost required: …"`` detail); |
| 22 | +3. builds the base trace row via the Slice-1 ``build_llm_call_row`` (the 18 |
| 23 | + ``LlmCallBase`` base fields + raw ``llmRequest``/``llmResponse``), folding in the |
| 24 | + derived status/error_detail by replacing the input; |
| 25 | +4. sets the typed ``voRequest`` (``input.llm_request``) and ``voResponse`` |
| 26 | + (``outcome.data`` — the extracted mirror dict) columns on the row → native jsonb; |
| 27 | +5. persists the row ONCE via the supplied recorder (the never-throwing Slice-1 |
| 28 | + ``persist_llm_call_row``, which redacts then records); |
| 29 | +6. returns a typed ``<Entity>TraceResult`` (``status``, ``error_detail``, |
| 30 | + ``vo_response``). |
| 31 | +
|
| 32 | +Skips (no helper emitted) when the entity is abstract, does not derive from |
| 33 | +``LlmCallBase``, has no nested ``template.prompt``, or that prompt carries NEITHER |
| 34 | +``@responseRef`` NOR ``@payloadRef`` (both gate the helper — matching the TS |
| 35 | +reference). The response VO is resolved via ``resolve_payload_vo`` (short-name OR |
| 36 | +FQN); a non-``object.value`` target is a hard generator error (matching Java). |
| 37 | +
|
| 38 | +STI/TPH discriminator handling is DEFERRED (matching the Java port's Slice 2 — a |
| 39 | +plain trace entity is the target). |
| 40 | +""" |
| 41 | +from __future__ import annotations |
| 42 | + |
| 43 | +from collections.abc import Callable |
| 44 | + |
| 45 | +from metaobjects.codegen import extract_schema_emitter as rse |
| 46 | +from metaobjects.codegen.constants import generated_header |
| 47 | +from metaobjects.codegen.format import ruff_format |
| 48 | +from metaobjects.codegen.generator import EmittedFile, GenContext, Generator |
| 49 | +from metaobjects.codegen.generators.payload_vo_generator import resolve_payload_vo |
| 50 | +from metaobjects.meta.core.object.meta_object import MetaObject |
| 51 | +from metaobjects.meta.core.object.object_constants import OBJECT_SUBTYPE_ENTITY |
| 52 | +from metaobjects.meta.meta_data import MetaData |
| 53 | +from metaobjects.meta.template import template_constants as tc |
| 54 | +from metaobjects.meta.template.meta_template import MetaTemplate |
| 55 | +from metaobjects.shared.base_types import TYPE_TEMPLATE |
| 56 | + |
| 57 | +_GENERATOR_NAME = "trace-helper" |
| 58 | + |
| 59 | +#: The abstract base entity a trace entity must (transitively) ``extends``. |
| 60 | +#: Cross-port constant — mirrors TS ``LLM_CALL_BASE`` / Java ``LLM_CALL_BASE``. |
| 61 | +LLM_CALL_BASE = "LlmCallBase" |
| 62 | + |
| 63 | + |
| 64 | +def _snake_case(name: str) -> str: |
| 65 | + """``GreetingCall`` → ``greeting_call``. PascalCase → snake_case with no |
| 66 | + acronym handling — matches the convention used by sibling generators |
| 67 | + (``output_parser_generator._snake_case`` etc.).""" |
| 68 | + out: list[str] = [] |
| 69 | + for i, ch in enumerate(name): |
| 70 | + if ch.isupper() and i > 0: |
| 71 | + out.append("_") |
| 72 | + out.append(ch.lower()) |
| 73 | + return "".join(out) |
| 74 | + |
| 75 | + |
| 76 | +def _extends_base(entity: MetaObject) -> bool: |
| 77 | + """Walk the resolved super chain looking for a node SHORT-named |
| 78 | + :data:`LLM_CALL_BASE`. ``MetaData.name`` holds the short name only (the package |
| 79 | + lives on ``MetaData.package``), so a plain ``name`` compare is the short-name |
| 80 | + test — mirrors the Java ``getShortName()`` walk and the TS ``superResolved`` |
| 81 | + walk.""" |
| 82 | + cur = entity.super_data |
| 83 | + visited: set[int] = set() |
| 84 | + while cur is not None and id(cur) not in visited: |
| 85 | + if cur.name == LLM_CALL_BASE: |
| 86 | + return True |
| 87 | + visited.add(id(cur)) |
| 88 | + cur = cur.super_data |
| 89 | + return False |
| 90 | + |
| 91 | + |
| 92 | +def _first_prompt(entity: MetaObject) -> MetaTemplate | None: |
| 93 | + """First OWN ``template.prompt`` child of *entity*, or ``None``. Own-only — |
| 94 | + the trace prompt is declared inline on the concrete trace entity (Slice 1/3 |
| 95 | + derive the typed columns from it).""" |
| 96 | + for child in entity.own_children(): |
| 97 | + if ( |
| 98 | + isinstance(child, MetaTemplate) |
| 99 | + and child.type == TYPE_TEMPLATE |
| 100 | + and child.sub_type == tc.TEMPLATE_SUBTYPE_PROMPT |
| 101 | + ): |
| 102 | + return child |
| 103 | + return None |
| 104 | + |
| 105 | + |
| 106 | +def render_trace_helper(entity: MetaObject, root: MetaData) -> str | None: |
| 107 | + """Render one ``record_<entity>.py`` for a concrete trace ``object.entity``. |
| 108 | +
|
| 109 | + Returns ``None`` when the entity is not a trace-helper target (abstract, not |
| 110 | + ``LlmCallBase``-derived, no nested ``template.prompt``, or that prompt carries |
| 111 | + neither ``@responseRef`` nor ``@payloadRef``). |
| 112 | +
|
| 113 | + Raises ``ValueError`` when the prompt's ``@responseRef`` does not resolve to an |
| 114 | + ``object.value`` (matching the Java port's hard ``GeneratorException``).""" |
| 115 | + if entity.is_abstract: |
| 116 | + return None |
| 117 | + if not _extends_base(entity): |
| 118 | + return None |
| 119 | + |
| 120 | + prompt = _first_prompt(entity) |
| 121 | + if prompt is None: |
| 122 | + return None |
| 123 | + |
| 124 | + response_ref = prompt.response_ref() |
| 125 | + payload_ref = prompt.payload_ref() |
| 126 | + # Both gate the helper — at least one of @responseRef / @payloadRef must be |
| 127 | + # present (matches the TS reference, which needs @responseRef to type the result |
| 128 | + # and @payloadRef to type the request). |
| 129 | + if response_ref is None and payload_ref is None: |
| 130 | + return None |
| 131 | + |
| 132 | + # The response VO drives the baked extract schema + the typed voResponse. When |
| 133 | + # only @payloadRef is set we still emit a helper (the request is typed); the |
| 134 | + # extract schema falls back to an empty descriptor (no response VO to shape it). |
| 135 | + response_vo = resolve_payload_vo(root, response_ref) if response_ref else None |
| 136 | + if response_ref is not None and response_vo is None: |
| 137 | + raise ValueError( |
| 138 | + f"{_GENERATOR_NAME}: entity {entity.name!r} prompt @responseRef " |
| 139 | + f"{response_ref!r} does not resolve to an object.value" |
| 140 | + ) |
| 141 | + |
| 142 | + entity_name = entity.name |
| 143 | + snake = _snake_case(entity_name) |
| 144 | + record_fn = f"record_{snake}" |
| 145 | + result_class = f"{entity_name}TraceResult" |
| 146 | + |
| 147 | + # Derive the parse format from the prompt's @format attr (default json) — the |
| 148 | + # SAME rule the output-parser / extractor generators use. |
| 149 | + fmt = prompt.attr(tc.TEMPLATE_ATTR_FORMAT) |
| 150 | + fmt_str = fmt if isinstance(fmt, str) else tc.TEMPLATE_FORMAT_DEFAULT |
| 151 | + |
| 152 | + fqn = entity.fqn() |
| 153 | + |
| 154 | + # Baked response-extract schema. REUSE the extract_schema_emitter exactly as the |
| 155 | + # output-parser generator does for its tolerant ``extract_lenient_*`` twin: |
| 156 | + # ``schema_literal(vo, fmt, root_name)`` emits an |
| 157 | + # ``ExtractSchema(Format.X, "<root>", [FieldSpec(...), …])`` literal, and |
| 158 | + # ``extract_map_imports(vo)`` returns the sorted/deduped ``extract_map`` accessor |
| 159 | + # names the literal's FieldSpec helpers reference. ``root_name`` is the response |
| 160 | + # VO's short name (the JSON/XML root the tolerant reader locates). When there is |
| 161 | + # no response VO (only @payloadRef set) the schema is an empty descriptor whose |
| 162 | + # extract yields ``{}``. |
| 163 | + if response_vo is not None: |
| 164 | + schema_literal = rse.schema_literal(response_vo, fmt_str, response_vo.name) |
| 165 | + helpers = rse.extract_map_imports(response_vo) |
| 166 | + else: |
| 167 | + schema_literal = f'ExtractSchema({_format_enum(fmt_str)}, "response", [])' |
| 168 | + helpers = [] |
| 169 | + |
| 170 | + request_doc = payload_ref if payload_ref else "the structured request object" |
| 171 | + |
| 172 | + lines: list[str] = [ |
| 173 | + generated_header(entity_name, fqn), |
| 174 | + "from __future__ import annotations", |
| 175 | + "", |
| 176 | + "import dataclasses", |
| 177 | + "from dataclasses import dataclass", |
| 178 | + "", |
| 179 | + "from metaobjects.render.extract import (", |
| 180 | + " ExtractSchema,", |
| 181 | + " FieldKind,", |
| 182 | + " FieldSpec,", |
| 183 | + " Format,", |
| 184 | + " extract,", |
| 185 | + ")", |
| 186 | + ] |
| 187 | + if helpers: |
| 188 | + lines.append("from metaobjects.render.extract.extract_map import (") |
| 189 | + for h in helpers: |
| 190 | + lines.append(f" {h},") |
| 191 | + lines.append(")") |
| 192 | + lines.extend( |
| 193 | + [ |
| 194 | + "from metaobjects.runtime import (", |
| 195 | + " LlmCallInput,", |
| 196 | + " LlmCallRecorder,", |
| 197 | + " LlmCallRow,", |
| 198 | + " STATUS_ERROR,", |
| 199 | + " STATUS_OK,", |
| 200 | + " build_llm_call_row,", |
| 201 | + " persist_llm_call_row,", |
| 202 | + ")", |
| 203 | + "", |
| 204 | + "from collections.abc import Callable", |
| 205 | + "", |
| 206 | + "", |
| 207 | + "# AI-trace baked response-extract descriptor — the format/root/field shape", |
| 208 | + "# the tolerant parser repairs the model's raw response text against.", |
| 209 | + f"_RESPONSE_SCHEMA: ExtractSchema = {schema_literal}", |
| 210 | + "", |
| 211 | + "", |
| 212 | + "@dataclass(frozen=True, slots=True)", |
| 213 | + f"class {result_class}:", |
| 214 | + f' """Typed result of ``{record_fn}``: the derived call outcome plus the', |
| 215 | + " best-effort extracted response VO.", |
| 216 | + "", |
| 217 | + " * ``status`` — ``STATUS_OK`` | ``STATUS_ERROR`` (a lost ``@required``", |
| 218 | + " response field → ``STATUS_ERROR``).", |
| 219 | + " * ``error_detail`` — a ``\"lost required: …\"`` summary when ``status`` is", |
| 220 | + " ``STATUS_ERROR``, else ``None``.", |
| 221 | + " * ``vo_response`` — the extracted response mirror dict (``None`` only when", |
| 222 | + ' extraction produced nothing)."""', |
| 223 | + " status: str", |
| 224 | + " error_detail: str | None", |
| 225 | + " vo_response: dict | None", |
| 226 | + "", |
| 227 | + "", |
| 228 | + f"def {record_fn}(", |
| 229 | + " recorder: LlmCallRecorder,", |
| 230 | + " input: LlmCallInput,", |
| 231 | + " redact: Callable[[LlmCallRow], LlmCallRow] | None = None,", |
| 232 | + f") -> {result_class}:", |
| 233 | + f' """Record a single ``{entity_name}`` LLM call: extract the typed response', |
| 234 | + " VO from ``input.llm_response_text`` and persist ONE trace row (the base", |
| 235 | + " envelope + raw I/O + typed voRequest/voResponse) via ``recorder`` —", |
| 236 | + " regardless of whether extraction succeeded.", |
| 237 | + "", |
| 238 | + " The tolerant ``extract`` NEVER raises; a lost ``@required`` response field", |
| 239 | + " drives ``status``/``error_detail`` (it does not abort the persist). The", |
| 240 | + f" request payload (``input.llm_request``, typed as ``{request_doc}`` at the", |
| 241 | + " call site) is threaded through as the typed ``voRequest`` column.", |
| 242 | + "", |
| 243 | + " :param recorder: the never-throwing write-side seam (Slice-1 recorder).", |
| 244 | + " :param input: the LLM call fields (envelope + raw request/response text).", |
| 245 | + " :param redact: optional row redaction applied before the single record.", |
| 246 | + ' """', |
| 247 | + " outcome = extract(input.llm_response_text, _RESPONSE_SCHEMA)", |
| 248 | + " failed = outcome.report.has_lost_required()", |
| 249 | + " status = STATUS_ERROR if failed else STATUS_OK", |
| 250 | + " error_detail = (", |
| 251 | + ' "lost required: " + ", ".join(outcome.report.lost_required())', |
| 252 | + " if failed", |
| 253 | + " else None", |
| 254 | + " )", |
| 255 | + " # Extraction owns the derived status/error_detail — fold them into the input", |
| 256 | + " # before building the base row (dataclasses.replace, leaving the original", |
| 257 | + " # input untouched).", |
| 258 | + " effective = dataclasses.replace(", |
| 259 | + " input, status=status, error_detail=error_detail", |
| 260 | + " )", |
| 261 | + " row = build_llm_call_row(effective)", |
| 262 | + " # Typed columns → native jsonb (pg8000 binds dict/list straight to jsonb).", |
| 263 | + " row[\"voResponse\"] = outcome.data", |
| 264 | + " row[\"voRequest\"] = input.llm_request", |
| 265 | + " # Persist ONCE (redact-then-record; the recorder never raises).", |
| 266 | + " persist_llm_call_row(recorder, row, redact)", |
| 267 | + f" return {result_class}(status, error_detail, outcome.data)", |
| 268 | + "", |
| 269 | + "", |
| 270 | + f'__all__ = ["{record_fn}", "{result_class}"]', |
| 271 | + "", |
| 272 | + ] |
| 273 | + ) |
| 274 | + return "\n".join(lines) |
| 275 | + |
| 276 | + |
| 277 | +def _format_enum(fmt: str) -> str: |
| 278 | + """``"xml"`` → ``Format.XML``; anything else → ``Format.JSON``. Matches the |
| 279 | + extract_schema_emitter's private ``_format_enum`` (kept local to avoid reaching |
| 280 | + into a private helper).""" |
| 281 | + return "Format.XML" if fmt.lower() == tc.TEMPLATE_FORMAT_XML else "Format.JSON" |
| 282 | + |
| 283 | + |
| 284 | +class TraceHelperGenerator: |
| 285 | + """Generator wrapping :func:`render_trace_helper`. Emits one |
| 286 | + ``record_<entity>.py`` per concrete ``object.entity`` that derives from |
| 287 | + ``LlmCallBase`` and nests a ``template.prompt`` with ``@responseRef`` / |
| 288 | + ``@payloadRef`` (skips everything else).""" |
| 289 | + |
| 290 | + name = _GENERATOR_NAME |
| 291 | + |
| 292 | + def __init__(self, *, filter: Callable[[MetaObject], bool] | None = None) -> None: |
| 293 | + # The ``filter`` arg matches the cross-generator contract. |
| 294 | + self.filter = filter |
| 295 | + |
| 296 | + def _render_module(self, entity: MetaObject, root: MetaData) -> str | None: |
| 297 | + """EXTENSION SEAM — render the whole ``record_<entity>`` module for one trace |
| 298 | + entity. Defaults to :func:`render_trace_helper`. Override to pre/post-process |
| 299 | + the emitted source or replace the render path entirely.""" |
| 300 | + return render_trace_helper(entity, root) |
| 301 | + |
| 302 | + def generate(self, ctx: GenContext) -> list[EmittedFile]: |
| 303 | + root = ctx.loaded_root |
| 304 | + if root is None: |
| 305 | + return [] |
| 306 | + files: list[EmittedFile] = [] |
| 307 | + # Stable name order — deterministic emission (matches the other generators). |
| 308 | + entities = sorted( |
| 309 | + ( |
| 310 | + c |
| 311 | + for c in root.own_children() |
| 312 | + if isinstance(c, MetaObject) and c.sub_type == OBJECT_SUBTYPE_ENTITY |
| 313 | + ), |
| 314 | + key=lambda c: c.name, |
| 315 | + ) |
| 316 | + for entity in entities: |
| 317 | + if self.filter is not None and not self.filter(entity): |
| 318 | + continue |
| 319 | + content = self._render_module(entity, root) |
| 320 | + if content is None: |
| 321 | + continue |
| 322 | + files.append( |
| 323 | + EmittedFile( |
| 324 | + path=f"record_{_snake_case(entity.name)}.py", |
| 325 | + content=ruff_format(content), |
| 326 | + ) |
| 327 | + ) |
| 328 | + return files |
| 329 | + |
| 330 | + |
| 331 | +def trace_helper_generator( |
| 332 | + *, filter: Callable[[MetaObject], bool] | None = None |
| 333 | +) -> Generator: |
| 334 | + """Factory mirroring the TS ``traceHelperFile()`` and the Java |
| 335 | + ``LlmTraceHelperGenerator``. Stable cross-port name ``trace-helper``.""" |
| 336 | + return TraceHelperGenerator(filter=filter) |
0 commit comments