Update Postgres helpers by KaQuMiQ · Pull Request #467 · miquido/draive

KaQuMiQ · 2025-11-17T10:45:03Z

No description provided.

coderabbitai · 2025-11-17T10:45:13Z

Walkthrough

The PR replaces many ctx.record(ObservabilityLevel.X, ...) calls with specialized helpers (ctx.record_info / record_debug / record_warning / record_error) and enriches observability payloads (kind, unit, attributes like provider/model/type, batch_size, dimensions). It introduces a concrete ModelMemory abstraction with protocols (ModelMemoryRecalling/Remembering/Maintaining) and related factory/methods, updates memory handling to use ModelMemory across call sites, changes Postgres memory APIs to use UUID identifiers and adds meta, removes the local PostgresConfigurationRepository implementation, and bumps the project/version and several dependency constraints.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60–90 minutes

Areas requiring extra attention:

src/draive/models/types.py: new ModelMemory class, runtime-checkable protocols, @statemethod usage, and changes to all.
src/draive/postgres/memory.py: public signature changes to use UUID, SQL adjustments, recall/remember/maintenance semantics, and added meta.
Removal of src/draive/postgres/configuration.PostgresConfigurationRepository and the re-export change in src/draive/postgres/init.py.
Widespread observability changes: ensure correct mapping from ObservabilityLevel→ctx.record_* and consistency of added metric fields (kind, unit, attributes) across embedding, messaging, and token-usage sites.
API/type ripple effects: Stage.execute and other call sites switched to ModelMemory — verify callers and tests for compatibility.
tests/test_conversation_completion.py: updated fixtures and wiring for ModelMemory/META_EMPTY/ModelMemoryRecall.

Possibly related PRs

Update telemetry and error handling #460 — Overlapping observability instrumentation changes (same ctx.record → ctx.record_* migration and metric metadata adjustments).
Add Postgres helpers #431 — Modifies the same Postgres modules and exports (PostgresConfigurationRepository / PostgresModelMemory).
Rework LMM to GenerativeModel #394 — Related memory abstraction refactor touching ModelMemory / ModelMemoryRecall and conversation/realtime memory normalization.

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings, 1 inconclusive)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The pull request has no description provided by the author, making it impossible to assess whether the changeset is properly documented or explained.	Add a detailed pull request description explaining the objectives, changes made, and reasoning behind the refactoring (especially regarding ObservabilityLevel deprecation and ModelMemory updates).
Docstring Coverage	⚠️ Warning	Docstring coverage is 4.76% which is insufficient. The required threshold is 80.00%.	You can run `@coderabbitai generate docstrings` to improve docstring coverage.
Title check	❓ Inconclusive	The title 'Update Postgres helpers' is too vague and generic, failing to convey specific details about the extensive changes made across 40+ files affecting observability APIs, memory handling, and dependencies.	Provide a more descriptive title that summarizes the primary changes, such as 'Refactor observability API: replace ObservabilityLevel with ctx.record_* methods and update memory model' or similar.

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment
Commit unit tests in branch feature/postgre

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

KaQuMiQ · 2025-11-17T11:16:30Z

@coderabbitai check

coderabbitai · 2025-11-17T11:16:46Z

@KaQuMiQ I'll review the changes in this PR for you.

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

coderabbitai

Actionable comments posted: 3

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/draive/postgres/memory.py (1)
322-345: Critical: SQL query uses wrong column name in WHERE clause

Line 335 uses WHERE identifier = $1::UUID but should use WHERE memories = $1::UUID to match the schema where memories is the foreign key column referencing the memories table.

Apply this diff:
         WHERE
-            identifier = $1::UUID
+            memories = $1::UUID

📜 Review details

Configuration used: CodeRabbit UI

Review profile: ASSERTIVE

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 09017dc and 99bd2e0.

⛔ Files ignored due to path filters (1)

uv.lock is excluded by !**/*.lock

📒 Files selected for processing (40)

docs/guides/Basics.md (2 hunks)
llms.txt (1 hunks)
pyproject.toml (3 hunks)
src/draive/anthropic/messages.py (7 hunks)
src/draive/bedrock/converse.py (5 hunks)
src/draive/bedrock/guardrails.py (2 hunks)
src/draive/cohere/embedding.py (7 hunks)
src/draive/conversation/completion/state.py (3 hunks)
src/draive/conversation/realtime/state.py (3 hunks)
src/draive/evaluation/evaluator.py (1 hunks)
src/draive/evaluation/scenario.py (1 hunks)
src/draive/gemini/embedding.py (4 hunks)
src/draive/gemini/generating.py (12 hunks)
src/draive/generation/audio/state.py (2 hunks)
src/draive/generation/image/state.py (2 hunks)
src/draive/generation/model/state.py (3 hunks)
src/draive/generation/text/state.py (2 hunks)
src/draive/mistral/completions.py (8 hunks)
src/draive/mistral/embedding.py (4 hunks)
src/draive/mistral/moderation.py (2 hunks)
src/draive/models/generative.py (1 hunks)
src/draive/models/tools/function.py (4 hunks)
src/draive/models/types.py (3 hunks)
src/draive/ollama/chat.py (3 hunks)
src/draive/ollama/embedding.py (4 hunks)
src/draive/openai/embedding.py (4 hunks)
src/draive/openai/moderation.py (2 hunks)
src/draive/openai/realtime.py (4 hunks)
src/draive/openai/responses.py (9 hunks)
src/draive/postgres/__init__.py (1 hunks)
src/draive/postgres/configuration.py (0 hunks)
src/draive/postgres/memory.py (11 hunks)
src/draive/postgres/templates.py (3 hunks)
src/draive/postgres/vector_index.py (2 hunks)
src/draive/stages/stage.py (6 hunks)
src/draive/utils/memory.py (4 hunks)
src/draive/utils/vector_index.py (3 hunks)
src/draive/vllm/embedding.py (4 hunks)
src/draive/vllm/messages.py (6 hunks)
tests/test_conversation_completion.py (3 hunks)

💤 Files with no reviewable changes (1)

src/draive/postgres/configuration.py

🧰 Additional context used

🧠 Learnings (2)

📚 Learning: 2025-06-16T10:28:07.434Z

Learnt from: KaQuMiQ
Repo: miquido/draive PR: 338
File: src/draive/lmm/__init__.py:1-2
Timestamp: 2025-06-16T10:28:07.434Z
Learning: The draive project requires Python 3.12+ as specified in pyproject.toml with "requires-python = ">=3.12"" and uses Python 3.12+ specific features like PEP 695 type aliases and generic syntax extensively throughout the codebase.

Applied to files:

pyproject.toml

📚 Learning: 2025-06-16T10:09:51.508Z

Learnt from: KaQuMiQ
Repo: miquido/draive PR: 338
File: src/draive/conversation/types.py:123-144
Timestamp: 2025-06-16T10:09:51.508Z
Learning: In the draive.commons.metadata.Meta class, properties like `identifier`, `origin_identifier`, `predecessor_identifier`, and `successor_identifier` automatically parse stored hex strings into UUID objects when accessed, returning None if the string is missing or invalid. This means accessing `meta.identifier` returns a UUID object, not a string, even though it's stored internally as a hex string.

Applied to files:

src/draive/postgres/memory.py

🔇 Additional comments (57)

src/draive/openai/moderation.py (1)

3-3: LGTM! Clean observability API migration.

The migration from ObservabilityLevel.INFO to ctx.record_info(...) is correct and consistent with the broader refactoring across the codebase. The attributes are preserved correctly.

Also applies to: 29-34

src/draive/mistral/moderation.py (1)

4-4: LGTM! Consistent observability migration.

The changes align with the observability API standardization. Import and usage updates are correct.

Also applies to: 31-36

src/draive/generation/audio/state.py (1)

3-3: LGTM! Template observability updates are correct.

Both template identifier recordings have been migrated to ctx.record_info(...) appropriately. The changes are consistent with the broader refactoring.

Also applies to: 44-51

src/draive/generation/text/state.py (1)

4-4: LGTM! Observability migration applied correctly.

Template recording updates match the pattern used across other generation modules.

Also applies to: 50-57

src/draive/generation/image/state.py (1)

3-3: LGTM! Consistent with other generation modules.

The observability migration follows the same pattern as audio and text generation.

Also applies to: 44-51

src/draive/vllm/messages.py (1)

7-7: LGTM! Enhanced observability with richer metrics.

The migration to specialized recording helpers is correct. The addition of unit="tokens" and kind="counter" for token metrics significantly improves observability by providing structured metadata. All attributes are preserved correctly, and the distinction between info-level and debug-level logging is appropriate.

Also applies to: 124-142, 239-252, 273-291

src/draive/bedrock/guardrails.py (1)

4-4: LGTM! Guardrails observability correctly updated.

The migration to ctx.record_info(...) is consistent with the broader refactoring. All guardrail attributes are preserved.

Also applies to: 83-89

src/draive/bedrock/converse.py (1)

5-5: LGTM! Observability enhancements align with vLLM changes.

The specialized recording helpers are used appropriately, and the enriched token metrics with unit="tokens" and kind="counter" improve observability consistency across model providers. All attributes are correctly preserved.

Also applies to: 123-142, 194-213

src/draive/postgres/__init__.py (1)

1-9: LGTM! Clean dependency migration.

The import source change from a local module to haiway.postgres aligns with the haiway dependency bump in pyproject.toml. The public API remains unchanged through the re-export.

src/draive/openai/embedding.py (1)

6-6: LGTM! Observability standardization looks excellent.

The migration from ObservabilityLevel-based logging to dedicated ctx.record_info calls with structured metrics (unit, kind, attributes) significantly improves observability instrumentation. The enriched metadata (provider, model, dimensions, batch_size, type) will enable better monitoring and analysis.

Also applies to: 33-77

src/draive/postgres/templates.py (1)

32-51: Helpful documentation addition.

The Notes section with example schema and index creation provides valuable setup guidance for users of this repository.

src/draive/mistral/embedding.py (1)

6-6: LGTM! Consistent observability standardization.

The changes mirror the pattern in src/draive/openai/embedding.py, providing consistent observability instrumentation across embedding providers with enriched metrics and structured attributes.

Also applies to: 32-74

src/draive/anthropic/messages.py (1)

144-307: LGTM! Comprehensive observability instrumentation.

The migration demonstrates excellent use of the new observability helpers:

ctx.record_info for configuration and metrics with structured attributes

ctx.record_debug for detailed context (instructions, tools, conversation history)

ctx.record_warning for rate limit events with retry metadata

Token metrics properly structured with unit="tokens", kind="counter", and provider/model tags

This provides a solid foundation for monitoring and debugging Anthropic API interactions.

src/draive/mistral/completions.py (1)

6-6: LGTM! Well-structured observability implementation.

The changes demonstrate consistent observability patterns across both streaming and non-streaming completion paths. The dedicated _record_usage_metrics helper function promotes code reuse and maintainability. Token metrics are properly structured with units, kinds, and provider/model attributes.

Also applies to: 129-147, 180-199, 311-329, 549-562

src/draive/ollama/chat.py (1)

8-8: LGTM! Consistent observability pattern.

The changes align with the observability standardization across other model provider implementations, using ctx.record_info for configuration metadata and ctx.record_debug for detailed context.

Also applies to: 114-131

pyproject.toml (1)

8-8: Dependency versions verified as available and compatible.

The version bumps are confirmed: haiway 0.38.0 and openai 2.8.0 both exist on PyPI and are compatible with the project's Python ≥3.13 requirement.

src/draive/postgres/vector_index.py (3)

7-7: LGTM: Clean import update

The removal of ObservabilityLevel aligns with the project-wide migration to specialized observability helpers.

40-59: Excellent documentation addition!

The schema example is comprehensive and correct, clearly demonstrating the required table structure with proper types (UUID, VECTOR, JSONB, TIMESTAMPTZ) and the recommended ivfflat index for cosine similarity.

167-167: Consider whether metrics should still be recorded for index operations

The observability recording was removed from the index method. While ctx.log_info provides basic logging, consider whether structured metrics (e.g., item counts, batch sizes, timing) should still be recorded via ctx.record_info to maintain visibility into indexing performance and usage patterns.

This same consideration applies to the search (line 268) and delete (lines 300, 317) methods. Should these operations emit metrics for operational observability?

src/draive/ollama/embedding.py (3)

6-6: LGTM: Import cleanup

Consistent with the project-wide refactor removing ObservabilityLevel.

32-42: Excellent observability enrichment

The addition of structured attributes (embedding.provider, embedding.model, embedding.batch_size, embedding.concurrent) and the debug-level config mapping provide much better visibility into embedding operations.

52-62: Well-structured metrics with proper metadata

Both metrics now include:

Proper unit="count" and kind="counter" declarations

Rich attributes for filtering/grouping (provider, model, type)

This follows observability best practices.

Also applies to: 67-78

src/draive/vllm/embedding.py (1)

6-6: LGTM: Consistent observability migration

The changes mirror the pattern in ollama/embedding.py with proper metric enrichment. The inclusion of embedding.dimensions is appropriate for VLLM's configurable embedding dimensions.

Also applies to: 33-40, 50-76

src/draive/cohere/embedding.py (1)

8-8: LGTM: Comprehensive observability updates for both text and image embeddings

Both embedding paths are properly updated with:

Specialized record helpers (ctx.record_info)

Rich metric attributes differentiating text vs. image types

Consistent structure with other embedding modules

Also applies to: 33-76, 126-167
src/draive/postgres/memory.py (5)
40-77: Excellent schema documentation with UUID integration

The schema example correctly demonstrates:

UUID columns with gen_random_uuid() defaults

Proper foreign key constraints with ON DELETE CASCADE

Appropriate indexes for query performance

79-79: Verify assertion change is correct

The assertion now allows recall_limit >= 0 instead of what was likely > 0 previously. With the default now being 0, this makes sense, but ensure that a limit of 0 has well-defined semantics (e.g., "load all" or "load none").

Looking at lines 238-277, it appears limit=0 means "no limit" (load all records), which is a valid design choice but differs from some conventions where 0 means "return nothing."

228-228: Good addition of metadata

The meta=Meta.of({"source": "postgres"}) provides useful traceability for memory instances.

86-107: Improved observability with structured events

The migration to ctx.record_info with structured attributes (identifier, limit, elements, variables) provides much better observability compared to the previous approach.

Also applies to: 168-175

5-5: Documentation example needs immediate update; breaking change to UUID type confirmed

The verification confirms the original concerns. The documentation at docs/guides/Postgres.md:30 contains a broken example: PostgresModelMemory("demo-session") passes a string literal but the current function signature requires UUID. This would fail at runtime or type-check.

Additionally:

The recall_limit semantic change from int | None = None to int = 0 alters behavior; the docs at line 145 still references it as "optional" which is no longer accurate

No callers of PostgresModelMemory exist in the codebase beyond the broken docs example

Required fixes:

Update docs/guides/Postgres.md:30 to pass a UUID object instead of string, e.g., PostgresModelMemory(UUID("..."))

Update docs text (line 145) to reflect that recall_limit is no longer optional and has a default of 0
⛔ Skipped due to learnings
Learnt from: KaQuMiQ
Repo: miquido/draive PR: 338
File: src/draive/conversation/types.py:123-144
Timestamp: 2025-06-16T10:09:51.508Z
Learning: In the draive.commons.metadata.Meta class, properties like `identifier`, `origin_identifier`, `predecessor_identifier`, and `successor_identifier` automatically parse stored hex strings into UUID objects when accessed, returning None if the string is missing or invalid. This means accessing `meta.identifier` returns a UUID object, not a string, even though it's stored internally as a hex string.
llms.txt (1)

215-215: LGTM: Documentation correctly updated

The observability example now reflects the new ctx.record_info API.

src/draive/utils/memory.py (1)

3-3: LGTM: Clean observability migration

The migration to ctx.record_info is consistent across all memory operations. The addition of "items": len(items) to the remember event (line 114) is a useful enhancement.

Also applies to: 86-88, 112-115, 136-138

src/draive/generation/model/state.py (1)

4-4: LGTM: Consistent observability updates in generation flow

The migration properly tracks:

Generated model metadata (qualname, schema_injection)

Template identifiers when templates are used

This maintains good observability for model generation operations.

Also applies to: 62-91

src/draive/evaluation/evaluator.py (1)

677-687: LGTM! Observability enhancement is well-implemented.

The migration from ObservabilityLevel-based logging to ctx.record_info() with enriched metadata is excellent. The addition of kind="histogram" is appropriate for performance metrics, and the structured attributes (passed, threshold, score) provide valuable context for analysis.

src/draive/models/tools/function.py (2)

275-277: LGTM! Proper use of dedicated recording helper.

The migration to ctx.record_info() is appropriate for recording tool invocation metadata. Including the call_id in attributes provides good traceability.

285-291: LGTM! Consistent error recording.

Both exception handlers now use ctx.record_error() with structured attributes capturing exception type and message. This provides better categorization than the previous approach and aligns with the observability API migration.

Also applies to: 305-311

docs/guides/Basics.md (1)

189-189: LGTM! Documentation updated to reflect new API.

The documentation examples have been correctly updated to use ctx.record_info() instead of the deprecated ctx.record(). This ensures users learn the current best practices.

Also applies to: 228-228, 235-235

src/draive/utils/vector_index.py (3)

94-100: LGTM! Enhanced indexing metrics.

The migration to ctx.record_info() with structured attributes is good. Adding the values count (length of values being indexed) provides valuable insight into batch sizes.

170-187: LGTM! Improved search metrics capture.

The reordering to execute the search first, then record metrics is beneficial because it allows capturing the results count in the recorded event. This provides more complete observability by showing both the query conditions and the result size.

219-225: LGTM! Cleaner requirements representation.

Using a boolean (requirements is not None) instead of a string representation is cleaner and more suitable for structured metrics.

src/draive/evaluation/scenario.py (1)

401-410: LGTM! Consistent observability enhancement.

The migration to ctx.record_info() with histogram kind and enriched attributes (passed status, evaluator names) provides excellent observability for scenario evaluation. The structured approach is consistent with the changes in evaluator.py.

src/draive/openai/responses.py (3)

179-198: LGTM! Enhanced model completion observability.

The migration introduces proper separation between info-level metrics (model configuration) and debug-level details (instructions, tools, context). The structured attributes provide comprehensive context for filtering and analysis.

272-291: LGTM! Structured token usage metrics.

The token usage metrics now include:

Explicit unit="tokens" for clarity

kind="counter" for proper metric aggregation

Provider and model attributes for filtering

This structured approach significantly improves observability and enables better metrics aggregation and visualization.

Also applies to: 574-593

421-440: LGTM! Consistent streaming observability.

The streaming completion path follows the same observability pattern as the non-streaming version, ensuring consistency across both execution modes.

src/draive/openai/realtime.py (2)

109-117: LGTM! Comprehensive session observability.

The session opening now records comprehensive context including provider, model name, tools configuration, and output selection. This provides excellent visibility into realtime session initialization.

222-245: LGTM! Robust token metrics with defensive defaults.

The token usage recording follows the same structured pattern as responses.py and includes defensive handling by defaulting to 0 when token values are None. This prevents potential errors while maintaining metric consistency.

src/draive/models/generative.py (1)

680-685: No issues found - change is correct and safe.

Verification confirms that ModelMemory.constant() accepts variadic *elements: ModelContextElement arguments (defined in src/draive/models/types.py:826). The implementation internally creates the ModelMemoryRecall object from these elements, so the simplified pattern at line 682 (ModelMemory.constant(*memory)) is correct and functionally equivalent. This usage pattern is consistent throughout the codebase (e.g., lines 81 and 121 in conversation modules), and the behavior remains consistent with the previous implementation.

src/draive/gemini/generating.py (3)

35-164: LGTM! Observability refactoring looks clean.

The migration from ObservabilityLevel-based logging to dedicated helpers (ctx.record_info, ctx.record_debug) improves code clarity and adds richer structured attributes for telemetry.

186-232: LGTM! Rate limit and safety logging updates are appropriate.

The use of ctx.record_warning for rate limits and ctx.record_info for safety results aligns with their severity levels and improves observability.

291-467: LGTM! Streaming and metrics observability is well-structured.

The usage metrics now include proper unit, kind, and provider/model attributes, enabling better telemetry aggregation and analysis.

src/draive/gemini/embedding.py (1)

7-76: LGTM! Embedding observability follows the same pattern.

The observability updates are consistent with the broader refactoring, adding structured attributes like embedding.provider, embedding.model, embedding.type, unit, and kind for comprehensive telemetry.

src/draive/conversation/realtime/state.py (1)

4-85: LGTM! ModelMemory API simplification improves usability.

The updated ModelMemory.constant() API now accepts context elements directly as variadic arguments instead of requiring a ModelMemoryRecall wrapper, making the code more concise and readable.

src/draive/conversation/completion/state.py (1)

4-125: LGTM! Consistent ModelMemory API updates.

The changes align with the new ModelMemory API, simplifying memory construction and improving consistency across the codebase. The empty memory case (ModelMemory.constant()) is particularly clean.

src/draive/stages/stage.py (2)

349-624: LGTM! Observability updates are consistent.

The template-based instruction and input logging now uses ctx.record_info with structured attributes, maintaining consistency with the broader observability refactoring.

2001-2065: LGTM! Stage.execute() properly updated for ModelMemory.

The method signature and implementation correctly handle the new ModelMemory type, recalling context at the start and remembering new context elements after execution.

src/draive/models/types.py (3)

782-806: LGTM! Well-designed memory protocols.

The @runtime_checkable protocols provide clear contracts for memory components with proper separation of concerns:

ModelMemoryRecalling: retrieves memory

ModelMemoryRemembering: persists context elements

ModelMemoryMaintaining: handles maintenance tasks

809-820: LGTM! Clean no-op implementations.

The helper functions provide appropriate default behavior for immutable memory instances created via ModelMemory.constant().

823-923: LGTM! Robust ModelMemory implementation.

The ModelMemory class provides a clean, extensible API:

constant() factory simplifies creating immutable memory from context elements

@statemethod decorators enable both class and instance method usage

Proper overloads ensure type safety

Protocol-based attributes allow flexible implementations

src/draive/postgres/memory.py

src/draive/postgres/templates.py

tests/test_conversation_completion.py

coderabbitai

Actionable comments posted: 3

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

src/draive/postgres/templates.py (1)

5-62: Add structured observability events to match postgres repository pattern.

The file currently lacks ctx.record_info() calls despite using ctx.log_info(). The similar implementation in src/draive/postgres/memory.py emits structured observability events:

Line 99: ctx.record_info(event="postgres.memory.recall", attributes={...})

Line 168: ctx.record_info(event="postgres.memory.remember", attributes={...})

Add corresponding ctx.record_info() calls after the database operations in listing(), load(), and defining() functions to maintain consistent instrumentation across the codebase.

src/draive/postgres/vector_index.py (1)

7-167: I need to verify the full scope of the Postgres implementation to confirm observability is missing from all operations.Now let me search for the full Postgres implementation to verify all operations:

Let me search for the complete Postgres implementation including search and delete operations:

Based on the verification results, I can now provide the rewritten review comment:

Add observability recording to track vector index operations in production.

The Postgres implementation removes all structured observability (ctx.record_info()) from vector index operations, whereas the base implementation in src/draive/utils/vector_index.py consistently emits metrics for vector_index.index, vector_index.search, and vector_index.delete events. Production systems require observability on these operations to track latency, throughput, success rates, and usage patterns. Add ctx.record_info() calls that mirror the base implementation for all three operations—index, search, and delete.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: ASSERTIVE

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 09017dc and 99bd2e0.

⛔ Files ignored due to path filters (1)

uv.lock is excluded by !**/*.lock

📒 Files selected for processing (40)

docs/guides/Basics.md (2 hunks)
llms.txt (1 hunks)
pyproject.toml (3 hunks)
src/draive/anthropic/messages.py (7 hunks)
src/draive/bedrock/converse.py (5 hunks)
src/draive/bedrock/guardrails.py (2 hunks)
src/draive/cohere/embedding.py (7 hunks)
src/draive/conversation/completion/state.py (3 hunks)
src/draive/conversation/realtime/state.py (3 hunks)
src/draive/evaluation/evaluator.py (1 hunks)
src/draive/evaluation/scenario.py (1 hunks)
src/draive/gemini/embedding.py (4 hunks)
src/draive/gemini/generating.py (12 hunks)
src/draive/generation/audio/state.py (2 hunks)
src/draive/generation/image/state.py (2 hunks)
src/draive/generation/model/state.py (3 hunks)
src/draive/generation/text/state.py (2 hunks)
src/draive/mistral/completions.py (8 hunks)
src/draive/mistral/embedding.py (4 hunks)
src/draive/mistral/moderation.py (2 hunks)
src/draive/models/generative.py (1 hunks)
src/draive/models/tools/function.py (4 hunks)
src/draive/models/types.py (3 hunks)
src/draive/ollama/chat.py (3 hunks)
src/draive/ollama/embedding.py (4 hunks)
src/draive/openai/embedding.py (4 hunks)
src/draive/openai/moderation.py (2 hunks)
src/draive/openai/realtime.py (4 hunks)
src/draive/openai/responses.py (9 hunks)
src/draive/postgres/__init__.py (1 hunks)
src/draive/postgres/configuration.py (0 hunks)
src/draive/postgres/memory.py (11 hunks)
src/draive/postgres/templates.py (3 hunks)
src/draive/postgres/vector_index.py (2 hunks)
src/draive/stages/stage.py (6 hunks)
src/draive/utils/memory.py (4 hunks)
src/draive/utils/vector_index.py (3 hunks)
src/draive/vllm/embedding.py (4 hunks)
src/draive/vllm/messages.py (6 hunks)
tests/test_conversation_completion.py (3 hunks)

💤 Files with no reviewable changes (1)

src/draive/postgres/configuration.py

🧰 Additional context used

🧠 Learnings (1)

📚 Learning: 2025-06-16T10:28:07.434Z

Learnt from: KaQuMiQ
Repo: miquido/draive PR: 338
File: src/draive/lmm/__init__.py:1-2
Timestamp: 2025-06-16T10:28:07.434Z
Learning: The draive project requires Python 3.12+ as specified in pyproject.toml with "requires-python = ">=3.12"" and uses Python 3.12+ specific features like PEP 695 type aliases and generic syntax extensively throughout the codebase.

Applied to files:

pyproject.toml

🧬 Code graph analysis (12)

src/draive/conversation/completion/state.py (4)

src/draive/models/types.py (2)

ModelMemory (824-923)

constant (826-848)

src/draive/utils/memory.py (1)

constant (54-66)

src/draive/conversation/realtime/state.py (1)

model_context_elements (73-79)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/conversation/realtime/state.py (4)

src/draive/models/types.py (2)

ModelMemory (824-923)

constant (826-848)

src/draive/utils/memory.py (1)

constant (54-66)

src/draive/conversation/completion/state.py (1)

model_context_elements (113-119)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/generation/model/state.py (1)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/generation/text/state.py (1)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/generation/image/state.py (1)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/models/types.py (3)

src/draive/utils/memory.py (13)

_remembering_none (38-42)

_maintaining_noop (45-48)

constant (54-66)

recall (70-73)

recall (76-79)

recall (82-89)

recalling (58-61)

remember (93-97)

remember (100-104)

remember (107-116)

maintenance (120-123)

maintenance (126-129)

maintenance (132-139)

src/draive/postgres/memory.py (3)

recall (81-112)

remember (114-176)

maintenance (178-222)

tests/test_stage.py (4)

recall (65-66)

recall (441-442)

remember (68-70)

remember (444-446)

tests/test_conversation_completion.py (1)

src/draive/models/types.py (3)

ModelMemory (824-923)

ModelMemoryRecall (730-777)

recalling (838-841)

src/draive/generation/audio/state.py (1)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/vllm/messages.py (1)

src/draive/stages/stage.py (1)

completion (287-381)

src/draive/stages/stage.py (2)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/models/types.py (1)

ModelMemory (824-923)

src/draive/postgres/memory.py (2)

src/draive/models/types.py (15)

ModelMemory (824-923)

recall (852-855)

recall (858-861)

recall (864-868)

ModelMemoryRecall (730-777)

of (283-305)

of (326-358)

of (386-422)

of (455-478)

of (545-569)

of (598-622)

of (746-772)

of (941-967)

of (994-1021)

of (1086-1114)

src/draive/utils/memory.py (3)

recall (70-73)

recall (76-79)

recall (82-89)

src/draive/models/generative.py (1)

src/draive/models/types.py (2)

ModelMemory (824-923)

constant (826-848)

🔇 Additional comments (64)

src/draive/postgres/__init__.py (1)

3-3: Verify external API compatibility before merging.

Local cleanup verified: the draive.postgres.configuration module is removed and no stale references remain. However, the critical dependency—whether haiway.postgres v0.38.0 exports PostgresConfigurationRepository—could not be confirmed through available documentation sources.

Before merging, confirm that the haiway package exposes this export. Run the codebase's test suite or type checker to catch import failures.

src/draive/generation/text/state.py (1)

4-4: LGTM! Observability API migration is correct.

The migration from ctx.record(ObservabilityLevel.INFO, ...) to ctx.record_info(...) is consistent and preserves all attributes. The removal of the ObservabilityLevel import aligns with the new specialized helper methods.

Also applies to: 50-52, 55-57

src/draive/generation/image/state.py (1)

3-3: LGTM! Consistent with the observability API migration.

The changes mirror the pattern in text/state.py and correctly migrate to the new ctx.record_info(...) API.

Also applies to: 44-46, 49-51

src/draive/openai/moderation.py (1)

3-3: LGTM! Observability migration applied correctly.

The migration preserves the guardrails provider and model attributes while using the new ctx.record_info(...) helper.

Also applies to: 29-34

src/draive/bedrock/guardrails.py (1)

4-4: LGTM! Bedrock guardrails observability updated correctly.

The migration to ctx.record_info(...) preserves all guardrails metadata including provider, identifier, and version.

Also applies to: 83-89

src/draive/generation/audio/state.py (1)

3-3: LGTM! Audio generation observability migration is consistent.

The changes follow the same pattern as text and image generation modules, correctly migrating to ctx.record_info(...).

Also applies to: 44-46, 49-51

src/draive/mistral/moderation.py (1)

4-4: LGTM! Mistral moderation observability migration is correct.

The migration mirrors the OpenAI moderation pattern and correctly preserves guardrails provider and model attributes.

Also applies to: 31-36

src/draive/evaluation/evaluator.py (1)

677-687: LGTM! Enhanced observability for evaluator performance.

The migration to ctx.record_info(...) not only updates the API but also enhances the metric recording by:

Adding kind="histogram" for proper metric aggregation

Including evaluation details (passed, threshold, score) in attributes

This provides better observability for evaluation performance tracking.

src/draive/ollama/chat.py (1)

8-8: LGTM! Ollama chat observability correctly migrated with appropriate levels.

The migration correctly uses:

ctx.record_info(...) for model configuration metadata (provider, name, temperature, etc.)

ctx.record_debug(...) for detailed context information (instructions, tools, context messages)

This proper separation of observability levels improves debugging and monitoring capabilities.

Verification confirms that haiway~=0.38.0 supports the ctx.record_info() and ctx.record_debug() methods. Both are consistently used across 20+ files throughout the codebase, including in other providers (vllm, openai, anthropic, bedrock, mistral, gemini), confirming the API is available and properly integrated.

src/draive/generation/model/state.py (2)

4-4: LGTM: Import updated to match new observability API.

The removal of ObservabilityLevel from the import aligns with the migration to specialized ctx.record_info calls throughout this file.

62-91: LGTM: Clean migration to specialized observability helpers.

The three ctx.record_info(...) calls correctly replace ctx.record(ObservabilityLevel.INFO, ...) while preserving all attributes and metadata.

src/draive/evaluation/scenario.py (1)

401-410: LGTM: Observability API migration preserves all metrics.

The migration from ctx.record(ObservabilityLevel.INFO, ...) to ctx.record_info(...) maintains the same metric name, value, unit, kind, and attributes.

llms.txt (1)

215-215: LGTM: Documentation updated to reflect new observability API.

The documentation correctly updates the example from ctx.record(metric_name, value=...) to ctx.record_info(metric_name, value=...).

src/draive/anthropic/messages.py (3)

144-164: LGTM: Observability calls migrated to specialized helpers.

The migration correctly uses ctx.record_info(...) and ctx.record_debug(...) in place of ObservabilityLevel.INFO and ObservabilityLevel.DEBUG, preserving all attributes.

231-268: LGTM: Rate limit warnings migrated consistently.

Both rate limit warning paths now use ctx.record_warning(...) instead of ctx.record(ObservabilityLevel.WARNING, ...), maintaining the same event name and attributes.

277-307: LGTM: Token metrics migrated with full metadata.

All three token metrics (model.input_tokens, model.input_tokens.cached, model.output_tokens) correctly use ctx.record_info(...) with preserved metric names, values, units, kinds, and attributes.

src/draive/ollama/embedding.py (2)

32-42: LGTM: Observability API migration with clean separation.

The info and debug recording calls are now properly separated using ctx.record_info(...) and ctx.record_debug(...).

52-78: LGTM: Enhanced metrics with richer metadata.

The embedding metrics now include unit, kind, and comprehensive attributes (provider, model, type), providing better observability than the previous implementation.

src/draive/cohere/embedding.py (2)

33-76: LGTM: Text embedding observability enhanced.

The migration to ctx.record_info(...) includes richer metadata for both scope attributes and metrics, with unit, kind, and detailed provider/model/type information.

126-167: LGTM: Image embedding observability enhanced.

The image embedding path mirrors the text embedding improvements, using ctx.record_info(...) with consistent metadata structure.

src/draive/postgres/templates.py (1)

32-51: LGTM: Helpful schema documentation added.

The Notes section with the example schema and index definition provides clear guidance for users setting up the Postgres backend.

src/draive/postgres/vector_index.py (1)

41-59: LGTM: Comprehensive schema documentation added.

The Notes section provides a complete example schema with vector indexing setup, which is valuable for users implementing Postgres-backed vector storage.

docs/guides/Basics.md (1)

189-189: LGTM! Documentation examples correctly updated to new observability API.

The documentation examples have been properly updated to use ctx.record_info() instead of the previous ctx.record() pattern, which aligns with the broader API migration across the codebase.

Also applies to: 228-228, 235-235

src/draive/models/tools/function.py (1)

4-4: LGTM! Observability API migration correctly applied.

The changes properly migrate from ObservabilityLevel-based recording to specialized helper methods (ctx.record_info and ctx.record_error). The error recording maintains the same event and attribute structure for both Exception and BaseException handlers.

Also applies to: 275-316

src/draive/vllm/messages.py (1)

7-7: LGTM! Observability API migration correctly applied with enhanced metrics.

The changes migrate from ObservabilityLevel-based calls to specialized methods (ctx.record_info, ctx.record_debug). Token usage metrics are enhanced with kind="counter" and structured attributes for better observability.

Also applies to: 124-142, 239-252, 273-291

src/draive/bedrock/converse.py (1)

5-5: LGTM! Observability API migration correctly applied.

The changes properly migrate from ObservabilityLevel-based recording to specialized helper methods. Token usage metrics include proper structure with kind="counter" and provider/model attributes for observability systems.

Also applies to: 123-142, 194-213

src/draive/openai/responses.py (1)

7-7: LGTM! Observability API migration correctly applied with defensive coding.

The changes migrate to specialized observability methods consistently across both streaming and non-streaming paths. Token metrics use getattr with default values for robustness and include proper kind="counter" attributes.

Also applies to: 179-198, 272-291, 421-440, 574-593

src/draive/mistral/completions.py (1)

6-6: LGTM! Observability API migration correctly applied with helper function.

The changes migrate to specialized observability methods consistently. The _record_usage_metrics helper function promotes code reuse for token metrics. Defensive coding with or 0 defaults handles potential None values.

Also applies to: 129-147, 180-199, 311-329, 549-562

src/draive/openai/realtime.py (1)

30-30: LGTM! Observability API migration correctly applied with explicit None handling.

The changes migrate to ctx.record_info consistently. Token usage metrics use explicit if X is not None else 0 checks for robust None handling, which is appropriate for the realtime API context.

Also applies to: 109-117, 222-245

src/draive/gemini/embedding.py (1)

7-7: LGTM! Observability API migration correctly applied for embedding metrics.

The changes migrate to ctx.record_info with enhanced metric structure. Embedding-specific metrics include proper kind="counter" attributes and structured metadata (provider, model, type) for effective observability.

Also applies to: 33-76

src/draive/utils/vector_index.py (2)

94-100: LGTM! Observability migration is clean.

The migration from ctx.record(ObservabilityLevel.INFO, ...) to ctx.record_info(...) is correctly implemented with appropriate attributes for tracking model and value counts.

170-187: Well-structured search observability.

The search recording correctly captures results after the search operation and uses boolean flags to indicate the presence of query and requirements rather than exposing their content, which is appropriate for INFO-level logging.

src/draive/mistral/embedding.py (2)

6-6: LGTM! Clean import update.

Correctly removed ObservabilityLevel from imports as part of the migration to specialized recording methods.

32-74: Excellent observability enrichment.

The migration to ctx.record_info includes enhanced metadata with proper metric tracking (unit="count", kind="counter") and structured attributes for provider, model, and type. This provides much richer observability data compared to the previous implementation.

src/draive/gemini/generating.py (4)

35-35: LGTM! Import cleanup.

Correctly removed ObservabilityLevel and MISSING from the import, keeping only the necessary symbols.

145-164: Excellent use of logging levels.

Properly separates configuration details (recorded at INFO level) from sensitive context data like instructions and conversation history (recorded at DEBUG level). This is a security best practice.

186-197: Consistent rate limit handling.

Both streaming and non-streaming paths correctly use ctx.record_warning for rate limit events and apply randomized backoff (0.3-3.0 seconds), which is appropriate for avoiding thundering herd problems.

Also applies to: 401-413

438-467: Well-structured usage metrics.

The _record_usage_metrics function correctly tracks input tokens, cached tokens, and output tokens with proper metric metadata (unit, kind, attributes). This provides comprehensive observability for token usage.

src/draive/vllm/embedding.py (1)

6-6: LGTM! Consistent embedding observability pattern.

The observability migration matches the pattern used in other embedding implementations (Mistral, OpenAI), ensuring consistency across providers. Properly includes VLLM-specific configuration like dimensions.

Also applies to: 33-76

src/draive/openai/embedding.py (1)

6-6: LGTM! Consistent with other embedding implementations.

The observability migration is identical to the VLLM implementation and consistent with the Mistral implementation, ensuring a uniform observability experience across all embedding providers.

Also applies to: 33-77

pyproject.toml (2)

8-8: Appropriate version bump for API changes.

Bumping from 0.92.1 to 0.93.0 is appropriate for this PR, which introduces API surface changes (migration to specialized observability methods and ModelMemory refactoring).

27-27: The codebase extensively uses ctx.record_info(), ctx.record_debug(), ctx.record_warning(), and ctx.record_error() methods (89+, 11, 4, and 2 usages respectively). Let me verify these are provided by haiway 0.38.0:

Dependency versions are verified: haiway 0.38.0 is published and provides required logging APIs.

The codebase extensively uses the logging methods that haiway 0.38.0 should provide (100+ usages of ctx.record_info, ctx.record_debug, ctx.record_warning, and ctx.record_error). The fact that these methods are already integrated throughout multiple modules (vllm, openai, mistral, gemini, anthropic, bedrock, postgres, qdrant, etc.) confirms that haiway 0.38.0 provides these APIs as required.

src/draive/utils/memory.py (1)

3-3: LGTM! Memory observability enhanced.

The migration correctly removes ObservabilityLevel and adds structured observability for all memory operations (recall, remember, maintenance). The addition of an items count in the remember event (line 114) provides valuable metrics.

Also applies to: 86-139

src/draive/conversation/completion/state.py (5)

4-4: LGTM! Correct import updates.

Properly removed ObservabilityLevel and added the new ModelMemory and ModelMemoryRecall imports to support the memory API migration.

Also applies to: 12-13

102-104: Correct usage of ModelMemory.constant().

Creating an empty memory with ModelMemory.constant() (no arguments) aligns with the new API, which internally creates an empty ModelMemoryRecall.

106-109: Proper ModelMemory type check and recall.

The type check correctly uses isinstance(memory, ModelMemory) and properly awaits memory.recall() to extract memory variables.

113-122: Clean conversion from iterable to ModelMemory.

The generator function correctly converts ConversationMessage objects into ModelInput/ModelOutput elements, and ModelMemory.constant(*model_context_elements()) properly creates memory from these elements.

125-127: LGTM! Template observability.

Correctly uses ctx.record_info to record the template identifier when instructions are provided as a template.

src/draive/models/generative.py (1)

681-682: LGTM! Memory construction simplified.

The change correctly uses the new ModelMemory.constant API by directly unpacking the sequence elements, eliminating the need for the intermediate ModelMemoryRecall wrapper. This aligns with the broader memory system refactor.

src/draive/conversation/realtime/state.py (2)

66-82: LGTM! Memory handling correctly migrated to ModelMemory API.

The isinstance check and memory construction are properly updated to use the new ModelMemory class. The generator unpacking in ModelMemory.constant(*model_context_elements()) correctly passes the yielded ModelInput | ModelOutput elements as variadic arguments.

85-87: LGTM! Observability logging updated to new API.

The migration from ctx.record(ObservabilityLevel.INFO, ...) to ctx.record_info(...) correctly uses the new specialized helper method while maintaining the same functionality.

tests/test_conversation_completion.py (1)

8-8: LGTM! Imports updated for new memory API.

The imports correctly reference the new public API entities (META_EMPTY, ModelMemory, ModelMemoryRecall) that are used in the test memory construction.

Also applies to: 19-20

src/draive/models/types.py (4)

25-34: LGTM! Imports and exports updated for ModelMemory implementation.

The new imports provide the necessary infrastructure for the ModelMemory state class, and the __all__ exports correctly expose the new memory-related protocols to the public API.

Also applies to: 51-54

782-807: LGTM! Memory protocols properly defined.

The three protocols (ModelMemoryRecalling, ModelMemoryRemembering, ModelMemoryMaintaining) are correctly decorated with @runtime_checkable and define appropriate async callback signatures for the memory system.

809-820: LGTM! Noop helper functions correctly implemented.

The helper functions _remembering_none and _maintaining_noop provide proper no-op implementations matching their respective protocol signatures, used for constant memory instances that don't need to persist changes.

823-923: LGTM! ModelMemory class well-implemented.

The ModelMemory class provides a clean state-based abstraction for memory management:

The constant() factory method creates immutable memory with fixed recall results

All methods (recall, remember, maintenance) use @statemethod to support both class and instance invocation

Methods properly delegate to the protocol callables, maintaining separation of concerns

The implementation aligns with the broader codebase patterns using haiway.State

src/draive/stages/stage.py (2)

349-350: LGTM! Observability logging consistently migrated.

All instances of ctx.record(ObservabilityLevel.INFO, attributes={...}) have been correctly updated to use the new ctx.record_info(attributes={...}) helper method across multiple completion methods.

Also applies to: 354-355, 456-457, 553-554, 624-625

2004-2004: LGTM! Memory handling correctly migrated to ModelMemory API.

The execute method properly updates:

Parameter type from Memory to ModelMemory

isinstance checks to use ModelMemory

Memory operations using the new recall() and remember() methods

The logic correctly computes and persists only the newly produced context elements (suffix after common prefix).

Also applies to: 2034-2063

src/draive/postgres/memory.py (6)

5-8: Imports updated correctly for UUID-based identifiers and Postgres helpers

The new imports (UUID, BasicValue, Map, Meta, ctx, Postgres, PostgresRow) line up with the UUID-based identifier API and the use of haiway.postgres throughout this module. No issues here.

81-112: Recall path and observability look consistent with ModelMemoryRecall

The recall flow now:

Loads variables via _load_variables(identifier=identifier).

Loads context via _load_context(identifier=identifier, limit=recall_limit).

Emits a postgres.memory.recall event with identifier, limit, and element/variable counts.

Returns a ModelMemoryRecall carrying both context and variables.

This is consistent with the ModelMemoryRecall type and the ctx.record_info‑based observability pattern used elsewhere. No functional issues spotted.

114-177: Remember path: transactional writes and metrics look sound

The remember implementation looks solid:

Early return when both elements and variables are empty avoids unnecessary DB work.

Both variable snapshots and context elements are written in a single transaction, with identifier cast to UUID and variables/content serialized as JSONB.

The microsecond offset per element on a single created_timestamp gives deterministic ordering for same-batch inserts.

The postgres.memory.remember event includes identifier and counts of elements/variables, which is useful for observability.

No correctness or reliability issues identified here.

178-222: Maintenance flow correctly ensures base row and optional variables

The maintenance routine:

Upserts a row into memories for the given identifier (using ON CONFLICT (identifier) DO NOTHING).

Optionally appends a variables snapshot to memories_variables within the same transaction.

Logs progress clearly via ctx.log_info.

Assuming your actual DDL references memories(identifier) from variables/elements, this is the right place to ensure the parent row exists. Implementation looks correct.

224-229: Setting ModelMemory.meta to {"source": "postgres"} is appropriate

Returning a ModelMemory with meta=Meta.of({"source": "postgres"}) mirrors the pattern used for other memory sources (e.g. the constant memory having "source": "constant"). This is a useful hook for downstream routing/diagnostics and fits the existing types.

322-350: Variable loading behaviour aligns with recall expectations

_load_variables:

Selects the latest row for a given identifier using DISTINCT ON (identifier) ... ORDER BY identifier, created DESC LIMIT 1, which correctly yields the newest snapshot for that identifier.

Casts identifier to UUID in the WHERE clause, in line with the UUID-based API.

Returns an empty Map() when no variables exist, ensuring recall always works with a mapping (and len(variables) in metrics stays safe).

This matches the Mapping[str, BasicValue] expectations in the rest of the module.

src/draive/postgres/memory.py

tests/test_conversation_completion.py

coderabbitai

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

src/draive/postgres/vector_index.py (1)
386-393: Critical: SQL syntax error in contains operator.

Line 390 is missing a closing parenthesis for jsonb_array_elements_text(). The AS element clause should come after the function call is closed, similar to the correct pattern in the contains_any case above (lines 380-383).

Apply this diff to fix the SQL syntax error:
         case "contains":
             resolved_arguments = [*arguments, requirement.rhs]
             return (
                 "EXISTS (SELECT 1 FROM jsonb_array_elements_text("  # nosec: B608
-                f"{_scalar_accessor(str(requirement.lhs))} AS element"
+                f"{_scalar_accessor(str(requirement.lhs))}) AS element"
                 f" WHERE element = ${len(resolved_arguments)})",
                 resolved_arguments,
             )
src/draive/postgres/memory.py (1)
178-222: Critical: Incorrect column name in maintenance INSERT statement.

The INSERT statement at lines 205-220 uses identifier column but should use memories column to match the schema documented in lines 50-62. The memories_variables table has:

identifier (primary key, auto-generated)

memories (foreign key to memories.identifier)

variables (data)

The INSERT should reference the foreign key column memories, not identifier.

Apply this diff to fix the column name:
                    await connection.execute(
                        """
                        INSERT INTO
                            memories_variables (
-                               identifier,
+                               memories,
                                variables
                            )

                        VALUES (
                            $1::UUID,
                            $2::JSONB
                        );
                        """,
                        identifier,
                        json.dumps(variables),
                    )

♻️ Duplicate comments (2)

src/draive/postgres/templates.py (1)

58-62: Observability pattern inconsistency remains.

Unlike other files in this PR that migrate to ctx.record_info/debug with enriched payloads, this file removes observability records entirely. Consider adding structured observability calls alongside the existing ctx.log_info to align with the broader instrumentation strategy.
tests/test_conversation_completion.py (1)
210-218: Fix typo in callback name.

The callback is named maintanance_callback but should be maintenance_callback for consistency with the maintaining field and standard spelling.

Apply this diff to fix the typo:
-    async def maintanance_callback(**_: Any):
+    async def maintenance_callback(**_: Any):
         pass
 
     memory = ModelMemory(
         recalling=recall_callback,
         remembering=remember_callback,
-        maintaining=maintanance_callback,
+        maintaining=maintenance_callback,
         meta=META_EMPTY,
     )

📜 Review details

Configuration used: CodeRabbit UI

Review profile: ASSERTIVE

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 99bd2e0 and 98c24b2.

⛔ Files ignored due to path filters (1)

uv.lock is excluded by !**/*.lock

📒 Files selected for processing (40)

docs/guides/Basics.md (2 hunks)
llms.txt (1 hunks)
pyproject.toml (3 hunks)
src/draive/anthropic/messages.py (7 hunks)
src/draive/bedrock/converse.py (5 hunks)
src/draive/bedrock/guardrails.py (2 hunks)
src/draive/cohere/embedding.py (7 hunks)
src/draive/conversation/completion/state.py (3 hunks)
src/draive/conversation/realtime/state.py (3 hunks)
src/draive/evaluation/evaluator.py (1 hunks)
src/draive/evaluation/scenario.py (1 hunks)
src/draive/gemini/embedding.py (4 hunks)
src/draive/gemini/generating.py (12 hunks)
src/draive/generation/audio/state.py (2 hunks)
src/draive/generation/image/state.py (2 hunks)
src/draive/generation/model/state.py (3 hunks)
src/draive/generation/text/state.py (2 hunks)
src/draive/mistral/completions.py (8 hunks)
src/draive/mistral/embedding.py (4 hunks)
src/draive/mistral/moderation.py (2 hunks)
src/draive/models/generative.py (1 hunks)
src/draive/models/tools/function.py (4 hunks)
src/draive/models/types.py (3 hunks)
src/draive/ollama/chat.py (3 hunks)
src/draive/ollama/embedding.py (4 hunks)
src/draive/openai/embedding.py (4 hunks)
src/draive/openai/moderation.py (2 hunks)
src/draive/openai/realtime.py (4 hunks)
src/draive/openai/responses.py (9 hunks)
src/draive/postgres/__init__.py (1 hunks)
src/draive/postgres/configuration.py (0 hunks)
src/draive/postgres/memory.py (11 hunks)
src/draive/postgres/templates.py (3 hunks)
src/draive/postgres/vector_index.py (2 hunks)
src/draive/stages/stage.py (6 hunks)
src/draive/utils/memory.py (4 hunks)
src/draive/utils/vector_index.py (3 hunks)
src/draive/vllm/embedding.py (4 hunks)
src/draive/vllm/messages.py (6 hunks)
tests/test_conversation_completion.py (3 hunks)

💤 Files with no reviewable changes (1)

src/draive/postgres/configuration.py

🧰 Additional context used

🧠 Learnings (2)

📚 Learning: 2025-06-16T10:28:07.434Z

Learnt from: KaQuMiQ
Repo: miquido/draive PR: 338
File: src/draive/lmm/__init__.py:1-2
Timestamp: 2025-06-16T10:28:07.434Z
Learning: The draive project requires Python 3.12+ as specified in pyproject.toml with "requires-python = ">=3.12"" and uses Python 3.12+ specific features like PEP 695 type aliases and generic syntax extensively throughout the codebase.

Applied to files:

pyproject.toml

📚 Learning: 2025-06-16T10:09:51.508Z

Learnt from: KaQuMiQ
Repo: miquido/draive PR: 338
File: src/draive/conversation/types.py:123-144
Timestamp: 2025-06-16T10:09:51.508Z
Learning: In the draive.commons.metadata.Meta class, properties like `identifier`, `origin_identifier`, `predecessor_identifier`, and `successor_identifier` automatically parse stored hex strings into UUID objects when accessed, returning None if the string is missing or invalid. This means accessing `meta.identifier` returns a UUID object, not a string, even though it's stored internally as a hex string.

Applied to files:

src/draive/postgres/memory.py

🧬 Code graph analysis (13)

src/draive/vllm/messages.py (1)

src/draive/stages/stage.py (1)

completion (287-381)

src/draive/generation/image/state.py (1)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/ollama/embedding.py (1)

src/draive/stages/stage.py (1)

concurrent (1219-1276)

tests/test_conversation_completion.py (1)

src/draive/models/types.py (3)

ModelMemory (824-923)

ModelMemoryRecall (730-777)

recalling (838-841)

src/draive/generation/text/state.py (1)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/generation/audio/state.py (1)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/conversation/realtime/state.py (2)

src/draive/models/types.py (2)

ModelMemory (824-923)

constant (826-848)

src/draive/conversation/completion/state.py (1)

model_context_elements (113-119)

src/draive/models/types.py (3)

src/draive/utils/memory.py (13)

_remembering_none (38-42)

_maintaining_noop (45-48)

constant (54-66)

recall (70-73)

recall (76-79)

recall (82-89)

recalling (58-61)

remember (93-97)

remember (100-104)

remember (107-116)

maintenance (120-123)

maintenance (126-129)

maintenance (132-139)

src/draive/postgres/memory.py (3)

recall (81-112)

remember (114-176)

maintenance (178-222)

tests/test_stage.py (4)

recall (65-66)

recall (441-442)

remember (68-70)

remember (444-446)

src/draive/conversation/completion/state.py (4)

src/draive/models/types.py (2)

ModelMemory (824-923)

constant (826-848)

src/draive/utils/memory.py (1)

constant (54-66)

src/draive/conversation/realtime/state.py (1)

model_context_elements (73-79)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/models/generative.py (2)

src/draive/models/types.py (2)

ModelMemory (824-923)

constant (826-848)

src/draive/utils/memory.py (1)

constant (54-66)

src/draive/stages/stage.py (2)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/models/types.py (1)

ModelMemory (824-923)

src/draive/generation/model/state.py (1)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/postgres/memory.py (1)

src/draive/models/types.py (15)

ModelMemory (824-923)

recall (852-855)

recall (858-861)

recall (864-868)

ModelMemoryRecall (730-777)

of (283-305)

of (326-358)

of (386-422)

of (455-478)

of (545-569)

of (598-622)

of (746-772)

of (941-967)

of (994-1021)

of (1086-1114)

🔇 Additional comments (52)

src/draive/vllm/messages.py (4)

7-7: LGTM: Import updated correctly.

The addition of ctx to the imports is appropriate for the refactored observability calls throughout this file.

124-142: LGTM: Observability refactoring is clean and correct.

The replacement of ctx.record(ObservabilityLevel.INFO/DEBUG, ...) with ctx.record_info(...) and ctx.record_debug(...) improves code readability while preserving all recorded metadata. The separation of high-level configuration (INFO) from detailed context (DEBUG) is appropriate.

239-252: LGTM: Token usage metrics correctly recorded.

The refactored token usage recording maintains the same metric names, values, and attributes while using the cleaner ctx.record_info(...) API.

273-291: LGTM: Streaming observability consistently refactored.

The observability calls in the streaming path mirror the non-streaming implementation, correctly recording the same metadata with model.stream: True. The refactoring maintains consistency across both code paths.

src/draive/utils/vector_index.py (1)

94-100: LGTM!

The migration to ctx.record_info is clean, and recording the count of values (len(values)) instead of the full collection is more efficient for observability.

src/draive/postgres/vector_index.py (2)

41-59: Excellent documentation addition.

The example schema and index definition provide clear guidance for users setting up the database. The schema correctly matches what the code expects (JSONB for payload and meta, VECTOR for embedding), and the recommended ivfflat index with vector_cosine_ops aligns with the cosine distance operator (<#>) used in the search queries.

7-7: Observability instrumentation was removed, not replaced.

The git diff confirms the original review flagged a real inconsistency. The actual changes show that ctx.record() calls were removed entirely from the index(), search(), and delete() functions—they were not replaced with ctx.record_info calls as the AI summary claimed. The ObservabilityLevel import removal is consistent with this removal.

Verify this was intentional and that the loss of observability telemetry for these three functions is acceptable for your use case.

src/draive/openai/moderation.py (1)

3-3: LGTM! Observability API migration applied correctly.

The import and telemetry call have been successfully migrated from ObservabilityLevel-based recording to the specialized ctx.record_info helper, preserving the same attributes and control flow.

Also applies to: 29-34

src/draive/mistral/moderation.py (1)

4-4: LGTM! Observability API migration applied correctly.

The migration from ObservabilityLevel.INFO to ctx.record_info is consistent with the PR-wide refactoring, preserving the moderation flow and telemetry attributes.

Also applies to: 31-36

src/draive/openai/realtime.py (2)

30-30: LGTM! Observability API migration applied correctly.

The import changes and telemetry call have been successfully migrated to the specialized ctx.record_info helper, preserving the session initialization attributes.

Also applies to: 109-117

222-245: LGTM! Defensive default value handling for token metrics.

The migration to ctx.record_info includes defensive coding that defaults token counts to 0 when None, preventing potential issues with metric recording systems that may not handle None values gracefully.

src/draive/ollama/chat.py (1)

8-8: LGTM! Observability API migration applied correctly.

The import changes and telemetry calls have been successfully migrated to the specialized ctx.record_info and ctx.record_debug helpers, preserving all attributes and the completion flow logic.

Also applies to: 114-131

src/draive/bedrock/guardrails.py (1)

4-4: LGTM! Observability API migration applied correctly.

The migration from ObservabilityLevel.INFO to ctx.record_info is consistent with the PR-wide refactoring, preserving the guardrails verification flow and telemetry attributes.

Also applies to: 83-89

src/draive/generation/audio/state.py (1)

3-3: LGTM! Observability API migration applied correctly.

The import and template identifier telemetry calls have been successfully migrated to the specialized ctx.record_info helper, preserving the audio generation flow.

Also applies to: 44-51

src/draive/generation/text/state.py (1)

4-4: LGTM! Observability API migration applied correctly.

The import and template identifier telemetry calls have been successfully migrated to the specialized ctx.record_info helper, preserving the text generation flow.

Also applies to: 50-57

src/draive/evaluation/scenario.py (1)

5-13: LGTM! Observability API migration applied correctly.

The import changes and performance metric telemetry call have been successfully migrated to the specialized ctx.record_info helper, preserving all metric attributes and the evaluation scenario flow.

Also applies to: 401-410

src/draive/evaluation/evaluator.py (1)

677-687: LGTM! Observability API migration correctly applied.

The migration from ObservabilityLevel-based logging to ctx.record_info is implemented correctly with enriched metric payloads including histogram kind and properly structured attributes.

src/draive/ollama/embedding.py (4)

6-6: LGTM! Import statement correctly updated.

The removal of ObservabilityLevel import aligns with the migration to specialized ctx.record_* methods.

32-42: LGTM! Telemetry instrumentation properly structured.

The separation of key parameters in record_info and full configuration in record_debug provides appropriate observability granularity.

52-62: LGTM! Embedding items metric correctly structured.

The metric properly tracks item counts with appropriate counter kind and enriched attributes for filtering by provider, model, and type.

67-78: LGTM! Batch count metric correctly implemented.

The batch calculation and metric recording are correct, with proper counter semantics and consistent attribute structure.

src/draive/generation/model/state.py (2)

62-67: LGTM! Model generation telemetry correctly structured.

The observability call properly captures the generated model type and schema injection strategy.

84-91: LGTM! Template usage tracking properly implemented.

The conditional telemetry correctly captures template identifiers when templates are used for instructions or input.

src/draive/generation/image/state.py (1)

44-51: LGTM! Template tracking consistent with model generation.

The observability instrumentation follows the same pattern as model generation, ensuring consistent telemetry across generation types.

llms.txt (1)

215-215: LGTM! Documentation correctly updated.

The API documentation accurately reflects the migration from ctx.record to ctx.record_info.

src/draive/mistral/completions.py (4)

129-147: LGTM! Stream completion telemetry properly instrumented.

The observability calls correctly separate high-level parameters (info) from detailed context (debug), with the stream flag properly set.

180-199: LGTM! Token usage metrics correctly structured.

The token metrics properly use counter semantics with appropriate fallbacks for null values and consistent attribute structure.

311-329: LGTM! Non-stream telemetry consistent with stream variant.

The telemetry structure correctly mirrors the streaming version with the appropriate stream flag value.

549-562: LGTM! Usage metrics helper function correctly extracted.

The extracted function properly encapsulates token recording logic with consistent structure.

docs/guides/Basics.md (2)

189-189: LGTM! Documentation examples correctly updated.

The examples properly demonstrate the new ctx.record_info API usage.

228-235: LGTM! Observability examples demonstrate proper usage.

The nested scope examples correctly show how ctx.record_info works with LoggerObservability and nested contexts.

src/draive/postgres/__init__.py (1)

3-3: No issues found—review comment can be resolved.

Verification confirms the relocation is intentional and safe:

haiway~=0.38.0 is a declared core dependency

PostgresConfigurationRepository is properly imported from haiway.postgres and correctly exported in all

Other classes from haiway.postgres (Postgres, PostgresRow, PostgresValue) are actively integrated throughout the codebase without issues

No internal code depends on the old location—this is purely a public API change

Git history shows intentional consolidation ("Update haiway with Postgres proxy")

The interface is compatible and the public API surface remains unchanged for consumers.

src/draive/bedrock/converse.py (1)

123-213: LGTM! Observability API migration is clean and consistent.

The migration from ObservabilityLevel-based calls to specialized ctx.record_info and ctx.record_debug helpers is well-executed. The enriched token metrics now include provider and model attributes, which improves observability granularity.

src/draive/models/tools/function.py (1)

275-316: LGTM! Error handling instrumentation is properly maintained.

The migration to ctx.record_info and ctx.record_error preserves all error attributes and maintains the existing exception handling logic.

src/draive/postgres/templates.py (1)

31-51: Helpful SQL schema documentation.

The added Notes section with example schema and index definition improves developer experience and clarifies table structure expectations.

src/draive/openai/responses.py (2)

179-291: LGTM! Comprehensive observability migration for completion path.

Both the observability API migration and token metrics enrichment are well-implemented. The provider and model attributes add valuable context for monitoring.

421-593: LGTM! Streaming path observability is consistent.

The streaming path follows the same instrumentation pattern as the non-streaming path, maintaining consistency across the module.

src/draive/gemini/embedding.py (1)

33-76: LGTM! Embedding instrumentation is comprehensive and consistent.

The enriched observability payloads include dimensions, batch_size, provider, model, and type attributes. The use of kind="counter" for items and batches metrics is appropriate.

src/draive/vllm/embedding.py (1)

33-76: LGTM! VLLM embedding instrumentation matches the standard pattern.

The observability implementation is consistent with other embedding modules in this PR.

src/draive/cohere/embedding.py (2)

33-76: LGTM! Text embedding instrumentation is consistent.

The observability migration follows the standard pattern with enriched metadata including batch_size and provider/model/type attributes.

126-167: LGTM! Image embedding instrumentation mirrors text pattern.

The image embedding path uses the same instrumentation strategy, ensuring consistency across embedding types.

src/draive/openai/embedding.py (1)

33-77: LGTM! OpenAI embedding instrumentation is clean and consistent.

The observability migration aligns with the standard pattern established across all embedding modules.

src/draive/mistral/embedding.py (1)

6-74: LGTM: Observability API migration is consistent.

The migration from ObservabilityLevel-based ctx.record(...) to specialized ctx.record_info(...) is correctly implemented. The addition of structured metrics with unit, kind, and attributes enriches the telemetry data without altering functional behavior.

pyproject.toml (1)

8-44: LGTM: Version bumps are appropriate.

The version bump to 0.93.0 and dependency updates (haiway ~0.38.0, openai ~2.8) align with the observability API migration and broader codebase changes.

src/draive/gemini/generating.py (1)

35-467: LGTM: Observability API migration is thorough.

The migration from ObservabilityLevel-based logging to specialized ctx.record_info(...), ctx.record_debug(...), and ctx.record_warning(...) is correctly applied across all code paths (non-streaming, streaming, error handling, and usage metrics).

src/draive/anthropic/messages.py (1)

144-307: LGTM: Observability migration is consistent.

The observability surface has been correctly updated to use ctx.record_info(...), ctx.record_debug(...), and ctx.record_warning(...). Token metrics now include structured attributes for better observability.

src/draive/utils/memory.py (1)

3-139: LGTM: Memory observability updated correctly.

The migration to ctx.record_info(...) is consistent, and the addition of item count in the memory.remember event provides useful telemetry data.

src/draive/models/generative.py (1)

680-683: LGTM: Memory normalization updated correctly.

The change from wrapping with ModelMemoryRecall.empty to directly expanding the sequence into ModelMemory.constant(*memory) aligns with the new ModelMemory.constant(...) signature that accepts variadic ModelContextElement arguments.

tests/test_conversation_completion.py (1)

8-21: LGTM: Test imports updated for new memory API.

The addition of META_EMPTY, ModelMemory, and ModelMemoryRecall imports correctly reflects the updated memory abstraction surface in the codebase.

src/draive/conversation/realtime/state.py (1)

4-103: LGTM: Memory handling and observability updated correctly.

The memory handling logic now properly differentiates between ModelMemory instances and message iterables, using memory.recall() for the former and expanding a generator into ModelMemory.constant(...) for the latter. The observability migration to ctx.record_info(...) is consistent with the broader refactor.

src/draive/models/types.py (1)

823-923: LGTM! Well-structured ModelMemory implementation.

The ModelMemory class is well-designed with:

Clean separation of concerns through protocol-based callables

Proper use of @statemethod for dual class/instance method support

Consistent factory pattern with constant() for immutable memory

Type-safe overloads for recall(), remember(), and maintenance()

src/draive/stages/stage.py (1)

2001-2065: LGTM! Memory type migration correctly implemented.

The execute method properly handles both ModelMemory and ModelContext:

Recalls initial context from ModelMemory when provided

Persists only newly produced context elements (suffix after common prefix)

Maintains backward compatibility with direct ModelContext usage

src/draive/utils/vector_index.py

coderabbitai

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/draive/postgres/memory.py (1)
178-222: Fix maintenance function column reference to match schema and remember function.

The maintenance function (lines 208-210) inserts into identifier instead of memories column. The schema defines memories_variables.memories as a foreign key to the memories table, while identifier is the primary key. The remember function (lines 131-132) correctly uses memories.

Update the INSERT statement to use the correct foreign key column:
                 if variables is not None:
                     ctx.log_info(f"...remembering {len(variables)} variables...")
                     await connection.execute(
                         """
                         INSERT INTO
                             memories_variables (
-                                identifier,
+                                memories,
                                 variables
                             )
 
                         VALUES (
                             $1::UUID,
                             $2::JSONB
                         );
                         """,
                         identifier,
                         json.dumps(variables),
                     )

♻️ Duplicate comments (1)

src/draive/postgres/templates.py (1)

5-5: Observability pattern inconsistency (previously flagged).

As noted in the previous review, this file removes observability records entirely instead of migrating them to ctx.record_info/debug/warning like other modules in this PR. For consistency with the observability improvements across the codebase, consider adding structured metrics alongside the existing ctx.log_info calls.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: ASSERTIVE

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 3fdddae and d7d042b.

⛔ Files ignored due to path filters (1)

uv.lock is excluded by !**/*.lock

📒 Files selected for processing (40)

docs/guides/Basics.md (2 hunks)
llms.txt (1 hunks)
pyproject.toml (3 hunks)
src/draive/anthropic/messages.py (7 hunks)
src/draive/bedrock/converse.py (5 hunks)
src/draive/bedrock/guardrails.py (2 hunks)
src/draive/cohere/embedding.py (7 hunks)
src/draive/conversation/completion/state.py (3 hunks)
src/draive/conversation/realtime/state.py (3 hunks)
src/draive/evaluation/evaluator.py (1 hunks)
src/draive/evaluation/scenario.py (1 hunks)
src/draive/gemini/embedding.py (4 hunks)
src/draive/gemini/generating.py (12 hunks)
src/draive/generation/audio/state.py (2 hunks)
src/draive/generation/image/state.py (2 hunks)
src/draive/generation/model/state.py (3 hunks)
src/draive/generation/text/state.py (2 hunks)
src/draive/mistral/completions.py (8 hunks)
src/draive/mistral/embedding.py (4 hunks)
src/draive/mistral/moderation.py (2 hunks)
src/draive/models/generative.py (1 hunks)
src/draive/models/tools/function.py (4 hunks)
src/draive/models/types.py (3 hunks)
src/draive/ollama/chat.py (3 hunks)
src/draive/ollama/embedding.py (4 hunks)
src/draive/openai/embedding.py (4 hunks)
src/draive/openai/moderation.py (2 hunks)
src/draive/openai/realtime.py (4 hunks)
src/draive/openai/responses.py (9 hunks)
src/draive/postgres/__init__.py (1 hunks)
src/draive/postgres/configuration.py (0 hunks)
src/draive/postgres/memory.py (11 hunks)
src/draive/postgres/templates.py (3 hunks)
src/draive/postgres/vector_index.py (2 hunks)
src/draive/stages/stage.py (6 hunks)
src/draive/utils/memory.py (4 hunks)
src/draive/utils/vector_index.py (4 hunks)
src/draive/vllm/embedding.py (4 hunks)
src/draive/vllm/messages.py (6 hunks)
tests/test_conversation_completion.py (3 hunks)

💤 Files with no reviewable changes (1)

src/draive/postgres/configuration.py

🧰 Additional context used

🧠 Learnings (1)

📚 Learning: 2025-06-16T10:28:07.434Z

Learnt from: KaQuMiQ
Repo: miquido/draive PR: 338
File: src/draive/lmm/__init__.py:1-2
Timestamp: 2025-06-16T10:28:07.434Z
Learning: The draive project requires Python 3.12+ as specified in pyproject.toml with "requires-python = ">=3.12"" and uses Python 3.12+ specific features like PEP 695 type aliases and generic syntax extensively throughout the codebase.

Applied to files:

pyproject.toml

🧬 Code graph analysis (12)

src/draive/models/generative.py (2)

src/draive/models/types.py (2)

ModelMemory (824-923)

constant (826-848)

src/draive/utils/memory.py (1)

constant (54-66)

src/draive/generation/model/state.py (1)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/vllm/messages.py (1)

src/draive/stages/stage.py (1)

completion (287-381)

src/draive/conversation/completion/state.py (4)

src/draive/models/types.py (2)

ModelMemory (824-923)

constant (826-848)

src/draive/utils/memory.py (1)

constant (54-66)

src/draive/conversation/realtime/state.py (1)

model_context_elements (73-79)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/generation/text/state.py (1)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

tests/test_conversation_completion.py (1)

src/draive/models/types.py (3)

ModelMemory (824-923)

ModelMemoryRecall (730-777)

recalling (838-841)

src/draive/stages/stage.py (2)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/models/types.py (1)

ModelMemory (824-923)

src/draive/generation/image/state.py (1)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/conversation/realtime/state.py (4)

src/draive/models/types.py (2)

ModelMemory (824-923)

constant (826-848)

src/draive/utils/memory.py (1)

constant (54-66)

src/draive/conversation/completion/state.py (1)

model_context_elements (113-119)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/generation/audio/state.py (1)

src/draive/multimodal/templates/types.py (1)

Template (39-139)

src/draive/postgres/memory.py (1)

src/draive/models/types.py (15)

ModelMemory (824-923)

recall (852-855)

recall (858-861)

recall (864-868)

ModelMemoryRecall (730-777)

of (283-305)

of (326-358)

of (386-422)

of (455-478)

of (545-569)

of (598-622)

of (746-772)

of (941-967)

of (994-1021)

of (1086-1114)

src/draive/models/types.py (2)

src/draive/utils/memory.py (13)

_remembering_none (38-42)

_maintaining_noop (45-48)

constant (54-66)

recall (70-73)

recall (76-79)

recall (82-89)

recalling (58-61)

remember (93-97)

remember (100-104)

remember (107-116)

maintenance (120-123)

maintenance (126-129)

maintenance (132-139)

src/draive/postgres/memory.py (3)

recall (81-112)

remember (114-176)

maintenance (178-222)

🔇 Additional comments (50)

docs/guides/Basics.md (1)

189-189: LGTM! Documentation updated correctly.

The documentation examples have been updated to reflect the new observability API surface, replacing ctx.record(...) with ctx.record_info(...).

Also applies to: 228-228, 235-235

src/draive/utils/vector_index.py (3)

94-100: LGTM! Observability enhanced with value count.

The observability call has been updated to use ctx.record_info and now includes the count of values being indexed, which provides useful telemetry.

170-183: LGTM! Past review concerns addressed.

The observability refactor resolves previous concerns:

Timing fixed: The event is now recorded before self.searching() (line 174), ensuring search failures remain observable.

Attributes simplified: Misleading boolean attributes ("query", "requirements") have been removed; only essential "model" attribute remains.

215-218: LGTM! Attributes simplified.

The observability payload has been streamlined to include only the essential "model" attribute, addressing previous concerns about misleading attribute names.

llms.txt (1)

215-215: LGTM! Documentation updated correctly.

The LLM usage guide has been updated to reflect the new observability API surface.

src/draive/ollama/chat.py (1)

8-8: LGTM! Observability API migration complete.

The code has been updated to use the new specialized observability helpers (ctx.record_info and ctx.record_debug) instead of ObservabilityLevel-based calls, with no changes to the underlying logic.

Also applies to: 114-131

src/draive/openai/moderation.py (1)

3-3: LGTM! Observability API migration complete.

The code has been updated to use ctx.record_info instead of ObservabilityLevel-based calls, with no changes to the moderation logic.

Also applies to: 29-34

src/draive/generation/audio/state.py (1)

3-3: LGTM! Observability API migration complete.

The code has been updated to use ctx.record_info for template tracking instead of ObservabilityLevel-based calls, with no changes to the audio generation logic.

Also applies to: 44-51

src/draive/mistral/moderation.py (1)

4-4: LGTM! Observability API migration complete.

The code has been updated to use ctx.record_info instead of ObservabilityLevel-based calls, with no changes to the content verification logic.

Also applies to: 31-36

src/draive/utils/memory.py (1)

86-88: LGTM! Clean observability upgrade with useful metadata enrichment.

The migration from ObservabilityLevel.INFO to ctx.record_info improves code clarity. The addition of attributes={"items": len(items)} in the remember operation (line 114) provides valuable operational insight for monitoring memory operations.

Also applies to: 112-115, 136-138

src/draive/evaluation/scenario.py (1)

401-410: LGTM! Excellent metric metadata enrichment.

The migration to ctx.record_info with kind="histogram" and unit="%" properly types this performance metric. The additional attributes (passed, evaluators) provide valuable context for observability dashboards and analysis.

src/draive/generation/model/state.py (1)

62-67: LGTM! Granular observability for model generation.

The three separate ctx.record_info calls provide excellent visibility into model generation operations: configuration attributes, template usage for instructions, and template usage for input. This granular approach enables targeted monitoring and debugging.

Also applies to: 84-91

src/draive/generation/image/state.py (1)

44-51: LGTM! Consistent template tracking.

The conditional template tracking for instructions and input mirrors the pattern used in text and model generation, ensuring consistent observability across all generation modules.

src/draive/generation/text/state.py (1)

50-57: LGTM! Consistent implementation across generation modules.

Template tracking implementation is consistent with image and model generation, maintaining uniform observability patterns across the codebase.

src/draive/evaluation/evaluator.py (1)

677-687: LGTM! Proper metric typing for evaluator performance.

The addition of kind="histogram" and unit="%" correctly types this performance metric. The comprehensive attributes (passed, threshold, score) enable rich observability and analysis of evaluation results.

src/draive/openai/responses.py (2)

179-198: LGTM! Excellent separation of info and debug levels.

The split between ctx.record_info for operational metadata and ctx.record_debug for detailed context (instructions, tools, context) follows best practices for log level separation. This enables appropriate filtering in production environments.

272-291: LGTM! Proper token metric typing with cost tracking attributes.

Token metrics are correctly typed as kind="counter" with unit="tokens", and include provider/model attributes. This enables accurate cost tracking and usage analysis across different models.

Also applies to: 574-593

src/draive/bedrock/guardrails.py (1)

83-89: LGTM! Essential guardrails configuration tracking.

The ctx.record_info call captures key guardrails configuration (provider, identifier, version) for operational visibility and debugging.

src/draive/mistral/completions.py (1)

6-6: LGTM! Observability refactoring is consistent.

The changes correctly migrate from ObservabilityLevel-based logging to the specialized ctx.record_info and ctx.record_debug helpers while preserving all observability data payloads. The pattern is applied consistently across both streaming and non-streaming completion paths, as well as usage metrics recording.

Also applies to: 129-147, 180-199, 311-329, 549-562

src/draive/openai/realtime.py (1)

30-30: LGTM! Observability migration is correct.

The migration from ObservabilityLevel-based calls to ctx.record_info is consistent and correct. The addition of default values (if usage.input_tokens is not None else 0) for token metrics ensures robust telemetry even when usage data is incomplete.

Also applies to: 109-117, 222-245

src/draive/postgres/vector_index.py (1)

40-59: Good addition of schema documentation.

The example schema and index definition in the Notes section provides helpful guidance for users implementing Postgres vector storage.

src/draive/postgres/templates.py (1)

31-51: Good addition of schema documentation.

The example schema and index definition in the Notes section provides helpful guidance for users implementing Postgres template storage.

src/draive/gemini/generating.py (1)

35-35: LGTM! Complete and consistent observability refactoring.

The changes correctly migrate from ObservabilityLevel-based logging to specialized helpers (ctx.record_info, ctx.record_debug, ctx.record_warning) throughout both streaming and non-streaming completion paths. Rate limit handling appropriately uses ctx.record_warning, and token metrics include proper metadata (kind="counter", unit="tokens").

Also applies to: 145-164, 186-197, 221-232, 291-310, 349-360, 401-413, 438-467

src/draive/anthropic/messages.py (1)

144-164: LGTM! Observability refactoring is thorough and correct.

The migration from ObservabilityLevel-based calls to specialized helpers is complete and consistent. Rate limit handling appropriately uses ctx.record_warning with retry_after tracking, and token metrics now include kind="counter" for better observability. The addition of cached input token tracking (model.input_tokens.cached) is a valuable enhancement.

Also applies to: 231-238, 256-268, 277-307

src/draive/bedrock/converse.py (1)

5-5: LGTM! Observability refactoring is consistent.

The migration from ObservabilityLevel-based calls to ctx.record_info and ctx.record_debug is complete and correct. Token metrics now include proper metadata (unit="tokens", kind="counter", provider/model attributes) for enhanced observability.

Also applies to: 123-142, 194-213

src/draive/models/generative.py (1)

680-683: Original review comment cannot be verified and appears inaccurate.

The current code at lines 681-682 is correct and type-safe. When memory is a Sequence, it's correctly identified as ModelContext (which is Sequence[ModelContextElement]), and unpacking with *memory properly passes each element to ModelMemory.constant(*elements: ModelContextElement).

However, the review comment's central claim—that this represents "a change from the previous double-wrapping pattern"—cannot be verified. No git diff shows recent changes to these lines, and no evidence of a previous "double-wrapping pattern" exists in the codebase. The comment makes assertions about what was changed without supporting evidence.

Likely an incorrect or invalid review comment.

src/draive/vllm/embedding.py (1)

6-6: LGTM! Observability refactoring is consistent and well-structured.

The migration from ObservabilityLevel-based logging to specialized ctx.record_info(...) calls enriches observability data with structured metrics (counter semantics, units) and attributes (provider, model, type, dimensions, batch_size). This pattern is consistent across the embedding module.

Also applies to: 33-76

src/draive/gemini/embedding.py (1)

7-7: LGTM! Consistent observability improvements.

The changes mirror the pattern across other embedding modules, removing ObservabilityLevel and adopting structured ctx.record_info(...) with richer metadata.

Also applies to: 33-76

src/draive/cohere/embedding.py (1)

8-8: LGTM! Observability refactoring covers both text and image embeddings.

The changes systematically apply the new observability pattern to both text and image embedding flows, maintaining consistency with appropriate embedding.type attributes ("text" vs "image").

Also applies to: 33-76, 126-167

src/draive/conversation/completion/state.py (2)

103-121: Memory abstraction update is correct.

The changes properly adapt to the ModelMemory API:

Line 103: ModelMemory.constant() correctly creates an empty memory (the old code passing ModelMemoryRecall.empty was incorrect for this signature)

Line 121: Passing elements directly via *model_context_elements() aligns with ModelMemory.constant(*elements, ...) signature

125-127: LGTM! Observability logging updated.

Consistent migration to ctx.record_info(...) for template identifier tracking.

src/draive/vllm/messages.py (2)

124-142: LGTM! Well-structured observability with appropriate log levels.

The changes properly distinguish between:

ctx.record_info(...) for high-level metadata and token metrics with counter semantics

ctx.record_debug(...) for detailed context (instructions, tools, context strings)

This provides better observability granularity.

Also applies to: 239-252

273-291: LGTM! Streaming observability matches completion pattern.

Consistent instrumentation for the streaming code path.

src/draive/models/tools/function.py (1)

275-316: LGTM! Proper log level usage for tool execution.

The changes appropriately use:

ctx.record_info(...) for successful tool invocations

ctx.record_error(...) for exception cases with proper event and exception details

src/draive/mistral/embedding.py (1)

6-6: LGTM! Consistent observability pattern.

The changes follow the same systematic refactoring applied across all embedding modules.

Also applies to: 32-74

src/draive/stages/stage.py (2)

349-356: LGTM! Template logging updated consistently.

All template identifier logging migrated to ctx.record_info(...) across completion, prompting_completion, loopback_completion, and result_completion methods.

Also applies to: 456-458, 553-555, 624-626

2001-2065: Memory abstraction updated in execute() method.

The changes properly migrate from Memory to ModelMemory:

Line 2004: Parameter type updated to ModelMemory | ModelContext

Lines 2034, 2049: Type checks updated to isinstance(memory, ModelMemory)

This aligns with the broader memory abstraction refactoring across the codebase.

Note: This is a breaking change to the public API of Stage.execute(). Callers passing Memory instances will need to migrate to ModelMemory.

src/draive/ollama/embedding.py (1)

32-78: LGTM! Observability API migration looks clean.

The migration from ObservabilityLevel-based logging to ctx.record_info and ctx.record_debug is well-structured, and the enriched telemetry (provider, model, type, batch_size, dimensions) will improve observability.

src/draive/conversation/realtime/state.py (2)

66-82: LGTM! Memory handling correctly migrated to ModelMemory abstraction.

The branching logic correctly distinguishes between ModelMemory instances (which are recalled) and iterable message sequences (which are converted to constant memory). The use of ModelMemory.constant(...) aligns with the new memory API.

85-87: LGTM! Observability migration is consistent.

The switch to ctx.record_info with template attributes follows the same pattern used across the codebase.

src/draive/openai/embedding.py (1)

33-77: LGTM! Observability API migration is consistent with other embedding modules.

The enriched telemetry follows the same pattern as src/draive/ollama/embedding.py, ensuring uniform observability across embedding backends.

tests/test_conversation_completion.py (2)

8-8: LGTM! Test imports correctly reference the new public API.

The addition of META_EMPTY, ModelMemory, and ModelMemoryRecall imports aligns with the memory abstraction refactor.

Also applies to: 19-21

210-217: LGTM! Typo fix applied and memory construction is correct.

The maintenance callback is now correctly named, and the ModelMemory construction properly wires the recalling, remembering, and maintaining fields.

src/draive/models/types.py (2)

783-821: LGTM! Protocol definitions and noop helpers are well-structured.

The runtime-checkable protocols clearly define the memory operation contracts, and the noop helpers (_remembering_none, _maintaining_noop) provide sensible defaults for constant/read-only memory scenarios.

823-923: LGTM! ModelMemory abstraction is well-designed.

The implementation provides:

A clean factory method (constant) for static memory

Proper overloading via @statemethod for both class and instance usage

Clear separation of concerns (recalling, remembering, maintaining)

Metadata support via the meta field

The design aligns well with the State pattern and supports both constant and mutable memory backends.

src/draive/postgres/memory.py (5)

22-24: LGTM! Type safety improvements with UUID and clearer recall_limit semantics.

Switching identifier from str to UUID improves type safety and aligns with the database schema. The recall_limit: int = 0 default with the assertion on line 79 makes the "unbounded recall" behavior explicit.

42-77: LGTM! Schema documentation now matches implementation.

The updated example schema correctly uses memories as the foreign key column in both memories_variables and memories_elements, matching how the SQL queries reference it throughout the module.

86-112: LGTM! Recall flow enhanced with logging and metadata.

The addition of detailed logging and ctx.record_info improves observability. Including variables in the ModelMemoryRecall construction (line 111) properly surfaces memory variables to callers.

237-278: LGTM! Context loading correctly uses UUID and memories FK.

The SQL queries properly use memories = $1::UUID to filter by the foreign key, and the conditional branching on limit handles both bounded and unbounded recall correctly.

322-350: LGTM! Variables loading correctly updated to use UUID and memories FK.

The query properly filters by memories = $1::UUID (line 335) to fetch the latest variables for the given memory identifier.

src/draive/postgres/__init__.py

src/draive/postgres/vector_index.py

KaQuMiQ force-pushed the feature/postgre branch from d61b191 to 99bd2e0 Compare November 17, 2025 11:00

coderabbitai bot reviewed Nov 17, 2025

View reviewed changes

src/draive/postgres/memory.py Show resolved Hide resolved

src/draive/postgres/templates.py Show resolved Hide resolved

tests/test_conversation_completion.py Outdated Show resolved Hide resolved

coderabbitai bot reviewed Nov 17, 2025

View reviewed changes

src/draive/postgres/memory.py Show resolved Hide resolved

src/draive/postgres/memory.py Show resolved Hide resolved

tests/test_conversation_completion.py Outdated Show resolved Hide resolved

KaQuMiQ force-pushed the feature/postgre branch from 99bd2e0 to 98c24b2 Compare November 17, 2025 11:30

coderabbitai bot reviewed Nov 17, 2025

View reviewed changes

src/draive/utils/vector_index.py Show resolved Hide resolved

src/draive/utils/vector_index.py Show resolved Hide resolved

KaQuMiQ force-pushed the feature/postgre branch from 98c24b2 to 3fdddae Compare November 17, 2025 12:02

Update Postgres helpers

d7d042b

KaQuMiQ force-pushed the feature/postgre branch from 3fdddae to d7d042b Compare November 17, 2025 12:27

coderabbitai bot reviewed Nov 17, 2025

View reviewed changes

src/draive/postgres/__init__.py Show resolved Hide resolved

src/draive/postgres/vector_index.py Show resolved Hide resolved

KaQuMiQ merged commit cfa0c55 into main Nov 17, 2025
3 checks passed

KaQuMiQ deleted the feature/postgre branch November 17, 2025 12:41

Comments

Conversation

KaQuMiQ commented Nov 17, 2025

Uh oh!

coderabbitai bot commented Nov 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Estimated code review effort

Possibly related PRs

Pre-merge checks and finishing touches

Uh oh!

KaQuMiQ commented Nov 17, 2025

Uh oh!

coderabbitai bot commented Nov 17, 2025

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

coderabbitai bot commented Nov 17, 2025 •

edited

Loading