Repository navigation
refactor(llm): one LLM module behind every call; old entry points delegate - #78
Conversation
…egate dokimi.llm gets one interface: LLM(model, provider, api_key, max_tokens, temperature, seed) with complete() -> Reply(text, usage, finish_reason; .truncated), ask(), chat(messages, tools=) (never response-cached) and metrics(); ask_parsed() with Unparseable; one PROVIDERS table plus aliases, resolve_model (Bedrock bare names via defaults.bedrock_profile), api_key_env, resolve_api_key and is_retryable. Behind it sits one adapter seam, (Request) -> Reply: LiteLLMAdapter in production, FakeAdapter in tests through the new fake_llm fixture. Cost keeps the _safe_completion_cost rule (an unpriced model is None, not $0). llm_call, llm_call_with_usage, chat_completion, the ChatOpenAI and ChatAnthropic shims and utils.get_llm now delegate to LLM; no caller changes. get_llm keeps its own provider table, cache keys and usage labels through a private LLM subclass until the judge moves. Response-cache keys are unchanged: tests/test_llm.py replays entries whose file names main's llm_call_with_usage and get_llm wrote. Differences from main, none in a request litellm sends: - google-vertex reads GEMINI_API_KEY on the llm_call path too (it read OPENAI_API_KEY); litellm's vertex_ai route does not use the key. - get_llm's retries and cut replies are logged as llm_call's are; chat_completion's retry log reads "LLM call failed". Part of intel#4. Signed-off-by: Nikolay Petrov <nikolay.a.petrov@intel.com>
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #78 +/- ##
==========================================
+ Coverage 89.26% 89.52% +0.26%
==========================================
Files 204 204
Lines 23626 23758 +132
Branches 3796 3800 +4
==========================================
+ Hits 21089 21270 +181
+ Misses 1840 1803 -37
+ Partials 697 685 -12
🚀 New features to boost your workflow:
|
antonwolfy
left a comment
There was a problem hiding this comment.
Review at a367770. The delegation is faithful: I found no behavioural difference from main on any path real litellm responses reach. The problems are in the tests. They do not pin the delegates' argument forwarding, the retry schedule or the stream path, so a refactor could lose those without a failure. Three test gaps are worth closing before merge, and there are four nits.
What I ran, on a Linux review host at a367770:
make ci-local: 18 passed, 1 failed, 4 skipped. The failure is the changelog-fragment check, which the PR covers withskip-changelog. 4,759 tests passed. Changed-line coverage is 99%.- A differential run of
mainand the head against a stubbedlitellm.completion, with the transport mocked. It coveredllm_call,llm_call_with_usage, streaming,chat_completion,get_llm(...).invokeand the shims. It used 11 providers,file:keys,max_retriesof 0, 1, 3 and 5, cache hits withcache_ifanddraw, and 18 response shapes. Sent kwargs, cache file names and contents, returned text and usage, sleep delays and raised exceptions were identical, apart from the differences the PR lists. - A mutation run on the new code: 131 single-line mutants. 32 pass the relevant test files, and I re-ran 8 of them on the full suite. All 8 pass (4,758 passed each time). They are listed in the inline comments.
Checked, no issues
- Cache keys: the keys for plain, system + prefix, pinned +
max_tokens+ draw, stream and the twoget_llmcases match the onesmainwrote. Mutating any key field (draw, stream,max_tokens, temperature) fails the suite. - Retries: attempt counts and sleep delays match
mainfor all four paths. Permanent errors are raised on the first attempt. - Renamed symbols:
_build_litellm_model,_is_retryableand_stream_callare kept. Nothing outsidellm.pyandtest_llm.pystill uses_resolve_api_keyor the old_warn_if_truncatedsignature. - Schemas and persisted output: no schema,
SCHEMA_VERSIONSormkdocs.ymlchange. The persistedUsageRecordcomes from the adapter, not from themetrics()total. - Tests: offline and deterministic.
time.sleepis patched, the fixture restores the default adapter, and 84 tests pass under several seeds.
| def test_chat_completion_is_llm_chat_returning_the_raw_response(fake_llm): | ||
| raw = object() | ||
| fake_llm.push(Reply(None, UsageRecord(), response=raw)) | ||
| tools = [{"type": "function", "function": {"name": "search"}}] | ||
| out = chat_completion([{"role": "user", "content": "q"}], model="m", provider="openai", | ||
| tools=tools, api_key="k", max_tokens=0, seed=1) | ||
| assert out is raw | ||
| (req,) = fake_llm.requests | ||
| assert (req.tools, req.tool_choice, req.max_tokens, req.seed, req.api_key) == ( | ||
| tools, None, 0, 1, "k") |
There was a problem hiding this comment.
The tests don't pin what the legacy delegates forward, so any of these can be dropped and the suite stays green.
test_chat_completion_is_llm_chat_returning_the_raw_response asserts tool_choice is None, so a forwarded value is never checked. Nothing checks that llm_call_with_usage forwards max_retries or retry_delay, or that the shims forward draw. I applied these mutants to a copy of the head and ran the full suite (4,758 passed, 2 skipped):
llm_call_with_usagebuildsLLM(...)withoutmax_retries=max_retries, retry_delay=retry_delay(llm.py:1010-1011).chat_completioncallsllm.chat(messages, tools=tools)withouttool_choice=tool_choice(llm.py:1052)._Shim.invokeomitsdraw=draw(llm.py:949-953).get_llmskips_read_key(api_key)on an explicit key (utils.py:141), soget_llm(..., api_key="file:/path")sends the literal string to the provider.
All four pass. With the first one, llm_call(..., max_retries=0) would retry three times and nothing notices. The PR's evidence says the delegates keep their contracts, but only the parameters that the old stubs happened to exercise are checked.
Suggestion: extend test_llm_call_sends_what_llm_complete_sends to pass non-default max_retries, retry_delay, max_tokens, draw and stream and assert the sleep schedule and the request. In the chat_completion test, pass tool_choice="required" and assert req.tool_choice == "required". Add a draw=1 case through ChatOpenAI.invoke, and a get_llm(..., api_key=f"file:{path}") case.
| def test_transient_errors_are_retried_with_doubling_backoff(error, fake_llm, no_cache, sleeps): | ||
| fake_llm.push(error, error, "ok") | ||
| assert LLM("m.x", "amazon-bedrock").ask("p") == "ok" | ||
| assert len(fake_llm.requests) == 3 | ||
| assert sleeps == [2.0, 4.0] |
There was a problem hiding this comment.
The backoff test can't tell doubling from adding 2, and never sets retry_delay.
It runs two retries with the default delay and asserts sleeps == [2.0, 4.0]. Doubling and delay += 2 give the same two values. Mutation check, full suite, 4,758 passed: delay *= 2 -> delay += 2 (llm.py:792) passes. A third retry would sleep 6s instead of 8s. No retry test passes a non-default retry_delay, so the value is never observed.
Suggestion: parametrize with LLM(..., retry_delay=0.5, max_retries=3) and three transient errors, and assert sleeps == [0.5, 1.0, 2.0].
| def test_the_litellm_adapter_streams_and_measures_ttft(monkeypatch): | ||
| import litellm | ||
|
|
||
| def chunk(content, finish_reason=None, usage=None): | ||
| delta = SimpleNamespace(content=content) | ||
| return SimpleNamespace(choices=[SimpleNamespace(delta=delta, finish_reason=finish_reason)], | ||
| usage=usage) | ||
|
|
||
| sent = [] | ||
| chunks = [chunk("Hello "), chunk("world", finish_reason="stop"), | ||
| chunk(None, usage=SimpleNamespace(prompt_tokens=5, completion_tokens=2, | ||
| total_tokens=7))] | ||
| monkeypatch.setattr(litellm, "completion", lambda **kw: sent.append(kw) or iter(chunks)) | ||
| monkeypatch.setattr(llm_mod, "_safe_completion_cost", lambda r, m, **kw: None) | ||
| reply = LiteLLMAdapter()(Request("m", "openai", "m", [], max_tokens=0, stream=True)) | ||
| assert "max_tokens" not in sent[0] and sent[0]["stream_options"] == {"include_usage": True} | ||
| assert (reply.text, reply.finish_reason, reply.usage.total_tokens) == ("Hello world", "stop", 7) | ||
| assert reply.usage.ttft_sec is not None |
There was a problem hiding this comment.
The stream test checks only that ttft_sec is set, so the stream request and TTFT can regress unnoticed.
It sends no key, temperature or seed. Mutation check, full suite, 4,758 passed:
_streamdropsapi_key=request.api_key(llm.py:478): a streamed call loses its key.if ttft is None:->if True:(llm.py:494): TTFT is stamped at the last delta, not the first._sampling_kwargsusesif seed:instead ofif seed is not None:(llm.py:173):seed=0is dropped from both the request and the cache key. No test passesseed=0throughLLM.
All three pass.
Suggestion: run the stream with api_key="k", temperature=0.0, seed=0 and several content chunks. Assert those kwargs in sent[0] and that ttft_sec is the time of the first chunk, using a patched time.time.
| logger.warning( | ||
| f"LLM reply hit the output ceiling ({model}, max_tokens=" | ||
| f"{max_tokens or 'provider default'}, " | ||
| f"{reply.usage.completion_tokens} tokens); " | ||
| "pass a higher --max-tokens if this repeats" |
There was a problem hiding this comment.
Nit: the truncation warning advises --max-tokens, which doesn't control the get_llm calls that now reach it.
get_llm("openrouter", "openai/gpt-oss-120b", "k").invoke("q") with a reply cut at the ceiling now logs LLM reply hit the output ceiling (openai/gpt-oss-120b, max_tokens=2000, 2000 tokens); pass a higher --max-tokens if this repeats. Here the ceiling is the hardcoded 2000 of the OpenRouter route, and the judge, persona and validator paths have no such flag. main never warned on this path. The PR description lists the new log, but not that its advice doesn't apply.
Suggestion: drop the flag from the message, or keep it only when the caller passed max_tokens.
| def metrics(self) -> UsageRecord: | ||
| """The sum of every reply this object returned, cache replays included. | ||
|
|
||
| ``n_calls`` counts replies; ``cost_usd`` is ``None`` when none was priced. | ||
| """ | ||
| with self._lock: | ||
| return self._total |
There was a problem hiding this comment.
Nit: metrics() says cache replays are included, but the total loses the replayed flag and counts a replay at its stored cost.
With the response cache on, two identical complete("p") calls give one adapter request, the second reply has usage.replayed == True, and metrics() returns n_calls=2, total_tokens=30, replayed=False. The seed record is replayed=False, and UsageRecord.__add__ computes self.replayed and other.replayed, so the total is never marked. A caller that reads metrics().cost_usd as spend, as the spend cap in novelty/client.py does by skipping replayed records, will overcount on a warm cache and can't tell why.
No caller reads it yet. Suggestion: say in the docstring that replays are summed at their stored cost, or exclude them from cost_usd and latency_sec.
| Usage: | ||
| from dokimi.llm import llm_call | ||
| text = llm_call("Your prompt", model="gpt-6.1-sol", provider="openai") | ||
| """Every LLM call in the framework goes through this module. |
There was a problem hiding this comment.
Nit: the module docstring describes the end state of the series, not this step.
- Line 4, "Every LLM call in the framework goes through this module": the key and route tables in
utils.get_llm,commands/benchmark.py(_credential_env_var) andeval/ragas_eval.py(_resolve_key) are still separate. The judge, persona and validator redraw loops don't useask_parsed. The same applies to "the one model-name map" and "the one key-variable map" at lines 23-24, and toarchitecture.md's "every LLM call". - Lines 46-47: "A live reply cut at the output ceiling ... is logged as a warning" holds for
complete/ask, butchat()never calls_warn_if_truncated. Achat()reply withfinish_reason="length"logs nothing. That matchesmain'schat_completion, but the docstring says otherwise.
Suggestion: word it as "the entry point for" and name what hasn't moved yet, and add "complete/ask" to the truncation sentence.
| def _stream_call(completion: Any, litellm_model: str, messages: List[Dict[str, Any]], | ||
| api_key: Optional[str], model: str, provider: str, | ||
| max_tokens: Optional[int] = None, temperature: Optional[float] = None, | ||
| seed: Optional[int] = None) -> Tuple[str, UsageRecord]: | ||
| """One streamed completion as ``(text, usage)``; the adapter's streaming path.""" | ||
| request = Request(model=model, provider=provider, litellm_model=litellm_model, | ||
| messages=messages, api_key=api_key, max_tokens=max_tokens, | ||
| temperature=temperature, seed=seed, stream=True) | ||
| reply = _stream(completion, request) | ||
| _warn_if_truncated(reply, model, max_tokens) | ||
| return reply.text or "", reply.usage |
There was a problem hiding this comment.
Nit: _stream_call is dead in production and its docstring calls it "the adapter's streaming path".
The adapter uses _stream (llm.py:519). Its only caller is tests/test_metrics_usage.py:226. It's the one place the old streaming test still runs, while the new code in _stream is covered by a weaker test (see the stream comment on test_llm.py).
Suggestion: point test_stream_call_measures_ttft at LiteLLMAdapter and delete _stream_call, or mark it as kept only for that test.
Signed-off-by: Nikolay Petrov <nikolay.a.petrov@intel.com>
Review follow-ups: - tests: retry_delay doubling ([0.5, 1.0, 2.0]); llm_call/_with_usage forward max_retries, retry_delay, max_tokens, stream; draw through the shims and llm_call; chat_completion forwards tool_choice; get_llm reads an explicit file: key; a streamed LLM call sends key/temperature/seed=0 and stamps TTFT at the first chunk. - remove the unused _stream_call; its TTFT test drives LiteLLMAdapter. - truncation warning no longer points at a --max-tokens flag. - metrics() says replays are summed at their stored cost. - module docstring names what has not moved to llm.py yet, and that chat() does not log truncation. Signed-off-by: Nikolay Petrov <nikolay.a.petrov@intel.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RW8mAicb3ibXdzHMV83wDn
79de0b1 to
106d65f
Compare
On main chat_completion returned the raw response, so the tool probe priced a reply with an empty choices list and its spend cap counted it. The adapter read choices[0] first and the usage was lost. Read the choice defensively; complete() still raises IndexError and caches nothing. Signed-off-by: Alexander Rybkin <alexander.rybkin@intel.com>
Empty tool list, chat_completion retry budget, get_llm's four attempts, a content-less reply stored as null, dict tool calls without arguments, file:~ keys and the openai-codex route, each as on main. Signed-off-by: Alexander Rybkin <alexander.rybkin@intel.com>
|
Ran a differential at the
The description's "same suite with no test edits" and " |
xaleryb
left a comment
There was a problem hiding this comment.
Approving at 7b93ffa: requests, cache keys, retries, usage and cost match main at the litellm.completion boundary over 2677 scenarios and live Bedrock calls, no key leaks, the billed-reply-without-choices case is priced again by its caller, the delegate behaviours the suite left open are pinned, and all checks pass. Per issue #4, #80 should land right after this.
Resolve conflicts with intel#78/intel#69/intel#43: - arms: keep main's questions.v2 loader; keep the id-type check before the run. - cross-judge: keep main's judge_materials gate; stored scores and the pairwise/checklist gate read records without a join, and a RecordError exits 1 naming the record (tests updated from exit 2). - arm_runner: digest lookup uses Record.question_id. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RW8mAicb3ibXdzHMV83wDn Signed-off-by: Nikolay Petrov <nikolay.a.petrov@intel.com>
Part of #4.
Summary
Step 1 of the single LLM-call module: the module, both adapters and their tests. Every old entry point now delegates to
LLM. No caller changes.Next to it:
PROVIDERS(canonical name → key variable) plusPROVIDER_ALIASES,resolve_model(bare Bedrock names go throughdefaults.bedrock_profile, the mapping isn't copied),api_key_env,resolve_api_key, andask_parsed(ask, prompt, parse, reprompt=).ask_parseduses a new draw per attempt, caches only replies that parse, and raisesUnparseable(ValueError). Cost still goes through_safe_completion_cost, so an unpriced model costsNone, not$0. The truncation warning is now a check onReply.truncated, made once inLLM.completefor every live reply.Evidence
main. After: the same suite with no test edits, which keeps the old entry points' contracts. That includes thelitellm.completionstubs behindtest_judge_transport_retry,test_judge_response_cache,test_sampling_forwarding,test_get_llm_bedrockand the truncation-log tests intest_novelty_silent_defaults.tests/test_llm.pyalso checks each delegate directly againstLLM:llm_call_with_usageandutils.get_llmagainst a stubbedlitellm.completionand recorded the six cache file names they wrote: plain; system + cache_prefix; pinned temperature/seed + max_tokens + draw 2; stream; get_llm on Bedrock with a bare name; get_llm on OpenRouter, pinned, draw 1.max_tokens: Noneto0. Three of the six cases failed.tests/test_llm.py, 84 tests) cover routes and keys per provider, retry and backoff ([2.0, 4.0]), sampling sent only when set, draws,cache_ifon store and replay,chatnever cached, the truncation log on a live call but not on a replay,metrics()stayingNonewhen unpriced,ask_parsed, and theLiteLLMAdapterrequest, tool-call and stream shapes.ruff,mypy(210 files,llm.pyincluded, no new exemption),pytestwith--randomly-seed=1234and98765(4758 passed each), diff-cover 99% (354 lines; the 1 missing line is unchangedextract_json_arraycode),make ci-local --offline: everything passes except the changelog fragment (see below).skip-changeloglabel.Merge Danger
Door: two-way. A revert restores the two old transports. Nothing persisted changes: cache keys are byte-identical and no schema is touched.
Blast Radius: internal
LLM. Requests sent to litellm are unchanged. Two small differences:google-vertexon thellm_callpath readsGEMINI_API_KEY, asget_llmand thegoogleandgeminispellings already did. Before, it readOPENAI_API_KEY. litellm'svertex_airoute doesn't read the key.get_llm's retries and cut replies are now logged the wayllm_call's are.chat_completion's retry log line now reads "LLM call failed".chat_completionnow builds aReplywith usage for every agent turn. The agent loop still reads the raw response.skip-changelog. This is an internal refactor behind delegates, and no request, score, artifact or flag changes.get_llmontoLLM+ask_parsed. That step decidesget_llm's divergences: the$ANTHROPIC_DEFAULT_SONNET_MODELfallback for bare Bedrock names, OpenRouter's hiddenmax_tokens=2000, and its stream-less cache keys. Those change behaviour and cache keys, so that PR gets a changelog entry.LLM.complete/LLM.chat, andchat_completionand the shims go.commands/benchmark.pyandeval/ragas_eval.pyreadapi_key_env, and the CLI provider lists readPROVIDERS.Checklist
pytestpassesruff,mypy, and schema checks pass.github/.licenserc.yaml)🤖 Generated with Claude Code