Conversation
WalkthroughAdds a Postgres-backed vector index implementation and exports it as PostgresVectorIndex. The new module implements index, search, and delete operations that convert inputs to embeddings, persist per-model vectors and JSON payloads into per-model pgvector-enabled tables, translate AttributeRequirement filters to SQL, support optional score_threshold and MMR-based reranking, and include helpers for JSON path access and table-name resolution. Updates package exports, adds tests that fake Postgres connections and stub TextEmbedding to validate inserts, searches (similarity, requirements, thresholds, limits, rerank), and deletes, and adds Postgres/pgvector documentation. Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Pre-merge checks and finishing touches❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (4)
docs/guides/Postgres.md(1 hunks)src/draive/postgres/__init__.py(2 hunks)src/draive/postgres/vector_index.py(1 hunks)tests/test_postgres_vector_index.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (5)
**/*.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/postgres/__init__.pytests/test_postgres_vector_index.pysrc/draive/postgres/vector_index.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/postgres/__init__.pysrc/draive/postgres/vector_index.py
src/draive/{httpx,mcp,postgres,opentelemetry}/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Place integrations under draive/httpx, draive/mcp, draive/postgres, draive/opentelemetry
Files:
src/draive/postgres/__init__.pysrc/draive/postgres/vector_index.py
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
tests/**/*.py: Do not perform real network I/O in unit tests; mock providers/HTTP
Keep tests fast and focused on changed code; start with unit tests around new types/functions/adapters
Use fixtures from tests/ or add focused ones; avoid heavy integration scaffolding
Use pytest-asyncio for coroutine tests (@pytest.mark.asyncio)
Prefer scoping with ctx.scope(...) in async tests and bind required State instances explicitly
Avoid real I/O and network in async tests; stub provider calls and HTTP
Files:
tests/test_postgres_vector_index.py
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
🧬 Code graph analysis (3)
src/draive/postgres/__init__.py (1)
src/draive/postgres/vector_index.py (1)
PostgresVectorIndex(41-328)
tests/test_postgres_vector_index.py (4)
src/draive/embedding/state.py (1)
TextEmbedding(15-127)src/draive/embedding/types.py (1)
Embedded(14-17)src/draive/parameters/model.py (3)
DataModel(418-768)from_json(705-720)to_json(752-768)src/draive/postgres/vector_index.py (4)
PostgresVectorIndex(41-328)index(71-171)search(173-290)delete(292-322)
src/draive/postgres/vector_index.py (6)
src/draive/embedding/types.py (1)
Embedded(14-17)src/draive/embedding/state.py (2)
ImageEmbedding(130-242)TextEmbedding(15-127)src/draive/multimodal/text.py (1)
TextContent(11-82)src/draive/parameters/model.py (3)
DataModel(418-768)to_json(752-768)from_json(705-720)src/draive/resources/types.py (1)
ResourceContent(126-212)src/draive/utils/vector_index.py (1)
VectorIndex(54-204)
KaQuMiQ
force-pushed
the
feature/pgvector
branch
from
8d3d2e3 to
701d4d8
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (5)
src/draive/postgres/vector_index.py (4)
81-89: Support literal attribute values.Lines 81-89 funnel every non-callable through the
AttributePathassertion, so passing a literalValue(allowed by theVectorIndexcontract) now explodes. TheVectorIndex.indexsignature permitsattribute: Callable[[Model], Value] | AttributePath[Model, Value] | Value, but the current match only handles callable vs AttributePath.Apply this diff to handle all three cases:
- match attribute: - case Callable() as selector: - value_selector = selector - - case path: - assert isinstance( # nosec: B101 - path, AttributePath - ), "Prepare parameter path by using Self._.path.to.property" - value_selector = cast(AttributePath[Model, Value], path).__call__ + if callable(attribute): + value_selector = cast(Callable[[Model], Value], attribute) + elif isinstance(attribute, AttributePath): + value_selector = cast(AttributePath[Model, Value], attribute).__call__ + else: + literal_value: Value = cast(Value, attribute) + value_selector = lambda _model: literal_value
101-105: FeedImageEmbeddingwith bytes, not base64 strings.Line 105 appends
resource_content.data, which is the base64 string. The subsequentall(isinstance(value, bytes))check on line 129 never passes, so image embeddings take the text path and fail. Decode to bytes so the image branch runs.Apply this diff:
case ResourceContent() as resource_content: if not resource_content.mime_type.startswith("image"): raise ValueError(f"{resource_content.mime_type} embedding is not supported") - selected_values.append(resource_content.data) + selected_values.append(resource_content.to_bytes())
226-239: Do not raise after handlingResourceContentqueries.The unconditional
raiseon line 239 fires even when the mime type matched "image" or "text", so everyResourceContentquery fails. Wrap it in anelseblock.Apply this diff:
case ResourceContent() as resource_content: if resource_content.mime_type.startswith("image"): embedded_image: Embedded[bytes] = await ImageEmbedding.embed( b64decode(resource_content.data) ) query_vector = embedded_image.vector elif resource_content.mime_type.startswith("text"): embedded_query: Embedded[str] = await TextEmbedding.embed( b64decode(resource_content.data).decode() ) query_vector = embedded_query.vector - raise ValueError(f"{resource_content.mime_type} embedding is not supported") + else: + raise ValueError(f"{resource_content.mime_type} embedding is not supported")
395-402: Close the JSONBcontainsSQL properly.The
containsbranch emitsjsonb_array_elements_text(without closing parentheses and never terminates the EXISTS expression, yielding invalid SQL. Add the missing closing parentheses.Apply this diff:
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" WHERE element = ${len(resolved_arguments)}", + f"{_scalar_accessor(str(requirement.lhs))}) AS element" + f" WHERE element = ${len(resolved_arguments)})", resolved_arguments, )docs/guides/Postgres.md (1)
188-188: Fix constructor signature and remove non-existent parameters.Line 188 shows
PostgresVectorIndex(embedding_dimensions=1536), but the actual signature only acceptsmmr_multiplier. Additionally, lines 202, 208, and 214 reference anamespaceparameter that doesn't exist in the implementation.Apply this diff to correct the constructor example:
-vector_index = PostgresVectorIndex(embedding_dimensions=1536) +vector_index = PostgresVectorIndex(mmr_multiplier=8)And remove all
namespace="docs"arguments from the index and search calls in the example, as the current implementation doesn't support namespacing.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (5)
docs/guides/Postgres.md(1 hunks)pyproject.toml(1 hunks)src/draive/postgres/__init__.py(2 hunks)src/draive/postgres/vector_index.py(1 hunks)tests/test_postgres_vector_index.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (6)
{pyproject.toml,pyrightconfig.json}
📄 CodeRabbit inference engine (AGENTS.md)
Use Ruff, Bandit, and Pyright (strict) via make lint
Files:
pyproject.toml
**/*.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/postgres/vector_index.pysrc/draive/postgres/__init__.pytests/test_postgres_vector_index.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/postgres/vector_index.pysrc/draive/postgres/__init__.py
src/draive/{httpx,mcp,postgres,opentelemetry}/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Place integrations under draive/httpx, draive/mcp, draive/postgres, draive/opentelemetry
Files:
src/draive/postgres/vector_index.pysrc/draive/postgres/__init__.py
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
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
tests/**/*.py: Do not perform real network I/O in unit tests; mock providers/HTTP
Keep tests fast and focused on changed code; start with unit tests around new types/functions/adapters
Use fixtures from tests/ or add focused ones; avoid heavy integration scaffolding
Use pytest-asyncio for coroutine tests (@pytest.mark.asyncio)
Prefer scoping with ctx.scope(...) in async tests and bind required State instances explicitly
Avoid real I/O and network in async tests; stub provider calls and HTTP
Files:
tests/test_postgres_vector_index.py
🧬 Code graph analysis (3)
src/draive/postgres/vector_index.py (5)
src/draive/embedding/types.py (1)
Embedded(14-17)src/draive/embedding/state.py (2)
ImageEmbedding(130-242)TextEmbedding(15-127)src/draive/multimodal/text.py (1)
TextContent(11-82)src/draive/resources/types.py (1)
ResourceContent(126-212)src/draive/utils/vector_index.py (1)
VectorIndex(54-204)
src/draive/postgres/__init__.py (1)
src/draive/postgres/vector_index.py (1)
PostgresVectorIndex(42-332)
tests/test_postgres_vector_index.py (5)
src/draive/embedding/state.py (1)
TextEmbedding(15-127)src/draive/embedding/types.py (1)
Embedded(14-17)src/draive/parameters/model.py (3)
DataModel(418-768)from_json(705-720)to_json(752-768)src/draive/utils/vector_index.py (1)
VectorIndex(54-204)src/draive/postgres/vector_index.py (4)
PostgresVectorIndex(42-332)index(72-175)search(177-294)delete(296-326)
🪛 GitHub Actions: CI
tests/test_postgres_vector_index.py
[error] 120-120: AssertionError: isinstance(params[2], datetime) failed; expected datetime but got '{}' in test_index_persists_entries.
🔇 Additional comments (2)
pyproject.toml (1)
8-8: LGTM!Version bump to 0.88.0 appropriately reflects the addition of the new
PostgresVectorIndexfeature.src/draive/postgres/__init__.py (1)
13-13: LGTM!The import and export of
PostgresVectorIndexcorrectly expose the new public API.Also applies to: 25-25
KaQuMiQ
force-pushed
the
feature/pgvector
branch
from
701d4d8 to
6c93004
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
src/draive/postgres/vector_index.py (2)
80-110: Literal attributes break indexing
Passing a constant value (allowed byVectorIndex.index) hits the fall-through branch and leavesvalue_selectorundefined, raising at runtime. Add an explicit branch for literal values instead of assertingAttributePath.- match attribute: - case Callable() as selector: - value_selector = selector - - case path: - assert isinstance( # nosec: B101 - path, AttributePath - ), "Prepare parameter path by using Self._.path.to.property" - value_selector = cast(AttributePath[Model, Value], path).__call__ + if callable(attribute): + value_selector = cast(Callable[[Model], Value], attribute) + elif isinstance(attribute, AttributePath): + value_selector = cast(AttributePath[Model, Value], attribute).__call__ + else: + literal_value: Value = cast(Value, attribute) + value_selector = lambda _model: literal_value
388-403: Fix JSON array predicates (contains_any/contains)
Both branches feedjsonb_array_elements_textwithpayload #>> …, which returns TEXT and raisesfunction jsonb_array_elements_text(text) does not exist. ThecontainsSQL is also missing the closing parenthesis beforeAS element. Use a JSONB accessor (payload #> …) and ensure the function call is properly closed.- case "contains_any": - 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" WHERE element = ANY(${len(resolved_arguments)}))", - resolved_arguments, - ) + case "contains_any": + resolved_arguments = [*arguments, requirement.rhs] + accessor = _json_accessor(str(requirement.lhs)) + return ( + "EXISTS (SELECT 1 FROM jsonb_array_elements_text(" # nosec: B608 + f"{accessor}) AS element WHERE element = ANY(${len(resolved_arguments)}))", + resolved_arguments, + ) ... - 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" WHERE element = ${len(resolved_arguments)})", - resolved_arguments, - ) + case "contains": + resolved_arguments = [*arguments, requirement.rhs] + accessor = _json_accessor(str(requirement.lhs)) + return ( + "EXISTS (SELECT 1 FROM jsonb_array_elements_text(" # nosec: B608 + f"{accessor}) AS element WHERE element = ${len(resolved_arguments)})", + resolved_arguments, + )Add:
def _scalar_accessor(path: str) -> str: return f"payload #>> {_path_literal(path)}" + + +def _json_accessor(path: str) -> str: + return f"payload #> {_path_literal(path)}"
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (5)
docs/guides/Postgres.md(1 hunks)pyproject.toml(1 hunks)src/draive/postgres/__init__.py(2 hunks)src/draive/postgres/vector_index.py(1 hunks)tests/test_postgres_vector_index.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (6)
{pyproject.toml,pyrightconfig.json}
📄 CodeRabbit inference engine (AGENTS.md)
Use Ruff, Bandit, and Pyright (strict) via make lint
Files:
pyproject.toml
**/*.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/postgres/__init__.pytests/test_postgres_vector_index.pysrc/draive/postgres/vector_index.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/postgres/__init__.pysrc/draive/postgres/vector_index.py
src/draive/{httpx,mcp,postgres,opentelemetry}/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Place integrations under draive/httpx, draive/mcp, draive/postgres, draive/opentelemetry
Files:
src/draive/postgres/__init__.pysrc/draive/postgres/vector_index.py
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
tests/**/*.py: Do not perform real network I/O in unit tests; mock providers/HTTP
Keep tests fast and focused on changed code; start with unit tests around new types/functions/adapters
Use fixtures from tests/ or add focused ones; avoid heavy integration scaffolding
Use pytest-asyncio for coroutine tests (@pytest.mark.asyncio)
Prefer scoping with ctx.scope(...) in async tests and bind required State instances explicitly
Avoid real I/O and network in async tests; stub provider calls and HTTP
Files:
tests/test_postgres_vector_index.py
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
🧬 Code graph analysis (3)
src/draive/postgres/__init__.py (1)
src/draive/postgres/vector_index.py (1)
PostgresVectorIndex(42-333)
tests/test_postgres_vector_index.py (5)
src/draive/embedding/state.py (1)
TextEmbedding(15-127)src/draive/embedding/types.py (1)
Embedded(14-17)src/draive/parameters/model.py (3)
DataModel(418-768)from_json(705-720)to_json(752-768)src/draive/utils/vector_index.py (1)
VectorIndex(54-204)src/draive/postgres/vector_index.py (4)
PostgresVectorIndex(42-333)index(72-175)search(177-295)delete(297-327)
src/draive/postgres/vector_index.py (7)
src/draive/embedding/types.py (1)
Embedded(14-17)src/draive/embedding/state.py (2)
ImageEmbedding(130-242)TextEmbedding(15-127)src/draive/multimodal/text.py (1)
TextContent(11-82)src/draive/parameters/model.py (3)
DataModel(418-768)to_json(752-768)from_json(705-720)src/draive/resources/types.py (1)
ResourceContent(126-212)src/draive/utils/vector_index.py (1)
VectorIndex(54-204)tests/test_postgres_vector_index.py (2)
transaction(50-51)execute(47-48)
| parameters: Sequence[Sequence[PostgresValue] | PostgresValue] = [ | ||
| *arguments, | ||
| limit or 8, | ||
| ] | ||
| results: Sequence[PostgresRow] = await Postgres.fetch( | ||
| f""" | ||
| SELECT | ||
| payload | ||
|
|
||
| FROM {resolve_table_name(model)} | ||
|
|
||
| {where_clause} | ||
| ORDER BY created DESC | ||
| LIMIT ${len(parameters)}; | ||
| """, # nosec: B608 | ||
| *parameters, | ||
| ) | ||
|
|
||
| return tuple(model.from_json(cast(str, result["payload"])) for result in results) | ||
|
|
||
| query_vector: Sequence[float] | ||
| match query: | ||
| case str() as text: | ||
| embedded_query: Embedded[str] = await TextEmbedding.embed(text) | ||
| query_vector = embedded_query.vector | ||
|
|
||
| case TextContent() as text_content: | ||
| embedded_query: Embedded[str] = await TextEmbedding.embed(text_content.text) | ||
| query_vector = embedded_query.vector | ||
|
|
||
| case ResourceContent() as resource_content: | ||
| if resource_content.mime_type.startswith("image"): | ||
| embedded_image: Embedded[bytes] = await ImageEmbedding.embed( | ||
| b64decode(resource_content.data) | ||
| ) | ||
| query_vector = embedded_image.vector | ||
|
|
||
| elif resource_content.mime_type.startswith("text"): | ||
| embedded_query: Embedded[str] = await TextEmbedding.embed( | ||
| b64decode(resource_content.data).decode() | ||
| ) | ||
| query_vector = embedded_query.vector | ||
|
|
||
| else: | ||
| raise ValueError(f"{resource_content.mime_type} embedding is not supported") | ||
|
|
||
| case vector: | ||
| query_vector = vector | ||
|
|
||
| arguments: Sequence[Sequence[PostgresValue] | PostgresValue] = (query_vector,) | ||
| similarity_expression: str = f"embedding <#> ${len(arguments)}" | ||
|
|
||
| where_clause, arguments = resolve_requirements(requirements, arguments=arguments) | ||
|
|
||
| if score_threshold is not None: | ||
| arguments = (*arguments, 1.0 - float(score_threshold)) | ||
| threshold_clause: str = f"{similarity_expression} <= ${len(arguments)}" | ||
| if where_clause: | ||
| where_clause = f"WHERE {threshold_clause} AND ({where_clause})" | ||
|
|
||
| else: | ||
| where_clause = f"WHERE {threshold_clause}" | ||
|
|
||
| elif where_clause: | ||
| where_clause = f"WHERE {where_clause}" | ||
|
|
||
| arguments = (*arguments, (limit or 8) * mmr_multiplier if rerank else (limit or 8)) | ||
| results: Sequence[PostgresRow] = await Postgres.fetch( | ||
| f""" | ||
| SELECT | ||
| embedding, | ||
| payload | ||
|
|
||
| FROM {resolve_table_name(model)} | ||
|
|
||
| {where_clause} | ||
| ORDER BY {similarity_expression} | ||
| LIMIT ${len(arguments)}; | ||
| """, # nosec: B608 | ||
| *arguments, | ||
| ) |
There was a problem hiding this comment.
Respect explicit limit values (including zero)
Both the no-query and similarity paths coerce limit via limit or 8, so limit=0 (or any falsy value) unexpectedly expands to 8. Use an explicit if limit is None fallback instead.
- parameters: Sequence[Sequence[PostgresValue] | PostgresValue] = [
- *arguments,
- limit or 8,
- ]
+ resolved_limit = limit if limit is not None else 8
+ parameters: Sequence[Sequence[PostgresValue] | PostgresValue] = [
+ *arguments,
+ resolved_limit,
+ ]
...
- arguments = (*arguments, (limit or 8) * mmr_multiplier if rerank else (limit or 8))
+ resolved_limit = limit if limit is not None else 8
+ arguments = (
+ *arguments,
+ resolved_limit * mmr_multiplier if rerank else resolved_limit,
+ )📝 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.
| parameters: Sequence[Sequence[PostgresValue] | PostgresValue] = [ | |
| *arguments, | |
| limit or 8, | |
| ] | |
| results: Sequence[PostgresRow] = await Postgres.fetch( | |
| f""" | |
| SELECT | |
| payload | |
| FROM {resolve_table_name(model)} | |
| {where_clause} | |
| ORDER BY created DESC | |
| LIMIT ${len(parameters)}; | |
| """, # nosec: B608 | |
| *parameters, | |
| ) | |
| return tuple(model.from_json(cast(str, result["payload"])) for result in results) | |
| query_vector: Sequence[float] | |
| match query: | |
| case str() as text: | |
| embedded_query: Embedded[str] = await TextEmbedding.embed(text) | |
| query_vector = embedded_query.vector | |
| case TextContent() as text_content: | |
| embedded_query: Embedded[str] = await TextEmbedding.embed(text_content.text) | |
| query_vector = embedded_query.vector | |
| case ResourceContent() as resource_content: | |
| if resource_content.mime_type.startswith("image"): | |
| embedded_image: Embedded[bytes] = await ImageEmbedding.embed( | |
| b64decode(resource_content.data) | |
| ) | |
| query_vector = embedded_image.vector | |
| elif resource_content.mime_type.startswith("text"): | |
| embedded_query: Embedded[str] = await TextEmbedding.embed( | |
| b64decode(resource_content.data).decode() | |
| ) | |
| query_vector = embedded_query.vector | |
| else: | |
| raise ValueError(f"{resource_content.mime_type} embedding is not supported") | |
| case vector: | |
| query_vector = vector | |
| arguments: Sequence[Sequence[PostgresValue] | PostgresValue] = (query_vector,) | |
| similarity_expression: str = f"embedding <#> ${len(arguments)}" | |
| where_clause, arguments = resolve_requirements(requirements, arguments=arguments) | |
| if score_threshold is not None: | |
| arguments = (*arguments, 1.0 - float(score_threshold)) | |
| threshold_clause: str = f"{similarity_expression} <= ${len(arguments)}" | |
| if where_clause: | |
| where_clause = f"WHERE {threshold_clause} AND ({where_clause})" | |
| else: | |
| where_clause = f"WHERE {threshold_clause}" | |
| elif where_clause: | |
| where_clause = f"WHERE {where_clause}" | |
| arguments = (*arguments, (limit or 8) * mmr_multiplier if rerank else (limit or 8)) | |
| results: Sequence[PostgresRow] = await Postgres.fetch( | |
| f""" | |
| SELECT | |
| embedding, | |
| payload | |
| FROM {resolve_table_name(model)} | |
| {where_clause} | |
| ORDER BY {similarity_expression} | |
| LIMIT ${len(arguments)}; | |
| """, # nosec: B608 | |
| *arguments, | |
| ) | |
| resolved_limit = limit if limit is not None else 8 | |
| parameters: Sequence[Sequence[PostgresValue] | PostgresValue] = [ | |
| *arguments, | |
| resolved_limit, | |
| ] | |
| results: Sequence[PostgresRow] = await Postgres.fetch( | |
| f""" | |
| SELECT | |
| payload | |
| FROM {resolve_table_name(model)} | |
| {where_clause} | |
| ORDER BY created DESC | |
| LIMIT ${len(parameters)}; | |
| """, # nosec: B608 | |
| *parameters, | |
| ) | |
| return tuple(model.from_json(cast(str, result["payload"])) for result in results) | |
| query_vector: Sequence[float] | |
| match query: | |
| case str() as text: | |
| embedded_query: Embedded[str] = await TextEmbedding.embed(text) | |
| query_vector = embedded_query.vector | |
| case TextContent() as text_content: | |
| embedded_query: Embedded[str] = await TextEmbedding.embed(text_content.text) | |
| query_vector = embedded_query.vector | |
| case ResourceContent() as resource_content: | |
| if resource_content.mime_type.startswith("image"): | |
| embedded_image: Embedded[bytes] = await ImageEmbedding.embed( | |
| b64decode(resource_content.data) | |
| ) | |
| query_vector = embedded_image.vector | |
| elif resource_content.mime_type.startswith("text"): | |
| embedded_query: Embedded[str] = await TextEmbedding.embed( | |
| b64decode(resource_content.data).decode() | |
| ) | |
| query_vector = embedded_query.vector | |
| else: | |
| raise ValueError(f"{resource_content.mime_type} embedding is not supported") | |
| case vector: | |
| query_vector = vector | |
| arguments: Sequence[Sequence[PostgresValue] | PostgresValue] = (query_vector,) | |
| similarity_expression: str = f"embedding <#> ${len(arguments)}" | |
| where_clause, arguments = resolve_requirements(requirements, arguments=arguments) | |
| if score_threshold is not None: | |
| arguments = (*arguments, 1.0 - float(score_threshold)) | |
| threshold_clause: str = f"{similarity_expression} <= ${len(arguments)}" | |
| if where_clause: | |
| where_clause = f"WHERE {threshold_clause} AND ({where_clause})" | |
| else: | |
| where_clause = f"WHERE {threshold_clause}" | |
| elif where_clause: | |
| where_clause = f"WHERE {where_clause}" | |
| resolved_limit = limit if limit is not None else 8 | |
| arguments = ( | |
| *arguments, | |
| resolved_limit * mmr_multiplier if rerank else resolved_limit, | |
| ) | |
| results: Sequence[PostgresRow] = await Postgres.fetch( | |
| f""" | |
| SELECT | |
| embedding, | |
| payload | |
| FROM {resolve_table_name(model)} | |
| {where_clause} | |
| ORDER BY {similarity_expression} | |
| LIMIT ${len(arguments)}; | |
| """, # nosec: B608 | |
| *arguments, | |
| ) |
🤖 Prompt for AI Agents
In src/draive/postgres/vector_index.py around lines 196-276, the code uses
"limit or 8" which treats limit=0 (and other falsy values) as 8; change to
explicitly default only when limit is None. Introduce a local limit_val = 8 if
limit is None else limit (ensure it's an int), then replace all occurrences of
"(limit or 8)" with "limit_val" and "(limit or 8) * mmr_multiplier" with
"limit_val * mmr_multiplier" so explicit zero or other falsy limits are
respected.
1 participant
No description provided.