Conversation
Walkthrough
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Pre-merge checks and finishing touches❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests
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. Comment |
There was a problem hiding this comment.
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/__init__.py (1)
335-350: Deduplicate repeated names in all.
There are duplicate entries: Resource, ResourceContent, ResourceMissing, ResourceReference, ResourceTemplate, resource. Clean to avoid confusing users and docs.Apply minimal removals:
- "Resource", - "ResourceContent", - "ResourceMissing", - "ResourceReference", - "ResourceTemplate", - "resource",Also applies to: 399-400
src/draive/gemini/generating.py (1)
326-356: Minor: log safety once per response, not per chunk.
The current loop may emit many safety logs; consider logging only on the first candidate that provides ratings to reduce noise in DEBUG/INFO.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (4)
pyproject.toml(3 hunks)src/draive/__init__.py(2 hunks)src/draive/gemini/generating.py(4 hunks)src/draive/postgres/__init__.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (3)
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: Use absolute imports from the draive package (e.g., from draive import ...), not relative imports
Follow Ruff import ordering: standard library, third-party, then local imports
Use Python 3.12+ typing features (e.g., | unions, PEP 695 generics)
Prefer abstract/base collection types (Sequence, Mapping, Set, Iterable) over concrete list, dict, set in type hints
Define and raise custom exception types for specific error cases
Use NumPy docstring convention for all functions, classes, and methods
Skip module-level docstrings unless explicitly requested
Docstrings should include Parameters, Returns, Raises, and Notes (if needed) sectionsFormat all Python code with Ruff formatter (run via
make format); do not use any other formatter
Files:
src/draive/postgres/__init__.pysrc/draive/__init__.pysrc/draive/gemini/generating.py
src/draive/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/draive/**/*.py: Import Haiway symbols directly asfrom haiway import State, ctx
Usectx.scope(...)to bind activeStateinstances; avoid global state
Route all logging throughctx.log_*and never use
Enforce strict typing on public APIs; avoid untyped definitions andAnyexcept at verified third‑party boundaries
Prefer explicit attribute access with static types; avoid dynamicgetattrexcept at narrow boundaries
Use abstract immutable protocols (Mapping,Sequence,Iterable) instead of concretedict/list/setin public types
Usefinalwhere applicable; prefer composition over complex inheritance
Use precise unions (|), narrow withmatch/isinstance, and avoidcastunless provably safe and localized
Model immutable configuration/data ashaiway.State; avoid in‑place mutation and preferState.updated(...)/functional builders
Access active state viahaiway.ctxwithin async scopes (ctx.scope(...))
Public state methods that dispatch on the active instance must use@statemethod
Add structured, concise logs around generation calls, tool dispatch, and provider requests/responses (without secrets); add metrics viactx.recordwhere applicable
All I/O is async; keep boundaries async, usectx.spawnfor detached tasks, and avoid custom threading
Translate provider/SDK errors into typed exceptions; never raise bareException; preserve context
Compose multimodal content withMultimodalContent.of(...); useResourceContent/Referencefor blobs; wrap custom data inArtifactContent(use hidden when needed)
Public symbols require NumPy-style docstrings with Parameters/Returns/Raises and rationale when non-obvious
Avoid docstrings for internal helpers; prefer self-explanatory names; skip module-level docstrings
Never log secrets or full request bodies containing keys/tokens
Read credentials from environment variables; resolve via helpers likegetenv_str
Files:
src/draive/postgres/__init__.pysrc/draive/__init__.pysrc/draive/gemini/generating.py
src/draive/__init__.py
📄 CodeRabbit inference engine (AGENTS.md)
Centralize and update all public exports in
src/draive/__init__.pywhen API surface changes
Files:
src/draive/__init__.py
🧠 Learnings (4)
📓 Common learnings
Learnt from: CR
PR: miquido/draive#0
File: AGENTS.md:0-0
Timestamp: 2025-09-11T15:22:15.967Z
Learning: Applies to src/draive/__init__.py : Centralize and update all public exports in `src/draive/__init__.py` when API surface changes
📚 Learning: 2025-09-11T15:22:15.967Z
Learnt from: CR
PR: miquido/draive#0
File: AGENTS.md:0-0
Timestamp: 2025-09-11T15:22:15.967Z
Learning: Applies to src/draive/__init__.py : Centralize and update all public exports in `src/draive/__init__.py` when API surface changes
Applied to files:
src/draive/postgres/__init__.pysrc/draive/__init__.py
📚 Learning: 2025-09-11T15:22:15.967Z
Learnt from: CR
PR: miquido/draive#0
File: AGENTS.md:0-0
Timestamp: 2025-09-11T15:22:15.967Z
Learning: Applies to src/draive/**/*.py : Import Haiway symbols directly as `from haiway import State, ctx`
Applied to files:
src/draive/__init__.py
📚 Learning: 2025-09-08T10:32:32.520Z
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-09-08T10:32:32.520Z
Learning: Applies to draive/**/__init__.py : Put exported public symbols into package __init__.py files
Applied to files:
src/draive/__init__.py
🧬 Code graph analysis (1)
src/draive/gemini/generating.py (2)
src/draive/gemini/config.py (2)
GeminiConfig(25-36)GeminiSafetyConfig(16-22)src/draive/models/types.py (4)
content(458-462)content(597-601)ModelOutputFailed(189-212)ModelOutputLimit(654-682)
🔇 Additional comments (5)
src/draive/__init__.py (2)
28-28: Re-exporting Immutable looks good.
Import from haiway is consistent with our public API pattern.
264-264: Added "Immutable" to all — OK.
Public surface stays aligned with the new import.src/draive/gemini/generating.py (1)
29-32: Import split for GeminiConfig/SafetyConfig — OK.
Ordering complies with Ruff and local-import rules.src/draive/postgres/__init__.py (1)
1-17: Straightforward re-exports — looks good.
API is concise and matches the optional dependency group.pyproject.toml (1)
24-24: No action needed — haiway~=0.32.0 is available on PyPI.
PyPI releases include 0.32.0, so the constraint will resolve as written.
| "opentelemetry-exporter-otlp-proto-grpc", | ||
| ] | ||
| httpx = ["haiway[httpx]", "httpx"] | ||
| postgres = ["haiway[postgres]", "asyncpg"] |
There was a problem hiding this comment.
🛠️ Refactor suggestion
🧩 Analysis chain
Optional group “postgres”: confirm duplication and pin asyncpg.
If haiway[postgres] already depends on asyncpg, listing asyncpg here duplicates it. If not, please pin a compatible range (e.g., asyncpg>=0.30,<1) to avoid resolver drift. Latest stable of asyncpg is 0.30.0. (pypi.org)
Option A (if haiway extra includes driver):
-postgres = ["haiway[postgres]", "asyncpg"]
+postgres = ["haiway[postgres]"]Option B (if not included):
-postgres = ["haiway[postgres]", "asyncpg"]
+postgres = ["haiway[postgres]", "asyncpg>=0.30,<1"]🌐 Web query:
Does the haiway package define a "postgres" extra that installs asyncpg?
💡 Result:
No — the haiway package does not expose a "postgres" extra that installs asyncpg. (haiway's PyPI project page and repository show no extras for a postgres/asyncpg extra). [1][2]
Sources:
[1] haiway on PyPI. [2] miquido/haiway repository (GitHub).
Pin asyncpg in pyproject.toml — haiway has no postgres extra
haiway does not expose a "postgres" extra that installs asyncpg — pin asyncpg to avoid resolver drift.
-postgres = ["haiway[postgres]", "asyncpg"]
+postgres = ["haiway[postgres]", "asyncpg>=0.30,<1"]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| postgres = ["haiway[postgres]", "asyncpg"] | |
| postgres = ["haiway[postgres]", "asyncpg>=0.30,<1"] |
🤖 Prompt for AI Agents
In pyproject.toml at line 50, the dependency list uses "asyncpg" unpinned while
haiway does not provide a "postgres" extra; update that entry to pin asyncpg to
a specific, compatible version range (for example a maintained minor range like
>=0.26.0,<0.28.0) so the resolver won’t drift — edit the postgres line to
replace the bare "asyncpg" with the chosen version specifier and run poetry
lock/resolve to verify.
| if finish_reason == FinishReason.SAFETY: | ||
| raise ModelOutputFailed( | ||
| provider="gemini", | ||
| model=config.model, | ||
| reason=f"Safety filtering: {finish_message or ''}", | ||
| ) | ||
|
|
||
| if finish_reason == FinishReason.MAX_TOKENS: | ||
| raise ModelOutputLimit( | ||
| provider="gemini", | ||
| model=config.model, | ||
| max_output_tokens=unwrap_missing( | ||
| config.max_output_tokens, | ||
| default=0, | ||
| ), | ||
| content=tuple(collected_blocks), | ||
| ) | ||
|
|
||
| if finish_reason not in (None, FinishReason.STOP): | ||
| raise ModelOutputFailed( | ||
| provider="gemini", | ||
| model=config.model, | ||
| reason=( | ||
| f"Completion error: {finish_message}" | ||
| if finish_message | ||
| else "Completion error" | ||
| ), | ||
| ) | ||
|
|
There was a problem hiding this comment.
Don’t swallow ModelOutputLimit/ModelOutputFailed raised post-stream.
The broad except Exception will catch the deliberate raises after the stream ends, converting them into ModelOutputFailed and losing semantics. Let our control‑flow exceptions pass through.
Apply a focused guard before the broad catch:
except ResourceExhausted as exc:
ctx.record(
ObservabilityLevel.WARNING,
event="model.rate_limit",
attributes={
"model.provider": "gemini",
"model.name": config.model,
},
)
# Propagate as ModelRateLimit with randomized backoff window
raise ModelRateLimit(
provider="gemini",
model=config.model,
retry_after=random.uniform(*RATE_LIMIT_RETRY_RANGE), # nosec: B311
) from exc
+ except (ModelOutputFailed, ModelOutputLimit):
+ # Preserve intended control-flow exceptions from end-of-stream checks
+ raise
+
except Exception as exc:
# Convert to ModelOutputFailed for consistency
raise ModelOutputFailed(
provider="gemini",
model=config.model,
reason=str(exc),
) from excAlso applies to: 412-419
🤖 Prompt for AI Agents
In src/draive/gemini/generating.py around lines 367-395 (and similarly 412-419),
the current broad except Exception block swallows deliberate
ModelOutputLimit/ModelOutputFailed exceptions raised after the stream; update
the exception handler so that it re-raises those control-flow exceptions instead
of converting them. Concretely, change the generic except to capture the
exception as e and add an immediate guard like: if isinstance(e,
(ModelOutputLimit, ModelOutputFailed)): raise; otherwise handle/log the
exception as before; ensure the names are in scope or imported as needed.
No description provided.