mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 14:42:28 +02:00
Adds a few of preliminaries: 1. Makes the returned "score" actually the result of the requested operation (cosine, inner_product, l2) 2. Sorts asc, etc. so that if you were to add an HNSW index (and not have any WHERE filters), it would be used 3. Drop the inner WHERE statement if no namespace or other filters are provided. See (2) for why. I don't yet add an index to the migrations since I think we need to agree on the right balance to ensure it's actually used in common query patterns.
785 lines
26 KiB
Python
785 lines
26 KiB
Python
# type: ignore
|
|
|
|
from contextlib import contextmanager
|
|
from typing import Any, Optional
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from langchain_core.embeddings import Embeddings
|
|
from psycopg import Connection
|
|
|
|
from langgraph.store.base import (
|
|
GetOp,
|
|
Item,
|
|
ListNamespacesOp,
|
|
MatchCondition,
|
|
PutOp,
|
|
SearchOp,
|
|
)
|
|
from langgraph.store.postgres import PostgresStore
|
|
from tests.conftest import (
|
|
DEFAULT_URI,
|
|
VECTOR_TYPES,
|
|
CharacterEmbeddings,
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="function", params=["default", "pipe", "pool"])
|
|
def store(request) -> PostgresStore:
|
|
database = f"test_{uuid4().hex[:16]}"
|
|
uri_parts = DEFAULT_URI.split("/")
|
|
uri_base = "/".join(uri_parts[:-1])
|
|
query_params = ""
|
|
if "?" in uri_parts[-1]:
|
|
db_name, query_params = uri_parts[-1].split("?", 1)
|
|
query_params = "?" + query_params
|
|
|
|
conn_string = f"{uri_base}/{database}{query_params}"
|
|
admin_conn_string = DEFAULT_URI
|
|
|
|
with Connection.connect(admin_conn_string, autocommit=True) as conn:
|
|
conn.execute(f"CREATE DATABASE {database}")
|
|
try:
|
|
with PostgresStore.from_conn_string(conn_string) as store:
|
|
store.setup()
|
|
|
|
if request.param == "pipe":
|
|
with PostgresStore.from_conn_string(conn_string, pipeline=True) as store:
|
|
yield store
|
|
elif request.param == "pool":
|
|
with PostgresStore.from_conn_string(
|
|
conn_string, pool_config={"min_size": 1, "max_size": 10}
|
|
) as store:
|
|
yield store
|
|
else: # default
|
|
with PostgresStore.from_conn_string(conn_string) as store:
|
|
yield store
|
|
finally:
|
|
with Connection.connect(admin_conn_string, autocommit=True) as conn:
|
|
conn.execute(f"DROP DATABASE {database}")
|
|
|
|
|
|
def test_batch_order(store: PostgresStore) -> None:
|
|
# Setup test data
|
|
store.put(("test", "foo"), "key1", {"data": "value1"})
|
|
store.put(("test", "bar"), "key2", {"data": "value2"})
|
|
|
|
ops = [
|
|
GetOp(namespace=("test", "foo"), key="key1"),
|
|
PutOp(namespace=("test", "bar"), key="key2", value={"data": "value2"}),
|
|
SearchOp(
|
|
namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
|
|
),
|
|
ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
|
|
GetOp(namespace=("test",), key="key3"),
|
|
]
|
|
|
|
results = store.batch(ops)
|
|
assert len(results) == 5
|
|
assert isinstance(results[0], Item)
|
|
assert isinstance(results[0].value, dict)
|
|
assert results[0].value == {"data": "value1"}
|
|
assert results[0].key == "key1"
|
|
assert results[1] is None # Put operation returns None
|
|
assert isinstance(results[2], list)
|
|
assert len(results[2]) == 1
|
|
assert isinstance(results[3], list)
|
|
assert len(results[3]) > 0 # Should contain at least our test namespaces
|
|
assert results[4] is None # Non-existent key returns None
|
|
|
|
# Test reordered operations
|
|
ops_reordered = [
|
|
SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
|
|
GetOp(namespace=("test", "bar"), key="key2"),
|
|
ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0),
|
|
PutOp(namespace=("test",), key="key3", value={"data": "value3"}),
|
|
GetOp(namespace=("test", "foo"), key="key1"),
|
|
]
|
|
|
|
results_reordered = store.batch(ops_reordered)
|
|
assert len(results_reordered) == 5
|
|
assert isinstance(results_reordered[0], list)
|
|
assert len(results_reordered[0]) >= 2 # Should find at least our two test items
|
|
assert isinstance(results_reordered[1], Item)
|
|
assert results_reordered[1].value == {"data": "value2"}
|
|
assert results_reordered[1].key == "key2"
|
|
assert isinstance(results_reordered[2], list)
|
|
assert len(results_reordered[2]) > 0
|
|
assert results_reordered[3] is None # Put operation returns None
|
|
assert isinstance(results_reordered[4], Item)
|
|
assert results_reordered[4].value == {"data": "value1"}
|
|
assert results_reordered[4].key == "key1"
|
|
|
|
|
|
def test_batch_get_ops(store: PostgresStore) -> None:
|
|
# Setup test data
|
|
store.put(("test",), "key1", {"data": "value1"})
|
|
store.put(("test",), "key2", {"data": "value2"})
|
|
|
|
ops = [
|
|
GetOp(namespace=("test",), key="key1"),
|
|
GetOp(namespace=("test",), key="key2"),
|
|
GetOp(namespace=("test",), key="key3"), # Non-existent key
|
|
]
|
|
|
|
results = store.batch(ops)
|
|
|
|
assert len(results) == 3
|
|
assert results[0] is not None
|
|
assert results[1] is not None
|
|
assert results[2] is None
|
|
assert results[0].key == "key1"
|
|
assert results[1].key == "key2"
|
|
|
|
|
|
def test_batch_put_ops(store: PostgresStore) -> None:
|
|
ops = [
|
|
PutOp(namespace=("test",), key="key1", value={"data": "value1"}),
|
|
PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
|
|
PutOp(namespace=("test",), key="key3", value=None), # Delete operation
|
|
]
|
|
|
|
results = store.batch(ops)
|
|
assert len(results) == 3
|
|
assert all(result is None for result in results)
|
|
|
|
# Verify the puts worked
|
|
item1 = store.get(("test",), "key1")
|
|
item2 = store.get(("test",), "key2")
|
|
item3 = store.get(("test",), "key3")
|
|
|
|
assert item1 and item1.value == {"data": "value1"}
|
|
assert item2 and item2.value == {"data": "value2"}
|
|
assert item3 is None
|
|
|
|
|
|
def test_batch_search_ops(store: PostgresStore) -> None:
|
|
# Setup test data
|
|
test_data = [
|
|
(("test", "foo"), "key1", {"data": "value1", "tag": "a"}),
|
|
(("test", "bar"), "key2", {"data": "value2", "tag": "a"}),
|
|
(("test", "baz"), "key3", {"data": "value3", "tag": "b"}),
|
|
]
|
|
for namespace, key, value in test_data:
|
|
store.put(namespace, key, value)
|
|
|
|
ops = [
|
|
SearchOp(namespace_prefix=("test",), filter={"tag": "a"}, limit=10, offset=0),
|
|
SearchOp(namespace_prefix=("test",), filter=None, limit=2, offset=0),
|
|
SearchOp(namespace_prefix=("test", "foo"), filter=None, limit=10, offset=0),
|
|
]
|
|
|
|
results = store.batch(ops)
|
|
assert len(results) == 3
|
|
|
|
# First search should find items with tag "a"
|
|
assert len(results[0]) == 2
|
|
assert all(item.value["tag"] == "a" for item in results[0])
|
|
|
|
# Second search should return first 2 items
|
|
assert len(results[1]) == 2
|
|
|
|
# Third search should only find items in test/foo namespace
|
|
assert len(results[2]) == 1
|
|
assert results[2][0].namespace == ("test", "foo")
|
|
|
|
|
|
def test_batch_list_namespaces_ops(store: PostgresStore) -> None:
|
|
# Setup test data with various namespaces
|
|
test_data = [
|
|
(("test", "documents", "public"), "doc1", {"content": "public doc"}),
|
|
(("test", "documents", "private"), "doc2", {"content": "private doc"}),
|
|
(("test", "images", "public"), "img1", {"content": "public image"}),
|
|
(("prod", "documents", "public"), "doc3", {"content": "prod doc"}),
|
|
]
|
|
for namespace, key, value in test_data:
|
|
store.put(namespace, key, value)
|
|
|
|
ops = [
|
|
ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
|
|
ListNamespacesOp(match_conditions=None, max_depth=2, limit=10, offset=0),
|
|
ListNamespacesOp(
|
|
match_conditions=[MatchCondition("suffix", "public")],
|
|
max_depth=None,
|
|
limit=10,
|
|
offset=0,
|
|
),
|
|
]
|
|
|
|
results = store.batch(ops)
|
|
assert len(results) == 3
|
|
|
|
# First operation should list all namespaces
|
|
assert len(results[0]) == len(test_data)
|
|
|
|
# Second operation should only return namespaces up to depth 2
|
|
assert all(len(ns) <= 2 for ns in results[1])
|
|
|
|
# Third operation should only return namespaces ending with "public"
|
|
assert all(ns[-1] == "public" for ns in results[2])
|
|
|
|
|
|
class TestPostgresStore:
|
|
@pytest.fixture(autouse=True)
|
|
def setup(self) -> None:
|
|
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
|
|
store.setup()
|
|
|
|
def test_basic_store_ops(self) -> None:
|
|
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
|
|
namespace = ("test", "documents")
|
|
item_id = "doc1"
|
|
item_value = {"title": "Test Document", "content": "Hello, World!"}
|
|
|
|
store.put(namespace, item_id, item_value)
|
|
item = store.get(namespace, item_id)
|
|
|
|
assert item
|
|
assert item.namespace == namespace
|
|
assert item.key == item_id
|
|
assert item.value == item_value
|
|
|
|
# Test update
|
|
updated_value = {"title": "Updated Document", "content": "Hello, Updated!"}
|
|
store.put(namespace, item_id, updated_value)
|
|
updated_item = store.get(namespace, item_id)
|
|
|
|
assert updated_item.value == updated_value
|
|
assert updated_item.updated_at > item.updated_at
|
|
|
|
# Test get from non-existent namespace
|
|
different_namespace = ("test", "other_documents")
|
|
item_in_different_namespace = store.get(different_namespace, item_id)
|
|
assert item_in_different_namespace is None
|
|
|
|
# Test delete
|
|
store.delete(namespace, item_id)
|
|
deleted_item = store.get(namespace, item_id)
|
|
assert deleted_item is None
|
|
|
|
def test_list_namespaces(self) -> None:
|
|
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
|
|
# Create test data with various namespaces
|
|
test_namespaces = [
|
|
("test", "documents", "public"),
|
|
("test", "documents", "private"),
|
|
("test", "images", "public"),
|
|
("test", "images", "private"),
|
|
("prod", "documents", "public"),
|
|
("prod", "documents", "private"),
|
|
]
|
|
|
|
# Insert test data
|
|
for namespace in test_namespaces:
|
|
store.put(namespace, "dummy", {"content": "dummy"})
|
|
|
|
# Test listing with various filters
|
|
all_namespaces = store.list_namespaces()
|
|
assert len(all_namespaces) == len(test_namespaces)
|
|
|
|
# Test prefix filtering
|
|
test_prefix_namespaces = store.list_namespaces(prefix=["test"])
|
|
assert len(test_prefix_namespaces) == 4
|
|
assert all(ns[0] == "test" for ns in test_prefix_namespaces)
|
|
|
|
# Test suffix filtering
|
|
public_namespaces = store.list_namespaces(suffix=["public"])
|
|
assert len(public_namespaces) == 3
|
|
assert all(ns[-1] == "public" for ns in public_namespaces)
|
|
|
|
# Test max depth
|
|
depth_2_namespaces = store.list_namespaces(max_depth=2)
|
|
assert all(len(ns) <= 2 for ns in depth_2_namespaces)
|
|
|
|
# Test pagination
|
|
paginated_namespaces = store.list_namespaces(limit=3)
|
|
assert len(paginated_namespaces) == 3
|
|
|
|
# Cleanup
|
|
for namespace in test_namespaces:
|
|
store.delete(namespace, "dummy")
|
|
|
|
def test_search(self) -> None:
|
|
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
|
|
# Create test data
|
|
test_data = [
|
|
(
|
|
("test", "docs"),
|
|
"doc1",
|
|
{"title": "First Doc", "author": "Alice", "tags": ["important"]},
|
|
),
|
|
(
|
|
("test", "docs"),
|
|
"doc2",
|
|
{"title": "Second Doc", "author": "Bob", "tags": ["draft"]},
|
|
),
|
|
(
|
|
("test", "images"),
|
|
"img1",
|
|
{"title": "Image 1", "author": "Alice", "tags": ["final"]},
|
|
),
|
|
]
|
|
|
|
for namespace, key, value in test_data:
|
|
store.put(namespace, key, value)
|
|
|
|
# Test basic search
|
|
all_items = store.search(["test"])
|
|
assert len(all_items) == 3
|
|
|
|
# Test namespace filtering
|
|
docs_items = store.search(["test", "docs"])
|
|
assert len(docs_items) == 2
|
|
assert all(item.namespace == ("test", "docs") for item in docs_items)
|
|
|
|
# Test value filtering
|
|
alice_items = store.search(["test"], filter={"author": "Alice"})
|
|
assert len(alice_items) == 2
|
|
assert all(item.value["author"] == "Alice" for item in alice_items)
|
|
|
|
# Test pagination
|
|
paginated_items = store.search(["test"], limit=2)
|
|
assert len(paginated_items) == 2
|
|
|
|
offset_items = store.search(["test"], offset=2)
|
|
assert len(offset_items) == 1
|
|
|
|
# Cleanup
|
|
for namespace, key, _ in test_data:
|
|
store.delete(namespace, key)
|
|
|
|
|
|
@contextmanager
|
|
def _create_vector_store(
|
|
vector_type: str,
|
|
distance_type: str,
|
|
fake_embeddings: Embeddings,
|
|
text_fields: Optional[list[str]] = None,
|
|
) -> PostgresStore:
|
|
"""Create a store with vector search enabled."""
|
|
database = f"test_{uuid4().hex[:16]}"
|
|
uri_parts = DEFAULT_URI.split("/")
|
|
uri_base = "/".join(uri_parts[:-1])
|
|
query_params = ""
|
|
if "?" in uri_parts[-1]:
|
|
db_name, query_params = uri_parts[-1].split("?", 1)
|
|
query_params = "?" + query_params
|
|
|
|
conn_string = f"{uri_base}/{database}{query_params}"
|
|
admin_conn_string = DEFAULT_URI
|
|
|
|
index_config = {
|
|
"dims": fake_embeddings.dims,
|
|
"embed": fake_embeddings,
|
|
"ann_index_config": {
|
|
"vector_type": vector_type,
|
|
},
|
|
"distance_type": distance_type,
|
|
"text_fields": text_fields,
|
|
}
|
|
|
|
with Connection.connect(admin_conn_string, autocommit=True) as conn:
|
|
conn.execute(f"CREATE DATABASE {database}")
|
|
try:
|
|
with PostgresStore.from_conn_string(
|
|
conn_string,
|
|
index=index_config,
|
|
) as store:
|
|
store.setup()
|
|
yield store
|
|
finally:
|
|
with Connection.connect(admin_conn_string, autocommit=True) as conn:
|
|
conn.execute(f"DROP DATABASE {database}")
|
|
|
|
|
|
@pytest.fixture(
|
|
scope="function",
|
|
params=[
|
|
(vector_type, distance_type)
|
|
for vector_type in VECTOR_TYPES
|
|
for distance_type in (
|
|
["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
|
|
)
|
|
],
|
|
ids=lambda p: f"{p[0]}_{p[1]}",
|
|
)
|
|
def vector_store(
|
|
request,
|
|
fake_embeddings: Embeddings,
|
|
) -> PostgresStore:
|
|
"""Create a store with vector search enabled."""
|
|
vector_type, distance_type = request.param
|
|
with _create_vector_store(vector_type, distance_type, fake_embeddings) as store:
|
|
yield store
|
|
|
|
|
|
def test_vector_store_initialization(
|
|
vector_store: PostgresStore, fake_embeddings: CharacterEmbeddings
|
|
) -> None:
|
|
"""Test store initialization with embedding config."""
|
|
# Store should be initialized with embedding config
|
|
assert vector_store.index_config is not None
|
|
assert vector_store.index_config["dims"] == fake_embeddings.dims
|
|
assert vector_store.index_config["embed"] == fake_embeddings
|
|
|
|
|
|
def test_vector_insert_with_auto_embedding(vector_store: PostgresStore) -> None:
|
|
"""Test inserting items that get auto-embedded."""
|
|
docs = [
|
|
("doc1", {"text": "short text"}),
|
|
("doc2", {"text": "longer text document"}),
|
|
("doc3", {"text": "longest text document here"}),
|
|
("doc4", {"description": "text in description field"}),
|
|
("doc5", {"content": "text in content field"}),
|
|
("doc6", {"body": "text in body field"}),
|
|
]
|
|
|
|
for key, value in docs:
|
|
vector_store.put(("test",), key, value)
|
|
|
|
results = vector_store.search(("test",), query="long text")
|
|
assert len(results) > 0
|
|
|
|
doc_order = [r.key for r in results]
|
|
assert "doc2" in doc_order
|
|
assert "doc3" in doc_order
|
|
|
|
|
|
def test_vector_update_with_embedding(vector_store: PostgresStore) -> None:
|
|
"""Test that updating items properly updates their embeddings."""
|
|
vector_store.put(("test",), "doc1", {"text": "zany zebra Xerxes"})
|
|
vector_store.put(("test",), "doc2", {"text": "something about dogs"})
|
|
vector_store.put(("test",), "doc3", {"text": "text about birds"})
|
|
|
|
results_initial = vector_store.search(("test",), query="Zany Xerxes")
|
|
assert len(results_initial) > 0
|
|
assert results_initial[0].key == "doc1"
|
|
initial_score = results_initial[0].score
|
|
|
|
vector_store.put(("test",), "doc1", {"text": "new text about dogs"})
|
|
|
|
results_after = vector_store.search(("test",), query="Zany Xerxes")
|
|
after_score = next((r.score for r in results_after if r.key == "doc1"), 0.0)
|
|
assert after_score < initial_score
|
|
|
|
results_new = vector_store.search(("test",), query="new text about dogs")
|
|
for r in results_new:
|
|
if r.key == "doc1":
|
|
assert r.score > after_score
|
|
|
|
# Don't index this one
|
|
vector_store.put(("test",), "doc4", {"text": "new text about dogs"}, index=False)
|
|
results_new = vector_store.search(("test",), query="new text about dogs", limit=3)
|
|
assert not any(r.key == "doc4" for r in results_new)
|
|
|
|
|
|
def test_vector_search_with_filters(vector_store: PostgresStore) -> None:
|
|
"""Test combining vector search with filters."""
|
|
# Insert test documents
|
|
docs = [
|
|
("doc1", {"text": "red apple", "color": "red", "score": 4.5}),
|
|
("doc2", {"text": "red car", "color": "red", "score": 3.0}),
|
|
("doc3", {"text": "green apple", "color": "green", "score": 4.0}),
|
|
("doc4", {"text": "blue car", "color": "blue", "score": 3.5}),
|
|
]
|
|
|
|
for key, value in docs:
|
|
vector_store.put(("test",), key, value)
|
|
|
|
results = vector_store.search(("test",), query="apple", filter={"color": "red"})
|
|
assert len(results) == 2
|
|
assert results[0].key == "doc1"
|
|
|
|
results = vector_store.search(("test",), query="car", filter={"color": "red"})
|
|
assert len(results) == 2
|
|
assert results[0].key == "doc2"
|
|
|
|
results = vector_store.search(
|
|
("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
|
|
)
|
|
assert len(results) == 3
|
|
assert results[0].key == "doc4"
|
|
|
|
# Multiple filters
|
|
results = vector_store.search(
|
|
("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"}
|
|
)
|
|
assert len(results) == 1
|
|
assert results[0].key == "doc3"
|
|
|
|
|
|
def test_vector_search_pagination(vector_store: PostgresStore) -> None:
|
|
"""Test pagination with vector search."""
|
|
# Insert multiple similar documents
|
|
for i in range(5):
|
|
vector_store.put(("test",), f"doc{i}", {"text": f"test document number {i}"})
|
|
|
|
# Test with different page sizes
|
|
results_page1 = vector_store.search(("test",), query="test", limit=2)
|
|
results_page2 = vector_store.search(("test",), query="test", limit=2, offset=2)
|
|
|
|
assert len(results_page1) == 2
|
|
assert len(results_page2) == 2
|
|
assert results_page1[0].key != results_page2[0].key
|
|
|
|
# Get all results
|
|
all_results = vector_store.search(("test",), query="test", limit=10)
|
|
assert len(all_results) == 5
|
|
|
|
|
|
def test_vector_search_edge_cases(vector_store: PostgresStore) -> None:
|
|
"""Test edge cases in vector search."""
|
|
vector_store.put(("test",), "doc1", {"text": "test document"})
|
|
|
|
results = vector_store.search(("test",), query="")
|
|
assert len(results) == 1
|
|
|
|
results = vector_store.search(("test",), query=None)
|
|
assert len(results) == 1
|
|
|
|
long_query = "test " * 100
|
|
results = vector_store.search(("test",), query=long_query)
|
|
assert len(results) == 1
|
|
|
|
special_query = "test!@#$%^&*()"
|
|
results = vector_store.search(("test",), query=special_query)
|
|
assert len(results) == 1
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"vector_type,distance_type",
|
|
[
|
|
("vector", "cosine"),
|
|
("vector", "inner_product"),
|
|
("halfvec", "cosine"),
|
|
("halfvec", "inner_product"),
|
|
],
|
|
)
|
|
def test_embed_with_path_sync(
|
|
request: Any,
|
|
fake_embeddings: CharacterEmbeddings,
|
|
vector_type: str,
|
|
distance_type: str,
|
|
) -> None:
|
|
"""Test vector search with specific text fields in Postgres store."""
|
|
with _create_vector_store(
|
|
vector_type,
|
|
distance_type,
|
|
fake_embeddings,
|
|
text_fields=["key0", "key1", "key3"],
|
|
) as store:
|
|
# This will have 2 vectors representing it
|
|
doc1 = {
|
|
# Omit key0 - check it doesn't raise an error
|
|
"key1": "xxx",
|
|
"key2": "yyy",
|
|
"key3": "zzz",
|
|
}
|
|
# This will have 3 vectors representing it
|
|
doc2 = {
|
|
"key0": "uuu",
|
|
"key1": "vvv",
|
|
"key2": "www",
|
|
"key3": "xxx",
|
|
}
|
|
store.put(("test",), "doc1", doc1)
|
|
store.put(("test",), "doc2", doc2)
|
|
|
|
# doc2.key3 and doc1.key1 both would have the highest score
|
|
results = store.search(("test",), query="xxx")
|
|
assert len(results) == 2
|
|
assert results[0].key != results[1].key
|
|
ascore = results[0].score
|
|
bscore = results[1].score
|
|
assert ascore == pytest.approx(bscore, abs=1e-3)
|
|
|
|
# ~Only match doc2
|
|
results = store.search(("test",), query="uuu")
|
|
assert len(results) == 2
|
|
assert results[0].key != results[1].key
|
|
assert results[0].key == "doc2"
|
|
assert results[0].score > results[1].score
|
|
assert ascore == pytest.approx(results[0].score, abs=1e-3)
|
|
|
|
# ~Only match doc1
|
|
results = store.search(("test",), query="zzz")
|
|
assert len(results) == 2
|
|
assert results[0].key != results[1].key
|
|
assert results[0].key == "doc1"
|
|
assert results[0].score > results[1].score
|
|
assert ascore == pytest.approx(results[0].score, abs=1e-3)
|
|
|
|
# Un-indexed - will have low results for both. Not zero (because we're projecting)
|
|
# but less than the above.
|
|
results = store.search(("test",), query="www")
|
|
assert len(results) == 2
|
|
assert results[0].key != results[1].key
|
|
assert results[0].score < ascore
|
|
assert results[1].score < ascore
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"vector_type,distance_type",
|
|
[
|
|
("vector", "cosine"),
|
|
("vector", "inner_product"),
|
|
("halfvec", "cosine"),
|
|
("halfvec", "inner_product"),
|
|
],
|
|
)
|
|
def test_embed_with_path_operation_config(
|
|
request: Any,
|
|
fake_embeddings: CharacterEmbeddings,
|
|
vector_type: str,
|
|
distance_type: str,
|
|
) -> None:
|
|
"""Test operation-level field configuration for vector search."""
|
|
|
|
with _create_vector_store(
|
|
vector_type,
|
|
distance_type,
|
|
fake_embeddings,
|
|
text_fields=["key17"], # Default fields that won't match our test data
|
|
) as store:
|
|
doc3 = {
|
|
"key0": "aaa",
|
|
"key1": "bbb",
|
|
"key2": "ccc",
|
|
"key3": "ddd",
|
|
}
|
|
doc4 = {
|
|
"key0": "eee",
|
|
"key1": "bbb", # Same as doc3.key1
|
|
"key2": "fff",
|
|
"key3": "ggg",
|
|
}
|
|
|
|
store.put(("test",), "doc3", doc3, index=["key0", "key1"])
|
|
store.put(("test",), "doc4", doc4, index=["key1", "key3"])
|
|
|
|
results = store.search(("test",), query="aaa")
|
|
assert len(results) == 2
|
|
assert results[0].key == "doc3"
|
|
assert len(set(r.key for r in results)) == 2
|
|
assert results[0].score > results[1].score
|
|
|
|
results = store.search(("test",), query="ggg")
|
|
assert len(results) == 2
|
|
assert results[0].key == "doc4"
|
|
assert results[0].score > results[1].score
|
|
|
|
results = store.search(("test",), query="bbb")
|
|
assert len(results) == 2
|
|
assert results[0].key != results[1].key
|
|
assert results[0].score == pytest.approx(results[1].score, abs=1e-3)
|
|
|
|
results = store.search(("test",), query="ccc")
|
|
assert len(results) == 2
|
|
assert all(
|
|
r.score < 0.9 for r in results
|
|
) # Unindexed field should have low scores
|
|
|
|
# Test index=False behavior
|
|
doc5 = {
|
|
"key0": "hhh",
|
|
"key1": "iii",
|
|
}
|
|
store.put(("test",), "doc5", doc5, index=False)
|
|
results = store.search(("test",))
|
|
assert len(results) == 3
|
|
assert all(r.score is None for r in results)
|
|
assert any(r.key == "doc5" for r in results)
|
|
|
|
results = store.search(("test",), query="hhh")
|
|
# TODO: We don't currently fill in additional results if there are not enough
|
|
# returned during vector search.
|
|
# assert len(results) == 3
|
|
# doc5_result = next(r for r in results if r.key == "doc5")
|
|
# assert doc5_result.score is None
|
|
|
|
|
|
def _cosine_similarity(X: list[float], Y: list[list[float]]) -> list[float]:
|
|
"""
|
|
Compute cosine similarity between a vector X and a matrix Y.
|
|
Lazy import numpy for efficiency.
|
|
"""
|
|
|
|
similarities = []
|
|
for y in Y:
|
|
dot_product = sum(a * b for a, b in zip(X, y))
|
|
norm1 = sum(a * a for a in X) ** 0.5
|
|
norm2 = sum(a * a for a in y) ** 0.5
|
|
similarity = dot_product / (norm1 * norm2) if norm1 > 0 and norm2 > 0 else 0.0
|
|
similarities.append(similarity)
|
|
|
|
return similarities
|
|
|
|
|
|
def _inner_product(X: list[float], Y: list[list[float]]) -> list[float]:
|
|
"""
|
|
Compute inner product between a vector X and a matrix Y.
|
|
Lazy import numpy for efficiency.
|
|
"""
|
|
|
|
similarities = []
|
|
for y in Y:
|
|
similarity = sum(a * b for a, b in zip(X, y))
|
|
similarities.append(similarity)
|
|
|
|
return similarities
|
|
|
|
|
|
def _neg_l2_distance(X: list[float], Y: list[list[float]]) -> list[float]:
|
|
"""
|
|
Compute l2 distance between a vector X and a matrix Y.
|
|
Lazy import numpy for efficiency.
|
|
"""
|
|
|
|
similarities = []
|
|
for y in Y:
|
|
similarity = sum((a - b) ** 2 for a, b in zip(X, y)) ** 0.5
|
|
similarities.append(-similarity)
|
|
|
|
return similarities
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"vector_type,distance_type",
|
|
[
|
|
("vector", "cosine"),
|
|
("vector", "inner_product"),
|
|
("halfvec", "l2"),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("query", ["aaa", "bbb", "ccc", "abcd", "poisson"])
|
|
def test_scores(
|
|
fake_embeddings: CharacterEmbeddings,
|
|
vector_type: str,
|
|
distance_type: str,
|
|
query: str,
|
|
) -> None:
|
|
"""Test operation-level field configuration for vector search."""
|
|
with _create_vector_store(
|
|
vector_type,
|
|
distance_type,
|
|
fake_embeddings,
|
|
text_fields=["key0"],
|
|
) as store:
|
|
doc = {
|
|
"key0": "aaa",
|
|
}
|
|
store.put(("test",), "doc", doc, index=["key0", "key1"])
|
|
|
|
results = store.search((), query=query)
|
|
vec0 = fake_embeddings.embed_query(doc["key0"])
|
|
vec1 = fake_embeddings.embed_query(query)
|
|
if distance_type == "cosine":
|
|
similarities = _cosine_similarity(vec1, [vec0])
|
|
elif distance_type == "inner_product":
|
|
similarities = _inner_product(vec1, [vec0])
|
|
elif distance_type == "l2":
|
|
similarities = _neg_l2_distance(vec1, [vec0])
|
|
|
|
assert len(results) == 1
|
|
assert results[0].score == pytest.approx(similarities[0], abs=1e-3)
|