Skip to content

refactor(themefinder): migrate LLM layer to Pydantic AI - #1428

Open
jjuritzno10 wants to merge 6 commits into
mainfrom
themefinder/refactor-pydantic-ai
Open

refactor(themefinder): migrate LLM layer to Pydantic AI#1428
jjuritzno10 wants to merge 6 commits into
mainfrom
themefinder/refactor-pydantic-ai

Conversation

@jjuritzno10

@jjuritzno10 jjuritzno10 commented Jun 19, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Context

Every themefinder LLM call went through a single OpenAI-SDK class (OpenAILLM), making provider switching hard. This replaces it with one Pydantic AI-backed LLM where the model is a single configurable value — LLM("openai:gpt-4o") vs LLM("anthropic:claude-..."), or the LiteLLM gateway via base_url+api_key.

Changes proposed in this pull request

  • llm.py: OpenAILLM → Pydantic AI LLM (same ainvoke/invoke interface). Direct mode (provider-prefixed string), gateway mode (OpenAIChatModel+OpenAIProvider), or a Model instance. Structured output kept on OpenAI-native behaviour via NativeOutput.
  • llm_batch_processor.py: provider-agnostic error handling (ModelHTTPError/UnexpectedModelBehavior); tiktoken fallback for non-OpenAI model names.
  • tasks.py: optional stage_llms on find_themes for per-stage models (backward compatible).
  • Consumers: pipeline scripts + evals construct LLM(...).
  • Tests: unit coverage (FunctionModel/TestModel) + live gateway tests on gambling_XS data (skip without creds).
  • Deps: add pydantic-ai-slim[openai,anthropic,google].

Guidance to review

make test (themefinder) for unit tests. Live tests run automatically when LLM_GATEWAY_URL/LITELLM_CONSULT_OPENAI_API_KEY are set — currently 8 pass.

Outstanding (future PRs)

  • Clean up and standardise environment variables and secrets.
  • Standardise the model-selection registry.
  • Improve consistency of langfuse logging.

Things to check

  • I have added any new ENV vars in all deployed environments and updated the .env.test files in the repo

🤖 Generated with Claude Code

@jjuritzno10
jjuritzno10 changed the base branch from main to jordanjuritz/pro-274-pr-1-core-monorepo-migration-atomic June 19, 2026 16:01
@jjuritzno10

jjuritzno10 commented Jun 19, 2026

Copy link
Copy Markdown
Collaborator Author

Temporarily based on jordanjuritz/pro-274-pr-1-core-monorepo-migration-atomic (#1332) to keep the diff scoped to the Pydantic AI work. TODO: retarget to main once #1332 merges (gh pr edit 1428 --base main).

@jjuritzno10
jjuritzno10 requested a review from Nazehs June 25, 2026 14:19
@Nazehs
Nazehs force-pushed the jordanjuritz/pro-274-pr-1-core-monorepo-migration-atomic branch 3 times, most recently from 126a158 to bcf500a Compare July 6, 2026 14:48
@Nazehs
Nazehs force-pushed the jordanjuritz/pro-274-pr-1-core-monorepo-migration-atomic branch 2 times, most recently from 7c1f5ee to d7ec478 Compare July 13, 2026 09:10
Base automatically changed from jordanjuritz/pro-274-pr-1-core-monorepo-migration-atomic to main July 13, 2026 10:28
@github-actions

github-actions Bot commented Jul 14, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Claude finished @Nazehs's task in 2m 9s —— View job

PR Review: refactor(themefinder): migrate LLM layer to Pydantic AI

  • Gather context
  • Review key changed files
  • Post feedback

Clean refactor — replaces the bespoke OpenAILLM with a thin Pydantic AI wrapper, keeps the ainvoke/invoke interface intact, and adds good test coverage via FunctionModel.

Two things worth addressing:

🔵 CONSIDER (llm.py:66, tasks.py): The system_prompt is still injected as a {system_prompt} user-turn placeholder rather than set as the agent's actual system message. Pydantic AI supports Agent(system_prompt=...) natively — promoting it to a real system turn is more semantically correct and typically more effective for instruction-following models.

⚪️ NIT (llm_batch_processor.py:98, 156): Optional[str] is still used in two function signatures; the rest of the PR uses str | None (Python 3.10+ style, consistent with the codebase). Fix this →

@Nazehs
Nazehs force-pushed the themefinder/refactor-pydantic-ai branch from c60c805 to 4917ff4 Compare July 15, 2026 10:57
@github-actions

github-actions Bot commented Jul 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Claude finished @Nazehs's task in 1m 1s —— View job

  • Gather context
  • Review key changed files
  • Post feedback

PR Review: refactor(themefinder): migrate LLM layer to Pydantic AI

Clean migration. The LLM wrapper is minimal, the Agent is correctly constructed once and reused across concurrent calls, and the stage_llms override validation is a solid addition.

Two carry-forwards from the previous review and one new one:

🔵 CONSIDER (llm.py:66): Agent(self._model) still has no system prompt — system_prompt is baked into user-turn text via template {system_prompt} substitution. Promoting it to Agent(self._model, system_prompt=...) is more semantically correct for instruction-following models, though it would require threading the value to __init__ (or a per-call override). Low urgency, future-PR material.

⚪️ NIT (llm_batch_processor.py:98, 156): Optional[str] in partition_dataframe and batch_task_input_df signatures still uses the legacy import style; rest of the file uses str | None. Fix this →

⚪️ NIT (llm_batch_processor.py:334): calculate_string_token_length(input_text: str, model: str = None) — the default None makes this effectively str | None without declaring it. Should be model: str | None = None to match the codebase style. Fix this →

@github-actions

github-actions Bot commented Jul 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Claude finished @Nazehs's task in 48s —— View job

  • Gather context
  • Review key changed files
  • Post feedback

PR Review: refactor(themefinder): migrate LLM layer to Pydantic AI

NITs from prior rounds are resolved. Only one carry-forward remains:

🔵 CONSIDER (llm.py:66): Agent(self._model) still has no system prompt — system_prompt is baked into user-turn text via {system_prompt} substitution rather than set as a real Agent(self._model, system_prompt=...) turn. Low urgency, future-PR material.

@github-actions

github-actions Bot commented Jul 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Claude finished @Nazehs's task in 48s —— View job

  • Gather context
  • Review key changed files
  • Post feedback

PR Review: refactor(themefinder): migrate LLM layer to Pydantic AI

Clean and minimal. All prior NITs resolved.

🔵 CONSIDER (llm.py:66, prompts.py:7,42,74...): system_prompt is still injected as the first line of every user-turn template ({system_prompt}\n...) rather than set via Agent(self._model, system_prompt=...). Low urgency — future-PR material — but worth promoting for instruction-following correctness.

@github-actions

github-actions Bot commented Jul 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Claude finished @Nazehs's task in 50s —— View job

  • Gather context
  • Review key changed files
  • Post feedback

PR Review: refactor(themefinder): migrate LLM layer to Pydantic AI

Clean and minimal — prior NITs resolved. One carry-forward:

🔵 CONSIDER (llm.py:66, prompts.py): Agent(self._model) still has no system prompt — {system_prompt} is injected as the first line of every user-turn template rather than set via Agent(self._model, system_prompt=...). Low urgency, future-PR material.

@github-actions

github-actions Bot commented Jul 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Claude finished @Nazehs's task in 29s —— View job

  • Gather context
  • Review key changed files
  • Post feedback

PR Review: refactor(themefinder): migrate LLM layer to Pydantic AI

Clean and minimal — prior NITs all resolved.

🔵 CONSIDER (llm.py:42-46): _resolve_model will silently ignore a passed Model instance when base_url or api_key is also provided — it would wrap it into a new OpenAIChatModel, losing the instance. This is probably unreachable today but worth a guard or type narrowing.

🔵 CONSIDER (llm.py:66, prompts.py): Agent(self._model) still has no system prompt — {system_prompt} is injected as the first line of every user-turn template rather than set via Agent(self._model, system_prompt=...). Low urgency, future-PR material.

@github-actions

github-actions Bot commented Jul 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Claude finished @Nazehs's task in 39s —— View job

  • Gather context
  • Review key changed files
  • Post feedback

PR Review: refactor(themefinder): migrate LLM layer to Pydantic AI

Clean migration with all prior NITs resolved. Two carry-forwards remain as future-PR material:

🔵 CONSIDER (llm.py:42-45): _resolve_model passes model directly into OpenAIChatModel(model, ...) when base_url/api_key are set — if model is already a Model instance this will TypeError at runtime. A quick isinstance(model, str) guard or a type-narrowing check would close this off. Fix this →

🔵 CONSIDER (llm.py:66, prompts.py): Agent(self._model) has no system prompt — {system_prompt} is injected as the first line of every user-turn template rather than promoted to a real Agent(self._model, system_prompt=...) turn. Low urgency, future-PR material.

@github-actions

github-actions Bot commented Jul 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Claude finished @Nazehs's task in 34s —— View job

  • Gather context
  • Review key changed files
  • Post feedback

PR Review: refactor(themefinder): migrate LLM layer to Pydantic AI

Clean and minimal — all prior NITs resolved, _resolve_model guard is in place.

🔵 CONSIDER (llm.py:66, prompts.py): Agent(self._model) still has no system prompt — {system_prompt} is injected into every user-turn template rather than promoted to Agent(self._model, system_prompt=...). Low urgency, future-PR material.

@Nazehs
refactor(themefinder): migrate LLM layer to Pydantic AI
5a1d282 
Replace the OpenAI-specific OpenAILLM with a single Pydantic AI-backed LLM
class that supports provider-prefixed models (openai:/anthropic:/google:) and
OpenAI-compatible gateway routing via base_url/api_key. Wire per-stage model
routing through tasks.py and the batch processor. Add pydantic-ai-slim
(openai/anthropic/google extras) as a themefinder dependency.

_resolve_model narrows on str first, so a ready-made Model instance is always
used as-is regardless of gateway args (avoids re-wrapping it in
OpenAIChatModel). Batch-processor signatures use PEP 604 str | None throughout.
@Nazehs
Nazehs force-pushed the themefinder/refactor-pydantic-ai branch from d8954ef to 8f40062 Compare July 15, 2026 11:57
@github-actions

github-actions Bot commented Jul 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Claude finished @Nazehs's task in 43s —— View job

  • Gather context
  • Review key changed files
  • Post feedback

PR Review: refactor(themefinder): migrate LLM layer to Pydantic AI

Clean and merge-ready. All prior NITs resolved.

⚪️ NIT (llm.py:67-68): If both model_settings and temperature are passed, temperature is silently ignored — model_settings or (...) short-circuits. A guard or docstring note would prevent silent misconfiguration.

🔵 CONSIDER (llm.py:70, prompts.py): Agent(self._model) still has no system prompt — {system_prompt} is injected into every user-turn template. Low urgency, future-PR material.

Nazehs added 5 commits July 15, 2026 13:06
@Nazehs
test(themefinder): cover Pydantic AI LLM wrapper and per-stage routing
309a5fc 
Add unit tests for the new LLM wrapper (gateway vs direct-provider resolution,
Model-instance passthrough incl. with gateway args, structured-output modes)
and live-inference tests gated on gateway credentials. Extend task and
batch-processor tests for the Pydantic AI code path.
@Nazehs
refactor(themefinder/evals): migrate eval + pipeline scripts off Open…
1f43ace 
…AILLM

Point benchmark, metrics, and the four eval scripts at the new LLM class, and
update the sign-off/mapping pipeline scripts to match. benchmark.py gains
--dataset/--provider/--models/--runs for multi-model, multi-repeat runs.

Also fix pre-existing eval breakage the revived Pydantic AI path exposed:
- align Azure deployment names with the gateway catalog (gpt-4.1, gpt-4o-sweden)
  instead of stale dated names that 400 as 'Invalid model name';
- reuse the resolved task LLM for generation metrics in the local fallback
  rather than building a judge from the AUTO_EVAL_4_1_SWEDEN_DEPLOYMENT secret;
- parse the prompt's ternary decision format (STRONG/PARTIAL/NO) via
  evaluators._parse_evaluation_response instead of treating dicts as scores.
@Nazehs
fix(themefinder/evals): commit gambling_XS reference data for local e…
d2039df 
…val fallback

The eval's local-data path reads inputs/question_part_*/question.json and
outputs/mapping/2025-07-22/question_part_*/themes.json, but a blanket *.json
ignore rule (for GCP creds) hid every .json under the dataset, so only the
.jsonl files were tracked. When Langfuse is unavailable the eval falls back to
local data and failed with FileNotFoundError in CI.

Re-include the reference dataset's .json via a targeted .gitignore negation and
commit the four missing files, copied verbatim from upstream themefinder
(commit 09918eb2, per ADR-0009; sibling .jsonl verified byte-identical).
@Nazehs
ci(themefinder-eval): add configurable runs input for multi-repeat evals
90ce3ab 
Expose a 'runs' workflow input and pass it through to benchmark.py as --runs,
raise the job timeout to 180m to accommodate repeated runs.
@Nazehs
chore(themefinder): ignore dev dirs in images; lock pydantic-ai deps
6a9eae9 
Exclude themefinder dev-only directories (evals, docs, tests, caches) and
sql.log from Docker build context; regenerate uv.lock to resolve
pydantic-ai-slim and its transitive dependencies.
@Nazehs
Nazehs force-pushed the themefinder/refactor-pydantic-ai branch from 8f40062 to 6a9eae9 Compare July 15, 2026 12:07
@github-actions

github-actions Bot commented Jul 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Claude finished @Nazehs's task in 37s —— View job

  • Gather context
  • Review key changed files
  • Post feedback

PR Review: refactor(themefinder): migrate LLM layer to Pydantic AI

Clean migration — all prior NITs resolved. Two items remain:

⚪️ NIT (llm.py:67-68): model_settings or (ModelSettings(temperature=...) if ...) short-circuits when model_settings is truthy, silently ignoring a passed temperature. Should raise if both are provided. Fix this →

🔵 CONSIDER (llm.py:70, prompts.py): Agent(self._model) has no system prompt — {system_prompt} is injected as the first line of every user-turn template rather than set via Agent(self._model, system_prompt=...). Low urgency, future-PR material.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Nazehs Nazehs Awaiting requested review from Nazehs

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@jjuritzno10 @Nazehs