feat(vector): add direct pgvector adapter
This commit is contained in:
@@ -1,7 +1,7 @@
|
||||
import pytest
|
||||
|
||||
from tht.adapters.dwh import PostgresDwhAdapter, ThothRestDwhAdapter
|
||||
from tht.adapters.vector import LegacyDirectVectorStore, ThothHttpVectorStore
|
||||
from tht.adapters.vector import PgVectorStore, ThothHttpVectorStore
|
||||
from tht.adapters.factory import build_dwh, build_vector_store
|
||||
from tht.config import Config, ConfigError
|
||||
|
||||
@@ -28,7 +28,11 @@ def _config(*, dwh_type="thoth_rest", vector_type="thoth_vector_http", reader=Tr
|
||||
vectors = (
|
||||
{
|
||||
"type": "thoth_vector_http",
|
||||
**({"reader": {"base_url": "https://vectors.test/", "api_key": "reader"}} if reader else {}),
|
||||
**(
|
||||
{"reader": {"base_url": "https://vectors.test/", "api_key": "reader"}}
|
||||
if reader
|
||||
else {}
|
||||
),
|
||||
**(
|
||||
{"writer": {"base_url": "https://vectors.test/", "api_key": "writer"}}
|
||||
if writer
|
||||
@@ -38,13 +42,32 @@ def _config(*, dwh_type="thoth_rest", vector_type="thoth_vector_http", reader=Tr
|
||||
if vector_type == "thoth_vector_http"
|
||||
else {
|
||||
"type": "pgvector_direct",
|
||||
"connection": {
|
||||
"host": "vector-db",
|
||||
"database": "postgres",
|
||||
"schema": "vectors",
|
||||
"user": "reader",
|
||||
"password": "secret",
|
||||
},
|
||||
**(
|
||||
{
|
||||
"reader": {
|
||||
"host": "vector-db",
|
||||
"database": "postgres",
|
||||
"schema": "vectors",
|
||||
"user": "reader",
|
||||
"password": "secret",
|
||||
}
|
||||
}
|
||||
if reader
|
||||
else {}
|
||||
),
|
||||
**(
|
||||
{
|
||||
"writer": {
|
||||
"host": "vector-db",
|
||||
"database": "postgres",
|
||||
"schema": "vectors",
|
||||
"user": "writer",
|
||||
"password": "secret",
|
||||
}
|
||||
}
|
||||
if writer
|
||||
else {}
|
||||
),
|
||||
}
|
||||
)
|
||||
legacy_database = (
|
||||
@@ -57,9 +80,7 @@ def _config(*, dwh_type="thoth_rest", vector_type="thoth_vector_http", reader=Tr
|
||||
"transport": "rest",
|
||||
}
|
||||
)
|
||||
return Config.model_validate(
|
||||
{"dwh": dwh, "vectors": vectors, "database": legacy_database}
|
||||
)
|
||||
return Config.model_validate({"dwh": dwh, "vectors": vectors, "database": legacy_database})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -87,14 +108,23 @@ def test_factory_builds_writer_only_http_vector_when_write_is_required():
|
||||
assert store.capabilities.upsert is True
|
||||
|
||||
|
||||
def test_factory_selects_direct_vector_reader():
|
||||
config = _config(vector_type="pgvector_direct")
|
||||
def test_factory_selects_direct_vector_store_and_requires_writer():
|
||||
config = _config(vector_type="pgvector_direct", writer=False)
|
||||
|
||||
assert isinstance(build_vector_store(config), LegacyDirectVectorStore)
|
||||
assert isinstance(build_vector_store(config), PgVectorStore)
|
||||
with pytest.raises(ConfigError, match="writer"):
|
||||
build_vector_store(config, require_write=True)
|
||||
|
||||
|
||||
def test_factory_builds_writer_only_direct_vector_when_write_is_required():
|
||||
store = build_vector_store(
|
||||
_config(vector_type="pgvector_direct", reader=False), require_write=True
|
||||
)
|
||||
assert isinstance(store, PgVectorStore)
|
||||
assert store.capabilities.search is False
|
||||
assert store.capabilities.upsert is True
|
||||
|
||||
|
||||
def test_factory_propagates_non_default_statement_timeout():
|
||||
config = _config(dwh_type="postgres_direct")
|
||||
config.execution.statement_timeout_ms = 12_345
|
||||
|
||||
Reference in New Issue
Block a user