diff --git a/README.md b/README.md index 84e677a..226d375 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ The data model is simple: - **Entity** — a named node with a `type` and optional JSON `properties` - **Link** — a directed edge from a *subject* entity to an *object* entity, labelled with a *predicate* string and optional JSON `properties` +- **Embedding** — a vector attached to an entity, with a model label and optional JSON metadata Example: `(Experiment "SAXS run 42") --[produced]--> (Dataset "raw_001.h5")` @@ -103,6 +104,58 @@ mutation { } ``` +### Create an embedding + +```bash +MODEL_ID=$(curl -X POST http://localhost:8080/splash_links/embedding-models \ + -H 'Content-Type: application/json' \ + -d '{ + "name": "text-embedding-3-small", + "version": "1", + "description": "Example model metadata" + }' | jq -r '.id') + +curl -X POST http://localhost:8080/splash_links/embeddings \ + -H 'Content-Type: application/json' \ + -d '{ + "entityId": "", + "embeddingModelId": "'"$MODEL_ID"'", + "vector": [0.12, -0.03, 0.88], + "properties": {"chunk": 1} + }' +``` + +### Find nearest embeddings + +Embedding CRUD uses REST. Nearest-neighbor search stays in GraphQL and uses cosine distance. For PostgreSQL, embeddings are stored in a native `pgvector` column; for SQLite, they are stored as compact binary blobs and searched in-process. + +```graphql +query { + nearestEmbeddings( + vector: [0.11, -0.02, 0.90] + embeddingModelId: "" + limit: 5 + ) { + distance + embedding { + id + entityId + embeddingModel { name version } + entity { name } + } + } +} +``` + +### Fetch or delete embeddings + +```bash +curl http://localhost:8080/splash_links/embeddings/ +curl 'http://localhost:8080/splash_links/embedding-models?name=text-embedding-3-small' +curl 'http://localhost:8080/splash_links/embeddings?entityId=&embeddingModelId=' +curl -X DELETE http://localhost:8080/splash_links/embeddings/ +``` + ### Health check ``` @@ -132,6 +185,13 @@ pixi run links -- --subject # outgoing from a node pixi run links -- --object # incoming to a node ``` +### List embeddings + +```bash +pixi run embeddings -- --entity +splash-links embeddings --model-id --limit 10 +``` + ### Raw SQLite shell ```bash @@ -221,6 +281,7 @@ Tests require ≥ 90% coverage and will fail the build if that threshold is not | `docs` | `pixi run docs` | Serve MkDocs site locally | | `entities` | `pixi run entities` | List entities in the database | | `links` | `pixi run links` | List links in the database | +| `embeddings` | `pixi run embeddings` | List embeddings in the database | | `db` | `pixi run db` | Open raw SQLite interactive shell | Pass extra flags after `--`, e.g. `pixi run entities -- --type Experiment --limit 5`. @@ -258,6 +319,16 @@ SPLASH_LINKS_DB=links.sqlite pixi run serve SPLASH_LINKS_DB=/data/links.sqlite pixi run serve ``` +#### PostgreSQL with pgvector + +PostgreSQL nearest-neighbor search uses the `pgvector` extension. The Alembic migration will create the extension automatically when permissions allow it. + +```bash +SPLASH_LINKS_DB=postgresql+psycopg2://user:pass@host/dbname pixi run serve +``` + +Embeddings use dialect-specific storage. PostgreSQL stores them in a native `pgvector` column, while SQLite stores packed float32 data in a BLOB. Base64 is intentionally not used, since it would only increase storage size and parsing overhead. + #### PostgreSQL (recommended for production / multi-user deployments) Use PostgreSQL when you need concurrent writes, role-based access control, or want to run the service behind a load balancer. A `docker-compose.yml` is provided that starts a Postgres instance alongside the application: diff --git a/_tests/test_base_client.py b/_tests/test_base_client.py index a4db415..5ffb43f 100644 --- a/_tests/test_base_client.py +++ b/_tests/test_base_client.py @@ -4,7 +4,7 @@ from splash_links.client import base as base_module from splash_links.client import tiled as tiled_module -from splash_links.client.base import Entity, LinksClient, from_uri +from splash_links.client.base import EmbeddingMatch, Entity, LinksClient, from_uri from splash_links.client.tiled import TiledEntity, _node_name, _node_properties, _node_uri, from_entity from splash_links.client.tiled import get_or_create_entity as tiled_get_or_create @@ -195,6 +195,165 @@ def fake_execute(query: str, variables: dict | None = None) -> dict: } +def test_create_embedding_posts_expected_payload(monkeypatch): + seen: dict[str, object] = {} + + def fake_post(url: str, json: dict, timeout: float): + seen["url"] = url + seen["json"] = json + seen["timeout"] = timeout + return FakeResponse( + { + "id": "emb-1", + "entityId": "ent-1", + "embeddingModelId": "model-1", + "embeddingModel": { + "id": "model-1", + "name": "model-a", + "description": None, + "url": None, + "version": "1", + }, + "vector": [0.1, 0.2, 0.3], + "dimensions": 3, + "properties": {"chunk": 1}, + "createdAt": "2026-01-01T00:00:00Z", + } + ) + + monkeypatch.setattr(base_module.httpx, "post", fake_post) + + client = from_uri("splash://api:8080") + embedding = client.create_embedding( + entity_id="ent-1", + vector=[0.1, 0.2, 0.3], + embedding_model_id="model-1", + properties={"chunk": 1}, + ) + + assert embedding.id == "emb-1" + assert embedding.embedding_model.id == "model-1" + assert seen["url"] == "http://api:8080/splash_links/embeddings" + assert seen["timeout"] == 30.0 + assert seen["json"] == { + "entityId": "ent-1", + "embeddingModelId": "model-1", + "vector": [0.1, 0.2, 0.3], + "properties": {"chunk": 1}, + } + + +def test_create_embedding_model_posts_expected_payload(monkeypatch): + seen: dict[str, object] = {} + + def fake_post(url: str, json: dict, timeout: float): + seen["url"] = url + seen["json"] = json + seen["timeout"] = timeout + return FakeResponse( + { + "id": "model-1", + "name": "model-a", + "description": "Example model", + "url": "https://example.com/model-a", + "version": "1", + } + ) + + monkeypatch.setattr(base_module.httpx, "post", fake_post) + + client = from_uri("splash://api:8080") + model = client.create_embedding_model( + name="model-a", + version="1", + description="Example model", + url="https://example.com/model-a", + ) + + assert model.id == "model-1" + assert seen["url"] == "http://api:8080/splash_links/embedding-models" + assert seen["json"] == { + "name": "model-a", + "version": "1", + "description": "Example model", + "url": "https://example.com/model-a", + } + + +def test_find_nearest_embeddings_posts_expected_payload(monkeypatch): + seen: dict[str, object] = {} + + def fake_execute(query: str, variables: dict | None = None) -> dict: + seen["query"] = query + seen["variables"] = variables + return { + "nearestEmbeddings": [ + { + "distance": 0.01, + "embedding": { + "id": "emb-1", + "entityId": "ent-1", + "embeddingModelId": "model-1", + "embeddingModel": { + "id": "model-1", + "name": "model-a", + "description": None, + "url": None, + "version": "1", + }, + "vector": [0.1, 0.2], + "dimensions": 2, + "properties": None, + "createdAt": "2026-01-01T00:00:00Z", + }, + } + ] + } + + client = LinksClient("http://example.com") + monkeypatch.setattr(client, "_execute", fake_execute) + + matches = client.find_nearest_embeddings( + vector=[0.1, 0.2], + embedding_model_id="model-1", + entity_id="ent-1", + limit=5, + offset=1, + ) + + assert matches == [ + EmbeddingMatch.model_validate( + { + "distance": 0.01, + "embedding": { + "id": "emb-1", + "entityId": "ent-1", + "embeddingModelId": "model-1", + "embeddingModel": { + "id": "model-1", + "name": "model-a", + "description": None, + "url": None, + "version": "1", + }, + "vector": [0.1, 0.2], + "dimensions": 2, + "properties": None, + "createdAt": "2026-01-01T00:00:00Z", + }, + } + ) + ] + assert seen["query"] == base_module._NEAREST_EMBEDDINGS_QUERY + assert seen["variables"] == { + "vector": [0.1, 0.2], + "embeddingModelId": "model-1", + "entityId": "ent-1", + "limit": 5, + "offset": 1, + } + + # --------------------------------------------------------------------------- # Tiled integration helpers # --------------------------------------------------------------------------- diff --git a/_tests/test_cli.py b/_tests/test_cli.py index 55b8046..0dd3a6c 100644 --- a/_tests/test_cli.py +++ b/_tests/test_cli.py @@ -8,7 +8,7 @@ import splash_links.cli as cli_module from splash_links.cli import app -from splash_links.store import EntityRecord, LinkRecord +from splash_links.store import EmbeddingModelRecord, EmbeddingRecord, EntityRecord, LinkRecord runner = CliRunner() @@ -44,10 +44,32 @@ def _make_link(**kw) -> LinkRecord: return LinkRecord(**defaults) +def _make_embedding(**kw) -> EmbeddingRecord: + defaults = dict( + id="emb-1", + entity_id="ent-1", + embedding_model_id="model-1", + embedding_model=EmbeddingModelRecord( + id="model-1", + name="model-a", + description=None, + url=None, + version="1", + ), + vector=[0.1, 0.2, 0.3], + dimensions=3, + properties={}, + created_at=datetime.now(timezone.utc), + ) + defaults.update(kw) + return EmbeddingRecord(**defaults) + + class FakeStore: - def __init__(self, entities=None, links=None): + def __init__(self, entities=None, links=None, embeddings=None): self._entities = entities or [] self._links = links or [] + self._embeddings = embeddings or [] def list_entities(self, entity_type=None, limit=50, offset=0): if entity_type: @@ -57,6 +79,14 @@ def list_entities(self, entity_type=None, limit=50, offset=0): def find_links(self, subject_id=None, predicate=None, object_id=None, limit=50, offset=0): return self._links + def list_embeddings(self, entity_id=None, embedding_model_id=None, limit=50, offset=0): + rows = self._embeddings + if entity_id: + rows = [embedding for embedding in rows if embedding.entity_id == entity_id] + if embedding_model_id: + rows = [embedding for embedding in rows if embedding.embedding_model_id == embedding_model_id] + return rows + def close(self): pass @@ -124,6 +154,26 @@ def test_links_with_properties(monkeypatch): assert "confidence" in result.output +# --------------------------------------------------------------------------- +# embeddings command +# --------------------------------------------------------------------------- + + +def test_embeddings_shows_rows(monkeypatch): + fake = FakeStore(embeddings=[_make_embedding()]) + monkeypatch.setattr(cli_module, "_open_store", lambda: fake) + result = runner.invoke(app, ["embeddings"]) + assert result.exit_code == 0 + assert "model-a" in result.output + + +def test_embeddings_no_rows_prints_message(monkeypatch): + monkeypatch.setattr(cli_module, "_open_store", lambda: FakeStore()) + result = runner.invoke(app, ["embeddings"]) + assert result.exit_code == 0 + assert "No embeddings found" in result.output + + # --------------------------------------------------------------------------- # _open_store — missing DB # --------------------------------------------------------------------------- diff --git a/_tests/test_client_cli.py b/_tests/test_client_cli.py index b9da7be..c1a6379 100644 --- a/_tests/test_client_cli.py +++ b/_tests/test_client_cli.py @@ -6,7 +6,7 @@ from splash_links import cli as root_cli from splash_links.client import cli as client_cli -from splash_links.client.base import Entity, Link +from splash_links.client.base import Embedding, EmbeddingMatch, EmbeddingModel, Entity, Link runner = CliRunner() @@ -241,6 +241,177 @@ def find_links(self, **kw): assert "Failed to find links" in result.output +def test_create_embedding_command_outputs_json(monkeypatch): + seen: dict[str, object] = {} + + class FakeClient: + def create_embedding(self, entity_id, vector, embedding_model_id, properties=None): + seen["entity_id"] = entity_id + seen["vector"] = vector + seen["embedding_model_id"] = embedding_model_id + seen["properties"] = properties + return Embedding( + id="emb-1", + entity_id=entity_id, + embedding_model_id=embedding_model_id, + embedding_model=EmbeddingModel( + id=embedding_model_id, + name="model-a", + description=None, + url=None, + version="1", + ), + vector=vector, + dimensions=len(vector), + properties=properties, + created_at="2026-01-01T00:00:00Z", + ) + + monkeypatch.setattr(client_cli, "from_uri", lambda uri: FakeClient()) + + result = runner.invoke( + client_cli.app, + [ + "create-embedding", + "ent-1", + "--vector", + "[0.1, 0.2, 0.3]", + "--model-id", + "model-1", + "--properties", + '{"chunk": 1}', + ], + ) + + assert result.exit_code == 0, result.output + payload = json.loads(result.stdout) + assert payload["id"] == "emb-1" + assert payload["embedding_model"]["name"] == "model-a" + assert seen == { + "entity_id": "ent-1", + "vector": [0.1, 0.2, 0.3], + "embedding_model_id": "model-1", + "properties": {"chunk": 1}, + } + + +def test_create_embedding_model_command_outputs_json(monkeypatch): + seen: dict[str, object] = {} + + class FakeClient: + def create_embedding_model(self, name, version, description=None, url=None): + seen["name"] = name + seen["version"] = version + seen["description"] = description + seen["url"] = url + return EmbeddingModel( + id="model-1", + name=name, + description=description, + url=url, + version=version, + ) + + monkeypatch.setattr(client_cli, "from_uri", lambda uri: FakeClient()) + + result = runner.invoke( + client_cli.app, + [ + "create-embedding-model", + "--name", + "model-a", + "--version", + "1", + "--description", + "Example model", + "--url", + "https://example.com/model-a", + ], + ) + + assert result.exit_code == 0, result.output + payload = json.loads(result.stdout) + assert payload["id"] == "model-1" + assert seen == { + "name": "model-a", + "version": "1", + "description": "Example model", + "url": "https://example.com/model-a", + } + + +def test_create_embedding_invalid_vector_exits_2(): + result = runner.invoke( + client_cli.app, + ["create-embedding", "ent-1", "--model-id", "model-1", "--vector", "not-json"], + ) + assert result.exit_code == 2 + assert "Invalid JSON passed to --vector" in result.output + + +def test_nearest_embeddings_command_outputs_list(monkeypatch): + seen: dict[str, object] = {} + + class FakeClient: + def find_nearest_embeddings(self, vector, embedding_model_id=None, entity_id=None, limit=10, offset=0): + seen["vector"] = vector + seen["embedding_model_id"] = embedding_model_id + seen["entity_id"] = entity_id + seen["limit"] = limit + seen["offset"] = offset + return [ + EmbeddingMatch( + distance=0.01, + embedding=Embedding( + id="emb-1", + entity_id="ent-1", + embedding_model_id="model-1", + embedding_model=EmbeddingModel( + id="model-1", + name="model-a", + description=None, + url=None, + version="1", + ), + vector=[0.1, 0.2], + dimensions=2, + properties=None, + created_at="2026-01-01T00:00:00Z", + ), + ) + ] + + monkeypatch.setattr(client_cli, "from_uri", lambda uri: FakeClient()) + + result = runner.invoke( + client_cli.app, + [ + "nearest-embeddings", + "--vector", + "[0.1, 0.2]", + "--model-id", + "model-1", + "--entity-id", + "ent-1", + "--limit", + "5", + "--offset", + "1", + ], + ) + + assert result.exit_code == 0, result.output + payload = json.loads(result.stdout) + assert payload[0]["embedding"]["id"] == "emb-1" + assert seen == { + "vector": [0.1, 0.2], + "embedding_model_id": "model-1", + "entity_id": "ent-1", + "limit": 5, + "offset": 1, + } + + def test_client_cli_main(monkeypatch): from splash_links.client import cli as client_cli_module diff --git a/_tests/test_service.py b/_tests/test_service.py index b0ad228..5668479 100644 --- a/_tests/test_service.py +++ b/_tests/test_service.py @@ -9,14 +9,11 @@ import pytest from fastapi.testclient import TestClient +from sqlalchemy import text from splash_links.app import create_app from splash_links.store import SQLiteStore -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - @pytest.fixture() def store(): @@ -34,10 +31,6 @@ def client(): yield c -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - _GQL = "/splash_links/graphql" @@ -49,6 +42,48 @@ def gql(client: TestClient, query: str, variables: dict | None = None) -> dict: return body["data"] +def create_embedding_model( + client: TestClient, + *, + name: str, + version: str, + description: str | None = None, + url: str | None = None, +) -> dict: + resp = client.post( + "/splash_links/embedding-models", + json={ + "name": name, + "version": version, + "description": description, + "url": url, + }, + ) + assert resp.status_code == 201, resp.text + return resp.json() + + +def create_embedding( + client: TestClient, + *, + entity_id: str, + embedding_model_id: str, + vector: list[float], + properties: dict | None = None, +) -> dict: + resp = client.post( + "/splash_links/embeddings", + json={ + "entityId": entity_id, + "embeddingModelId": embedding_model_id, + "vector": vector, + "properties": properties, + }, + ) + assert resp.status_code == 201, resp.text + return resp.json() + + CREATE_ENTITY = """ mutation CreateEntity($input: CreateEntityInput!) { createEntity(input: $input) { @@ -77,22 +112,12 @@ def gql(client: TestClient, query: str, variables: dict | None = None) -> dict: """ -# --------------------------------------------------------------------------- -# Health check -# --------------------------------------------------------------------------- - - def test_health(client): resp = client.get("/splash_links/health") assert resp.status_code == 200 assert resp.json() == {"status": "ok"} -# --------------------------------------------------------------------------- -# Store unit tests (no HTTP) -# --------------------------------------------------------------------------- - - class TestSQLiteStore: def test_create_and_get_entity(self, store): e = store.create_entity("Dataset", "SAXS run 001", properties={"beamline": "12.3.1"}) @@ -128,7 +153,7 @@ def test_delete_entity_cascades_links(self, store): assert store.delete_entity(e1.id) is True assert store.get_entity(e1.id) is None - assert store.get_link(lnk.id) is None # cascade deleted + assert store.get_link(lnk.id) is None def test_delete_entity_not_found(self, store): assert store.delete_entity("ghost") is False @@ -180,7 +205,7 @@ def test_delete_link(self, store): assert store.delete_link(lnk.id) is True assert store.get_link(lnk.id) is None - assert store.delete_link(lnk.id) is False # already gone + assert store.delete_link(lnk.id) is False def test_update_entity_name(self, store): e = store.create_entity("Dataset", "original") @@ -221,10 +246,84 @@ def test_update_link_predicate(self, store): def test_update_link_not_found(self, store): assert store.update_link("ghost", "anything") is None + def test_create_and_get_embedding(self, store): + entity = store.create_entity("Dataset", "run-1") + model = store.create_embedding_model("text-embedding-3-small", "1", description="Test model") + embedding = store.create_embedding( + entity.id, + embedding_model_id=model.id, + vector=[0.1, 0.2, 0.3], + properties={"chunk": 1}, + ) + + assert embedding.entity_id == entity.id + assert embedding.embedding_model_id == model.id + assert embedding.embedding_model.name == "text-embedding-3-small" + assert embedding.vector == [0.1, 0.2, 0.3] + assert embedding.dimensions == 3 + assert embedding.properties == {"chunk": 1} -# --------------------------------------------------------------------------- -# GraphQL integration tests (via HTTP) -# --------------------------------------------------------------------------- + fetched = store.get_embedding(embedding.id) + assert fetched is not None + assert fetched.id == embedding.id + + def test_sqlite_embeddings_are_stored_as_blob(self, store): + entity = store.create_entity("Dataset", "run-1") + model = store.create_embedding_model("default", "1") + embedding = store.create_embedding(entity.id, embedding_model_id=model.id, vector=[0.1, 0.2, 0.3]) + + with store._engine.connect() as conn: + storage_type = conn.execute( + text("SELECT typeof(vector) FROM embeddings WHERE id = :id"), + {"id": embedding.id}, + ).scalar_one() + + assert storage_type == "blob" + + def test_create_embedding_missing_entity_raises(self, store): + model = store.create_embedding_model("default", "1") + with pytest.raises(ValueError, match="Entity"): + store.create_embedding("missing", embedding_model_id=model.id, vector=[0.1, 0.2, 0.3]) + + def test_create_embedding_rejects_zero_vector(self, store): + entity = store.create_entity("Dataset", "run-1") + model = store.create_embedding_model("default", "1") + with pytest.raises(ValueError, match="all zeros"): + store.create_embedding(entity.id, embedding_model_id=model.id, vector=[0.0, 0.0, 0.0]) + + def test_list_embeddings_filters(self, store): + entity = store.create_entity("Dataset", "run-1") + other = store.create_entity("Dataset", "run-2") + model_a = store.create_embedding_model("model-a", "1") + model_b = store.create_embedding_model("model-b", "1") + store.create_embedding(entity.id, embedding_model_id=model_a.id, vector=[0.1, 0.2]) + store.create_embedding(entity.id, embedding_model_id=model_b.id, vector=[0.2, 0.3]) + store.create_embedding(other.id, embedding_model_id=model_a.id, vector=[0.3, 0.4]) + + rows = store.list_embeddings(entity_id=entity.id, embedding_model_id=model_a.id) + assert len(rows) == 1 + assert rows[0].embedding_model.name == "model-a" + assert rows[0].entity_id == entity.id + + def test_find_nearest_embeddings_orders_by_distance(self, store): + entity = store.create_entity("Dataset", "run-1") + model = store.create_embedding_model("model-a", "1") + near = store.create_embedding(entity.id, embedding_model_id=model.id, vector=[1.0, 0.0]) + far = store.create_embedding(entity.id, embedding_model_id=model.id, vector=[0.0, 1.0]) + store.create_embedding(entity.id, embedding_model_id=model.id, vector=[1.0, 1.0, 0.0]) + + matches = store.find_nearest_embeddings([0.9, 0.1], embedding_model_id=model.id) + + assert [match.embedding.id for match in matches] == [near.id, far.id] + assert matches[0].distance < matches[1].distance + + def test_delete_entity_cascades_embeddings(self, store): + entity = store.create_entity("Dataset", "run-1") + model = store.create_embedding_model("default", "1") + embedding = store.create_embedding(entity.id, embedding_model_id=model.id, vector=[0.1, 0.2, 0.3]) + + assert store.delete_entity(entity.id) is True + assert store.get_embedding(embedding.id) is None class TestGraphQL: @@ -240,11 +339,9 @@ def test_create_entity(self, client): assert e["properties"] == {"energy": 7.0} def test_query_entity(self, client): - created = gql( - client, - CREATE_ENTITY, - {"input": {"entityType": "Experiment", "name": "exp-42"}}, - )["createEntity"] + created = gql(client, CREATE_ENTITY, {"input": {"entityType": "Experiment", "name": "exp-42"}})[ + "createEntity" + ] fetched = gql( client, @@ -255,10 +352,7 @@ def test_query_entity(self, client): assert fetched["name"] == "exp-42" def test_query_entity_not_found(self, client): - result = gql( - client, - '{ entity(id: "00000000-0000-0000-0000-000000000000") { id } }', - ) + result = gql(client, '{ entity(id: "00000000-0000-0000-0000-000000000000") { id } }') assert result["entity"] is None def test_list_entities(self, client): @@ -284,11 +378,7 @@ def test_create_and_traverse_link(self, client): def test_traverse_outgoing_and_incoming(self, client): exp = gql(client, CREATE_ENTITY, {"input": {"entityType": "Experiment", "name": "exp"}})["createEntity"] ds = gql(client, CREATE_ENTITY, {"input": {"entityType": "Dataset", "name": "ds"}})["createEntity"] - gql( - client, - CREATE_LINK, - {"input": {"subjectId": exp["id"], "predicate": "produced", "objectId": ds["id"]}}, - ) + gql(client, CREATE_LINK, {"input": {"subjectId": exp["id"], "predicate": "produced", "objectId": ds["id"]}}) out = gql( client, @@ -342,7 +432,6 @@ def test_delete_link(self, client): result = gql(client, "mutation D($id: ID!) { deleteLink(id: $id) }", {"id": lnk["id"]}) assert result["deleteLink"] is True - # Entities still exist still_there = gql(client, "query Q($id: ID!) { entity(id: $id) { id } }", {"id": e1["id"]}) assert still_there["entity"] is not None @@ -368,9 +457,9 @@ def test_update_entity_not_found_returns_null(self, client): def test_update_link_mutation(self, client): e1 = gql(client, CREATE_ENTITY, {"input": {"entityType": "A", "name": "a"}})["createEntity"] e2 = gql(client, CREATE_ENTITY, {"input": {"entityType": "B", "name": "b"}})["createEntity"] - lnk = gql( - client, CREATE_LINK, {"input": {"subjectId": e1["id"], "predicate": "old", "objectId": e2["id"]}} - )["createLink"] + lnk = gql(client, CREATE_LINK, {"input": {"subjectId": e1["id"], "predicate": "old", "objectId": e2["id"]}})[ + "createLink" + ] data = gql( client, """mutation U($id: ID!, $input: UpdateLinkInput!) { @@ -383,7 +472,126 @@ def test_update_link_mutation(self, client): def test_update_link_not_found_returns_null(self, client): data = gql( client, - 'mutation { updateLink(id: "00000000-0000-0000-0000-000000000000",' - ' input: { predicate: "x" }) { id } }', + 'mutation { updateLink(id: "00000000-0000-0000-0000-000000000000", input: { predicate: "x" }) { id } }', ) assert data["updateLink"] is None + + def test_nearest_embeddings(self, client): + entity = gql(client, CREATE_ENTITY, {"input": {"entityType": "Dataset", "name": "run-1"}})["createEntity"] + model = create_embedding_model(client, name="model-a", version="1") + create_embedding(client, entity_id=entity["id"], embedding_model_id=model["id"], vector=[1.0, 0.0]) + create_embedding(client, entity_id=entity["id"], embedding_model_id=model["id"], vector=[0.0, 1.0]) + create_embedding(client, entity_id=entity["id"], embedding_model_id=model["id"], vector=[1.0, 1.0, 0.0]) + + data = gql( + client, + """ + query Search($vector: [Float!]!, $model: ID) { + nearestEmbeddings(vector: $vector, embeddingModelId: $model) { + distance + embedding { + entityId + embeddingModelId + embeddingModel { name version } + vector + } + } + } + """, + {"vector": [0.9, 0.1], "model": model["id"]}, + )["nearestEmbeddings"] + + assert len(data) == 2 + assert data[0]["embedding"]["vector"] == [1.0, 0.0] + assert data[0]["embedding"]["embeddingModel"]["name"] == "model-a" + assert data[0]["distance"] < data[1]["distance"] + + +class TestEmbeddingModelRestAPI: + def test_create_and_list_embedding_models(self, client): + created = create_embedding_model( + client, + name="text-embedding-3-small", + version="1", + description="Test model", + url="https://example.com/model", + ) + + assert created["name"] == "text-embedding-3-small" + assert created["version"] == "1" + + listed = client.get("/splash_links/embedding-models", params={"name": "text-embedding-3-small"}) + assert listed.status_code == 200 + rows = listed.json() + assert len(rows) == 1 + assert rows[0]["id"] == created["id"] + + def test_delete_embedding_model_in_use_returns_conflict(self, client): + entity = gql(client, CREATE_ENTITY, {"input": {"entityType": "Dataset", "name": "run-1"}})["createEntity"] + model = create_embedding_model(client, name="text-embedding-3-small", version="1") + create_embedding(client, entity_id=entity["id"], embedding_model_id=model["id"], vector=[0.1, 0.2, 0.3]) + + response = client.delete(f"/splash_links/embedding-models/{model['id']}") + + assert response.status_code == 409 + assert "still referenced" in response.json()["detail"] + + +class TestEmbeddingRestAPI: + def test_create_embedding(self, client): + entity = gql(client, CREATE_ENTITY, {"input": {"entityType": "Dataset", "name": "run-1"}})["createEntity"] + model = create_embedding_model(client, name="text-embedding-3-small", version="1") + + payload = create_embedding( + client, + entity_id=entity["id"], + embedding_model_id=model["id"], + vector=[0.1, 0.2, 0.3], + properties={"chunk": 1}, + ) + + assert payload["entityId"] == entity["id"] + assert payload["embeddingModelId"] == model["id"] + assert payload["embeddingModel"]["name"] == "text-embedding-3-small" + assert payload["dimensions"] == 3 + assert payload["properties"] == {"chunk": 1} + + def test_get_and_list_embeddings(self, client): + entity = gql(client, CREATE_ENTITY, {"input": {"entityType": "Dataset", "name": "run-1"}})["createEntity"] + model = create_embedding_model(client, name="model-a", version="1") + created = create_embedding(client, entity_id=entity["id"], embedding_model_id=model["id"], vector=[0.1, 0.2]) + + fetched = client.get(f"/splash_links/embeddings/{created['id']}") + assert fetched.status_code == 200 + assert fetched.json()["id"] == created["id"] + + listed = client.get( + "/splash_links/embeddings", + params={"entityId": entity["id"], "embeddingModelId": model["id"]}, + ) + assert listed.status_code == 200 + rows = listed.json() + assert len(rows) == 1 + assert rows[0]["id"] == created["id"] + + def test_delete_embedding(self, client): + entity = gql(client, CREATE_ENTITY, {"input": {"entityType": "Dataset", "name": "run-1"}})["createEntity"] + model = create_embedding_model(client, name="default", version="1") + created = create_embedding(client, entity_id=entity["id"], embedding_model_id=model["id"], vector=[0.1, 0.2, 0.3]) + + deleted = client.delete(f"/splash_links/embeddings/{created['id']}") + assert deleted.status_code == 204 + + missing = client.get(f"/splash_links/embeddings/{created['id']}") + assert missing.status_code == 404 + + def test_delete_entity_cascades_embeddings(self, client): + entity = gql(client, CREATE_ENTITY, {"input": {"entityType": "Dataset", "name": "run-1"}})["createEntity"] + model = create_embedding_model(client, name="default", version="1") + created = create_embedding(client, entity_id=entity["id"], embedding_model_id=model["id"], vector=[0.1, 0.2, 0.3]) + + deleted = gql(client, "mutation D($id: ID!) { deleteEntity(id: $id) }", {"id": entity["id"]}) + assert deleted["deleteEntity"] is True + + missing = client.get(f"/splash_links/embeddings/{created['id']}") + assert missing.status_code == 404 diff --git a/alembic/versions/6a86ce49c068_add_embeddings_table.py b/alembic/versions/6a86ce49c068_add_embeddings_table.py new file mode 100644 index 0000000..51680d3 --- /dev/null +++ b/alembic/versions/6a86ce49c068_add_embeddings_table.py @@ -0,0 +1,85 @@ +"""add embeddings table + +Revision ID: 6a86ce49c068 +Revises: 8fa312b52756 +Create Date: 2026-04-05 00:00:00.000000 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa +from sqlalchemy.types import UserDefinedType + +from alembic import op + + +class PostgreSQLVectorType(UserDefinedType): + cache_ok = True + + def get_col_spec(self, **_kw) -> str: + return "vector" + +# revision identifiers, used by Alembic. +revision: str = "6a86ce49c068" +down_revision: Union[str, Sequence[str], None] = "8fa312b52756" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + bind = op.get_bind() + if bind.dialect.name == "postgresql": + op.execute("CREATE EXTENSION IF NOT EXISTS vector") + + vector_type: sa.types.TypeEngine + if bind.dialect.name == "postgresql": + vector_type = PostgreSQLVectorType() + else: + vector_type = sa.LargeBinary() + + op.create_table( + "embedding_models", + sa.Column("id", sa.String(), nullable=False), + sa.Column("name", sa.String(), nullable=False), + sa.Column("description", sa.String(), nullable=True), + sa.Column("url", sa.String(), nullable=True), + sa.Column("version", sa.String(), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + "embedding_models_name_version_idx", + "embedding_models", + ["name", "version"], + unique=True, + ) + + op.create_table( + "embeddings", + sa.Column("id", sa.String(), nullable=False), + sa.Column("entity_id", sa.String(), sa.ForeignKey("entities.id", ondelete="CASCADE"), nullable=False), + sa.Column("embedding_model_id", sa.String(), sa.ForeignKey("embedding_models.id", ondelete="RESTRICT"), nullable=False), + sa.Column("vector", vector_type, nullable=False), + sa.Column("dimensions", sa.Integer(), nullable=False), + sa.Column("properties", sa.JSON(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + "embeddings_entity_model_created_idx", + "embeddings", + ["entity_id", "embedding_model_id", "created_at"], + ) + op.create_index( + "embeddings_model_dimensions_idx", + "embeddings", + ["embedding_model_id", "dimensions"], + ) + + +def downgrade() -> None: + op.drop_index("embeddings_model_dimensions_idx", table_name="embeddings") + op.drop_index("embeddings_entity_model_created_idx", table_name="embeddings") + op.drop_table("embeddings") + op.drop_index("embedding_models_name_version_idx", table_name="embedding_models") + op.drop_table("embedding_models") \ No newline at end of file diff --git a/pixi.toml b/pixi.toml index 6779363..9189e57 100644 --- a/pixi.toml +++ b/pixi.toml @@ -46,6 +46,7 @@ docs = "mkdocs serve -f mkdocs/mkdocs.yml" db = "splash-links shell" entities = "splash-links entities" links = "splash-links links" +embeddings = "splash-links embeddings" migrate = "alembic upgrade head" frontend-install = "npm install --prefix frontend" frontend-dev = { cmd = "npm run dev --prefix frontend", depends-on = ["frontend-install"] } diff --git a/pyproject.toml b/pyproject.toml index 0e83763..17eedf7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ requires-python = ">=3.11" dependencies = [ "alembic>=1.13.0", "fastapi>=0.111.0", + "pgvector>=0.4.0", "pydantic>=2.0.0", "rich>=13.0.0", "sqlalchemy>=2.0.0", diff --git a/src/splash_links/app.py b/src/splash_links/app.py index a63c3d3..b7e5db3 100644 --- a/src/splash_links/app.py +++ b/src/splash_links/app.py @@ -18,17 +18,92 @@ from contextlib import asynccontextmanager from typing import AsyncGenerator, Optional -from fastapi import FastAPI, Request +from fastapi import FastAPI, HTTPException, Query, Request, Response, status from fastapi.staticfiles import StaticFiles +from pydantic import BaseModel, ConfigDict, Field from strawberry.fastapi import GraphQLRouter from .schema import schema from .store import SQLAlchemyStore as SQLiteStore -from .store import Store, _make_engine, _url_from_path +from .store import EmbeddingModelRecord, EmbeddingRecord, Store, _make_engine, _url_from_path logger = logging.getLogger(__name__) +class EmbeddingModelPayload(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + id: str + name: str + description: Optional[str] = None + url: Optional[str] = None + version: str + + +class EmbeddingPayload(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + id: str + entity_id: str = Field(alias="entityId") + embedding_model_id: str = Field(alias="embeddingModelId") + embedding_model: EmbeddingModelPayload = Field(alias="embeddingModel") + vector: list[float] + dimensions: int + properties: dict + created_at: str = Field(alias="createdAt") + + +class CreateEmbeddingModelPayload(BaseModel): + name: str + description: Optional[str] = None + url: Optional[str] = None + version: str + + +class CreateEmbeddingPayload(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + entity_id: str = Field(alias="entityId") + embedding_model_id: str = Field(alias="embeddingModelId") + vector: list[float] + properties: Optional[dict] = None + + +def _embedding_model_payload(record: EmbeddingModelRecord) -> EmbeddingModelPayload: + return EmbeddingModelPayload( + id=record.id, + name=record.name, + description=record.description, + url=record.url, + version=record.version, + ) + + +def _embedding_payload(record: EmbeddingRecord) -> EmbeddingPayload: + return EmbeddingPayload( + id=record.id, + entity_id=record.entity_id, + embedding_model_id=record.embedding_model_id, + embedding_model=_embedding_model_payload(record.embedding_model), + vector=record.vector, + dimensions=record.dimensions, + properties=record.properties, + created_at=record.created_at.isoformat(), + ) + + +def _embedding_error(exc: ValueError) -> HTTPException: + message = str(exc) + lowered = message.lower() + if "not found" in lowered: + status_code = status.HTTP_404_NOT_FOUND + elif "still referenced" in lowered: + status_code = status.HTTP_409_CONFLICT + else: + status_code = status.HTTP_400_BAD_REQUEST + return HTTPException(status_code=status_code, detail=message) + + def _run_migrations(db_url: str) -> None: """Stamp existing DBs and apply all pending Alembic migrations. @@ -122,6 +197,131 @@ async def get_context(request: Request) -> dict: def health() -> dict: return {"status": "ok"} + @app.post( + "/splash_links/embedding-models", + response_model=EmbeddingModelPayload, + status_code=status.HTTP_201_CREATED, + tags=["embeddings"], + summary="Create an embedding model", + ) + def create_embedding_model(payload: CreateEmbeddingModelPayload, request: Request) -> EmbeddingModelPayload: + try: + record = request.app.state.store.create_embedding_model( + name=payload.name, + description=payload.description, + url=payload.url, + version=payload.version, + ) + except ValueError as exc: + raise _embedding_error(exc) from exc + return _embedding_model_payload(record) + + @app.get( + "/splash_links/embedding-models/{model_id}", + response_model=EmbeddingModelPayload, + tags=["embeddings"], + summary="Fetch one embedding model", + ) + def get_embedding_model(model_id: str, request: Request) -> EmbeddingModelPayload: + record = request.app.state.store.get_embedding_model(model_id) + if record is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Embedding model not found") + return _embedding_model_payload(record) + + @app.get( + "/splash_links/embedding-models", + response_model=list[EmbeddingModelPayload], + tags=["embeddings"], + summary="List embedding models", + ) + def list_embedding_models( + request: Request, + name: Optional[str] = Query(None), + limit: int = Query(100, ge=1, le=1000), + offset: int = Query(0, ge=0), + ) -> list[EmbeddingModelPayload]: + records = request.app.state.store.list_embedding_models(name=name, limit=limit, offset=offset) + return [_embedding_model_payload(record) for record in records] + + @app.delete( + "/splash_links/embedding-models/{model_id}", + status_code=status.HTTP_204_NO_CONTENT, + tags=["embeddings"], + summary="Delete one embedding model", + ) + def delete_embedding_model(model_id: str, request: Request) -> Response: + try: + deleted = request.app.state.store.delete_embedding_model(model_id) + except ValueError as exc: + raise _embedding_error(exc) from exc + if not deleted: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Embedding model not found") + return Response(status_code=status.HTTP_204_NO_CONTENT) + + @app.post( + "/splash_links/embeddings", + response_model=EmbeddingPayload, + status_code=status.HTTP_201_CREATED, + tags=["embeddings"], + summary="Create an embedding", + ) + def create_embedding(payload: CreateEmbeddingPayload, request: Request) -> EmbeddingPayload: + try: + record = request.app.state.store.create_embedding( + entity_id=payload.entity_id, + vector=payload.vector, + embedding_model_id=payload.embedding_model_id, + properties=payload.properties, + ) + except ValueError as exc: + raise _embedding_error(exc) from exc + return _embedding_payload(record) + + @app.get( + "/splash_links/embeddings/{embedding_id}", + response_model=EmbeddingPayload, + tags=["embeddings"], + summary="Fetch one embedding", + ) + def get_embedding(embedding_id: str, request: Request) -> EmbeddingPayload: + record = request.app.state.store.get_embedding(embedding_id) + if record is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Embedding not found") + return _embedding_payload(record) + + @app.get( + "/splash_links/embeddings", + response_model=list[EmbeddingPayload], + tags=["embeddings"], + summary="List embeddings", + ) + def list_embeddings( + request: Request, + entity_id: Optional[str] = Query(None, alias="entityId"), + embedding_model_id: Optional[str] = Query(None, alias="embeddingModelId"), + limit: int = Query(100, ge=1, le=1000), + offset: int = Query(0, ge=0), + ) -> list[EmbeddingPayload]: + records = request.app.state.store.list_embeddings( + entity_id=entity_id, + embedding_model_id=embedding_model_id, + limit=limit, + offset=offset, + ) + return [_embedding_payload(record) for record in records] + + @app.delete( + "/splash_links/embeddings/{embedding_id}", + status_code=status.HTTP_204_NO_CONTENT, + tags=["embeddings"], + summary="Delete one embedding", + ) + def delete_embedding(embedding_id: str, request: Request) -> Response: + deleted = request.app.state.store.delete_embedding(embedding_id) + if not deleted: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Embedding not found") + return Response(status_code=status.HTTP_204_NO_CONTENT) + static_dir = os.environ.get("SPLASH_LINKS_STATIC_DIR", "") if static_dir and os.path.isdir(static_dir): app.mount("/splash_links", StaticFiles(directory=static_dir, html=True), name="static") diff --git a/src/splash_links/cli.py b/src/splash_links/cli.py index 013375d..757bfe1 100644 --- a/src/splash_links/cli.py +++ b/src/splash_links/cli.py @@ -4,6 +4,7 @@ Usage: splash-links entities [--type TYPE] [--limit N] splash-links links [--subject ID] [--predicate PRED] [--object ID] [--limit N] + splash-links embeddings [--entity ID] [--model-id ID] [--limit N] splash-links shell # drop into the raw SQLite CLI splash-links client --help @@ -120,6 +121,49 @@ def links( console.print(f"[dim]{len(rows)} row(s)[/dim]") +@app.command() +def embeddings( + entity: Optional[str] = typer.Option(None, "--entity", "-e", help="Filter by entity ID."), + model_id: Optional[str] = typer.Option(None, "--model-id", "-m", help="Filter by embedding model ID."), + limit: int = typer.Option(50, "--limit", "-n", help="Maximum rows to show."), +) -> None: + """List embeddings stored in the database.""" + store = _open_store() + try: + rows = store.list_embeddings(entity_id=entity, embedding_model_id=model_id, limit=limit) + finally: + store.close() + + if not rows: + console.print("[yellow]No embeddings found.[/yellow]") + return + + table = Table(box=box.SIMPLE_HEAVY, show_lines=False) + table.add_column("ID", style="dim", no_wrap=True) + table.add_column("Entity ID", style="dim", no_wrap=True) + table.add_column("Model", style="cyan") + table.add_column("Dims", justify="right") + table.add_column("Vector") + table.add_column("Properties") + table.add_column("Created", style="dim", no_wrap=True) + + for embedding in rows: + props = json.dumps(embedding.properties) if embedding.properties else "" + vector = json.dumps(embedding.vector) + table.add_row( + embedding.id[:8], + embedding.entity_id[:8], + embedding.embedding_model.name, + str(embedding.dimensions), + vector, + props, + embedding.created_at.isoformat(timespec="seconds"), + ) + + console.print(table) + console.print(f"[dim]{len(rows)} row(s)[/dim]") + + @app.command() def shell() -> None: """Open an interactive SQL shell against the database.""" diff --git a/src/splash_links/client/base.py b/src/splash_links/client/base.py index cdc145f..a8732d6 100644 --- a/src/splash_links/client/base.py +++ b/src/splash_links/client/base.py @@ -1,5 +1,5 @@ """ -HTTP client for the splash-links GraphQL service. +HTTP client for the splash-links service. Usage:: @@ -47,12 +47,47 @@ class Link(BaseModel): created_at: str = Field(alias="createdAt") +class EmbeddingModel(BaseModel): + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + id: str + name: str + description: Optional[str] = None + url: Optional[str] = None + version: str + + +class Embedding(BaseModel): + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + id: str + entity_id: str = Field(alias="entityId") + embedding_model_id: str = Field(alias="embeddingModelId") + embedding_model: EmbeddingModel = Field(alias="embeddingModel") + vector: list[float] + dimensions: int + properties: Optional[dict] + created_at: str = Field(alias="createdAt") + + +class EmbeddingMatch(BaseModel): + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + embedding: Embedding + distance: float + + # --------------------------------------------------------------------------- # GraphQL operations # --------------------------------------------------------------------------- _ENTITY_FIELDS = "id entityType name uri properties createdAt" _LINK_FIELDS = "id subjectId predicate objectId properties createdAt" +_EMBEDDING_MODEL_FIELDS = "id name description url version" +_EMBEDDING_FIELDS = ( + "id entityId embeddingModelId vector dimensions properties createdAt " + f"embeddingModel {{ {_EMBEDDING_MODEL_FIELDS} }}" +) _CREATE_ENTITY_MUTATION = f""" mutation CreateEntity($input: CreateEntityInput!) {{ @@ -77,6 +112,21 @@ class Link(BaseModel): }} """ +_NEAREST_EMBEDDINGS_QUERY = f""" +query NearestEmbeddings($vector: [Float!]!, $embeddingModelId: ID, $entityId: ID, $limit: Int, $offset: Int) {{ + nearestEmbeddings( + vector: $vector + embeddingModelId: $embeddingModelId + entityId: $entityId + limit: $limit + offset: $offset + ) {{ + distance + embedding {{ {_EMBEDDING_FIELDS} }} + }} +}} +""" + # --------------------------------------------------------------------------- # Record -> model helpers @@ -96,6 +146,14 @@ def _link_from_dict(d: dict) -> Link: return Link.model_validate(d) +def _embedding_from_dict(d: dict) -> Embedding: + return Embedding.model_validate(d) + + +def _embedding_match_from_dict(d: dict) -> EmbeddingMatch: + return EmbeddingMatch.model_validate(d) + + # --------------------------------------------------------------------------- # Client # --------------------------------------------------------------------------- @@ -103,13 +161,16 @@ def _link_from_dict(d: dict) -> Link: class LinksClient: """ - Synchronous HTTP client for the splash-links GraphQL API. + Synchronous HTTP client for the splash-links API. Instantiate via :func:`from_uri` rather than directly. """ def __init__(self, base_url: str) -> None: - self._gql_url = base_url.rstrip("/") + "/splash_links/graphql" + self._base_url = base_url.rstrip("/") + self._gql_url = self._base_url + "/splash_links/graphql" + self._embeddings_url = self._base_url + "/splash_links/embeddings" + self._embedding_models_url = self._base_url + "/splash_links/embedding-models" # Cache: tiled node URI -> Entity, avoids duplicate entity creation self._tiled_cache: dict[str, Entity] = {} @@ -234,6 +295,67 @@ def find_links( links.append(_link_from_dict(record)) return links + def create_embedding( + self, + entity_id: Any, + vector: list[float], + embedding_model_id: str, + properties: Optional[dict] = None, + ) -> Embedding: + resolved_entity_id = self._resolve(entity_id) + resp = httpx.post( + self._embeddings_url, + json={ + "entityId": resolved_entity_id, + "embeddingModelId": embedding_model_id, + "vector": vector, + "properties": properties, + }, + timeout=30.0, + ) + resp.raise_for_status() + return _embedding_from_dict(resp.json()) + + def create_embedding_model( + self, + name: str, + version: str, + description: Optional[str] = None, + url: Optional[str] = None, + ) -> EmbeddingModel: + resp = httpx.post( + self._embedding_models_url, + json={ + "name": name, + "version": version, + "description": description, + "url": url, + }, + timeout=30.0, + ) + resp.raise_for_status() + return EmbeddingModel.model_validate(resp.json()) + + def find_nearest_embeddings( + self, + vector: list[float], + embedding_model_id: Optional[str] = None, + entity_id: Optional[Any] = None, + limit: int = 10, + offset: int = 0, + ) -> list[EmbeddingMatch]: + data = self._execute( + _NEAREST_EMBEDDINGS_QUERY, + { + "vector": vector, + "embeddingModelId": embedding_model_id, + "entityId": self._resolve(entity_id) if entity_id is not None else None, + "limit": limit, + "offset": offset, + }, + ) + return [_embedding_match_from_dict(record) for record in data["nearestEmbeddings"]] + # --------------------------------------------------------------------------- # Factory diff --git a/src/splash_links/client/cli.py b/src/splash_links/client/cli.py index 194bdbc..d70df6e 100644 --- a/src/splash_links/client/cli.py +++ b/src/splash_links/client/cli.py @@ -7,7 +7,7 @@ import typer -from .base import Entity, Link, from_uri +from .base import Embedding, EmbeddingMatch, EmbeddingModel, Entity, Link, from_uri app = typer.Typer(help="Interact with a splash-links GraphQL service.") @@ -36,6 +36,20 @@ def _link_as_dict(link: Link) -> dict[str, Any]: return link.model_dump() +def _embedding_as_dict(embedding: Embedding) -> dict[str, Any]: + return embedding.model_dump() + + +def _embedding_model_as_dict(model: EmbeddingModel) -> dict[str, Any]: + return model.model_dump() + + +def _embedding_match_as_dict(match: EmbeddingMatch) -> dict[str, Any]: + payload = match.model_dump() + payload["embedding"] = _embedding_as_dict(match.embedding) + return payload + + def _emit_json(payload: Any) -> None: typer.echo(json.dumps(payload, indent=2, sort_keys=True)) @@ -128,5 +142,113 @@ def find_links( _emit_json([_link_as_dict(link) for link in links]) +@app.command("create-embedding") +def create_embedding( + entity_id: str = typer.Argument(..., help="Entity ID that owns the embedding."), + vector: str = typer.Option(..., "--vector", "-v", help="JSON array of numeric embedding values."), + embedding_model_id: str = typer.Option(..., "--model-id", "-m", help="Embedding model ID."), + properties: Optional[str] = typer.Option( + None, + "--properties", + "-p", + help='JSON object of metadata. Example: {"chunk": 3}', + ), + uri: str = typer.Option( + "splash://localhost:8080", + "--uri", + "-u", + envvar="SPLASH_LINKS_URI", + help="Service URI. Supports splash://, http://, or https://.", + ), +) -> None: + """Create an embedding for an entity.""" + props = _parse_json_option("properties", properties) + try: + raw_vector = json.loads(vector) + except json.JSONDecodeError as exc: + typer.echo(f"Invalid JSON passed to --vector: {exc}", err=True) + raise typer.Exit(code=2) from exc + if not isinstance(raw_vector, list): + typer.echo("--vector must decode to a JSON array.", err=True) + raise typer.Exit(code=2) + + client = from_uri(uri) + try: + embedding = client.create_embedding( + entity_id=entity_id, + vector=raw_vector, + embedding_model_id=embedding_model_id, + properties=props, + ) + except Exception as exc: + typer.echo(f"Failed to create embedding: {exc}", err=True) + raise typer.Exit(code=1) from exc + _emit_json(_embedding_as_dict(embedding)) + + +@app.command("create-embedding-model") +def create_embedding_model( + name: str = typer.Option(..., "--name", help="Embedding model name."), + version: str = typer.Option(..., "--version", help="Embedding model version string."), + description: Optional[str] = typer.Option(None, "--description", help="Optional model description."), + url: Optional[str] = typer.Option(None, "--url", help="Optional model URL."), + uri: str = typer.Option( + "splash://localhost:8080", + "--uri", + "-u", + envvar="SPLASH_LINKS_URI", + help="Service URI. Supports splash://, http://, or https://.", + ), +) -> None: + """Create an embedding model record.""" + client = from_uri(uri) + try: + model = client.create_embedding_model(name=name, version=version, description=description, url=url) + except Exception as exc: + typer.echo(f"Failed to create embedding model: {exc}", err=True) + raise typer.Exit(code=1) from exc + _emit_json(_embedding_model_as_dict(model)) + + +@app.command("nearest-embeddings") +def nearest_embeddings( + vector: str = typer.Option(..., "--vector", "-v", help="JSON array of numeric query values."), + embedding_model_id: Optional[str] = typer.Option(None, "--model-id", "-m", help="Optional embedding model filter."), + entity_id: Optional[str] = typer.Option(None, "--entity-id", "-e", help="Optional entity filter."), + limit: int = typer.Option(10, "--limit", "-n", help="Maximum number of matches to fetch."), + offset: int = typer.Option(0, "--offset", "-o", help="Pagination offset."), + uri: str = typer.Option( + "splash://localhost:8080", + "--uri", + "-u", + envvar="SPLASH_LINKS_URI", + help="Service URI. Supports splash://, http://, or https://.", + ), +) -> None: + """Find embeddings nearest to a query vector.""" + try: + raw_vector = json.loads(vector) + except json.JSONDecodeError as exc: + typer.echo(f"Invalid JSON passed to --vector: {exc}", err=True) + raise typer.Exit(code=2) from exc + if not isinstance(raw_vector, list): + typer.echo("--vector must decode to a JSON array.", err=True) + raise typer.Exit(code=2) + + client = from_uri(uri) + try: + matches = client.find_nearest_embeddings( + vector=raw_vector, + embedding_model_id=embedding_model_id, + entity_id=entity_id, + limit=limit, + offset=offset, + ) + except Exception as exc: + typer.echo(f"Failed to search embeddings: {exc}", err=True) + raise typer.Exit(code=1) from exc + _emit_json([_embedding_match_as_dict(match) for match in matches]) + + def main() -> None: app() diff --git a/src/splash_links/schema.py b/src/splash_links/schema.py index 047be4d..c3b744f 100644 --- a/src/splash_links/schema.py +++ b/src/splash_links/schema.py @@ -6,13 +6,14 @@ - Link — a directed, predicate-labeled edge between two entities Query highlights: - - entity / entities — fetch nodes - - link / links — fetch edges, filterable by subject, predicate, object - - Entity.outgoing_links / incoming_links — graph traversal from a node + - entity / entities — fetch nodes + - link / links — fetch edges, filterable by subject, predicate, object + - nearest_embeddings — cosine nearest-neighbor search over stored embeddings + - Entity.outgoing_links / incoming_links — graph traversal from a node Mutations: - - createEntity / createLink - - deleteEntity (cascades to attached links) / deleteLink + - createEntity / createLink + - deleteEntity (cascades to attached links) / deleteLink """ from __future__ import annotations @@ -24,7 +25,7 @@ from strawberry.scalars import JSON as StrawberryJSON from strawberry.types import Info -from .store import EntityRecord, LinkRecord, Store +from .store import EmbeddingMatchRecord, EmbeddingModelRecord, EmbeddingRecord, EntityRecord, LinkRecord, Store logger = logging.getLogger(__name__) @@ -101,6 +102,38 @@ def object(self, info: Info) -> Optional[Entity]: return _entity_from_record(record) if record else None +@strawberry.type +class EmbeddingModel: + id: strawberry.ID + name: str + description: Optional[str] + url: Optional[str] + version: str + + +@strawberry.type +class Embedding: + id: strawberry.ID + entity_id: strawberry.ID + embedding_model_id: strawberry.ID + embedding_model: EmbeddingModel + vector: list[float] + dimensions: int + properties: Optional[JSON] # type: ignore[valid-type] + created_at: str + + @strawberry.field + def entity(self, info: Info) -> Optional[Entity]: + record = _store(info).get_entity(str(self.entity_id)) + return _entity_from_record(record) if record else None + + +@strawberry.type +class EmbeddingMatch: + embedding: Embedding + distance: float + + # --------------------------------------------------------------------------- # Record -> GQL type converters # --------------------------------------------------------------------------- @@ -128,6 +161,32 @@ def _link_from_record(r: LinkRecord) -> Link: ) +def _embedding_from_record(r: EmbeddingRecord) -> Embedding: + return Embedding( + id=strawberry.ID(r.id), + entity_id=strawberry.ID(r.entity_id), + embedding_model_id=strawberry.ID(r.embedding_model_id), + embedding_model=EmbeddingModel( + id=strawberry.ID(r.embedding_model.id), + name=r.embedding_model.name, + description=r.embedding_model.description, + url=r.embedding_model.url, + version=r.embedding_model.version, + ), + vector=r.vector, + dimensions=r.dimensions, + properties=r.properties if r.properties else None, + created_at=r.created_at.isoformat(), + ) + + +def _embedding_match_from_record(r: EmbeddingMatchRecord) -> EmbeddingMatch: + return EmbeddingMatch( + embedding=_embedding_from_record(r.embedding), + distance=r.distance, + ) + + # --------------------------------------------------------------------------- # Input types # --------------------------------------------------------------------------- @@ -208,6 +267,25 @@ def links( ) return [_link_from_record(r) for r in records] + @strawberry.field(description="Find embeddings nearest to a query vector using cosine distance.") + def nearest_embeddings( + self, + info: Info, + vector: list[float], + embedding_model_id: Optional[strawberry.ID] = None, + entity_id: Optional[strawberry.ID] = None, + limit: int = 10, + offset: int = 0, + ) -> list[EmbeddingMatch]: + records = _store(info).find_nearest_embeddings( + query_vector=vector, + embedding_model_id=str(embedding_model_id) if embedding_model_id else None, + entity_id=str(entity_id) if entity_id else None, + limit=limit, + offset=offset, + ) + return [_embedding_match_from_record(r) for r in records] + @strawberry.type class Mutation: diff --git a/src/splash_links/store.py b/src/splash_links/store.py index 8449a49..00135d2 100644 --- a/src/splash_links/store.py +++ b/src/splash_links/store.py @@ -18,6 +18,9 @@ from __future__ import annotations import abc +import json +import math +import struct import uuid from datetime import datetime, timezone from typing import Optional @@ -29,6 +32,8 @@ DateTime, ForeignKey, Index, + Integer, + LargeBinary, MetaData, String, Table, @@ -37,10 +42,18 @@ event, insert, select, + text, update, ) from sqlalchemy.engine import Engine +from sqlalchemy.exc import IntegrityError from sqlalchemy.pool import StaticPool +from sqlalchemy.types import TypeDecorator, UserDefinedType + +try: + from pgvector.psycopg2 import register_vector as register_pgvector_psycopg2 +except ImportError: # pragma: no cover - dependency may not be installed in minimal environments + register_pgvector_psycopg2 = None # --------------------------------------------------------------------------- # Data records @@ -69,6 +82,71 @@ class LinkRecord(BaseModel): created_at: datetime +class EmbeddingModelRecord(BaseModel): + model_config = ConfigDict(extra="forbid") + + id: str + name: str + description: Optional[str] + url: Optional[str] + version: str + + +class EmbeddingRecord(BaseModel): + model_config = ConfigDict(extra="forbid") + + id: str + entity_id: str + embedding_model_id: str + embedding_model: EmbeddingModelRecord + vector: list[float] + dimensions: int + properties: dict + created_at: datetime + + +class EmbeddingMatchRecord(BaseModel): + model_config = ConfigDict(extra="forbid") + + embedding: EmbeddingRecord + distance: float + + +class PostgreSQLVectorType(UserDefinedType): + cache_ok = True + + def get_col_spec(self, **_kw) -> str: + return "vector" + + +class EmbeddingVectorType(TypeDecorator): + cache_ok = True + impl = LargeBinary + + def load_dialect_impl(self, dialect): + if dialect.name == "postgresql": + return dialect.type_descriptor(PostgreSQLVectorType()) + return dialect.type_descriptor(LargeBinary()) + + def process_bind_param(self, value, dialect): + if value is None: + return None + if dialect.name == "postgresql": + return _serialize_vector_postgresql(value) + return _serialize_vector_blob(value) + + def process_result_value(self, value, dialect): + if value is None: + return None + if dialect.name == "postgresql": + if isinstance(value, str): + return _deserialize_vector_postgresql(value) + return [float(item) for item in value] + if isinstance(value, memoryview): + value = value.tobytes() + return _deserialize_vector_blob(value) + + # --------------------------------------------------------------------------- # Abstract interface # --------------------------------------------------------------------------- @@ -137,6 +215,63 @@ def delete_link(self, id: str) -> bool: ... @abc.abstractmethod def update_link(self, id: str, predicate: str) -> Optional[LinkRecord]: ... + @abc.abstractmethod + def create_embedding_model( + self, + name: str, + version: str, + description: Optional[str] = None, + url: Optional[str] = None, + ) -> EmbeddingModelRecord: ... + + @abc.abstractmethod + def get_embedding_model(self, id: str) -> Optional[EmbeddingModelRecord]: ... + + @abc.abstractmethod + def list_embedding_models( + self, + name: Optional[str] = None, + limit: int = 100, + offset: int = 0, + ) -> list[EmbeddingModelRecord]: ... + + @abc.abstractmethod + def delete_embedding_model(self, id: str) -> bool: ... + + @abc.abstractmethod + def create_embedding( + self, + entity_id: str, + vector: list[float], + embedding_model_id: str, + properties: Optional[dict] = None, + ) -> EmbeddingRecord: ... + + @abc.abstractmethod + def get_embedding(self, id: str) -> Optional[EmbeddingRecord]: ... + + @abc.abstractmethod + def list_embeddings( + self, + entity_id: Optional[str] = None, + embedding_model_id: Optional[str] = None, + limit: int = 100, + offset: int = 0, + ) -> list[EmbeddingRecord]: ... + + @abc.abstractmethod + def find_nearest_embeddings( + self, + query_vector: list[float], + embedding_model_id: Optional[str] = None, + entity_id: Optional[str] = None, + limit: int = 10, + offset: int = 0, + ) -> list[EmbeddingMatchRecord]: ... + + @abc.abstractmethod + def delete_embedding(self, id: str) -> bool: ... + @abc.abstractmethod def close(self) -> None: ... @@ -174,6 +309,108 @@ def close(self) -> None: ... Index("links_triple_idx", "subject_id", "predicate", "object_id"), ) +_embedding_models = Table( + "embedding_models", + _metadata, + Column("id", String, primary_key=True), + Column("name", String, nullable=False), + Column("description", String, nullable=True), + Column("url", String, nullable=True), + Column("version", String, nullable=False), + Index("embedding_models_name_version_idx", "name", "version", unique=True), +) + +_embeddings = Table( + "embeddings", + _metadata, + Column("id", String, primary_key=True), + Column("entity_id", String, ForeignKey("entities.id", ondelete="CASCADE"), nullable=False), + Column("embedding_model_id", String, ForeignKey("embedding_models.id", ondelete="RESTRICT"), nullable=False), + Column("vector", EmbeddingVectorType(), nullable=False), + Column("dimensions", Integer, nullable=False), + Column("properties", JSON, nullable=False), + Column("created_at", DateTime(timezone=True), nullable=False), + Index("embeddings_entity_model_created_idx", "entity_id", "embedding_model_id", "created_at"), + Index("embeddings_model_dimensions_idx", "embedding_model_id", "dimensions"), +) + + +def _normalize_vector(vector: list[float]) -> list[float]: + if not vector: + raise ValueError("Embedding vectors must contain at least one value") + + normalized: list[float] = [] + for value in vector: + if isinstance(value, bool): + raise ValueError("Embedding vectors must contain only numeric values") + try: + numeric = float(value) + except (TypeError, ValueError) as exc: + raise ValueError("Embedding vectors must contain only numeric values") from exc + if not math.isfinite(numeric): + raise ValueError("Embedding vectors must contain only finite values") + normalized.append(numeric) + + magnitude = math.sqrt(sum(value * value for value in normalized)) + if magnitude == 0.0: + raise ValueError("Embedding vectors must not be all zeros") + + return normalized + + +def _serialize_vector(vector: list[float]) -> str: + return json.dumps(vector, separators=(",", ":")) + + +def _deserialize_vector(value: str) -> list[float]: + parsed = json.loads(value) + return [float(item) for item in parsed] + + +def _serialize_vector_postgresql(vector: list[float]) -> str: + return _serialize_vector(vector) + + +def _deserialize_vector_postgresql(value: str) -> list[float]: + return _deserialize_vector(value) + + +def _serialize_vector_blob(vector: list[float]) -> bytes: + return struct.pack(f"<{len(vector)}f", *vector) + + +def _deserialize_vector_blob(value: bytes) -> list[float]: + if len(value) % 4 != 0: + raise ValueError("Stored embedding blob has an invalid length") + dimensions = len(value) // 4 + return list(struct.unpack(f"<{dimensions}f", value)) + + +def _cosine_distance(left: list[float], right: list[float]) -> float: + if len(left) != len(right): + raise ValueError("Embedding vectors must have the same number of dimensions") + + dot = sum(lhs * rhs for lhs, rhs in zip(left, right, strict=True)) + left_norm = math.sqrt(sum(value * value for value in left)) + right_norm = math.sqrt(sum(value * value for value in right)) + return 1.0 - (dot / (left_norm * right_norm)) + + +def _embedding_select(): + return select( + _embeddings.c.id, + _embeddings.c.entity_id, + _embeddings.c.embedding_model_id, + _embeddings.c.vector, + _embeddings.c.dimensions, + _embeddings.c.properties, + _embeddings.c.created_at, + _embedding_models.c.name.label("embedding_model_name"), + _embedding_models.c.description.label("embedding_model_description"), + _embedding_models.c.url.label("embedding_model_url"), + _embedding_models.c.version.label("embedding_model_version"), + ).select_from(_embeddings.join(_embedding_models, _embeddings.c.embedding_model_id == _embedding_models.c.id)) + def _make_engine(db_url: str) -> Engine: """Create a SQLAlchemy engine from a URL, applying dialect-specific tuning.""" @@ -194,6 +431,11 @@ def _make_engine(db_url: str) -> Engine: def _set_sqlite_pragma(conn, _record): conn.execute("PRAGMA foreign_keys=ON") + if engine.dialect.name == "postgresql" and register_pgvector_psycopg2 is not None: + @event.listens_for(engine, "connect") + def _register_pgvector(conn, _record): + register_pgvector_psycopg2(conn) + return engine @@ -250,6 +492,45 @@ def _to_link(row) -> LinkRecord: created_at=row.created_at, ) + @staticmethod + def _to_embedding_model(row) -> EmbeddingModelRecord: + mapping = row._mapping if hasattr(row, "_mapping") else row + return EmbeddingModelRecord( + id=mapping["id"], + name=mapping["name"], + description=mapping["description"], + url=mapping["url"], + version=mapping["version"], + ) + + @staticmethod + def _to_embedding(row) -> EmbeddingRecord: + mapping = row._mapping if hasattr(row, "_mapping") else row + vector = mapping["vector"] + if isinstance(vector, str): + vector = _deserialize_vector_postgresql(vector) + elif isinstance(vector, (bytes, memoryview)): + blob = vector.tobytes() if isinstance(vector, memoryview) else vector + vector = _deserialize_vector_blob(blob) + else: + vector = [float(item) for item in vector] + return EmbeddingRecord( + id=mapping["id"], + entity_id=mapping["entity_id"], + embedding_model_id=mapping["embedding_model_id"], + embedding_model=EmbeddingModelRecord( + id=mapping["embedding_model_id"], + name=mapping["embedding_model_name"], + description=mapping["embedding_model_description"], + url=mapping["embedding_model_url"], + version=mapping["embedding_model_version"], + ), + vector=vector, + dimensions=mapping["dimensions"], + properties=mapping["properties"] or {}, + created_at=mapping["created_at"], + ) + # ------------------------------------------------------------------ # Entity operations # ------------------------------------------------------------------ @@ -387,6 +668,235 @@ def update_link(self, id: str, predicate: str) -> Optional[LinkRecord]: row = conn.execute(select(_links).where(_links.c.id == id)).one_or_none() return self._to_link(row) if row else None + # ------------------------------------------------------------------ + # Embedding model operations + # ------------------------------------------------------------------ + + def create_embedding_model( + self, + name: str, + version: str, + description: Optional[str] = None, + url: Optional[str] = None, + ) -> EmbeddingModelRecord: + stmt = select(_embedding_models).where( + _embedding_models.c.name == name, + _embedding_models.c.version == version, + ) + with self._engine.begin() as conn: + existing = conn.execute(stmt).one_or_none() + if existing is not None: + raise ValueError(f"Embedding model '{name}' version '{version}' already exists") + + id_ = str(uuid.uuid4()) + try: + conn.execute( + insert(_embedding_models).values( + id=id_, + name=name, + description=description, + url=url, + version=version, + ) + ) + except IntegrityError as exc: + raise ValueError(f"Embedding model '{name}' version '{version}' already exists") from exc + row = conn.execute(select(_embedding_models).where(_embedding_models.c.id == id_)).one() + return self._to_embedding_model(row) + + def get_embedding_model(self, id: str) -> Optional[EmbeddingModelRecord]: + with self._engine.connect() as conn: + row = conn.execute(select(_embedding_models).where(_embedding_models.c.id == id)).one_or_none() + return self._to_embedding_model(row) if row else None + + def list_embedding_models( + self, + name: Optional[str] = None, + limit: int = 100, + offset: int = 0, + ) -> list[EmbeddingModelRecord]: + stmt = select(_embedding_models).order_by(_embedding_models.c.name, _embedding_models.c.version).limit(limit).offset(offset) + if name is not None: + stmt = stmt.where(_embedding_models.c.name == name) + with self._engine.connect() as conn: + rows = conn.execute(stmt).all() + return [self._to_embedding_model(row) for row in rows] + + def delete_embedding_model(self, id: str) -> bool: + try: + with self._engine.begin() as conn: + result = conn.execute(delete(_embedding_models).where(_embedding_models.c.id == id)) + except IntegrityError as exc: + raise ValueError(f"Embedding model '{id}' is still referenced by embeddings") from exc + return result.rowcount > 0 + + # ------------------------------------------------------------------ + # Embedding operations + # ------------------------------------------------------------------ + + def create_embedding( + self, + entity_id: str, + vector: list[float], + embedding_model_id: str, + properties: Optional[dict] = None, + ) -> EmbeddingRecord: + if not self.get_entity(entity_id): + raise ValueError(f"Entity '{entity_id}' not found") + if not self.get_embedding_model(embedding_model_id): + raise ValueError(f"Embedding model '{embedding_model_id}' not found") + + normalized_vector = _normalize_vector(vector) + id_ = str(uuid.uuid4()) + now = datetime.now(timezone.utc) + with self._engine.begin() as conn: + conn.execute( + insert(_embeddings).values( + id=id_, + entity_id=entity_id, + embedding_model_id=embedding_model_id, + vector=normalized_vector, + dimensions=len(normalized_vector), + properties=properties or {}, + created_at=now, + ) + ) + row = conn.execute(_embedding_select().where(_embeddings.c.id == id_)).one() + return self._to_embedding(row) + + def get_embedding(self, id: str) -> Optional[EmbeddingRecord]: + with self._engine.connect() as conn: + row = conn.execute(_embedding_select().where(_embeddings.c.id == id)).one_or_none() + return self._to_embedding(row) if row else None + + def list_embeddings( + self, + entity_id: Optional[str] = None, + embedding_model_id: Optional[str] = None, + limit: int = 100, + offset: int = 0, + ) -> list[EmbeddingRecord]: + stmt = _embedding_select().order_by(_embeddings.c.created_at).limit(limit).offset(offset) + if entity_id is not None: + stmt = stmt.where(_embeddings.c.entity_id == entity_id) + if embedding_model_id is not None: + stmt = stmt.where(_embeddings.c.embedding_model_id == embedding_model_id) + with self._engine.connect() as conn: + rows = conn.execute(stmt).all() + return [self._to_embedding(row) for row in rows] + + def find_nearest_embeddings( + self, + query_vector: list[float], + embedding_model_id: Optional[str] = None, + entity_id: Optional[str] = None, + limit: int = 10, + offset: int = 0, + ) -> list[EmbeddingMatchRecord]: + normalized_query = _normalize_vector(query_vector) + if self._engine.dialect.name == "postgresql": + return self._find_nearest_embeddings_postgresql( + query_vector=normalized_query, + embedding_model_id=embedding_model_id, + entity_id=entity_id, + limit=limit, + offset=offset, + ) + return self._find_nearest_embeddings_python( + query_vector=normalized_query, + embedding_model_id=embedding_model_id, + entity_id=entity_id, + limit=limit, + offset=offset, + ) + + def _find_nearest_embeddings_python( + self, + query_vector: list[float], + embedding_model_id: Optional[str], + entity_id: Optional[str], + limit: int, + offset: int, + ) -> list[EmbeddingMatchRecord]: + stmt = _embedding_select().where(_embeddings.c.dimensions == len(query_vector)) + if entity_id is not None: + stmt = stmt.where(_embeddings.c.entity_id == entity_id) + if embedding_model_id is not None: + stmt = stmt.where(_embeddings.c.embedding_model_id == embedding_model_id) + + with self._engine.connect() as conn: + rows = conn.execute(stmt).all() + + matches = [ + EmbeddingMatchRecord( + embedding=record, + distance=_cosine_distance(record.vector, query_vector), + ) + for record in (self._to_embedding(row) for row in rows) + ] + matches.sort(key=lambda match: (match.distance, match.embedding.created_at, match.embedding.id)) + return matches[offset : offset + limit] + + def _find_nearest_embeddings_postgresql( + self, + query_vector: list[float], + embedding_model_id: Optional[str], + entity_id: Optional[str], + limit: int, + offset: int, + ) -> list[EmbeddingMatchRecord]: + conditions = ["e.dimensions = :dimensions"] + params: dict[str, object] = { + "dimensions": len(query_vector), + "query_vector": _serialize_vector(query_vector), + "limit": limit, + "offset": offset, + } + if entity_id is not None: + conditions.append("e.entity_id = :entity_id") + params["entity_id"] = entity_id + if embedding_model_id is not None: + conditions.append("e.embedding_model_id = :embedding_model_id") + params["embedding_model_id"] = embedding_model_id + + sql = f""" + SELECT + e.id, + e.entity_id, + e.embedding_model_id, + e.vector, + e.dimensions, + e.properties, + e.created_at, + m.name AS embedding_model_name, + m.description AS embedding_model_description, + m.url AS embedding_model_url, + m.version AS embedding_model_version, + e.vector <=> CAST(:query_vector AS vector) AS distance + FROM embeddings e + JOIN embedding_models m ON m.id = e.embedding_model_id + WHERE {' AND '.join(conditions)} + ORDER BY distance, created_at, id + LIMIT :limit + OFFSET :offset + """ + + with self._engine.connect() as conn: + rows = conn.execute(text(sql), params).mappings().all() + + return [ + EmbeddingMatchRecord( + embedding=self._to_embedding(row), + distance=float(row["distance"]), + ) + for row in rows + ] + + def delete_embedding(self, id: str) -> bool: + with self._engine.begin() as conn: + result = conn.execute(delete(_embeddings).where(_embeddings.c.id == id)) + return result.rowcount > 0 + def close(self) -> None: self._engine.dispose()