Add dedicated pgvector connection extension

[KaQuMiQ]

No description provided.

[coderabbitai]

Walkthrough

Adds pgvector integration guidance to the Postgres guide, introducing an initialize_pgvector(connection: Connection) -> None initializer that calls register_vector(connection, schema="public"/"embeddings") and shows PostgresConnectionPool.of(..., initialize=initialize_pgvector) usage. Bumps project version from 0.91.0 to 0.91.1 and updates dependency haiway to ~=0.37.3 in pyproject.toml. Refactors GeminiGenerating._completion_stream to accumulate chunk.usage_metadata into a last_usage_meta and record usage metrics once in a finally block instead of emitting per-chunk.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

• Review items requiring extra attention:
  • src/draive/gemini/generating.py: verify streaming control flow, ensure last_usage_meta correctly aggregates and that metrics are always recorded in all early-exit paths.
  • docs/guides/Postgres.md: confirm example imports/types (Connection, register_vector) match public API and that initializer usage is accurate for PostgresConnectionPool.of.
  • pyproject.toml: verify dependency bump (haiway) is compatible with project and version bump is intentional.

Possibly related PRs

• Add Postgres VectorIndex #442: Adds pgvector integration and initializer to the Postgres guide — likely directly related to the docs changes in this PR.
• Update dependencies #385: Updates pyproject.toml version and the haiway dependency — directly related to the dependency/version changes.
• Update haiway with Postgres proxy #417: Modifies the same _completion_stream implementation in src/draive/gemini/generating.py — related code-level change to streaming/usage-metric handling.

Pre-merge checks and finishing touches

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.
Description Check	❓ Inconclusive	No pull request description was provided by the author. An empty or missing description conveys no information about the changeset and cannot be evaluated for relevance. While the pre-merge check is intentionally lenient and passes as long as the description is not off-topic, a completely absent description falls into a gray area—it is neither actively misleading nor does it communicate anything about the changes. This makes a definitive pass or fail determination inconclusive based on the evaluation criteria provided.	Consider adding a brief pull request description that explains the motivation and scope of these changes. This would help reviewers understand the intent behind bundling the pgvector documentation update, version bump, and Gemini changes together in a single PR.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The PR title "Add dedicated pgvector connection extension" is clearly related to the actual changes in the pull request. The majority of the documented changes in docs/guides/Postgres.md focus on adding pgvector integration guidance, including a new initialize_pgvector function and examples of how to use it with PostgresConnectionPool. While the PR also includes a version bump and unrelated Gemini changes, the title accurately captures the primary focus of the documentation changes. The title is specific, concise, and helps a teammate scanning history understand the main contribution.

✨ 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/pgvector

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: ASSERTIVE

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between ba4619d and e7b62cb.

⛔ Files ignored due to path filters (1)

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

📒 Files selected for processing (3)

• docs/guides/Postgres.md (3 hunks)
• pyproject.toml (2 hunks)
• src/draive/gemini/generating.py (4 hunks)

🧰 Additional context used

📓 Path-based instructions (5)

{pyproject.toml,pyrightconfig.json}

📄 CodeRabbit inference engine (AGENTS.md)

Use Ruff, Bandit, and Pyright (strict) via make lint

Files:

• pyproject.toml

docs/**/*

📄 CodeRabbit inference engine (AGENTS.md)

docs/**/*: When behavior/API changes, update relevant docs under docs/ and examples as applicable
When adding public APIs, update examples/guides and ensure cross-links render

Files:

• docs/guides/Postgres.md

**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: Use Python 3.12+ features and syntax across the codebase
Format code exclusively with Ruff (make format); do not use other formatters
Skip module-level docstrings

Files:

• src/draive/gemini/generating.py

src/draive/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

src/draive/**/*.py: Import Haiway symbols directly (from haiway import State, ctx)
Use ctx.scope(...) to bind scoped Disposables and active State; avoid global state
Route all logs through ctx.log_debug/info/warn/error; do not use print
Use latest, most strict typing syntax (Python 3.12+), with strict typing only for public APIs
Avoid loose Any except at explicit third‑party boundaries
Prefer explicit attribute access with static types; avoid dynamic getattr except at narrow boundaries
Prefer Mapping/Sequence/Iterable in public types over dict/list/set
Use final where applicable; avoid inheritance and prefer composition
Use precise unions (|) and narrow with match/isinstance; avoid cast unless provably safe and localized
Model immutable data/config and facades with haiway.State; provide ergonomic classmethods like .of(...)
Avoid in-place mutation; use State.updated(...) or functional builders to produce new instances
Access active state via haiway.ctx inside async scopes (ctx.scope(...))
Use @statemethod for public state methods that dispatch on the active instance
Log around generation calls, tool dispatch, and provider requests/responses without leaking secrets; prefer structured/concise messages
Add metrics via ctx.record where applicable
All I/O is async; keep boundaries async and use ctx.spawn for detached tasks
Use structured concurrency and valid coroutine usage; rely on haiway/asyncio; avoid custom threading
Construct multimodal content with MultimodalContent.of(...) and compose blocks explicitly
Use ResourceContent/ResourceReference for media/data blobs
Wrap custom types/data within ArtifactContent; use hidden when needed
Add NumPy-style docstrings for public symbols with Parameters/Returns/Raises and rationale when non-obvious
Avoid docstrings on internal helpers; keep names self-explanatory
Keep docstrings high-quality; mkdocstrings pulls them into API reference
Never log secrets or full request bodies containing keys/tokens

Files:

• src/draive/gemini/generating.py

src/draive/{openai,anthropic,mistral,gemini,vllm,ollama,bedrock,cohere}/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

src/draive/{openai,anthropic,mistral,gemini,vllm,ollama,bedrock,cohere}/**/*.py: Provider-specific feature modules live under their respective provider directories
Translate provider/SDK errors into typed exceptions; do not raise bare Exception and preserve context
Use environment variables for credentials and resolve via helper functions like getenv_str

Files:

• src/draive/gemini/generating.py

🧠 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 (1)

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

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

• _record_usage_metrics (545-564)

🔇 Additional comments (6)

pyproject.toml (1)

8-8: LGTM!

The patch version bump is appropriate for this enhancement.

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

284-284: LGTM!

The additional PLR0912 annotation appropriately suppresses the "too many branches" lint warning for this complex streaming function.

---

329-345: LGTM! Correctly accumulates cumulative usage metadata.

The change properly defers usage metrics recording to the end of the stream. Since Gemini's usage_metadata is cumulative (not per-chunk deltas), storing only the last chunk's metadata captures the complete usage totals.

---

431-436: LGTM! Finally block ensures usage metrics are always recorded.

The finally block guarantees usage metrics are emitted at stream completion regardless of success or failure, which improves observability. This aligns with similar patterns in other provider implementations.

docs/guides/Postgres.md (2)

41-53: LGTM! Clear and correct pgvector initialization example.

The code properly demonstrates the initialize callback pattern with correct imports, function signature, and usage of register_vector.

---

200-223: LGTM! Practical pgvector integration example.

The example demonstrates the initialize callback in a realistic context, showing how to integrate pgvector registration with the connection pool and vector index usage.

---

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

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 2

📜 Review details

Configuration used: CodeRabbit UI

Review profile: ASSERTIVE

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between e83f649 and ba4619d.

⛔ Files ignored due to path filters (1)

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

📒 Files selected for processing (3)

• docs/guides/Postgres.md (3 hunks)
• pyproject.toml (2 hunks)
• src/draive/gemini/generating.py (4 hunks)

🧰 Additional context used

📓 Path-based instructions (5)

{pyproject.toml,pyrightconfig.json}

📄 CodeRabbit inference engine (AGENTS.md)

Use Ruff, Bandit, and Pyright (strict) via make lint

Files:

• pyproject.toml

docs/**/*

📄 CodeRabbit inference engine (AGENTS.md)

docs/**/*: When behavior/API changes, update relevant docs under docs/ and examples as applicable
When adding public APIs, update examples/guides and ensure cross-links render

Files:

• docs/guides/Postgres.md

**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: Use Python 3.12+ features and syntax across the codebase
Format code exclusively with Ruff (make format); do not use other formatters
Skip module-level docstrings

Files:

• src/draive/gemini/generating.py

src/draive/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

src/draive/**/*.py: Import Haiway symbols directly (from haiway import State, ctx)
Use ctx.scope(...) to bind scoped Disposables and active State; avoid global state
Route all logs through ctx.log_debug/info/warn/error; do not use print
Use latest, most strict typing syntax (Python 3.12+), with strict typing only for public APIs
Avoid loose Any except at explicit third‑party boundaries
Prefer explicit attribute access with static types; avoid dynamic getattr except at narrow boundaries
Prefer Mapping/Sequence/Iterable in public types over dict/list/set
Use final where applicable; avoid inheritance and prefer composition
Use precise unions (|) and narrow with match/isinstance; avoid cast unless provably safe and localized
Model immutable data/config and facades with haiway.State; provide ergonomic classmethods like .of(...)
Avoid in-place mutation; use State.updated(...) or functional builders to produce new instances
Access active state via haiway.ctx inside async scopes (ctx.scope(...))
Use @statemethod for public state methods that dispatch on the active instance
Log around generation calls, tool dispatch, and provider requests/responses without leaking secrets; prefer structured/concise messages
Add metrics via ctx.record where applicable
All I/O is async; keep boundaries async and use ctx.spawn for detached tasks
Use structured concurrency and valid coroutine usage; rely on haiway/asyncio; avoid custom threading
Construct multimodal content with MultimodalContent.of(...) and compose blocks explicitly
Use ResourceContent/ResourceReference for media/data blobs
Wrap custom types/data within ArtifactContent; use hidden when needed
Add NumPy-style docstrings for public symbols with Parameters/Returns/Raises and rationale when non-obvious
Avoid docstrings on internal helpers; keep names self-explanatory
Keep docstrings high-quality; mkdocstrings pulls them into API reference
Never log secrets or full request bodies containing keys/tokens

Files:

• src/draive/gemini/generating.py

src/draive/{openai,anthropic,mistral,gemini,vllm,ollama,bedrock,cohere}/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

src/draive/{openai,anthropic,mistral,gemini,vllm,ollama,bedrock,cohere}/**/*.py: Provider-specific feature modules live under their respective provider directories
Translate provider/SDK errors into typed exceptions; do not raise bare Exception and preserve context
Use environment variables for credentials and resolve via helper functions like getenv_str

Files:

• src/draive/gemini/generating.py

🧠 Learnings (1)

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

Learnt from: KaQuMiQ
PR: miquido/draive#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 (1)

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

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

• _record_usage_metrics (545-564)

🔇 Additional comments (7)

pyproject.toml (1)

8-8: LGTM: Appropriate patch version increment.

The version bump from 0.91.0 to 0.91.1 is consistent with the scope of changes—pgvector documentation additions and streaming metrics refactoring.

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

284-284: LGTM: Appropriate lint annotation.

Adding PLR0912 to suppress "too many branches" warnings is reasonable given the function's inherent complexity from handling multiple streaming edge cases.

---

329-329: LGTM: Proper accumulator for streaming usage metadata.

Introducing last_usage_meta to capture the final usage summary from streaming chunks is correctly typed and appropriately scoped.

---

344-345: LGTM: Efficient usage metadata collection for streaming.

Since chunks provide cumulative usage summaries rather than deltas (as noted in the inline comment), capturing only the last chunk's metadata is correct and more efficient than recording per-chunk.

---

431-435: LGTM: Proper finally block placement ensures metrics recording.

The finally block guarantees that usage metrics are recorded after streaming completes or fails, and correctly handles the case where last_usage_meta remains None (no chunks received) via the guard at line 443.

docs/guides/Postgres.md (2)

41-53: LGTM: Correct pgvector initialization pattern.

The code example properly demonstrates how to register the pgvector codec using the initialize callback. The pattern is reusable and follows asyncpg conventions.

---

200-223: LGTM: Complete working example with pgvector integration.

The full example correctly demonstrates the entire flow from initialization through indexing and search. The duplication of the initialize_pgvector function is appropriate for a self-contained example.