Repository navigation
feat(OMN-2318): integrate SPI 0.9.0 LLM cost tracking contracts - #345
Merged
Merged
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR extends Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~40 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@src/omnibase_infra/nodes/effects/models/converter_llm_usage_to_contract.py`:
- Around line 1-30: Rename the module file from
converter_llm_usage_to_contract.py to an approved pattern (e.g.,
adapter_llm_usage_to_contract.py or model_llm_usage_to_contract.py) and update
every import that references the old module name; specifically update imports
that reference symbols exported from the module such as ModelLlmUsage,
ContractLlmCallMetrics, ContractLlmUsageRaw, ContractLlmUsageNormalized and
ContractEnumUsageSource where the current import path points to
converter_llm_usage_to_contract, as well as any package __init__.py exports,
tests, and CI/deployment references; ensure the module-level docstring and
internal references remain unchanged and run tests to confirm no import errors.
- Around line 90-133: The function to_call_metrics should validate that model_id
is a non-empty string before building the ContractLlmCallMetrics; add an early
guard in to_call_metrics that checks model_id (e.g., if not isinstance(model_id,
str) or not model_id.strip():) and raise a ValueError with a clear message like
"model_id must be a non-empty string" so invalid/empty model identifiers are
rejected before calling ContractLlmCallMetrics.
In
`@src/omnibase_infra/nodes/node_llm_embedding_effect/handlers/handler_embedding_openai_compatible.py`:
- Line 43: In handler_embedding_openai_compatible.py update the usage-extraction
logic (the function that builds the usage object / sets ContractEnumUsageSource
to API) so that when prompt_tokens is missing or not an int but total_tokens is
present and valid you fall back to using total_tokens for prompt_tokens (and for
any other missing token fields as appropriate) rather than reporting zero;
ensure you cast/validate total_tokens as int before assigning, keep
usage_source=API, and apply the same fallback change in the second
usage-extraction block around lines 220-258 so both code paths preserve token
counts when only total_tokens is provided.
Comment on lines
+1
to
+30
| # SPDX-License-Identifier: MIT | ||
| # Copyright (c) 2025 OmniNode Team | ||
| """Converter bridging ModelLlmUsage to SPI LLM cost tracking contracts. | ||
|
|
||
| This module provides pure functions that translate between the infra-layer | ||
| ``ModelLlmUsage`` value object and the SPI measurement contracts: | ||
|
|
||
| - ``ContractLlmCallMetrics`` | ||
| - ``ContractLlmUsageRaw`` | ||
| - ``ContractLlmUsageNormalized`` | ||
| - ``ContractEnumUsageSource`` | ||
|
|
||
| All functions are stateless and produce frozen Pydantic models suitable for | ||
| downstream measurement pipeline ingestion. | ||
|
|
||
| Related: | ||
| - OMN-2318: Integrate SPI 0.9.0 LLM cost tracking contracts | ||
| - ModelLlmUsage: Source infra-layer usage model | ||
| - ContractLlmCallMetrics: Target SPI per-call metrics contract | ||
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| from omnibase_infra.nodes.effects.models.model_llm_usage import ModelLlmUsage | ||
| from omnibase_spi.contracts.measurement import ( | ||
| ContractEnumUsageSource, | ||
| ContractLlmCallMetrics, | ||
| ContractLlmUsageNormalized, | ||
| ContractLlmUsageRaw, | ||
| ) |
There was a problem hiding this comment.
Rename module to match required Python file naming patterns.
The current filename converter_llm_usage_to_contract.py does not match the allowed prefixes. Please rename it (and update imports) to an approved pattern (e.g., adapter_llm_usage_to_contract.py or model_llm_usage_to_contract.py, depending on intent).
As per coding guidelines, File naming must follow pattern: model_.py, adapter_.py, dispatcher_.py, enum_.py, mixin_.py, protocol_.py, service_.py, store_.py, validator_.py, registry_infra_.py.
Also applies to: 136-140
🤖 Prompt for AI Agents
In `@src/omnibase_infra/nodes/effects/models/converter_llm_usage_to_contract.py`
around lines 1 - 30, Rename the module file from
converter_llm_usage_to_contract.py to an approved pattern (e.g.,
adapter_llm_usage_to_contract.py or model_llm_usage_to_contract.py) and update
every import that references the old module name; specifically update imports
that reference symbols exported from the module such as ModelLlmUsage,
ContractLlmCallMetrics, ContractLlmUsageRaw, ContractLlmUsageNormalized and
ContractEnumUsageSource where the current import path points to
converter_llm_usage_to_contract, as well as any package __init__.py exports,
tests, and CI/deployment references; ensure the module-level docstring and
internal references remain unchanged and run tests to confirm no import errors.
| from omnibase_infra.nodes.node_llm_embedding_effect.models.model_llm_embedding_response import ( | ||
| ModelLlmEmbeddingResponse, | ||
| ) | ||
| from omnibase_spi.contracts.measurement import ContractEnumUsageSource |
There was a problem hiding this comment.
Preserve token counts when only total_tokens is provided.
If prompt_tokens is missing or non-int but total_tokens is valid, the current logic reports zero tokens while still setting usage_source=API. Consider falling back to total_tokens for embeddings to avoid under-reporting.
Proposed fix
- prompt_tokens = usage_raw.get("prompt_tokens", 0)
- if not isinstance(prompt_tokens, int):
- prompt_tokens = 0
+ prompt_tokens = usage_raw.get("prompt_tokens", 0)
+ if not isinstance(prompt_tokens, int):
+ prompt_tokens = 0
- total_tokens = usage_raw.get("total_tokens", 0)
- if not isinstance(total_tokens, int):
- total_tokens = 0
+ total_tokens = usage_raw.get("total_tokens", 0)
+ if not isinstance(total_tokens, int):
+ total_tokens = 0
+
+ if prompt_tokens == 0 and total_tokens > 0:
+ prompt_tokens = total_tokensAlso applies to: 220-258
🤖 Prompt for AI Agents
In
`@src/omnibase_infra/nodes/node_llm_embedding_effect/handlers/handler_embedding_openai_compatible.py`
at line 43, In handler_embedding_openai_compatible.py update the
usage-extraction logic (the function that builds the usage object / sets
ContractEnumUsageSource to API) so that when prompt_tokens is missing or not an
int but total_tokens is present and valid you fall back to using total_tokens
for prompt_tokens (and for any other missing token fields as appropriate) rather
than reporting zero; ensure you cast/validate total_tokens as int before
assigning, keep usage_source=API, and apply the same fallback change in the
second usage-extraction block around lines 220-258 so both code paths preserve
token counts when only total_tokens is provided.
…infra Bridge ModelLlmUsage to the new SPI measurement contracts (ContractLlmCallMetrics, ContractLlmUsageRaw, ContractLlmUsageNormalized) introduced in omnibase_spi 0.9.0. - Add usage_source (ContractEnumUsageSource) and raw_provider_usage fields to ModelLlmUsage for provenance tracking (API/ESTIMATED/MISSING) - Create converter module (to_call_metrics, to_usage_normalized, to_usage_raw) that maps infra-layer ModelLlmUsage to SPI contract models - Update all 4 LLM handlers to populate provenance and raw usage data: handler_llm_openai_compatible, handler_llm_ollama, handler_embedding_openai_compatible, handler_embedding_ollama - Update dependencies: omnibase-core ^0.18.0, omnibase-spi ^0.9.0 - Add 51 new unit tests for provenance fields and converters
…cases Major: - Fix has_usage always True in Ollama inference handler: add > 0 checks so usage_source correctly falls back to MISSING when Ollama omits usage Minor: - Guard against IndexError on empty embeddings in both OpenAI and Ollama embedding handlers - Fix tokens_total fallback in converter to compute from input + output instead of defaulting to 0 - Change latency_ms column from INTEGER to NUMERIC(10,2) to preserve sub-millisecond precision matching SPI contract type
…onstraints - OpenAI inference/embedding handlers now check has_usage before setting usage_source=API, matching the Ollama handler pattern (all-zero tokens → MISSING) - Migration 031: pg_column_size → octet_length for predictable 64KB limit - Migration 031: session_id nullable until write path integrated - Migration 031: latency_ms nullable for failed LLM calls
Remove dead total-token fallback in to_usage_normalized and to_call_metrics; ModelLlmUsage.model_validator guarantees tokens_total is always populated.
- Use `or 0` fallback for `tokens_total` in converter to satisfy mypy when field type is `int | None` (model_validator guarantees non-None at runtime, but static analysis cannot see that) - Remove duplicated model_validator guarantee comments - Simplify Ollama handler `has_usage` isinstance to `int` only (drop redundant `float`) matching embedding handler pattern
jonahgabriel
force-pushed
the
feat/integrate-spi-llm-contracts
branch
from
February 16, 2026 18:20
874ddfe to
444f42e
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In
`@src/omnibase_infra/nodes/node_llm_inference_effect/handlers/handler_llm_ollama.py`:
- Around line 416-443: The has_usage check in handler_llm_ollama.py currently
tests isinstance(tokens_input/output, int) which can miss numeric usage when
values are floats; update the logic to determine has_usage from the resolved
rounded values (resolved_input, resolved_output) — e.g., consider has_usage true
when resolved_input > 0 or resolved_output > 0 (and include any
tokens_total-equivalent if present) before constructing ModelLlmUsage; adjust
the has_usage reference used in the ModelLlmUsage constructor so provenance
matches the OpenAI handler's behavior.
🧹 Nitpick comments (1)
src/omnibase_infra/nodes/node_llm_inference_effect/handlers/handler_llm_openai_compatible.py (1)
137-192: Missinghandler_typeandhandler_categoryproperties.The
HandlerLlmOpenaiCompatibleclass does not implement the@property handler_typereturningEnumHandlerTypeor@property handler_categoryreturningEnumHandlerTypeCategory, which are required for handler classes. However, the class docstring explicitly notes this is a "documented deviation from the standard ONEX handler pattern" since it operates at the infrastructure layer with typed request/response models rather than envelope-based dispatch.If this deviation is intentional and approved, consider adding a comment or annotation to suppress this guideline check. Otherwise, add the required properties similar to
HandlerLlmOllama.As per coding guidelines, handlers must include
@property handler_type returning EnumHandlerTypeand@property handler_category returning EnumHandlerTypeCategoryin all handler classes.
- [major] rename converter_llm_usage_to_contract → adapter_ to match approved file naming prefixes (git mv preserves history) - [minor] add model_id validation guard in to_call_metrics() + test - [minor] fallback to total_tokens when prompt_tokens missing in embedding OpenAI handler - [minor] use resolved integer values for has_usage in Ollama handler, aligning with OpenAI handler provenance logic - [nitpick] add handler_type/handler_category properties to HandlerLlmOpenaiCompatible + update test assertion
…set stale metrics - Widen input_hash VARCHAR(64) to VARCHAR(71) in migration 031 to fit sha256- prefix (7 chars) + 64 hex chars - Add SUFFIX_INTELLIGENCE_LLM_CALL_COMPLETED to provisioned topic specs so TopicProvisioner auto-creates the LLM call completed topic - Reset self.last_call_metrics = None at start of handle() to prevent stale metrics surviving if _build_usage_metrics throws
Some LLM providers include cached/reasoning tokens in total_tokens, making it exceed prompt_tokens + completion_tokens. The ModelLlmUsage model_validator would raise ValueError on mismatch, crashing handle(). Now _parse_usage falls back to auto-compute when the provider total doesn't match the sum. Raw provider data preserved in raw_provider_usage for auditing.
There was a problem hiding this comment.
🤖 Fix all issues with AI agents
Before applying any fix, first verify the finding against the current code and
decide whether a code change is actually needed. If the finding is not valid or
no change is required, do not modify code for that item and briefly explain why
it was skipped.
In
`@src/omnibase_infra/nodes/node_llm_embedding_effect/handlers/handler_embedding_openai_compatible.py`:
- Around line 150-160: The empty-embeddings guard in
handler_embedding_openai_compatible.py is redundant because
_parse_openai_embeddings already raises for empty/invalid data; remove the if
not embeddings block (the ModelInfraErrorContext.with_correlation +
InfraProtocolError raise) from the method that calls _parse_openai_embeddings,
or alternatively move the distinct error message into the
_parse_openai_embeddings implementation so a single failure path (inside
_parse_openai_embeddings) reports the desired context/message; update callers to
rely on _parse_openai_embeddings to raise instead of checking embeddings
themselves.
🧹 Nitpick comments (1)
🤖 Fix all nitpicks with AI agents
Before applying any fix, first verify the finding against the current code and decide whether a code change is actually needed. If the finding is not valid or no change is required, do not modify code for that item and briefly explain why it was skipped. In `@src/omnibase_infra/nodes/node_llm_embedding_effect/handlers/handler_embedding_openai_compatible.py`: - Around line 150-160: The empty-embeddings guard in handler_embedding_openai_compatible.py is redundant because _parse_openai_embeddings already raises for empty/invalid data; remove the if not embeddings block (the ModelInfraErrorContext.with_correlation + InfraProtocolError raise) from the method that calls _parse_openai_embeddings, or alternatively move the distinct error message into the _parse_openai_embeddings implementation so a single failure path (inside _parse_openai_embeddings) reports the desired context/message; update callers to rely on _parse_openai_embeddings to raise instead of checking embeddings themselves.src/omnibase_infra/nodes/node_llm_embedding_effect/handlers/handler_embedding_openai_compatible.py (1)
150-160: Consider removing the redundant empty-embeddings guard.
_parse_openai_embeddingsalready raises on empty/invalid data, so this branch looks unreachable. If you want a distinct message, consider moving that logic into the parser instead.♻️ Optional simplification
- if not embeddings: - ctx = ModelInfraErrorContext.with_correlation( - correlation_id=request.correlation_id, - transport_type=EnumInfraTransportType.HTTP, - operation="parse_openai_embeddings", - target_name=self._llm_target_name, - ) - raise InfraProtocolError( - "OpenAI embedding response returned no embeddings", - context=ctx, - )🤖 Prompt for AI Agents
Before applying any fix, first verify the finding against the current code and decide whether a code change is actually needed. If the finding is not valid or no change is required, do not modify code for that item and briefly explain why it was skipped. In `@src/omnibase_infra/nodes/node_llm_embedding_effect/handlers/handler_embedding_openai_compatible.py` around lines 150 - 160, The empty-embeddings guard in handler_embedding_openai_compatible.py is redundant because _parse_openai_embeddings already raises for empty/invalid data; remove the if not embeddings block (the ModelInfraErrorContext.with_correlation + InfraProtocolError raise) from the method that calls _parse_openai_embeddings, or alternatively move the distinct error message into the _parse_openai_embeddings implementation so a single failure path (inside _parse_openai_embeddings) reports the desired context/message; update callers to rely on _parse_openai_embeddings to raise instead of checking embeddings themselves.
4 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Integrates the new SPI 0.9.0 LLM cost tracking contracts (
ContractLlmCallMetrics,ContractLlmUsageRaw,ContractLlmUsageNormalized,ContractEnumUsageSource) into omnibase_infra's LLM effect handlers.Ticket: OMN-2318
Pipeline Run: omn2318a
Changes
converter_llm_usage_to_contract.py): Three pure functions bridgingModelLlmUsageto SPI contractsto_usage_raw()→ContractLlmUsageRawto_usage_normalized()→ContractLlmUsageNormalizedto_call_metrics()→ContractLlmCallMetricsusage_source: ContractEnumUsageSourcefield toModelLlmUsage(default:MISSING)raw_provider_usage: dict | Nonefield toModelLlmUsageomnibase-coreto>=0.18.0andomnibase-spito>=0.9.0Test Plan
Summary by CodeRabbit
New Features
Bug Fixes
Chores