wehub-resource-sync
260 lines
9.4 KiB
Python
260 lines
9.4 KiB
Python
# Copyright (c) Microsoft. All rights reserved.
|
|
|
|
import uuid
|
|
from collections.abc import AsyncGenerator, Sequence
|
|
from contextlib import asynccontextmanager
|
|
from typing import Annotated, Any
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
import pytest_asyncio
|
|
from pydantic import BaseModel
|
|
|
|
from semantic_kernel.connectors.postgres import PostgresCollection, PostgresSettings, PostgresStore
|
|
from semantic_kernel.data.vector import (
|
|
DistanceFunction,
|
|
IndexKind,
|
|
VectorStoreCollectionDefinition,
|
|
VectorStoreField,
|
|
vectorstoremodel,
|
|
)
|
|
from semantic_kernel.exceptions.memory_connector_exceptions import (
|
|
MemoryConnectorConnectionException,
|
|
MemoryConnectorInitializationError,
|
|
)
|
|
|
|
try:
|
|
import psycopg # noqa: F401
|
|
import psycopg_pool # noqa: F401
|
|
|
|
psycopg_pool_installed = True
|
|
except ImportError:
|
|
psycopg_pool_installed = False
|
|
|
|
pg_settings: PostgresSettings = PostgresSettings()
|
|
try:
|
|
connection_params_present = any(pg_settings.get_connection_args().values())
|
|
except MemoryConnectorInitializationError:
|
|
connection_params_present = False
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
not (psycopg_pool_installed or connection_params_present),
|
|
reason="psycopg_pool is not installed" if not psycopg_pool_installed else "No connection parameters provided",
|
|
)
|
|
|
|
|
|
@vectorstoremodel
|
|
class SimpleDataModel(BaseModel):
|
|
id: Annotated[int, VectorStoreField("key")]
|
|
embedding: Annotated[
|
|
list[float] | str | None,
|
|
VectorStoreField(
|
|
"vector",
|
|
index_kind=IndexKind.HNSW,
|
|
dimensions=3,
|
|
distance_function=DistanceFunction.COSINE_SIMILARITY,
|
|
),
|
|
] = None
|
|
data: Annotated[
|
|
dict[str, Any],
|
|
VectorStoreField("data", type="JSONB"),
|
|
]
|
|
|
|
def model_post_init(self, context: Any) -> None:
|
|
if self.embedding is None:
|
|
self.embedding = self.data
|
|
|
|
|
|
def DataModelPandas(record) -> tuple:
|
|
definition = VectorStoreCollectionDefinition(
|
|
fields=[
|
|
VectorStoreField(
|
|
"vector",
|
|
name="embedding",
|
|
index_kind="hnsw",
|
|
dimensions=3,
|
|
distance_function="cosine_similarity",
|
|
type="float",
|
|
),
|
|
VectorStoreField("key", name="id", type="int"),
|
|
VectorStoreField("data", name="data", type="dict"),
|
|
],
|
|
container_mode=True,
|
|
to_dict=lambda x: x.to_dict(orient="records"),
|
|
from_dict=lambda x, **_: pd.DataFrame(x),
|
|
)
|
|
df = pd.DataFrame([record])
|
|
return definition, df
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def vector_store() -> AsyncGenerator[PostgresStore, None]:
|
|
try:
|
|
async with await pg_settings.create_connection_pool() as pool:
|
|
yield PostgresStore(connection_pool=pool)
|
|
except MemoryConnectorConnectionException:
|
|
pytest.skip("Postgres connection not available")
|
|
yield None
|
|
return
|
|
|
|
|
|
@asynccontextmanager
|
|
async def create_simple_collection(
|
|
vector_store: PostgresStore,
|
|
) -> AsyncGenerator[PostgresCollection[int, SimpleDataModel], None]:
|
|
"""Returns a collection with a unique name that is deleted after the context.
|
|
|
|
This can be moved to use a fixture with scope=function and loop_scope=session
|
|
after upgrade to pytest-asyncio 0.24. With the current version, the fixture
|
|
would both cache and use the event loop of the declared scope.
|
|
"""
|
|
suffix = str(uuid.uuid4()).replace("-", "")[:8]
|
|
collection_id = f"test_collection_{suffix}"
|
|
collection = vector_store.get_collection(collection_name=collection_id, record_type=SimpleDataModel)
|
|
assert isinstance(collection, PostgresCollection)
|
|
await collection.ensure_collection_exists()
|
|
try:
|
|
yield collection
|
|
finally:
|
|
await collection.ensure_collection_deleted()
|
|
|
|
|
|
def test_create_store(vector_store):
|
|
assert vector_store is not None
|
|
assert vector_store.connection_pool is not None
|
|
|
|
|
|
async def test_ensure_collection_exists_exists_and_delete(vector_store: PostgresStore):
|
|
suffix = str(uuid.uuid4()).replace("-", "")[:8]
|
|
|
|
collection = vector_store.get_collection(collection_name=f"test_collection_{suffix}", record_type=SimpleDataModel)
|
|
|
|
does_exist_1 = await collection.collection_exists()
|
|
assert does_exist_1 is False
|
|
|
|
await collection.ensure_collection_exists()
|
|
does_exist_2 = await collection.collection_exists()
|
|
assert does_exist_2 is True
|
|
|
|
await collection.ensure_collection_deleted()
|
|
does_exist_3 = await collection.collection_exists()
|
|
assert does_exist_3 is False
|
|
|
|
|
|
async def test_list_collection_names(vector_store):
|
|
async with create_simple_collection(vector_store) as simple_collection:
|
|
simple_collection_id = simple_collection.collection_name
|
|
result = await vector_store.list_collection_names()
|
|
assert simple_collection_id in result
|
|
|
|
|
|
async def test_upsert_get_and_delete(vector_store: PostgresStore):
|
|
record = SimpleDataModel(id=1, embedding=[1.1, 2.2, 3.3], data={"key": "value"})
|
|
async with create_simple_collection(vector_store) as simple_collection:
|
|
result_before_upsert = await simple_collection.get(1)
|
|
assert result_before_upsert is None
|
|
|
|
await simple_collection.upsert(record)
|
|
result = await simple_collection.get(1)
|
|
assert result is not None
|
|
assert result.id == record.id
|
|
assert result.embedding == record.embedding
|
|
assert result.data == record.data
|
|
|
|
# Check that the table has an index
|
|
connection_pool = simple_collection.connection_pool
|
|
async with connection_pool.connection() as conn, conn.cursor() as cur:
|
|
await cur.execute(
|
|
"SELECT indexname FROM pg_indexes WHERE tablename = %s", (simple_collection.collection_name,)
|
|
)
|
|
rows = await cur.fetchall()
|
|
index_names = [index[0] for index in rows]
|
|
assert any("embedding_idx" in index_name for index_name in index_names)
|
|
|
|
await simple_collection.delete(1)
|
|
result_after_delete = await simple_collection.get(1)
|
|
assert result_after_delete is None
|
|
|
|
|
|
async def test_upsert_get_and_delete_pandas(vector_store):
|
|
record = SimpleDataModel(id=1, embedding=[1.1, 2.2, 3.3], data={"key": "value"})
|
|
definition, df = DataModelPandas(record.model_dump())
|
|
|
|
suffix = str(uuid.uuid4()).replace("-", "")[:8]
|
|
collection = vector_store.get_collection(
|
|
collection_name=f"test_collection_{suffix}",
|
|
record_type=pd.DataFrame,
|
|
definition=definition,
|
|
)
|
|
await collection.ensure_collection_exists()
|
|
|
|
try:
|
|
result_before_upsert = await collection.get(1)
|
|
assert result_before_upsert is None
|
|
|
|
await collection.upsert(df)
|
|
result: pd.DataFrame = await collection.get(1)
|
|
assert result is not None
|
|
row = result.iloc[0]
|
|
assert row.id == record.id
|
|
assert row.embedding == record.embedding
|
|
assert row.data == record.data
|
|
|
|
await collection.delete(1)
|
|
result_after_delete = await collection.get(1)
|
|
assert result_after_delete is None
|
|
finally:
|
|
await collection.ensure_collection_deleted()
|
|
|
|
|
|
async def test_upsert_get_and_delete_multiple(vector_store: PostgresStore):
|
|
async with create_simple_collection(vector_store) as simple_collection:
|
|
record1 = SimpleDataModel(id=1, embedding=[1.1, 2.2, 3.3], data={"key": "value"})
|
|
record2 = SimpleDataModel(id=2, embedding=[4.4, 5.5, 6.6], data={"key": "value"})
|
|
|
|
result_before_upsert = await simple_collection.get([1, 2])
|
|
assert result_before_upsert is None
|
|
|
|
await simple_collection.upsert([record1, record2])
|
|
# Test get for the two existing keys and one non-existing key;
|
|
# this should return only the two existing records.
|
|
result = await simple_collection.get([1, 2, 3])
|
|
assert result is not None
|
|
assert isinstance(result, Sequence)
|
|
assert len(result) == 2
|
|
assert result[0] is not None
|
|
assert result[0].id == record1.id
|
|
assert result[0].embedding == record1.embedding
|
|
assert result[0].data == record1.data
|
|
assert result[1] is not None
|
|
assert result[1].id == record2.id
|
|
assert result[1].embedding == record2.embedding
|
|
assert result[1].data == record2.data
|
|
|
|
await simple_collection.delete([1, 2])
|
|
result_after_delete = await simple_collection.get([1, 2])
|
|
assert result_after_delete is None
|
|
|
|
|
|
async def test_search(vector_store: PostgresStore):
|
|
async with create_simple_collection(vector_store) as simple_collection:
|
|
records = [
|
|
SimpleDataModel(id=1, embedding=[1.0, 0.0, 0.0], data={"key": "value1"}),
|
|
SimpleDataModel(id=2, embedding=[0.8, 0.2, 0.0], data={"key": "value2"}),
|
|
SimpleDataModel(id=3, embedding=[0.6, 0.0, 0.4], data={"key": "value3"}),
|
|
SimpleDataModel(id=4, embedding=[1.0, 1.0, 0.0], data={"key": "value4"}),
|
|
SimpleDataModel(id=5, embedding=[0.0, 1.0, 1.0], data={"key": "value5"}),
|
|
SimpleDataModel(id=6, embedding=[1.0, 0.0, 1.0], data={"key": "value6"}),
|
|
]
|
|
|
|
await simple_collection.upsert(records)
|
|
|
|
try:
|
|
search_results = await simple_collection.search(vector=[1.0, 0.0, 0.0], top=3, include_total_count=True)
|
|
assert search_results is not None
|
|
assert search_results.total_count == 3
|
|
assert {result.record.id async for result in search_results.results} == {1, 2, 3}
|
|
|
|
finally:
|
|
await simple_collection.delete([r.id for r in records])
|