From aeaa0e3360ad731c5c97086297ea4dd128bd1a4b Mon Sep 17 00:00:00 2001 From: Dylan McReynolds Date: Sun, 5 Apr 2026 18:18:40 -0700 Subject: [PATCH 1/2] Add vector storage and search --- README.md | 61 +++ _tests/test_base_client.py | 102 +++- _tests/test_cli.py | 47 +- _tests/test_client_cli.py | 114 ++++- _tests/test_service.py | 175 ++++++- .../6a86ce49c068_add_embeddings_table.py | 67 +++ pixi.toml | 1 + pyproject.toml | 1 + src/splash_links/app.py | 108 ++++- src/splash_links/cli.py | 44 ++ src/splash_links/client/base.py | 92 +++- src/splash_links/client/cli.py | 96 +++- src/splash_links/schema.py | 73 ++- src/splash_links/store.py | 440 ++++++++++++++++++ 14 files changed, 1404 insertions(+), 17 deletions(-) create mode 100644 alembic/versions/6a86ce49c068_add_embeddings_table.py diff --git a/README.md b/README.md index 84e677a..aac1751 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,48 @@ mutation { } ``` +### Create an embedding + +```bash +curl -X POST http://localhost:8080/splash_links/embeddings \ + -H 'Content-Type: application/json' \ + -d '{ + "entityId": "", + "embeddingModel": "text-embedding-3-small", + "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] + embeddingModel: "text-embedding-3-small" + limit: 5 + ) { + distance + embedding { + id + entityId + entity { name } + } + } +} +``` + +### Fetch or delete embeddings + +```bash +curl http://localhost:8080/splash_links/embeddings/ +curl 'http://localhost:8080/splash_links/embeddings?entityId=&embeddingModel=text-embedding-3-small' +curl -X DELETE http://localhost:8080/splash_links/embeddings/ +``` + ### Health check ``` @@ -132,6 +175,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 text-embedding-3-small --limit 10 +``` + ### Raw SQLite shell ```bash @@ -221,6 +271,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 +309,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..45f71e0 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,106 @@ 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", + "embeddingModel": "model-a", + "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="model-a", + properties={"chunk": 1}, + ) + + assert embedding.id == "emb-1" + assert seen["url"] == "http://api:8080/splash_links/embeddings" + assert seen["timeout"] == 30.0 + assert seen["json"] == { + "entityId": "ent-1", + "vector": [0.1, 0.2, 0.3], + "embeddingModel": "model-a", + "properties": {"chunk": 1}, + } + + +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", + "embeddingModel": "model-a", + "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="model-a", + entity_id="ent-1", + limit=5, + offset=1, + ) + + assert matches == [ + EmbeddingMatch.model_validate( + { + "distance": 0.01, + "embedding": { + "id": "emb-1", + "entityId": "ent-1", + "embeddingModel": "model-a", + "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], + "embeddingModel": "model-a", + "entityId": "ent-1", + "limit": 5, + "offset": 1, + } + + # --------------------------------------------------------------------------- # Tiled integration helpers # --------------------------------------------------------------------------- diff --git a/_tests/test_cli.py b/_tests/test_cli.py index 55b8046..7d51e02 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 EmbeddingRecord, EntityRecord, LinkRecord runner = CliRunner() @@ -44,10 +44,25 @@ def _make_link(**kw) -> LinkRecord: return LinkRecord(**defaults) +def _make_embedding(**kw) -> EmbeddingRecord: + defaults = dict( + id="emb-1", + entity_id="ent-1", + embedding_model="model-a", + 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 +72,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=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: + rows = [embedding for embedding in rows if embedding.embedding_model == embedding_model] + return rows + def close(self): pass @@ -124,6 +147,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..f7f1336 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, Entity, Link runner = CliRunner() @@ -241,6 +241,118 @@ 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="default", properties=None): + seen["entity_id"] = entity_id + seen["vector"] = vector + seen["embedding_model"] = embedding_model + seen["properties"] = properties + return Embedding( + id="emb-1", + entity_id=entity_id, + embedding_model=embedding_model, + 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", + "model-a", + "--properties", + '{"chunk": 1}', + ], + ) + + assert result.exit_code == 0, result.output + payload = json.loads(result.stdout) + assert payload["id"] == "emb-1" + assert payload["embedding_model"] == "model-a" + assert seen == { + "entity_id": "ent-1", + "vector": [0.1, 0.2, 0.3], + "embedding_model": "model-a", + "properties": {"chunk": 1}, + } + + +def test_create_embedding_invalid_vector_exits_2(): + result = runner.invoke( + client_cli.app, + ["create-embedding", "ent-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=None, entity_id=None, limit=10, offset=0): + seen["vector"] = vector + seen["embedding_model"] = embedding_model + 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="model-a", + 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", + "model-a", + "--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": "model-a", + "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..d29ccc2 100644 --- a/_tests/test_service.py +++ b/_tests/test_service.py @@ -8,6 +8,7 @@ from __future__ import annotations import pytest +from sqlalchemy import text from fastapi.testclient import TestClient from splash_links.app import create_app @@ -49,6 +50,27 @@ def gql(client: TestClient, query: str, variables: dict | None = None) -> dict: return body["data"] +def create_embedding( + client: TestClient, + *, + entity_id: str, + vector: list[float], + embedding_model: str = "default", + properties: dict | None = None, +) -> dict: + resp = client.post( + "/splash_links/embeddings", + json={ + "entityId": entity_id, + "vector": vector, + "embeddingModel": embedding_model, + "properties": properties, + }, + ) + assert resp.status_code == 201, resp.text + return resp.json() + + CREATE_ENTITY = """ mutation CreateEntity($input: CreateEntityInput!) { createEntity(input: $input) { @@ -76,7 +98,6 @@ def gql(client: TestClient, query: str, variables: dict | None = None) -> dict: } """ - # --------------------------------------------------------------------------- # Health check # --------------------------------------------------------------------------- @@ -221,6 +242,76 @@ 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") + embedding = store.create_embedding( + entity.id, + [0.1, 0.2, 0.3], + embedding_model="text-embedding-3-small", + properties={"chunk": 1}, + ) + + assert embedding.entity_id == entity.id + assert embedding.embedding_model == "text-embedding-3-small" + assert embedding.vector == [0.1, 0.2, 0.3] + assert embedding.dimensions == 3 + assert embedding.properties == {"chunk": 1} + + 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") + embedding = store.create_embedding(entity.id, [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): + with pytest.raises(ValueError, match="Entity"): + store.create_embedding("missing", [0.1, 0.2, 0.3]) + + def test_create_embedding_rejects_zero_vector(self, store): + entity = store.create_entity("Dataset", "run-1") + with pytest.raises(ValueError, match="all zeros"): + store.create_embedding(entity.id, [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") + store.create_embedding(entity.id, [0.1, 0.2], embedding_model="model-a") + store.create_embedding(entity.id, [0.2, 0.3], embedding_model="model-b") + store.create_embedding(other.id, [0.3, 0.4], embedding_model="model-a") + + rows = store.list_embeddings(entity_id=entity.id, embedding_model="model-a") + assert len(rows) == 1 + assert rows[0].embedding_model == "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") + near = store.create_embedding(entity.id, [1.0, 0.0], embedding_model="model-a") + far = store.create_embedding(entity.id, [0.0, 1.0], embedding_model="model-a") + store.create_embedding(entity.id, [1.0, 1.0, 0.0], embedding_model="model-a") + + matches = store.find_nearest_embeddings([0.9, 0.1], embedding_model="model-a") + + 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") + embedding = store.create_embedding(entity.id, [0.1, 0.2, 0.3]) + + assert store.delete_entity(entity.id) is True + assert store.get_embedding(embedding.id) is None + # --------------------------------------------------------------------------- # GraphQL integration tests (via HTTP) @@ -387,3 +478,85 @@ def test_update_link_not_found_returns_null(self, client): ' 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"] + create_embedding(client, entity_id=entity["id"], embedding_model="model-a", vector=[1.0, 0.0]) + create_embedding(client, entity_id=entity["id"], embedding_model="model-a", vector=[0.0, 1.0]) + create_embedding(client, entity_id=entity["id"], embedding_model="model-a", vector=[1.0, 1.0, 0.0]) + + data = gql( + client, + """ + query Search($vector: [Float!]!, $model: String) { + nearestEmbeddings(vector: $vector, embeddingModel: $model) { + distance + embedding { + entityId + embeddingModel + vector + } + } + } + + + class TestEmbeddingRestAPI: + def test_create_embedding(self, client): + entity = gql(client, CREATE_ENTITY, {"input": {"entityType": "Dataset", "name": "run-1"}})["createEntity"] + + payload = create_embedding( + client, + entity_id=entity["id"], + embedding_model="text-embedding-3-small", + vector=[0.1, 0.2, 0.3], + properties={"chunk": 1}, + ) + + assert payload["entityId"] == entity["id"] + assert payload["embeddingModel"] == "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"] + created = create_embedding(client, entity_id=entity["id"], embedding_model="model-a", 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"], "embeddingModel": "model-a"}, + ) + 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"] + created = create_embedding(client, entity_id=entity["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"] + created = create_embedding(client, entity_id=entity["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 + )["createEmbedding"] + + deleted = gql(client, "mutation D($id: ID!) { deleteEntity(id: $id) }", {"id": entity["id"]}) + assert deleted["deleteEntity"] is True + + gone = gql(client, "query Q($id: ID!) { embedding(id: $id) { id } }", {"id": embedding["id"]}) + assert gone["embedding"] is None diff --git a/alembic/versions/6a86ce49c068_add_embeddings_table.py b/alembic/versions/6a86ce49c068_add_embeddings_table.py new file mode 100644 index 0000000..b84dccb --- /dev/null +++ b/alembic/versions/6a86ce49c068_add_embeddings_table.py @@ -0,0 +1,67 @@ +"""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( + "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", sa.String(), 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", "created_at"], + ) + op.create_index( + "embeddings_model_dimensions_idx", + "embeddings", + ["embedding_model", "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") \ 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..90f8700 100644 --- a/src/splash_links/app.py +++ b/src/splash_links/app.py @@ -18,17 +18,57 @@ 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 EmbeddingRecord, Store, _make_engine, _url_from_path logger = logging.getLogger(__name__) +class EmbeddingPayload(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + id: str + entity_id: str = Field(alias="entityId") + embedding_model: str = Field(alias="embeddingModel") + vector: list[float] + dimensions: int + properties: dict + created_at: str = Field(alias="createdAt") + + +class CreateEmbeddingPayload(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + entity_id: str = Field(alias="entityId") + vector: list[float] + embedding_model: str = Field(default="default", alias="embeddingModel") + properties: Optional[dict] = None + + +def _embedding_payload(record: EmbeddingRecord) -> EmbeddingPayload: + return EmbeddingPayload( + id=record.id, + entity_id=record.entity_id, + embedding_model=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) + status_code = status.HTTP_404_NOT_FOUND if "not found" in message.lower() else 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 +162,70 @@ async def get_context(request: Request) -> dict: def health() -> dict: return {"status": "ok"} + @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=payload.embedding_model, + 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: Optional[str] = Query(None, alias="embeddingModel"), + 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=embedding_model, + 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..d89e373 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 NAME] [--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: Optional[str] = typer.Option(None, "--model", "-m", help="Filter by embedding model."), + 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=model, 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, + 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..38a06aa 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,32 @@ class Link(BaseModel): created_at: str = Field(alias="createdAt") +class Embedding(BaseModel): + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + id: str + entity_id: str = Field(alias="entityId") + embedding_model: str = 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_FIELDS = "id entityId embeddingModel vector dimensions properties createdAt" _CREATE_ENTITY_MUTATION = f""" mutation CreateEntity($input: CreateEntityInput!) {{ @@ -77,6 +97,21 @@ class Link(BaseModel): }} """ +_NEAREST_EMBEDDINGS_QUERY = f""" +query NearestEmbeddings($vector: [Float!]!, $embeddingModel: String, $entityId: ID, $limit: Int, $offset: Int) {{ + nearestEmbeddings( + vector: $vector + embeddingModel: $embeddingModel + entityId: $entityId + limit: $limit + offset: $offset + ) {{ + distance + embedding {{ {_EMBEDDING_FIELDS} }} + }} +}} +""" + # --------------------------------------------------------------------------- # Record -> model helpers @@ -96,6 +131,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 +146,15 @@ 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" # Cache: tiled node URI -> Entity, avoids duplicate entity creation self._tiled_cache: dict[str, Entity] = {} @@ -234,6 +279,47 @@ def find_links( links.append(_link_from_dict(record)) return links + def create_embedding( + self, + entity_id: Any, + vector: list[float], + embedding_model: str = "default", + properties: Optional[dict] = None, + ) -> Embedding: + resolved_entity_id = self._resolve(entity_id) + resp = httpx.post( + self._embeddings_url, + json={ + "entityId": resolved_entity_id, + "vector": vector, + "embeddingModel": embedding_model, + "properties": properties, + }, + timeout=30.0, + ) + resp.raise_for_status() + return _embedding_from_dict(resp.json()) + + def find_nearest_embeddings( + self, + vector: list[float], + embedding_model: Optional[str] = None, + entity_id: Optional[Any] = None, + limit: int = 10, + offset: int = 0, + ) -> list[EmbeddingMatch]: + data = self._execute( + _NEAREST_EMBEDDINGS_QUERY, + { + "vector": vector, + "embeddingModel": embedding_model, + "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..4f3d2cd 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, Entity, Link, from_uri app = typer.Typer(help="Interact with a splash-links GraphQL service.") @@ -36,6 +36,16 @@ 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_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 +138,89 @@ 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: str = typer.Option("default", "--model", "-m", help="Embedding model label."), + 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=embedding_model, + 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("nearest-embeddings") +def nearest_embeddings( + vector: str = typer.Option(..., "--vector", "-v", help="JSON array of numeric query values."), + embedding_model: Optional[str] = typer.Option(None, "--model", "-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=embedding_model, + 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..5d6906c 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, EmbeddingRecord, EntityRecord, LinkRecord, Store logger = logging.getLogger(__name__) @@ -101,6 +102,28 @@ def object(self, info: Info) -> Optional[Entity]: return _entity_from_record(record) if record else None +@strawberry.type +class Embedding: + id: strawberry.ID + entity_id: strawberry.ID + embedding_model: str + 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 +151,25 @@ 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=r.embedding_model, + 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 +250,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: Optional[str] = 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=embedding_model, + 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..a328f8a 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,17 @@ event, insert, select, + text, update, ) from sqlalchemy.engine import Engine 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 +81,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 +214,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 +308,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 +430,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 +491,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 +667,166 @@ 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 operations + # ------------------------------------------------------------------ + + def create_embedding( + self, + entity_id: str, + vector: list[float], + embedding_model: str = "default", + properties: Optional[dict] = None, + ) -> EmbeddingRecord: + if not self.get_entity(entity_id): + raise ValueError(f"Entity '{entity_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=embedding_model, + vector=normalized_vector, + dimensions=len(normalized_vector), + properties=properties or {}, + created_at=now, + ) + ) + row = conn.execute(select(_embeddings).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(select(_embeddings).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: Optional[str] = None, + limit: int = 100, + offset: int = 0, + ) -> list[EmbeddingRecord]: + stmt = select(_embeddings).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 is not None: + stmt = stmt.where(_embeddings.c.embedding_model == embedding_model) + 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: 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=embedding_model, + entity_id=entity_id, + limit=limit, + offset=offset, + ) + return self._find_nearest_embeddings_python( + query_vector=normalized_query, + embedding_model=embedding_model, + entity_id=entity_id, + limit=limit, + offset=offset, + ) + + def _find_nearest_embeddings_python( + self, + query_vector: list[float], + embedding_model: Optional[str], + entity_id: Optional[str], + limit: int, + offset: int, + ) -> list[EmbeddingMatchRecord]: + stmt = select(_embeddings).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 is not None: + stmt = stmt.where(_embeddings.c.embedding_model == embedding_model) + + 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: Optional[str], + entity_id: Optional[str], + limit: int, + offset: int, + ) -> list[EmbeddingMatchRecord]: + conditions = ["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("entity_id = :entity_id") + params["entity_id"] = entity_id + if embedding_model is not None: + conditions.append("embedding_model = :embedding_model") + params["embedding_model"] = embedding_model + + sql = f""" + SELECT + id, + entity_id, + embedding_model, + vector, + dimensions, + properties, + created_at, + CAST(vector AS vector) <=> CAST(:query_vector AS vector) AS distance + FROM embeddings + 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() From e6029b79846730e730c5b028dc65f5c207c3bb0b Mon Sep 17 00:00:00 2001 From: Dylan McReynolds Date: Sun, 5 Apr 2026 18:30:35 -0700 Subject: [PATCH 2/2] add embedding_model table --- README.md | 18 +- _tests/test_base_client.py | 73 ++++- _tests/test_cli.py | 17 +- _tests/test_client_cli.py | 89 +++++- _tests/test_service.py | 263 ++++++++++-------- .../6a86ce49c068_add_embeddings_table.py | 26 +- src/splash_links/app.py | 112 +++++++- src/splash_links/cli.py | 8 +- src/splash_links/client/base.py | 52 +++- src/splash_links/client/cli.py | 38 ++- src/splash_links/schema.py | 27 +- src/splash_links/store.py | 130 +++++++-- 12 files changed, 644 insertions(+), 209 deletions(-) diff --git a/README.md b/README.md index aac1751..226d375 100644 --- a/README.md +++ b/README.md @@ -107,11 +107,19 @@ 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": "", - "embeddingModel": "text-embedding-3-small", + "embeddingModelId": "'"$MODEL_ID"'", "vector": [0.12, -0.03, 0.88], "properties": {"chunk": 1} }' @@ -125,13 +133,14 @@ Embedding CRUD uses REST. Nearest-neighbor search stays in GraphQL and uses cosi query { nearestEmbeddings( vector: [0.11, -0.02, 0.90] - embeddingModel: "text-embedding-3-small" + embeddingModelId: "" limit: 5 ) { distance embedding { id entityId + embeddingModel { name version } entity { name } } } @@ -142,7 +151,8 @@ query { ```bash curl http://localhost:8080/splash_links/embeddings/ -curl 'http://localhost:8080/splash_links/embeddings?entityId=&embeddingModel=text-embedding-3-small' +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/ ``` @@ -179,7 +189,7 @@ pixi run links -- --object # incoming to a node ```bash pixi run embeddings -- --entity -splash-links embeddings --model text-embedding-3-small --limit 10 +splash-links embeddings --model-id --limit 10 ``` ### Raw SQLite shell diff --git a/_tests/test_base_client.py b/_tests/test_base_client.py index 45f71e0..5ffb43f 100644 --- a/_tests/test_base_client.py +++ b/_tests/test_base_client.py @@ -206,7 +206,14 @@ def fake_post(url: str, json: dict, timeout: float): { "id": "emb-1", "entityId": "ent-1", - "embeddingModel": "model-a", + "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}, @@ -220,21 +227,59 @@ def fake_post(url: str, json: dict, timeout: float): embedding = client.create_embedding( entity_id="ent-1", vector=[0.1, 0.2, 0.3], - embedding_model="model-a", + 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], - "embeddingModel": "model-a", "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] = {} @@ -248,7 +293,14 @@ def fake_execute(query: str, variables: dict | None = None) -> dict: "embedding": { "id": "emb-1", "entityId": "ent-1", - "embeddingModel": "model-a", + "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, @@ -263,7 +315,7 @@ def fake_execute(query: str, variables: dict | None = None) -> dict: matches = client.find_nearest_embeddings( vector=[0.1, 0.2], - embedding_model="model-a", + embedding_model_id="model-1", entity_id="ent-1", limit=5, offset=1, @@ -276,7 +328,14 @@ def fake_execute(query: str, variables: dict | None = None) -> dict: "embedding": { "id": "emb-1", "entityId": "ent-1", - "embeddingModel": "model-a", + "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, @@ -288,7 +347,7 @@ def fake_execute(query: str, variables: dict | None = None) -> dict: assert seen["query"] == base_module._NEAREST_EMBEDDINGS_QUERY assert seen["variables"] == { "vector": [0.1, 0.2], - "embeddingModel": "model-a", + "embeddingModelId": "model-1", "entityId": "ent-1", "limit": 5, "offset": 1, diff --git a/_tests/test_cli.py b/_tests/test_cli.py index 7d51e02..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 EmbeddingRecord, EntityRecord, LinkRecord +from splash_links.store import EmbeddingModelRecord, EmbeddingRecord, EntityRecord, LinkRecord runner = CliRunner() @@ -48,7 +48,14 @@ def _make_embedding(**kw) -> EmbeddingRecord: defaults = dict( id="emb-1", entity_id="ent-1", - embedding_model="model-a", + 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={}, @@ -72,12 +79,12 @@ 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=None, limit=50, offset=0): + 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: - rows = [embedding for embedding in rows if embedding.embedding_model == embedding_model] + if embedding_model_id: + rows = [embedding for embedding in rows if embedding.embedding_model_id == embedding_model_id] return rows def close(self): diff --git a/_tests/test_client_cli.py b/_tests/test_client_cli.py index f7f1336..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 Embedding, EmbeddingMatch, Entity, Link +from splash_links.client.base import Embedding, EmbeddingMatch, EmbeddingModel, Entity, Link runner = CliRunner() @@ -245,15 +245,22 @@ def test_create_embedding_command_outputs_json(monkeypatch): seen: dict[str, object] = {} class FakeClient: - def create_embedding(self, entity_id, vector, embedding_model="default", properties=None): + def create_embedding(self, entity_id, vector, embedding_model_id, properties=None): seen["entity_id"] = entity_id seen["vector"] = vector - seen["embedding_model"] = embedding_model + seen["embedding_model_id"] = embedding_model_id seen["properties"] = properties return Embedding( id="emb-1", entity_id=entity_id, - embedding_model=embedding_model, + 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, @@ -269,8 +276,8 @@ def create_embedding(self, entity_id, vector, embedding_model="default", propert "ent-1", "--vector", "[0.1, 0.2, 0.3]", - "--model", - "model-a", + "--model-id", + "model-1", "--properties", '{"chunk": 1}', ], @@ -279,19 +286,64 @@ def create_embedding(self, entity_id, vector, embedding_model="default", propert assert result.exit_code == 0, result.output payload = json.loads(result.stdout) assert payload["id"] == "emb-1" - assert payload["embedding_model"] == "model-a" + assert payload["embedding_model"]["name"] == "model-a" assert seen == { "entity_id": "ent-1", "vector": [0.1, 0.2, 0.3], - "embedding_model": "model-a", + "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", "--vector", "not-json"], + ["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 @@ -301,9 +353,9 @@ def test_nearest_embeddings_command_outputs_list(monkeypatch): seen: dict[str, object] = {} class FakeClient: - def find_nearest_embeddings(self, vector, embedding_model=None, entity_id=None, limit=10, offset=0): + def find_nearest_embeddings(self, vector, embedding_model_id=None, entity_id=None, limit=10, offset=0): seen["vector"] = vector - seen["embedding_model"] = embedding_model + seen["embedding_model_id"] = embedding_model_id seen["entity_id"] = entity_id seen["limit"] = limit seen["offset"] = offset @@ -313,7 +365,14 @@ def find_nearest_embeddings(self, vector, embedding_model=None, entity_id=None, embedding=Embedding( id="emb-1", entity_id="ent-1", - embedding_model="model-a", + 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, @@ -330,8 +389,8 @@ def find_nearest_embeddings(self, vector, embedding_model=None, entity_id=None, "nearest-embeddings", "--vector", "[0.1, 0.2]", - "--model", - "model-a", + "--model-id", + "model-1", "--entity-id", "ent-1", "--limit", @@ -346,7 +405,7 @@ def find_nearest_embeddings(self, vector, embedding_model=None, entity_id=None, assert payload[0]["embedding"]["id"] == "emb-1" assert seen == { "vector": [0.1, 0.2], - "embedding_model": "model-a", + "embedding_model_id": "model-1", "entity_id": "ent-1", "limit": 5, "offset": 1, diff --git a/_tests/test_service.py b/_tests/test_service.py index d29ccc2..5668479 100644 --- a/_tests/test_service.py +++ b/_tests/test_service.py @@ -8,16 +8,12 @@ from __future__ import annotations import pytest -from sqlalchemy import text 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(): @@ -35,10 +31,6 @@ def client(): yield c -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - _GQL = "/splash_links/graphql" @@ -50,20 +42,41 @@ 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], - embedding_model: str = "default", properties: dict | None = None, ) -> dict: resp = client.post( "/splash_links/embeddings", json={ "entityId": entity_id, + "embeddingModelId": embedding_model_id, "vector": vector, - "embeddingModel": embedding_model, "properties": properties, }, ) @@ -98,10 +111,6 @@ def create_embedding( } """ -# --------------------------------------------------------------------------- -# Health check -# --------------------------------------------------------------------------- - def test_health(client): resp = client.get("/splash_links/health") @@ -109,11 +118,6 @@ def test_health(client): 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"}) @@ -149,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 @@ -201,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") @@ -244,15 +248,17 @@ def test_update_link_not_found(self, store): 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, - [0.1, 0.2, 0.3], - embedding_model="text-embedding-3-small", + 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 == "text-embedding-3-small" + 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} @@ -263,7 +269,8 @@ def test_create_and_get_embedding(self, store): def test_sqlite_embeddings_are_stored_as_blob(self, store): entity = store.create_entity("Dataset", "run-1") - embedding = store.create_embedding(entity.id, [0.1, 0.2, 0.3]) + 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( @@ -274,50 +281,51 @@ def test_sqlite_embeddings_are_stored_as_blob(self, store): 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", [0.1, 0.2, 0.3]) + 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, [0.0, 0.0, 0.0]) + 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") - store.create_embedding(entity.id, [0.1, 0.2], embedding_model="model-a") - store.create_embedding(entity.id, [0.2, 0.3], embedding_model="model-b") - store.create_embedding(other.id, [0.3, 0.4], embedding_model="model-a") + 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="model-a") + rows = store.list_embeddings(entity_id=entity.id, embedding_model_id=model_a.id) assert len(rows) == 1 - assert rows[0].embedding_model == "model-a" + 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") - near = store.create_embedding(entity.id, [1.0, 0.0], embedding_model="model-a") - far = store.create_embedding(entity.id, [0.0, 1.0], embedding_model="model-a") - store.create_embedding(entity.id, [1.0, 1.0, 0.0], embedding_model="model-a") + 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="model-a") + 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") - embedding = store.create_embedding(entity.id, [0.1, 0.2, 0.3]) + 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 -# --------------------------------------------------------------------------- -# GraphQL integration tests (via HTTP) -# --------------------------------------------------------------------------- - - class TestGraphQL: def test_create_entity(self, client): data = gql( @@ -331,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, @@ -346,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): @@ -375,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, @@ -433,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 @@ -459,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!) { @@ -474,89 +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"] - create_embedding(client, entity_id=entity["id"], embedding_model="model-a", vector=[1.0, 0.0]) - create_embedding(client, entity_id=entity["id"], embedding_model="model-a", vector=[0.0, 1.0]) - create_embedding(client, entity_id=entity["id"], embedding_model="model-a", vector=[1.0, 1.0, 0.0]) + 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: String) { - nearestEmbeddings(vector: $vector, embeddingModel: $model) { + query Search($vector: [Float!]!, $model: ID) { + nearestEmbeddings(vector: $vector, embeddingModelId: $model) { distance embedding { entityId - embeddingModel + 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 TestEmbeddingRestAPI: - def test_create_embedding(self, client): - entity = gql(client, CREATE_ENTITY, {"input": {"entityType": "Dataset", "name": "run-1"}})["createEntity"] - payload = create_embedding( - client, - entity_id=entity["id"], - embedding_model="text-embedding-3-small", - vector=[0.1, 0.2, 0.3], - properties={"chunk": 1}, - ) +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 payload["entityId"] == entity["id"] - assert payload["embeddingModel"] == "text-embedding-3-small" - assert payload["dimensions"] == 3 - assert payload["properties"] == {"chunk": 1} + assert created["name"] == "text-embedding-3-small" + assert created["version"] == "1" - def test_get_and_list_embeddings(self, client): - entity = gql(client, CREATE_ENTITY, {"input": {"entityType": "Dataset", "name": "run-1"}})["createEntity"] - created = create_embedding(client, entity_id=entity["id"], embedding_model="model-a", vector=[0.1, 0.2]) + 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"] - fetched = client.get(f"/splash_links/embeddings/{created['id']}") - assert fetched.status_code == 200 - assert fetched.json()["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]) - listed = client.get( - "/splash_links/embeddings", - params={"entityId": entity["id"], "embeddingModel": "model-a"}, - ) - assert listed.status_code == 200 - rows = listed.json() - assert len(rows) == 1 - assert rows[0]["id"] == created["id"] + response = client.delete(f"/splash_links/embedding-models/{model['id']}") - def test_delete_embedding(self, client): - entity = gql(client, CREATE_ENTITY, {"input": {"entityType": "Dataset", "name": "run-1"}})["createEntity"] - created = create_embedding(client, entity_id=entity["id"], vector=[0.1, 0.2, 0.3]) + assert response.status_code == 409 + assert "still referenced" in response.json()["detail"] - 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 +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") - def test_delete_entity_cascades_embeddings(self, client): - entity = gql(client, CREATE_ENTITY, {"input": {"entityType": "Dataset", "name": "run-1"}})["createEntity"] - created = create_embedding(client, entity_id=entity["id"], vector=[0.1, 0.2, 0.3]) + 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 - 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 - missing = client.get(f"/splash_links/embeddings/{created['id']}") - assert missing.status_code == 404 - )["createEmbedding"] + 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 - gone = gql(client, "query Q($id: ID!) { embedding(id: $id) { id } }", {"id": embedding["id"]}) - assert gone["embedding"] is None + 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 index b84dccb..51680d3 100644 --- a/alembic/versions/6a86ce49c068_add_embeddings_table.py +++ b/alembic/versions/6a86ce49c068_add_embeddings_table.py @@ -38,11 +38,27 @@ def upgrade() -> None: 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", sa.String(), 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), @@ -52,16 +68,18 @@ def upgrade() -> None: op.create_index( "embeddings_entity_model_created_idx", "embeddings", - ["entity_id", "embedding_model", "created_at"], + ["entity_id", "embedding_model_id", "created_at"], ) op.create_index( "embeddings_model_dimensions_idx", "embeddings", - ["embedding_model", "dimensions"], + ["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") \ No newline at end of file + 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/src/splash_links/app.py b/src/splash_links/app.py index 90f8700..b7e5db3 100644 --- a/src/splash_links/app.py +++ b/src/splash_links/app.py @@ -25,37 +25,66 @@ from .schema import schema from .store import SQLAlchemyStore as SQLiteStore -from .store import EmbeddingRecord, 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: str = Field(alias="embeddingModel") + 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] - embedding_model: str = Field(default="default", alias="embeddingModel") 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=record.embedding_model, + embedding_model_id=record.embedding_model_id, + embedding_model=_embedding_model_payload(record.embedding_model), vector=record.vector, dimensions=record.dimensions, properties=record.properties, @@ -65,7 +94,13 @@ def _embedding_payload(record: EmbeddingRecord) -> EmbeddingPayload: def _embedding_error(exc: ValueError) -> HTTPException: message = str(exc) - status_code = status.HTTP_404_NOT_FOUND if "not found" in message.lower() else status.HTTP_400_BAD_REQUEST + 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) @@ -162,6 +197,67 @@ 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, @@ -174,7 +270,7 @@ def create_embedding(payload: CreateEmbeddingPayload, request: Request) -> Embed record = request.app.state.store.create_embedding( entity_id=payload.entity_id, vector=payload.vector, - embedding_model=payload.embedding_model, + embedding_model_id=payload.embedding_model_id, properties=payload.properties, ) except ValueError as exc: @@ -202,13 +298,13 @@ def get_embedding(embedding_id: str, request: Request) -> EmbeddingPayload: def list_embeddings( request: Request, entity_id: Optional[str] = Query(None, alias="entityId"), - embedding_model: Optional[str] = Query(None, alias="embeddingModel"), + 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=embedding_model, + embedding_model_id=embedding_model_id, limit=limit, offset=offset, ) diff --git a/src/splash_links/cli.py b/src/splash_links/cli.py index d89e373..757bfe1 100644 --- a/src/splash_links/cli.py +++ b/src/splash_links/cli.py @@ -4,7 +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 NAME] [--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 @@ -124,13 +124,13 @@ def links( @app.command() def embeddings( entity: Optional[str] = typer.Option(None, "--entity", "-e", help="Filter by entity ID."), - model: Optional[str] = typer.Option(None, "--model", "-m", help="Filter by embedding model."), + 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=model, limit=limit) + rows = store.list_embeddings(entity_id=entity, embedding_model_id=model_id, limit=limit) finally: store.close() @@ -153,7 +153,7 @@ def embeddings( table.add_row( embedding.id[:8], embedding.entity_id[:8], - embedding.embedding_model, + embedding.embedding_model.name, str(embedding.dimensions), vector, props, diff --git a/src/splash_links/client/base.py b/src/splash_links/client/base.py index 38a06aa..a8732d6 100644 --- a/src/splash_links/client/base.py +++ b/src/splash_links/client/base.py @@ -47,12 +47,23 @@ 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: str = Field(alias="embeddingModel") + embedding_model_id: str = Field(alias="embeddingModelId") + embedding_model: EmbeddingModel = Field(alias="embeddingModel") vector: list[float] dimensions: int properties: Optional[dict] @@ -72,7 +83,11 @@ class EmbeddingMatch(BaseModel): _ENTITY_FIELDS = "id entityType name uri properties createdAt" _LINK_FIELDS = "id subjectId predicate objectId properties createdAt" -_EMBEDDING_FIELDS = "id entityId embeddingModel vector dimensions 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!) {{ @@ -98,10 +113,10 @@ class EmbeddingMatch(BaseModel): """ _NEAREST_EMBEDDINGS_QUERY = f""" -query NearestEmbeddings($vector: [Float!]!, $embeddingModel: String, $entityId: ID, $limit: Int, $offset: Int) {{ +query NearestEmbeddings($vector: [Float!]!, $embeddingModelId: ID, $entityId: ID, $limit: Int, $offset: Int) {{ nearestEmbeddings( vector: $vector - embeddingModel: $embeddingModel + embeddingModelId: $embeddingModelId entityId: $entityId limit: $limit offset: $offset @@ -155,6 +170,7 @@ def __init__(self, base_url: str) -> None: 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] = {} @@ -283,7 +299,7 @@ def create_embedding( self, entity_id: Any, vector: list[float], - embedding_model: str = "default", + embedding_model_id: str, properties: Optional[dict] = None, ) -> Embedding: resolved_entity_id = self._resolve(entity_id) @@ -291,8 +307,8 @@ def create_embedding( self._embeddings_url, json={ "entityId": resolved_entity_id, + "embeddingModelId": embedding_model_id, "vector": vector, - "embeddingModel": embedding_model, "properties": properties, }, timeout=30.0, @@ -300,10 +316,30 @@ def create_embedding( 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: Optional[str] = None, + embedding_model_id: Optional[str] = None, entity_id: Optional[Any] = None, limit: int = 10, offset: int = 0, @@ -312,7 +348,7 @@ def find_nearest_embeddings( _NEAREST_EMBEDDINGS_QUERY, { "vector": vector, - "embeddingModel": embedding_model, + "embeddingModelId": embedding_model_id, "entityId": self._resolve(entity_id) if entity_id is not None else None, "limit": limit, "offset": offset, diff --git a/src/splash_links/client/cli.py b/src/splash_links/client/cli.py index 4f3d2cd..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 Embedding, EmbeddingMatch, 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.") @@ -40,6 +40,10 @@ 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) @@ -142,7 +146,7 @@ def find_links( 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: str = typer.Option("default", "--model", "-m", help="Embedding model label."), + embedding_model_id: str = typer.Option(..., "--model-id", "-m", help="Embedding model ID."), properties: Optional[str] = typer.Option( None, "--properties", @@ -173,7 +177,7 @@ def create_embedding( embedding = client.create_embedding( entity_id=entity_id, vector=raw_vector, - embedding_model=embedding_model, + embedding_model_id=embedding_model_id, properties=props, ) except Exception as exc: @@ -182,10 +186,34 @@ def create_embedding( _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: Optional[str] = typer.Option(None, "--model", "-m", help="Optional embedding model filter."), + 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."), @@ -211,7 +239,7 @@ def nearest_embeddings( try: matches = client.find_nearest_embeddings( vector=raw_vector, - embedding_model=embedding_model, + embedding_model_id=embedding_model_id, entity_id=entity_id, limit=limit, offset=offset, diff --git a/src/splash_links/schema.py b/src/splash_links/schema.py index 5d6906c..c3b744f 100644 --- a/src/splash_links/schema.py +++ b/src/splash_links/schema.py @@ -25,7 +25,7 @@ from strawberry.scalars import JSON as StrawberryJSON from strawberry.types import Info -from .store import EmbeddingMatchRecord, EmbeddingRecord, EntityRecord, LinkRecord, Store +from .store import EmbeddingMatchRecord, EmbeddingModelRecord, EmbeddingRecord, EntityRecord, LinkRecord, Store logger = logging.getLogger(__name__) @@ -102,11 +102,21 @@ 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: str + embedding_model_id: strawberry.ID + embedding_model: EmbeddingModel vector: list[float] dimensions: int properties: Optional[JSON] # type: ignore[valid-type] @@ -155,7 +165,14 @@ def _embedding_from_record(r: EmbeddingRecord) -> Embedding: return Embedding( id=strawberry.ID(r.id), entity_id=strawberry.ID(r.entity_id), - embedding_model=r.embedding_model, + 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, @@ -255,14 +272,14 @@ def nearest_embeddings( self, info: Info, vector: list[float], - embedding_model: Optional[str] = None, + 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=embedding_model, + 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, diff --git a/src/splash_links/store.py b/src/splash_links/store.py index a328f8a..00135d2 100644 --- a/src/splash_links/store.py +++ b/src/splash_links/store.py @@ -46,6 +46,7 @@ update, ) from sqlalchemy.engine import Engine +from sqlalchemy.exc import IntegrityError from sqlalchemy.pool import StaticPool from sqlalchemy.types import TypeDecorator, UserDefinedType @@ -667,6 +668,68 @@ 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 # ------------------------------------------------------------------ @@ -675,11 +738,13 @@ def create_embedding( self, entity_id: str, vector: list[float], - embedding_model: str = "default", + 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()) @@ -689,33 +754,33 @@ def create_embedding( insert(_embeddings).values( id=id_, entity_id=entity_id, - embedding_model=embedding_model, + embedding_model_id=embedding_model_id, vector=normalized_vector, dimensions=len(normalized_vector), properties=properties or {}, created_at=now, ) ) - row = conn.execute(select(_embeddings).where(_embeddings.c.id == id_)).one() + 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(select(_embeddings).where(_embeddings.c.id == id)).one_or_none() + 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: Optional[str] = None, + embedding_model_id: Optional[str] = None, limit: int = 100, offset: int = 0, ) -> list[EmbeddingRecord]: - stmt = select(_embeddings).order_by(_embeddings.c.created_at).limit(limit).offset(offset) + 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 is not None: - stmt = stmt.where(_embeddings.c.embedding_model == embedding_model) + 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] @@ -723,7 +788,7 @@ def list_embeddings( def find_nearest_embeddings( self, query_vector: list[float], - embedding_model: Optional[str] = None, + embedding_model_id: Optional[str] = None, entity_id: Optional[str] = None, limit: int = 10, offset: int = 0, @@ -732,14 +797,14 @@ def find_nearest_embeddings( if self._engine.dialect.name == "postgresql": return self._find_nearest_embeddings_postgresql( query_vector=normalized_query, - embedding_model=embedding_model, + 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=embedding_model, + embedding_model_id=embedding_model_id, entity_id=entity_id, limit=limit, offset=offset, @@ -748,16 +813,16 @@ def find_nearest_embeddings( def _find_nearest_embeddings_python( self, query_vector: list[float], - embedding_model: Optional[str], + embedding_model_id: Optional[str], entity_id: Optional[str], limit: int, offset: int, ) -> list[EmbeddingMatchRecord]: - stmt = select(_embeddings).where(_embeddings.c.dimensions == len(query_vector)) + 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 is not None: - stmt = stmt.where(_embeddings.c.embedding_model == embedding_model) + 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() @@ -775,12 +840,12 @@ def _find_nearest_embeddings_python( def _find_nearest_embeddings_postgresql( self, query_vector: list[float], - embedding_model: Optional[str], + embedding_model_id: Optional[str], entity_id: Optional[str], limit: int, offset: int, ) -> list[EmbeddingMatchRecord]: - conditions = ["dimensions = :dimensions"] + conditions = ["e.dimensions = :dimensions"] params: dict[str, object] = { "dimensions": len(query_vector), "query_vector": _serialize_vector(query_vector), @@ -788,23 +853,28 @@ def _find_nearest_embeddings_postgresql( "offset": offset, } if entity_id is not None: - conditions.append("entity_id = :entity_id") + conditions.append("e.entity_id = :entity_id") params["entity_id"] = entity_id - if embedding_model is not None: - conditions.append("embedding_model = :embedding_model") - params["embedding_model"] = embedding_model + 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 - id, - entity_id, - embedding_model, - vector, - dimensions, - properties, - created_at, - CAST(vector AS vector) <=> CAST(:query_vector AS vector) AS distance - FROM embeddings + 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