refactor(ai): back assert_response_judged with a JudgeAgent#185
Conversation
Replace RecordingAgent._judge_live's hand-rolled init_chat_model()/invoke() call with JudgeAgent, an Agent subclass. The judge now goes through the same provider/model resolution pipeline as any other agent, and is fakeable via JudgeAgent.fake() and replayable via JudgeAgent.record() instead of only being testable by mocking langchain directly. assert_response_judged() is now async (the judge call is a real Agent.prompt() under the hood) and accepts an optional provider kwarg, forwarded through to JudgeAgent and folded into the cached verdict's cache key.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
tmgbedu
left a comment
There was a problem hiding this comment.
Code Review Verdict: APPROVE ✅ (posting as comment — GH blocks self-approval)
Clean, well-scoped refactor. It replaces the hand-rolled init_chat_model().invoke() judge with a JudgeAgent(Agent) that runs through the standard Agent pipeline — and as a bonus it fixes the latent type bug I flagged on #183.
✅ Async signature change — no silent breakage (the key risk)
assert_response_judged() going sync→async is dangerous for a public assert API: an un-awaited call returns a truthy coroutine and the assertion silently never runs (false green). I grepped every caller in the repo — all 9 are awaited: the example test_router_agent.py and 8 sites in test_agent_record_fluent.py. No un-awaited caller exists anywhere (src, tests, example). The internal _judge/_judge_live chain is consistently async. ✓
✅ JudgeAgent reuses the standard Agent pipeline
JudgeAgent.judge() → await self.prompt(...) → normal _build_messages/_build_model (Ai().get_model_for) model+provider resolution. So it's fakeable via JudgeAgent.fake([...]) and replayable via JudgeAgent.record() by class name, exactly like any agent (covered by test_judge_is_usable_via_the_record_fluent_dsl). Constructor correctly overrides model/provider only when provided. _parse_verdict moved cleanly to judge.py (no duplicate left in testing.py). provider is folded into the verdict cache key (judge_provider), so verdicts don't collide across providers. ✓
✅ Fixes my earlier #183 Blocker #2
The old _judge_live passed chat_model.invoke(...).content (str | list[...]) into _parse_verdict(raw: str) — a reportArgumentType error + latent TypeError on list-style content. Now the judge returns an AgentResponse (whose .content is typed str), so _parse_verdict(str) is type-clean. pyright drops 660 (dev baseline) → 659 here; 0 errors in judge.py/testing.py.
✅ Verification
- ruff check + format clean; no
type: ignore/pyright: ignoreadded. - Full suite 1867 passed / 7 skipped locally (excluding env-only Postgres). Your reported 1874 differs by 7 — consistent with optional/API-key-gated tests that are skipped in my env; no failures either way.
- New
test_judge_agent.py(7 tests): verdict parse pass/fail, prose-tolerant JSON extraction, prompt contents, constructor model/provider, fluent-DSL integration. Good direct coverage.
Notes (non-blocking)
- Unplanned scope (as you flagged): this is a follow-on refactor, but a net improvement — removes hand-rolled langchain code, fixes a real latent bug, unifies the judge into the standard pipeline. No objection to it landing on merit.
- Mixed-async DSL ergonomics:
assert_response_judgedis nowasyncwhile the sibling asserts (assert_text_response,assert_tool_called,assert_response_time_lt) stay sync — so some asserts needawaitand some don't. Justified (only the judge does I/O) but worth a one-line docstring note so users don't forget theawait. - Carry-over (not this PR):
AgentSnapshotremains orphaned indev(from #183, merged directly with that unaddressed). Out of scope here; flagging so it's not forgotten for a future cleanup.
Not merging per instructions — approve only.
… JSON Give JudgeAgent a Verdict pydantic schema() and put the grading rubric in instructions() -- the JSON reply is now turned into a typed result through the same structured-output path any Agent gets from schema()/response.parsed. Drops the bespoke _build_prompt()/_parse_verdict() helpers.
…alls When an Agent declares schema() and has no tools, build the model with chat_model.with_structured_output(schema, include_raw=True) so the provider enforces the shape, instead of relying on prompt instructions + post-hoc JSON parsing. include_raw keeps the raw message so content/usage/cassettes still work; the Runner passes the structured result through without executing the synthetic tool call, and _to_agent_response unwraps it into response.parsed. Streaming opts out (needs raw token chunks), tools take precedence over a schema in a single call, and the fake/record paths are unchanged (they keep exercising the JSON-string parse path for deterministic replay). Also drops two stale docstrings from Ai.
When an agent declares both tools() and schema(), bind them together via bind_tools([*tools, schema]) so a single model call offers both. The model returns either a real tool call (which the Runner executes) or the schema as its structured answer (which the Runner parses into response.parsed). Schema without tools still uses with_structured_output() for enforcement. _apply_schema no longer coerces content when the agent has tools, so a tool's plain-text output is not force-parsed into the schema.
…ut branch schema() is just appended to tools() and the whole set is bound in one call. This collapses the schema-only special case: with_structured_output only added tool_choice="any" enforcement, which is redundant now that _apply_schema parses plain JSON text for tool-less agents -- so a schema-only agent gets its parsed result whether the model emits the schema tool call or replies in JSON text. Also drops the now-dead structured-dict passthrough in Runner.run.
…ng as-is
Bind tools exactly as before, then wrap the model with
with_structured_output(schema, include_raw=True) when a schema is declared.
The wrapped model returns {raw, parsed, parsing_error}; the Runner passes it
through and _to_agent_response unwraps it into response.parsed.
Reverts the schema-as-a-tool detection and the _apply_schema tools guard.
Summary
JudgeAgent(Agent)— the LLM-as-judge call now runs through the standard Agent provider/model-resolution pipeline instead ofRecordingAgent._judge_livehand-rolling its owninit_chat_model()/.invoke()call.JudgeAgentis fakeable viaJudgeAgent.fake()and replayable viaJudgeAgent.record(), same as any other agent.assert_response_judged()is nowasync(the judge call is a realAgent.prompt()under the hood) and takes an optionalproviderkwarg, forwarded toJudgeAgentand folded into the cached verdict's cache key.test_router_agent.pyfor theawait+providerchange.Test plan
uv run pytest --ignore=tests/masoniteorm/postgres— 1874 passed, 7 skippeduv run ruff check ./uv run ruff format --check .— cleantests/ai/test_judge_agent.pycoversJudgeAgentdirectly (verdict parsing, prose-tolerant JSON extraction, constructor model/provider, and usable via.record())