Add dedicated pgvector connection extension

docs/guides/Postgres.md

@@ -33,6 +33,27 @@ async with ctx.scope(
 Each adapter relies on the same connection scope, so you can freely mix them within a single
 context.
 
+When working with pgvector-backed components, add the Python package `pgvector` to your environment
+(`pip install pgvector`) and ensure every connection registers the codec. haiway ≥ 0.37.1 exposes an
+`initialize` callback on `PostgresConnectionPool.of(...)` so you can point to the schema where the
+extension is installed:
+
+```python
+from asyncpg.connection import Connection
+from pgvector.asyncpg import register_vector
+
+
+async def initialize_pgvector(connection: Connection) -> None:
+    await register_vector(connection, schema="public")
+
+PostgresConnectionPool.of(
+    dsn="postgresql://draive:secret@localhost:5432/draive",
+    initialize=initialize_pgvector,
+)
+```
+
+Reuse the same initializer across scopes to keep the codec registration consistent.
+
 ## ConfigurationRepository implementation
 
 `PostgresConfigurationRepository` persists configuration snapshots inside a `configurations` table
@@ -176,6 +197,9 @@ from collections.abc import Sequence
 from typing import Annotated
 
 from draive import Alias, DataModel, ctx
+from asyncpg.connection import Connection
+from pgvector.asyncpg import register_vector
+
 from draive.postgres import PostgresConnectionPool, PostgresVectorIndex
 from draive.utils import VectorIndex
 
@@ -185,11 +209,18 @@ class Chunk(DataModel):
     text: str
 
 
+async def initialize_pgvector(connection: Connection) -> None:
+    await register_vector(connection, schema="embeddings")
+
+
 async with ctx.scope(
     "pgvector-demo",
     PostgresVectorIndex(),
     disposables=(
-        PostgresConnectionPool.of(dsn="postgresql://draive:secret@localhost:5432/draive"),
+        PostgresConnectionPool.of(
+            dsn="postgresql://draive:secret@localhost:5432/draive",
+            initialize=initialize_pgvector,
+        ),
     ),
 ):
     await VectorIndex.index(

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "uv_build"
 [project]
 name = "draive"
 description = "Framework designed to simplify and accelerate the development of LLM-based applications."
-version = "0.91.0"
+version = "0.91.1"
 readme = "README.md"
 maintainers = [
   { name = "Kacper Kaliński", email = "kacper.kalinski@miquido.com" },
@@ -24,7 +24,7 @@ classifiers = [
   "Topic :: Software Development :: Libraries :: Application Frameworks",
 ]
 license = { file = "LICENSE" }
-dependencies = ["numpy~=2.3", "haiway~=0.37.0"]
+dependencies = ["numpy~=2.3", "haiway~=0.37.3"]
 
 [project.urls]
 Homepage = "https://miquido.com"

src/draive/gemini/generating.py

@@ -281,7 +281,7 @@ async def _completion(
                 },
             )
 
-    async def _completion_stream(  # noqa: C901
+    async def _completion_stream(  # noqa: C901, PLR0912
         self,
         *,
         instructions: ModelInstructions,
@@ -326,6 +326,7 @@ async def _completion_stream(  # noqa: C901
                 _context_element_as_content(element) for element in context
             ]
 
+            last_usage_meta: GenerateContentResponseUsageMetadata | None = None
             try:
                 response_stream: AsyncIterator[
                     GenerateContentResponse
@@ -340,10 +341,8 @@ async def _completion_stream(  # noqa: C901
                 finish_message: str | None = None
 
                 async for chunk in response_stream:
-                    _record_usage_metrics(
-                        chunk.usage_metadata,
-                        model=config.model,
-                    )
+                    if chunk.usage_metadata:  # chunks provide usage summary instead of delta
+                        last_usage_meta = chunk.usage_metadata
 
                     if not chunk.candidates:
                         continue
@@ -429,6 +428,12 @@ async def _completion_stream(  # noqa: C901
                     reason=str(exc),
                 ) from exc
 
+            finally:
+                _record_usage_metrics(
+                    last_usage_meta,
+                    model=config.model,
+                )
+
 
 def _record_usage_metrics(
     usage: GenerateContentResponseUsageMetadata | None,