From 1308d06d551e1385578606d728f744aaf849bfb4 Mon Sep 17 00:00:00 2001 From: Jason Shen Date: Wed, 26 Aug 2026 02:48:44 +1200 Subject: [PATCH] addressed issue 61 --- CONTRIBUTING.md | 10 ++ config.toml.example | 3 + docs/agent-runtime.md | 85 +++++++++++- docs/agent-runtime.zh-CN.md | 85 +++++++++++- docs/roadmap.md | 2 +- docs/roadmap.zh-CN.md | 2 +- internal/rag/dimensions.go | 114 +++++++++++++++ internal/rag/dimensions_test.go | 231 +++++++++++++++++++++++++++++++ internal/rag/embedding.go | 12 +- internal/rag/pgvector.go | 86 ++++++++++++ internal/rag/pgvector_pg_test.go | 224 ++++++++++++++++++++++++++++++ internal/rag/supabase.go | 107 +++++++++++++- 12 files changed, 946 insertions(+), 15 deletions(-) create mode 100644 internal/rag/dimensions.go create mode 100644 internal/rag/dimensions_test.go create mode 100644 internal/rag/pgvector_pg_test.go diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c87b9a0..4bfb8fb 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -58,6 +58,16 @@ the README badge always reflects the real state of `main`. The Go toolchain version comes from `go.mod` via `go-version-file`, so bumping the `go` directive is all it takes to move CI with it. +The RAG store tests skip unless there is a database to run against. If you are +touching `internal/rag`, give them one — they are what keeps the schema in this +repo and the one `streamcore-cli` writes from drifting apart: + +```bash +docker run -d --name pgvector-test -e POSTGRES_PASSWORD=test -e POSTGRES_DB=ragtest \ + -p 55433:5432 pgvector/pgvector:pg16 +STREAMCORE_TEST_PG=postgres://postgres:test@localhost:55433/ragtest go test -race ./internal/rag/ +``` + Formatting and vet failures are the single most common reason a PR sits. ## Where things live diff --git a/config.toml.example b/config.toml.example index 4bb34a4..a775c85 100644 --- a/config.toml.example +++ b/config.toml.example @@ -209,6 +209,9 @@ tts_url = "" # Optional; defaults to wss://openspeech.bytedance.com/a provider = "supabase" # or "pgvector", or omit entirely to disable top_k = 3 embedding_model = "text-embedding-3-small" # optional, this is the default +# The model fixes the column width: text-embedding-3-small and -ada-002 need +# vector(1536), text-embedding-3-large needs vector(3072). The server checks +# this against the store at startup and refuses to boot on a mismatch. [supabase] url = "https://xxx.supabase.co" diff --git a/docs/agent-runtime.md b/docs/agent-runtime.md index 8e4c6cb..da343a5 100644 --- a/docs/agent-runtime.md +++ b/docs/agent-runtime.md @@ -78,11 +78,17 @@ CREATE EXTENSION IF NOT EXISTS vector; CREATE TABLE documents ( id SERIAL PRIMARY KEY, content TEXT NOT NULL, - embedding vector(1536), - source TEXT + embedding vector(1536) NOT NULL, + embedding_model TEXT NOT NULL, + source TEXT, + created_at TIMESTAMP DEFAULT NOW() ); + +CREATE INDEX IF NOT EXISTS documents_embedding_model_idx ON documents (embedding_model); ``` +`streamcore-cli setup` creates all of this for you; the SQL is here for anyone who would rather run it themselves. + ```toml [rag] provider = "pgvector" @@ -99,11 +105,14 @@ CREATE EXTENSION IF NOT EXISTS vector; CREATE TABLE documents ( id SERIAL PRIMARY KEY, content TEXT NOT NULL, - embedding vector(1536), + embedding vector(1536) NOT NULL, + embedding_model TEXT NOT NULL, source TEXT, created_at TIMESTAMP DEFAULT NOW() ); +CREATE INDEX IF NOT EXISTS documents_embedding_model_idx ON documents (embedding_model); + CREATE OR REPLACE FUNCTION match_documents( query_embedding vector(1536), match_count int DEFAULT 3 @@ -142,10 +151,37 @@ function = "match_documents" table = "documents" ``` +### Embedding model and vector width + +Every row records the model that embedded it, and the vector column is sized for that model. Both halves are checked when the server boots, and a store it cannot use is a startup failure rather than retrieval that quietly returns the wrong chunks. + +| `embedding_model` | Column type | +|---|---| +| `text-embedding-3-small` (default) | `vector(1536)` | +| `text-embedding-ada-002` | `vector(1536)` | +| `text-embedding-3-large` | `vector(3072)` | + +The width alone is not enough. `ada-002` and `3-small` are both 1536 wide, so vectors written by one and searched with the other produce no error at all — just answers drawn from the wrong chunks. That is what `embedding_model` is for: the server looks for a single row disagreeing with `rag.embedding_model` and refuses to start if it finds one. + +For `pgvector` the width comes from the column definition and the model check is one indexed query. For Supabase there is no catalog to read over PostgREST, so both come from rows: an empty table stays unverified until something has been ingested, and a project that is unreachable at boot logs a warning instead of blocking startup. + +If you ingested before `embedding_model` existed, the server will tell you so and print this migration: + +```sql +ALTER TABLE documents ADD COLUMN embedding_model TEXT; +UPDATE documents SET embedding_model = ''; +ALTER TABLE documents ALTER COLUMN embedding_model SET NOT NULL; +CREATE INDEX IF NOT EXISTS documents_embedding_model_idx ON documents (embedding_model); +``` + +Only you know which model those rows came from, which is why the backfill is a placeholder. If you no longer know, re-ingest. + ### Ingesting documents The server handles query-time retrieval only. Populate your vector store with [`streamcore-cli`](https://github.com/streamcoreai/streamcore-cli), a separate Go binary. It stays separate so the PDF, docx and xlsx parsers never end up in the server image. +Prebuilt binaries for macOS and Linux, both architectures, are on the [releases page](https://github.com/streamcoreai/streamcore-cli/releases). With a Go toolchain: + ```bash go install github.com/streamcoreai/streamcore-cli@latest @@ -154,7 +190,9 @@ git clone https://github.com/streamcoreai/streamcore-cli cd streamcore-cli && go build -o streamcore-cli . ``` -`streamcore-cli setup` asks for your provider, OpenAI key and credentials, then writes `~/.streamcore/config.toml`. If you already have a server `config.toml`, skip it — the CLI reads the same format and falls back to the server's file, so nothing is configured twice. +`streamcore-cli setup` asks for your provider, OpenAI key and credentials, writes `~/.streamcore/config.toml`, and creates the table with the vector width your chosen model needs. For Supabase it asks for the project's direct Postgres connection string, since PostgREST can insert rows but cannot run DDL; leave that blank and it prints the SQL for the dashboard's editor instead. + +If you already have a server `config.toml`, the CLI reads the same format and falls back to the server's file, so credentials are never configured twice. ```bash streamcore-cli setup @@ -174,6 +212,43 @@ Config is looked up in order: `--config`, `~/.streamcore/config.toml`, `./config | `--chunk-size` | 512 | Target chunk size in words | | `--chunk-overlap` | 64 | Overlap between chunks in words | -Ingest and query must use the same `embedding_model`. Vectors written by one model and searched by another are not comparable, and the result is bad retrieval rather than an error — so if you change the model, re-ingest. +Ingest and query must use the same `embedding_model` — see [Embedding model and vector width](#embedding-model-and-vector-width) for which column each model needs. `ingest` checks the store against the configured model before it writes anything, so a mismatch costs one query rather than a table full of unusable vectors. Full command reference, supported formats and database DDL: [streamcore-cli README](https://github.com/streamcoreai/streamcore-cli#readme). + +### End to end + +Pick something the model cannot already know. A PDF of your own release notes works; so does a text file you write on the spot. + +```bash +cat > closing-hours.md <<'EOF' +The Wellington workshop closes at 3pm on the last Friday of every month +for maintenance. All other Fridays it closes at 6pm. +EOF + +streamcore-cli setup # provider, key, model, and the table +streamcore-cli ingest closing-hours.md +``` + +``` +Using config: /Users/you/.streamcore/config.toml +Processing closing-hours.md ... + Extracted 1 chunks + Uploaded 1/1 chunks +Done. 1 chunks uploaded to supabase (text-embedding-3-small). +``` + +Point the server at the same config and start it. It reads the table at boot and says nothing if the contract holds: + +``` +RAG enabled — provider: supabase +Voice agent server listening on :8080 +``` + +Then connect a client and ask *"when does the Wellington workshop close on the last Friday of the month?"* The answer should be 3pm. Ask before ingesting, or against a store built with a different model, and you get a generic non-answer instead — which is the failure this contract exists to make loud. + +### Why ingestion is a separate binary + +The parsers are the reason. PDF, docx and xlsx pull in dependencies that the server has no use for at query time, and the server image is something people deploy; the ingestion tool is something they run once from a laptop. Keeping them apart means the server image never carries a document parser it will never call. + +The cost is a second artifact with its own config, which is why the CLI reads the server's `config.toml` and falls back to it — one set of credentials, two binaries. If that stops holding, the argument for the split is worth revisiting. diff --git a/docs/agent-runtime.zh-CN.md b/docs/agent-runtime.zh-CN.md index 3938978..abe6dea 100644 --- a/docs/agent-runtime.zh-CN.md +++ b/docs/agent-runtime.zh-CN.md @@ -78,11 +78,17 @@ CREATE EXTENSION IF NOT EXISTS vector; CREATE TABLE documents ( id SERIAL PRIMARY KEY, content TEXT NOT NULL, - embedding vector(1536), - source TEXT + embedding vector(1536) NOT NULL, + embedding_model TEXT NOT NULL, + source TEXT, + created_at TIMESTAMP DEFAULT NOW() ); + +CREATE INDEX IF NOT EXISTS documents_embedding_model_idx ON documents (embedding_model); ``` +`streamcore-cli setup` 会替你建好这一切；这段 SQL 是留给想自己动手的人的。 + ```toml [rag] provider = "pgvector" @@ -99,11 +105,14 @@ CREATE EXTENSION IF NOT EXISTS vector; CREATE TABLE documents ( id SERIAL PRIMARY KEY, content TEXT NOT NULL, - embedding vector(1536), + embedding vector(1536) NOT NULL, + embedding_model TEXT NOT NULL, source TEXT, created_at TIMESTAMP DEFAULT NOW() ); +CREATE INDEX IF NOT EXISTS documents_embedding_model_idx ON documents (embedding_model); + CREATE OR REPLACE FUNCTION match_documents( query_embedding vector(1536), match_count int DEFAULT 3 @@ -142,10 +151,37 @@ function = "match_documents" table = "documents" ``` +### 嵌入模型与向量宽度 + +每一行都记录了给它做 embedding 的模型，而向量列的宽度也是按那个模型定的。服务端启动时两边都会检查，用不了的向量库会直接导致启动失败，而不是让检索悄悄返回错误的片段。 + +| `embedding_model` | 列类型 | +|---|---| +| `text-embedding-3-small`（默认） | `vector(1536)` | +| `text-embedding-ada-002` | `vector(1536)` | +| `text-embedding-3-large` | `vector(3072)` | + +光看宽度是不够的。`ada-002` 和 `3-small` 都是 1536 维，用一个写入、用另一个检索，不会报任何错，只会从错误的片段里给出答案。`embedding_model` 就是为此存在的：服务端会去找一行与 `rag.embedding_model` 不一致的记录，找到就拒绝启动。 + +`pgvector` 的宽度取自列定义，模型检查是一次走索引的查询。Supabase 那边没有可以通过 PostgREST 读到的系统目录，所以两项都靠取行来判断：空表在入库之前无法校验，启动时连不上的项目只记一条警告，不会阻塞启动。 + +如果你是在 `embedding_model` 这一列出现之前入的库，服务端会告诉你，并打印出这段迁移 SQL： + +```sql +ALTER TABLE documents ADD COLUMN embedding_model TEXT; +UPDATE documents SET embedding_model = '<你当初入库用的模型>'; +ALTER TABLE documents ALTER COLUMN embedding_model SET NOT NULL; +CREATE INDEX IF NOT EXISTS documents_embedding_model_idx ON documents (embedding_model); +``` + +那些行到底来自哪个模型，只有你自己知道，所以回填的值留成了占位符。要是已经想不起来了，就重新入库一遍。 + ### 文档入库 服务端只负责查询时的检索。向量库的内容由 [`streamcore-cli`](https://github.com/streamcoreai/streamcore-cli) 填充 —— 那是一个独立的 Go 二进制。之所以独立，是为了让 PDF、docx、xlsx 的解析依赖不会进到服务端镜像里。 +macOS 和 Linux 两种架构的预编译二进制都在[发布页](https://github.com/streamcoreai/streamcore-cli/releases)。如果本机有 Go 工具链： + ```bash go install github.com/streamcoreai/streamcore-cli@latest @@ -154,7 +190,9 @@ git clone https://github.com/streamcoreai/streamcore-cli cd streamcore-cli && go build -o streamcore-cli . ``` -`streamcore-cli setup` 会依次询问服务商、OpenAI key 和凭据，然后写入 `~/.streamcore/config.toml`。如果你已经有服务端的 `config.toml`，可以跳过这一步 —— CLI 读的是同一种格式，会回退到服务端那份文件，因此没有任何东西需要配置两遍。 +`streamcore-cli setup` 会依次询问服务商、OpenAI key 和凭据，写入 `~/.streamcore/config.toml`，并按你选的模型所需的向量宽度把表建好。Supabase 会额外问一个项目的 Postgres 直连串 —— PostgREST 能插入数据，但跑不了 DDL；这一项留空，它就把 SQL 打印出来给你贴到控制台的 SQL 编辑器里。 + +如果你已经有服务端的 `config.toml`，CLI 读的是同一种格式，会回退到服务端那份文件，因此凭据不需要配置两遍。 ```bash streamcore-cli setup @@ -174,6 +212,43 @@ streamcore-cli ingest --chunk-size 256 --chunk-overlap 32 manual.docx | `--chunk-size` | 512 | 目标分块大小（词数） | | `--chunk-overlap` | 64 | 分块之间的重叠（词数） | -入库和查询必须使用同一个 `embedding_model`。用一个模型写入、用另一个模型检索出来的向量之间没有可比性，而且不会报错，只会让召回变差 —— 所以换了模型就重新入库一遍。 +入库和查询必须使用同一个 `embedding_model` —— 每个模型对应哪种列，见[嵌入模型与向量宽度](#嵌入模型与向量宽度)。`ingest` 在写入任何数据之前，会先拿配置里的模型去校验向量库，所以对不上的代价是一次查询，而不是一整张写满了用不了的向量的表。 完整的命令说明、支持的格式和建表 SQL：[streamcore-cli README](https://github.com/streamcoreai/streamcore-cli/blob/main/README.zh-CN.md)。 + +### 端到端跑一遍 + +挑一份模型不可能已经知道的内容。你自己的发布说明 PDF 可以，现写一个文本文件也可以。 + +```bash +cat > closing-hours.md <<'EOF' +惠灵顿工坊每月最后一个周五下午 3 点闭店做维护。 +其余的周五都是 6 点闭店。 +EOF + +streamcore-cli setup # 服务商、key、模型，以及建表 +streamcore-cli ingest closing-hours.md +``` + +``` +Using config: /Users/you/.streamcore/config.toml +Processing closing-hours.md ... + Extracted 1 chunks + Uploaded 1/1 chunks +Done. 1 chunks uploaded to supabase (text-embedding-3-small). +``` + +让服务端指向同一份配置并启动。它会在启动时读一遍这张表，契约没问题就什么都不说： + +``` +RAG enabled — provider: supabase +Voice agent server listening on :8080 +``` + +然后接一个客户端上去，问「惠灵顿工坊每月最后一个周五几点闭店？」，答案应该是下午 3 点。入库之前问，或者对着用另一个模型建起来的库问，得到的就是一句没有信息量的场面话 —— 而这正是这套契约要让它变响的那种失败。 + +### 为什么入库是一个单独的二进制 + +原因在解析器。PDF、docx、xlsx 会拖进一堆依赖，而服务端在查询时根本用不上它们；服务端镜像是要部署出去的东西，入库工具则是在自己笔记本上跑一次的东西。分开之后，服务端镜像里就不会带着一个永远不会被调用的文档解析器。 + +代价是多了一个带自己配置的产物 —— 所以 CLI 会去读服务端的 `config.toml` 并回退到它：一份凭据，两个二进制。哪天这一点不再成立了，这个拆分就值得重新讨论。 diff --git a/docs/roadmap.md b/docs/roadmap.md index ad0dbbc..a05af99 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -33,7 +33,7 @@ Want one of these? Say so in [Discord](https://discord.gg/xKGFaGWawT) or open an - [ ] **Broader examples** proving the positioning: realtime translator, AI-hosted voice room, browser copilot, embedded device, SIP application, and a raw audio-processing app with no LLM at all. - [ ] **Embedded client hardening.** The ESP32-S3 firmware in [`esp32`](https://github.com/streamcoreai/esp32) connects over WHIP but is not production-ready. -- [x] **`streamcore-cli` release.** Shipped as [`streamcoreai/streamcore-cli`](https://github.com/streamcoreai/streamcore-cli) — RAG document ingestion, see [Agent runtime → Ingesting documents](./agent-runtime.md#ingesting-documents). Tagged releases with prebuilt binaries remain open. +- [x] **`streamcore-cli` release.** Shipped as [`streamcoreai/streamcore-cli`](https://github.com/streamcoreai/streamcore-cli) — RAG document ingestion, see [Agent runtime → Ingesting documents](./agent-runtime.md#ingesting-documents). Tagging a version publishes prebuilt binaries for macOS and Linux on both architectures. - [ ] **React Native SDK on npm.** `@streamcore/react-native-sdk` is built and usable from source, but unpublished. --- diff --git a/docs/roadmap.zh-CN.md b/docs/roadmap.zh-CN.md index dc6c413..3e833bb 100644 --- a/docs/roadmap.zh-CN.md +++ b/docs/roadmap.zh-CN.md @@ -33,7 +33,7 @@ - [ ] **更多示例**，以印证产品定位：实时翻译、AI 主持的语音房间、浏览器副驾、嵌入式设备、SIP 应用，以及一个完全不含 LLM 的纯音频处理应用。 - [ ] **嵌入式客户端加固。** [`esp32`](https://github.com/streamcoreai/esp32) 中的 ESP32-S3 固件已能通过 WHIP 连接，但尚未达到生产可用。 -- [x] **`streamcore-cli` 发布。** 已发布为 [`streamcoreai/streamcore-cli`](https://github.com/streamcoreai/streamcore-cli) —— 提供 RAG 文档入库，见[智能体运行时 → 文档入库](./agent-runtime.zh-CN.md#文档入库)。打 tag 的预编译二进制仍在计划中。 +- [x] **`streamcore-cli` 发布。** 已发布为 [`streamcoreai/streamcore-cli`](https://github.com/streamcoreai/streamcore-cli) —— 提供 RAG 文档入库，见[智能体运行时 → 文档入库](./agent-runtime.zh-CN.md#文档入库)。打一个版本 tag 即会发布 macOS 与 Linux 两种架构的预编译二进制。 - [ ] **React Native SDK 发布到 npm。** `@streamcore/react-native-sdk` 已完成、可从源码使用，但尚未发布。 --- diff --git a/internal/rag/dimensions.go b/internal/rag/dimensions.go new file mode 100644 index 0000000..25c9c20 --- /dev/null +++ b/internal/rag/dimensions.go @@ -0,0 +1,114 @@ +package rag + +import ( + "encoding/json" + "fmt" + "strconv" + "strings" +) + +const defaultEmbeddingModel = "text-embedding-3-small" + +// Output width of every embedding model we know about. Width catches half the +// contract: Postgres rejects a 3072-wide vector against a vector(1536) column +// on its own. The other half is the same-width case — ada-002 and 3-small are +// both 1536 and produce no error at all, only bad retrieval — which is why +// every row also carries the model that wrote it. +var modelDimensions = map[string]int{ + "text-embedding-3-small": 1536, + "text-embedding-3-large": 3072, + "text-embedding-ada-002": 1536, +} + +// dimensionsFor reports the vector width for model. ok is false for models +// released after this table was written; callers skip the width check instead +// of refusing to start on a model they simply haven't heard of. +func dimensionsFor(model string) (dims int, ok bool) { + if model == "" { + model = defaultEmbeddingModel + } + dims, ok = modelDimensions[model] + return dims, ok +} + +func normalizeModel(model string) string { + if model == "" { + return defaultEmbeddingModel + } + return model +} + +// dimensionMismatch names both ways out, because which one is right depends on +// whether the config or the store is the thing that changed. +func dimensionMismatch(table, model string, want, got int) error { + model = normalizeModel(model) + return fmt.Errorf( + "table %q holds vector(%d) but embedding_model %q produces %d dimensions: "+ + "re-ingest with streamcore-cli against a vector(%d) column, or set "+ + "rag.embedding_model back to the model the store was built with", + table, got, model, want, want) +} + +// modelMismatch is the failure the width check cannot see: same-size vectors +// from a different model, which retrieve nonsense without erroring anywhere. +func modelMismatch(table, want, got string) error { + if got == "" { + got = "(null)" + } + return fmt.Errorf( + "table %q holds rows embedded with %q but rag.embedding_model is %q: "+ + "vectors from different models are not comparable, so re-ingest with "+ + "streamcore-cli or set rag.embedding_model back to %[2]q", + table, got, normalizeModel(want)) +} + +// missingModelColumn carries the migration inline. Anyone hitting this ingested +// before the column existed, and the value to backfill is something only they +// know, so the UPDATE is left with a placeholder rather than a guess. +func missingModelColumn(table string) error { + return fmt.Errorf( + "table %q has no embedding_model column, so the model its vectors were "+ + "written with cannot be checked. Migrate it:\n"+ + " ALTER TABLE %[1]s ADD COLUMN embedding_model TEXT;\n"+ + " UPDATE %[1]s SET embedding_model = '';\n"+ + " ALTER TABLE %[1]s ALTER COLUMN embedding_model SET NOT NULL;\n"+ + " CREATE INDEX IF NOT EXISTS %[1]s_embedding_model_idx ON %[1]s (embedding_model);", + table) +} + +// vectorWidth counts the dimensions in an embedding column value as PostgREST +// returns it. pgvector has no JSON mapping there, so a row arrives as the +// quoted text form "[0.1,0.2]" rather than a JSON array — accept both, since +// the day PostgREST learns the type is the day this silently stops checking. +func vectorWidth(raw json.RawMessage) (int, bool) { + trimmed := strings.TrimSpace(string(raw)) + if trimmed == "" || trimmed == "null" { + return 0, false + } + + if trimmed[0] == '"' { + var s string + if err := json.Unmarshal(raw, &s); err != nil { + return 0, false + } + s = strings.TrimSpace(s) + s = strings.TrimPrefix(s, "[") + s = strings.TrimSuffix(s, "]") + if s == "" { + return 0, false + } + parts := strings.Split(s, ",") + for _, p := range parts { + if _, err := strconv.ParseFloat(strings.TrimSpace(p), 64); err != nil { + return 0, false + } + } + return len(parts), true + } + + var vec []float64 + if err := json.Unmarshal(raw, &vec); err != nil || len(vec) == 0 { + return 0, false + } + return len(vec), true +} diff --git a/internal/rag/dimensions_test.go b/internal/rag/dimensions_test.go new file mode 100644 index 0000000..a1795ad --- /dev/null +++ b/internal/rag/dimensions_test.go @@ -0,0 +1,231 @@ +package rag + +import ( + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestDimensionsFor(t *testing.T) { + tests := []struct { + model string + want int + ok bool + }{ + {"", 1536, true}, + {"text-embedding-3-small", 1536, true}, + {"text-embedding-3-large", 3072, true}, + {"text-embedding-ada-002", 1536, true}, + {"some-model-shipped-next-year", 0, false}, + } + + for _, tt := range tests { + got, ok := dimensionsFor(tt.model) + if got != tt.want || ok != tt.ok { + t.Errorf("dimensionsFor(%q) = (%d, %v), want (%d, %v)", tt.model, got, ok, tt.want, tt.ok) + } + } +} + +func TestNewEmbeddingClientTracksWidth(t *testing.T) { + if c := newEmbeddingClient("k", ""); c.model != defaultEmbeddingModel || c.dims != 1536 { + t.Errorf("default client = (%q, %d), want (%q, 1536)", c.model, c.dims, defaultEmbeddingModel) + } + // An unknown model disables the check rather than failing every Embed call. + if c := newEmbeddingClient("k", "mystery-model"); c.dims != 0 { + t.Errorf("unknown model dims = %d, want 0", c.dims) + } +} + +func TestVectorWidth(t *testing.T) { + tests := []struct { + name string + raw string + want int + ok bool + }{ + {"postgrest text form", `"[0.1,0.2,-0.3]"`, 3, true}, + {"text form with spaces", `"[0.1, 0.2]"`, 2, true}, + {"json array", `[0.1,0.2,0.3,0.4]`, 4, true}, + {"null", `null`, 0, false}, + {"empty vector", `"[]"`, 0, false}, + {"not numbers", `"[a,b]"`, 0, false}, + {"empty array", `[]`, 0, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := vectorWidth(json.RawMessage(tt.raw)) + if got != tt.want || ok != tt.ok { + t.Errorf("vectorWidth(%s) = (%d, %v), want (%d, %v)", tt.raw, got, ok, tt.want, tt.ok) + } + }) + } +} + +func TestDimensionMismatchMessage(t *testing.T) { + err := dimensionMismatch("documents", "text-embedding-3-large", 3072, 1536) + for _, want := range []string{"documents", "vector(1536)", "3072", "text-embedding-3-large", "streamcore-cli"} { + if !strings.Contains(err.Error(), want) { + t.Errorf("message %q missing %q", err.Error(), want) + } + } +} + +func TestModelMismatchMessage(t *testing.T) { + err := modelMismatch("documents", "text-embedding-3-small", "text-embedding-ada-002") + for _, want := range []string{"documents", "text-embedding-ada-002", "text-embedding-3-small", "streamcore-cli"} { + if !strings.Contains(err.Error(), want) { + t.Errorf("message %q missing %q", err.Error(), want) + } + } + + // A null model reads as an unmigrated row, not an empty string. + if got := modelMismatch("documents", "text-embedding-3-small", "").Error(); !strings.Contains(got, "(null)") { + t.Errorf("null model message = %q, want it to mention (null)", got) + } +} + +func TestMissingModelColumnCarriesMigration(t *testing.T) { + err := missingModelColumn("chunks") + for _, want := range []string{"ALTER TABLE chunks ADD COLUMN embedding_model", "SET NOT NULL", "chunks_embedding_model_idx"} { + if !strings.Contains(err.Error(), want) { + t.Errorf("migration %q missing %q", err.Error(), want) + } + } +} + +func TestIsMissingColumn(t *testing.T) { + tests := []struct { + err error + want bool + }{ + {errors.New(`returned 400: {"code":"42703","message":"column documents.embedding_model does not exist"}`), true}, + {errors.New(`returned 400: {"message":"column embedding_model does not exist"}`), true}, + {errors.New(`returned 401: {"message":"invalid api key"}`), false}, + {errors.New("connection refused"), false}, + } + for _, tt := range tests { + if got := isMissingColumn(tt.err); got != tt.want { + t.Errorf("isMissingColumn(%v) = %v, want %v", tt.err, got, tt.want) + } + } +} + +func TestVerifySupabaseStore(t *testing.T) { + vec := func(n int) string { + parts := make([]string, n) + for i := range parts { + parts[i] = "0.1" + } + return "[" + strings.Join(parts, ",") + "]" + } + row := func(width int, model string) string { + return fmt.Sprintf(`[{"embedding":"%s","embedding_model":"%s"}]`, vec(width), model) + } + + tests := []struct { + name string + model string + // mismatch* is the reply to the "any row with a different model" probe; + // sampleBody answers the width query that follows it. + mismatchStatus int + mismatchBody string + sampleBody string + wantErr string + }{ + { + name: "store agrees", + model: "text-embedding-3-small", + mismatchStatus: 200, + mismatchBody: `[]`, + sampleBody: row(1536, "text-embedding-3-small"), + }, + { + name: "same width, different model", + model: "text-embedding-3-small", + mismatchStatus: 200, + mismatchBody: row(1536, "text-embedding-ada-002"), + wantErr: "text-embedding-ada-002", + }, + { + name: "unmigrated rows read as null", + model: "text-embedding-3-small", + mismatchStatus: 200, + mismatchBody: `[{"embedding":"[0.1]","embedding_model":null}]`, + wantErr: "(null)", + }, + { + name: "column predates the contract", + model: "text-embedding-3-small", + mismatchStatus: 400, + mismatchBody: `{"code":"42703","message":"column documents.embedding_model does not exist"}`, + wantErr: "ALTER TABLE documents ADD COLUMN embedding_model", + }, + { + name: "empty table", + model: "text-embedding-3-small", + mismatchStatus: 200, + mismatchBody: `[]`, + sampleBody: `[]`, + }, + { + name: "width disagrees", + model: "text-embedding-3-large", + mismatchStatus: 200, + mismatchBody: `[]`, + sampleBody: row(1536, "text-embedding-3-large"), + wantErr: "vector(1536)", + }, + { + name: "unauthorized leaves the store unverified", + model: "text-embedding-3-small", + mismatchStatus: 401, + mismatchBody: `{"message":"invalid api key"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var calls int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + if calls == 1 { + if !strings.Contains(r.URL.RawQuery, "embedding_model.neq.") { + t.Errorf("first query = %q, want the model mismatch probe", r.URL.RawQuery) + } + w.WriteHeader(tt.mismatchStatus) + _, _ = w.Write([]byte(tt.mismatchBody)) + return + } + _, _ = w.Write([]byte(tt.sampleBody)) + })) + defer srv.Close() + + err := verifySupabaseStore(srv.Client(), srv.URL, "key", "documents", tt.model) + switch { + case tt.wantErr == "" && err != nil: + t.Fatalf("err = %v, want nil", err) + case tt.wantErr != "" && err == nil: + t.Fatalf("err = nil, want one mentioning %q", tt.wantErr) + case tt.wantErr != "" && !strings.Contains(err.Error(), tt.wantErr): + t.Fatalf("err = %v, want it to mention %q", err, tt.wantErr) + } + }) + } +} + +func TestVerifySupabaseStoreUnreachable(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) + client := srv.Client() + srv.Close() + + // A project that is down at boot must not take the server down with it. + if err := verifySupabaseStore(client, srv.URL, "key", "documents", "text-embedding-3-small"); err != nil { + t.Fatalf("unreachable project returned %v, want nil", err) + } +} diff --git a/internal/rag/embedding.go b/internal/rag/embedding.go index 8c21f38..fe1c5f5 100644 --- a/internal/rag/embedding.go +++ b/internal/rag/embedding.go @@ -14,16 +14,19 @@ import ( type embeddingClient struct { apiKey string model string + dims int // expected output width, 0 for models not in modelDimensions client *http.Client } func newEmbeddingClient(apiKey, model string) *embeddingClient { if model == "" { - model = "text-embedding-3-small" + model = defaultEmbeddingModel } + dims, _ := dimensionsFor(model) return &embeddingClient{ apiKey: apiKey, model: model, + dims: dims, client: &http.Client{Timeout: 10 * time.Second}, } } @@ -76,5 +79,10 @@ func (e *embeddingClient) Embed(ctx context.Context, text string) ([]float32, er return nil, fmt.Errorf("embedding API returned no data") } - return result.Data[0].Embedding, nil + vec := result.Data[0].Embedding + if e.dims > 0 && len(vec) != e.dims { + return nil, fmt.Errorf("embedding model %q returned %d dimensions, expected %d", e.model, len(vec), e.dims) + } + + return vec, nil } diff --git a/internal/rag/pgvector.go b/internal/rag/pgvector.go index a3284de..dbea700 100644 --- a/internal/rag/pgvector.go +++ b/internal/rag/pgvector.go @@ -2,9 +2,12 @@ package rag import ( "context" + "errors" "fmt" "strings" + "time" + "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" "github.com/streamcoreai/streamcore-server/internal/config" ) @@ -36,6 +39,11 @@ func NewPgvectorClient(cfg *config.Config) (Client, error) { table = "documents" } + if err := verifyStore(pool, table, cfg.RAG.EmbeddingModel); err != nil { + pool.Close() + return nil, fmt.Errorf("pgvector: %w", err) + } + topK := cfg.RAG.TopK if topK == 0 { topK = 3 @@ -85,6 +93,84 @@ func (c *pgvectorClient) Search(ctx context.Context, query string, topK int) ([] return chunks, rows.Err() } +// verifyStore refuses to start against a store whose vectors cannot be +// compared with the ones this server will produce. Both halves matter: the +// declared width, which pgvector keeps straight in atttypmod (-1 for a bare +// `vector` column, which takes anything), and the model recorded on the rows, +// which is the only way to catch two 1536-wide models being mixed. +func verifyStore(pool *pgxpool.Pool, table, model string) error { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + cols, err := describeColumns(ctx, pool, table) + if err != nil { + return err + } + if len(cols) == 0 { + return fmt.Errorf("table %q is missing or has no embedding column: run `streamcore-cli setup` to create it, then ingest", table) + } + + typmod, ok := cols["embedding"] + if !ok { + return fmt.Errorf("table %q has no embedding column: run `streamcore-cli setup` to create it, then ingest", table) + } + if want, known := dimensionsFor(model); known && typmod > 0 && int(typmod) != want { + return dimensionMismatch(table, model, want, int(typmod)) + } + + if _, ok := cols["embedding_model"]; !ok { + return missingModelColumn(table) + } + return verifyRowModels(ctx, pool, table, model) +} + +func describeColumns(ctx context.Context, pool *pgxpool.Pool, table string) (map[string]int32, error) { + rows, err := pool.Query(ctx, + `SELECT attname, atttypmod FROM pg_attribute + WHERE attrelid = to_regclass($1) AND attname IN ('embedding', 'embedding_model') + AND attnum > 0 AND NOT attisdropped`, table) + if err != nil { + return nil, fmt.Errorf("inspect %s: %w", table, err) + } + defer rows.Close() + + cols := make(map[string]int32, 2) + for rows.Next() { + var name string + var typmod int32 + if err := rows.Scan(&name, &typmod); err != nil { + return nil, fmt.Errorf("inspect %s: %w", table, err) + } + cols[name] = typmod + } + return cols, rows.Err() +} + +// verifyRowModels looks for one row that disagrees with the configured model. +// IS DISTINCT FROM rather than <> so a NULL counts as a disagreement, and the +// query stops at the first offender — which on a store built entirely with the +// wrong model is the first row it reads. +func verifyRowModels(ctx context.Context, pool *pgxpool.Pool, table, model string) error { + want := normalizeModel(model) + + var got *string + err := pool.QueryRow(ctx, fmt.Sprintf( + `SELECT embedding_model FROM %s WHERE embedding_model IS DISTINCT FROM $1 LIMIT 1`, table), + want).Scan(&got) + if errors.Is(err, pgx.ErrNoRows) { + return nil // empty table, or every row agrees + } + if err != nil { + return fmt.Errorf("inspect %s.embedding_model: %w", table, err) + } + + found := "" + if got != nil { + found = *got + } + return modelMismatch(table, want, found) +} + func formatVector(v []float32) string { parts := make([]string, len(v)) for i, f := range v { diff --git a/internal/rag/pgvector_pg_test.go b/internal/rag/pgvector_pg_test.go new file mode 100644 index 0000000..7798057 --- /dev/null +++ b/internal/rag/pgvector_pg_test.go @@ -0,0 +1,224 @@ +package rag + +import ( + "context" + "os" + "strings" + "testing" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/streamcoreai/streamcore-server/internal/config" +) + +// These run against a live pgvector database, because what they check is +// whether the store contract in docs/agent-runtime.md survives contact with +// Postgres — the shape of atttypmod for a vector column, and whether the DDL +// we publish is the DDL the startup check accepts. Without one they skip: +// +// docker run -d -e POSTGRES_PASSWORD=test -e POSTGRES_DB=ragtest \ +// -p 55433:5432 pgvector/pgvector:pg16 +// STREAMCORE_TEST_PG=postgres://postgres:test@localhost:55433/ragtest go test ./internal/rag/ +func testPool(t *testing.T) (*pgxpool.Pool, string) { + t.Helper() + connString := os.Getenv("STREAMCORE_TEST_PG") + if connString == "" { + t.Skip("set STREAMCORE_TEST_PG to a pgvector database to run this") + } + + pool, err := pgxpool.New(context.Background(), connString) + if err != nil { + t.Fatalf("connect: %v", err) + } + t.Cleanup(pool.Close) + return pool, connString +} + +// documentedDDL is the SQL from docs/agent-runtime.md, kept here so a change to +// one without the other fails a test rather than a deployment. +func documentedDDL(table string, dims int) string { + return strings.NewReplacer("$TABLE", table, "$DIMS", itoa(dims)).Replace(` +CREATE EXTENSION IF NOT EXISTS vector; + +CREATE TABLE $TABLE ( + id SERIAL PRIMARY KEY, + content TEXT NOT NULL, + embedding vector($DIMS) NOT NULL, + embedding_model TEXT NOT NULL, + source TEXT, + created_at TIMESTAMP DEFAULT NOW() +); + +CREATE INDEX IF NOT EXISTS $TABLE_embedding_model_idx ON $TABLE (embedding_model);`) +} + +func itoa(n int) string { + if n == 0 { + return "0" + } + digits := "" + for n > 0 { + digits = string(rune('0'+n%10)) + digits + n /= 10 + } + return digits +} + +func createTable(t *testing.T, pool *pgxpool.Pool, ddl, table string) { + t.Helper() + ctx := context.Background() + if _, err := pool.Exec(ctx, "DROP TABLE IF EXISTS "+table); err != nil { + t.Fatalf("drop: %v", err) + } + if _, err := pool.Exec(ctx, ddl); err != nil { + t.Fatalf("create %s: %v", table, err) + } + t.Cleanup(func() { + _, _ = pool.Exec(context.Background(), "DROP TABLE IF EXISTS "+table) + }) +} + +func insertRow(t *testing.T, pool *pgxpool.Pool, table string, dims int, model string) { + t.Helper() + vec := make([]float32, dims) + vec[0] = 0.5 + _, err := pool.Exec(context.Background(), + "INSERT INTO "+table+" (content, embedding, embedding_model, source) VALUES ($1, $2::vector, $3, $4)", + "the workshop closes at 3pm", formatVector(vec), model, "hours.md") + if err != nil { + t.Fatalf("insert: %v", err) + } +} + +func TestVerifyStoreAcceptsTheDocumentedSchema(t *testing.T) { + pool, _ := testPool(t) + const table = "rag_documented_docs" + createTable(t, pool, documentedDDL(table, 1536), table) + + if err := verifyStore(pool, table, "text-embedding-3-small"); err != nil { + t.Fatalf("empty documented table rejected: %v", err) + } + + insertRow(t, pool, table, 1536, "text-embedding-3-small") + if err := verifyStore(pool, table, "text-embedding-3-small"); err != nil { + t.Fatalf("populated documented table rejected: %v", err) + } + + // The default model has to be read as text-embedding-3-small, or every + // config that omits embedding_model would fail against its own store. + if err := verifyStore(pool, table, ""); err != nil { + t.Fatalf("empty model config rejected: %v", err) + } +} + +func TestVerifyStoreRejectsAMismatchedStore(t *testing.T) { + pool, _ := testPool(t) + const table = "rag_mismatch_docs" + createTable(t, pool, documentedDDL(table, 1536), table) + insertRow(t, pool, table, 1536, "text-embedding-ada-002") + + // Same width, so nothing but the recorded model can catch it. + err := verifyStore(pool, table, "text-embedding-3-small") + if err == nil { + t.Fatal("a store built with ada-002 was accepted for 3-small") + } + if !strings.Contains(err.Error(), "text-embedding-ada-002") { + t.Errorf("error does not name the model in the store: %v", err) + } + + // Different width, caught by the column definition. + if err := verifyStore(pool, table, "text-embedding-3-large"); err == nil { + t.Error("a vector(1536) store was accepted for a 3072-wide model") + } +} + +func TestVerifyStoreOnLegacyAndOddTables(t *testing.T) { + pool, _ := testPool(t) + ctx := context.Background() + + t.Run("missing table", func(t *testing.T) { + _, _ = pool.Exec(ctx, "DROP TABLE IF EXISTS rag_absent_docs") + if err := verifyStore(pool, "rag_absent_docs", "text-embedding-3-small"); err == nil { + t.Error("a table that does not exist verified clean") + } + }) + + t.Run("no embedding_model column", func(t *testing.T) { + const table = "rag_legacy_docs" + createTable(t, pool, `CREATE TABLE `+table+` ( + id SERIAL PRIMARY KEY, + content TEXT NOT NULL, + embedding vector(1536), + source TEXT + )`, table) + + err := verifyStore(pool, table, "text-embedding-3-small") + if err == nil { + t.Fatal("an unmigrated table verified clean") + } + if !strings.Contains(err.Error(), "ALTER TABLE") { + t.Errorf("error carries no migration: %v", err) + } + }) + + t.Run("bare vector column takes any width", func(t *testing.T) { + const table = "rag_bare_vector_docs" + createTable(t, pool, `CREATE TABLE `+table+` ( + id SERIAL PRIMARY KEY, + content TEXT NOT NULL, + embedding vector, + embedding_model TEXT NOT NULL + )`, table) + + // atttypmod is -1 here, so there is no declared width to disagree with. + if err := verifyStore(pool, table, "text-embedding-3-large"); err != nil { + t.Errorf("undeclared width rejected: %v", err) + } + }) + + t.Run("null model reads as unmigrated rows", func(t *testing.T) { + const table = "rag_null_model_docs" + createTable(t, pool, `CREATE TABLE `+table+` ( + id SERIAL PRIMARY KEY, + content TEXT NOT NULL, + embedding vector(1536), + embedding_model TEXT + )`, table) + + vec := make([]float32, 1536) + if _, err := pool.Exec(ctx, + "INSERT INTO "+table+" (content, embedding) VALUES ($1, $2::vector)", "chunk", formatVector(vec)); err != nil { + t.Fatalf("insert: %v", err) + } + + err := verifyStore(pool, table, "text-embedding-3-small") + if err == nil || !strings.Contains(err.Error(), "(null)") { + t.Errorf("null embedding_model = %v, want a mismatch naming (null)", err) + } + }) +} + +func TestNewPgvectorClientRefusesABadStore(t *testing.T) { + _, connString := testPool(t) + pool, _ := testPool(t) + const table = "rag_client_docs" + createTable(t, pool, documentedDDL(table, 1536), table) + insertRow(t, pool, table, 1536, "text-embedding-ada-002") + + cfg := &config.Config{} + cfg.Pgvector.ConnectionString = connString + cfg.Pgvector.Table = table + cfg.RAG.EmbeddingModel = "text-embedding-3-small" + + // main.go turns this error into log.Fatalf, so a mismatch here is the + // server declining to boot. + if _, err := NewPgvectorClient(cfg); err == nil { + t.Error("client constructed against a store built with another model") + } + + cfg.RAG.EmbeddingModel = "text-embedding-ada-002" + client, err := NewPgvectorClient(cfg) + if err != nil { + t.Fatalf("client rejected its own store: %v", err) + } + _ = client +} diff --git a/internal/rag/supabase.go b/internal/rag/supabase.go index 4cc6ba7..37f4c92 100644 --- a/internal/rag/supabase.go +++ b/internal/rag/supabase.go @@ -6,7 +6,10 @@ import ( "encoding/json" "fmt" "io" + "log" "net/http" + "net/url" + "strings" "time" "github.com/streamcoreai/streamcore-server/internal/config" @@ -59,16 +62,115 @@ func NewSupabaseClient(cfg *config.Config) (Client, error) { topK = 3 } + table := cfg.Supabase.Table + if table == "" { + table = "documents" + } + + httpClient := &http.Client{Timeout: 10 * time.Second} + if err := verifySupabaseStore(httpClient, cfg.Supabase.URL, cfg.Supabase.APIKey, table, cfg.RAG.EmbeddingModel); err != nil { + return nil, fmt.Errorf("supabase: %w", err) + } + return &supabaseClient{ url: cfg.Supabase.URL, apiKey: cfg.Supabase.APIKey, function: fn, embedder: newEmbeddingClient(cfg.OpenAI.APIKey, cfg.RAG.EmbeddingModel), topK: topK, - client: &http.Client{Timeout: 10 * time.Second}, + client: httpClient, }, nil } +// verifySupabaseStore runs the same two checks as the pgvector client, against +// the only surface PostgREST offers: rows. There is no catalog to introspect, +// so the width comes from a sampled embedding and the model check is a filter +// that asks the database for one disagreeing row. A project that is down or +// still empty leaves the store unverified rather than blocking startup; a +// mismatch, or a table predating the embedding_model column, is fatal. +func verifySupabaseStore(client *http.Client, baseURL, apiKey, table, model string) error { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + base := strings.TrimSuffix(baseURL, "/") + "/rest/v1/" + url.PathEscape(table) + want := normalizeModel(model) + + // A row whose model differs, or is null. PostgREST's neq skips nulls, so the + // null arm has to be spelled out. + query := fmt.Sprintf("?select=embedding,embedding_model&or=(embedding_model.neq.%s,embedding_model.is.null)&limit=1", + url.QueryEscape(want)) + rows, err := supabaseSelect(ctx, client, base+query, apiKey) + if err != nil { + if isMissingColumn(err) { + return missingModelColumn(table) + } + log.Printf("Warning: RAG store check skipped — %s: %v", table, err) + return nil + } + if len(rows) > 0 { + return modelMismatch(table, want, rows[0].EmbeddingModel) + } + + // Every row agrees on the model, so any row will do for the width. + rows, err = supabaseSelect(ctx, client, base+"?select=embedding,embedding_model&limit=1", apiKey) + if err != nil { + log.Printf("Warning: RAG dimension check skipped — %s: %v", table, err) + return nil + } + if len(rows) == 0 { + return nil // nothing ingested yet + } + + wantDims, known := dimensionsFor(model) + got, ok := vectorWidth(rows[0].Embedding) + if !known || !ok { + return nil + } + if got != wantDims { + return dimensionMismatch(table, model, wantDims, got) + } + return nil +} + +type supabaseRow struct { + Embedding json.RawMessage `json:"embedding"` + EmbeddingModel string `json:"embedding_model"` +} + +func supabaseSelect(ctx context.Context, client *http.Client, endpoint, apiKey string) ([]supabaseRow, error) { + req, err := http.NewRequestWithContext(ctx, "GET", endpoint, nil) + if err != nil { + return nil, err + } + req.Header.Set("apikey", apiKey) + req.Header.Set("Authorization", "Bearer "+apiKey) + + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 512)) + return nil, fmt.Errorf("returned %d: %s", resp.StatusCode, strings.TrimSpace(string(body))) + } + + var rows []supabaseRow + if err := json.NewDecoder(resp.Body).Decode(&rows); err != nil { + return nil, fmt.Errorf("decode response: %w", err) + } + return rows, nil +} + +// isMissingColumn separates "your table predates embedding_model" from every +// other reason a select can fail. PostgREST reports it as 42703 in the body. +func isMissingColumn(err error) bool { + msg := err.Error() + return strings.Contains(msg, "42703") || + (strings.Contains(msg, "embedding_model") && strings.Contains(msg, "does not exist")) +} + type supabaseRPCRequest struct { QueryEmbedding []float32 `json:"query_embedding"` MatchCount int `json:"match_count"` @@ -114,6 +216,9 @@ func (c *supabaseClient) Search(ctx context.Context, query string, topK int) ([] if resp.StatusCode != http.StatusOK { respBody, _ := io.ReadAll(resp.Body) + if strings.Contains(string(respBody), "dimensions") { + return nil, fmt.Errorf("supabase rpc returned %d: %s — %s stores vectors of a different width than embedding_model %q produces; re-ingest with streamcore-cli or change the model back", resp.StatusCode, string(respBody), c.function, c.embedder.model) + } return nil, fmt.Errorf("supabase rpc returned %d: %s", resp.StatusCode, string(respBody)) }