Add vector storage and search

Author: dylanmcreynolds

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": "<entity-id>",
+    "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/<embedding-id>
+curl 'http://localhost:8080/splash_links/embeddings?entityId=<entity-id>&embeddingModel=text-embedding-3-small'
+curl -X DELETE http://localhost:8080/splash_links/embeddings/<embedding-id>
+```
+
 ### Health check
 
 ```
@@ -132,6 +175,13 @@ pixi run links -- --subject <entity-id>     # outgoing from a node
 pixi run links -- --object  <entity-id>     # incoming to a node
 ```
 
+### List embeddings
+
+```bash
+pixi run embeddings -- --entity <entity-id>
+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:

_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
 # ---------------------------------------------------------------------------

_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
 # ---------------------------------------------------------------------------