feat: add kb/upsert and kb/query connectors for RAG#169
Conversation
Native pgvector convenience connectors that compose with ai/embed, as thin sugar over a bring-your-own Postgres database (the step credential). - kb/upsert: store document text + embedding (+ optional JSONB metadata) into a pgvector table. Takes ai/embed's pgvector literal directly, supports single or batch (chunk) inserts via unnest, and idempotent ON CONFLICT (conflict_target) DO NOTHING. - kb/query: distance-ordered nearest-neighbour search (cosine/l2/ inner_product), returning the closest rows + a distance field. - Table/column names are validated as SQL identifiers to prevent injection (they can't be bound as parameters). The SQL-building logic is factored into pure prepareUpsert/prepareQuery functions with thorough unit tests (happy paths, arg alignment, injection rejection) — no DB needed. - Reuses postgres.go's connection + exec/select helpers; single postgres credential per step (no engine-managed store). - RAG example (rag-ingest/rag-ask/rag-kb-schema) and guide rewritten to use kb/*; connector reference documents both. Both workflows pass `mantle validate`. They don't manage schema; native chunking and metadata filtering on kb/query remain follow-ups on #153. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79
|
Warning Review limit reached
Next review available in: 50 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (3)
📝 WalkthroughWalkthroughIntroduces two new pgvector-backed Postgres connectors, ChangesKB Connectors Implementation and Documentation
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant Workflow
participant KBUpsertConnector
participant prepareUpsert
participant Postgres
Workflow->>KBUpsertConnector: Execute(params with content, vector, metadata)
KBUpsertConnector->>prepareUpsert: build INSERT SQL and args
prepareUpsert-->>KBUpsertConnector: SQL, args
KBUpsertConnector->>Postgres: pgx connect and exec
Postgres-->>KBUpsertConnector: rows_affected
KBUpsertConnector-->>Workflow: result map
sequenceDiagram
participant Workflow
participant KBQueryConnector
participant prepareQuery
participant Postgres
Workflow->>KBQueryConnector: Execute(params with vector, top_k, metric)
KBQueryConnector->>prepareQuery: build SELECT SQL ordered by distance
prepareQuery-->>KBQueryConnector: SQL, args
KBQueryConnector->>Postgres: pgx connect and select
Postgres-->>KBQueryConnector: rows, row_count
KBQueryConnector-->>Workflow: result map
Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2297623fc7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if err != nil { | ||
| return nil, fmt.Errorf("kb/upsert: %w", err) | ||
| } | ||
| delete(params, "_credential") |
There was a problem hiding this comment.
When a kb/upsert or kb/query step has retry.max_attempts > 1, executeStepLogic resolves _credential once and then reuses the same resolvedParams map for every attempt (engine.go resolves at line 570 and calls Execute at line 663). Deleting _credential here mutates that shared map, so after any first attempt that reaches this line and then fails, such as a transient connect or query error, the next attempt fails immediately with a missing credential instead of retrying the database operation; avoid mutating params here, and apply the same fix to KBQueryConnector below.
Useful? React with 👍 / 👎.
Addresses PR review (Codex P2). The engine resolves a step's params once and reuses that map across retry attempts, so `delete(params, "_credential")` in kb/upsert and kb/query made attempt 2+ fail with a missing credential instead of retrying a transient DB error. Read the credential without mutating the shared map. Adds regression tests asserting _credential survives a failed Execute. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79
|
Good catch — fixed in Heads-up: Generated by Claude Code |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (5)
packages/engine/internal/connector/kb_test.go (2)
64-85: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the metadata/metadatas mutual-exclusivity error.
kbMetadataexplicitly errors when bothmetadataandmetadatasare set (if hasSingle && hasMany { return ... "set only one of metadata or metadatas" }), but no test case exercises this branch. Given the surrounding table already tests the analogouscontent/contentsconflict (line 71), consider adding a sibling case for the metadata param, and optionally one forvector/vectorsboth being set, since that ambiguity path isn't covered either.♻️ Suggested additional test cases
{"count mismatch", map[string]any{"table": "t", "contents": []any{"a", "b"}, "vectors": []any{"[1]"}}}, {"metadata mismatch", map[string]any{"table": "t", "content": "a", "vector": "[1]", "metadatas": []any{map[string]any{}, map[string]any{}}}}, + {"both metadata forms", map[string]any{"table": "t", "content": "a", "vector": "[1]", "metadata": map[string]any{}, "metadatas": []any{map[string]any{}}}}, + {"both vector forms", map[string]any{"table": "t", "content": "a", "vector": "[1]", "vectors": []any{"[1]"}}}, {"sql injection in table", map[string]any{"table": "kb; DROP TABLE users;--", "content": "x", "vector": "[1]"}},🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/engine/internal/connector/kb_test.go` around lines 64 - 85, Add test coverage in TestPrepareUpsert_Errors for the kbMetadata mutual-exclusivity branch by including a case where both metadata and metadatas are set, and assert prepareUpsert returns an error. Use the existing table-driven pattern alongside the content/contents conflict so the new case is easy to locate, and optionally add a similar case for vector and vectors if you want to cover that ambiguity path too.
137-156: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueApprove; optional boundary case for
top_k.Error cases look solid (missing params, bad metric, negative top_k, injection). Since
TestPrepareQuery_TopKCapverifies the upper clamp at 1000, consider also asserting behavior at the lower boundary (top_k: 0) to confirm whether it's treated as "use default" or rejected — currently untested.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/engine/internal/connector/kb_test.go` around lines 137 - 156, Add a boundary test for `prepareQuery` to cover `top_k: 0`, since `TestPrepareQuery_Errors` already checks negative values and `TestPrepareQuery_TopKCap` covers the upper cap. Extend the existing `TestPrepareQuery_Errors` table or add a dedicated case so it asserts the intended behavior for `top_k` zero in `prepareQuery`, confirming whether it is rejected or treated as the default.packages/engine/internal/connector/kb.go (2)
24-48: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winIdentifier regex works well against injection, but blocks legitimate composite conflict targets.
kbIdentReonly matches a single (optionally schema-qualified) identifier. Applied toconflict_targetat Line 165, this rejects any composite unique-key conflict target (e.g."doc_id, chunk_id"), which is a common real-world case for dedupe keys spanning multiple columns.Consider relaxing validation for
conflict_targetto accept a comma-separated list of valid identifiers (each individually checked againstkbIdentRe), similar to howcolumnsis handled at Lines 236-246.♻️ Proposed fix to support composite conflict targets
var conflict string if ct, ok := params["conflict_target"].(string); ok && ct != "" { - if !kbIdentRe.MatchString(ct) { - return "", nil, fmt.Errorf("conflict_target %q is not a valid SQL identifier", ct) - } - conflict = fmt.Sprintf(" ON CONFLICT (%s) DO NOTHING", ct) + parts := strings.Split(ct, ",") + for i, p := range parts { + p = strings.TrimSpace(p) + if !kbIdentRe.MatchString(p) { + return "", nil, fmt.Errorf("conflict_target %q is not a valid SQL identifier", ct) + } + parts[i] = p + } + conflict = fmt.Sprintf(" ON CONFLICT (%s) DO NOTHING", strings.Join(parts, ", ")) }Also applies to: 163-169, 242-246
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/engine/internal/connector/kb.go` around lines 24 - 48, The identifier validation in kbIdentParam is too strict for conflict_target and blocks valid composite unique keys. Update the conflict_target handling in the kb.go path that reads params so it accepts a comma-separated list of identifiers, validating each piece with kbIdentRe instead of treating the whole string as one identifier. Keep the existing single-identifier behavior for other params, and mirror the list-splitting approach already used for columns.
186-198: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueError message omits documented metric aliases.
kbDistanceOperatoracceptseuclideanandipas aliases forl2/inner_product, but the error message on Line 196 only listscosine, l2, or inner_product. Minor inconsistency for users who pass an invalid metric and read the error to self-correct.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/engine/internal/connector/kb.go` around lines 186 - 198, The kbDistanceOperator error message is missing the accepted alias names, so update the default branch in kbDistanceOperator to mention all documented valid metrics and aliases, including euclidean and ip alongside cosine, l2, and inner_product. Keep the existing switch behavior unchanged; just adjust the fmt.Errorf message so callers who pass an invalid metric can self-correct using the full set of accepted values.packages/site/src/content/examples/rag-kb-schema.sql (1)
32-33: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winOnly cosine has a matching index;
l2/inner_productwill sequential-scan.
kb/query'smetricparam supportsl2andinner_product, but this schema only builds avector_cosine_opsHNSW index. Queries using the other metrics will silently fall back to a sequential scan instead of using an index.📝 Suggested comment addition
CREATE INDEX IF NOT EXISTS idx_kb_documents_embedding ON kb_documents USING hnsw (embedding vector_cosine_ops); +-- Note: this indexes cosine distance only (kb/query's default metric). +-- If you use metric: l2 or metric: inner_product, add a matching index +-- (vector_l2_ops / vector_ip_ops) or those queries will sequential-scan.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/site/src/content/examples/rag-kb-schema.sql` around lines 32 - 33, The kb_documents HNSW index only covers vector_cosine_ops, so kb/query requests using metric values like l2 or inner_product can’t use an index. Update the rag-kb-schema.sql example to either add matching HNSW indexes for the other supported metrics or clearly note in the schema/comment that only cosine is indexed; use the kb/query metric parameter and the existing CREATE INDEX on kb_documents as the reference points.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/engine/internal/connector/kb.go`:
- Around line 297-301: Add a timeout around the pgx.Connect call in kb/upsert
and the matching kb/query path, since they currently pass the step context
through unchanged and can hang forever on a stalled Postgres handshake. Wrap the
existing ctx with a short connection timeout before calling pgx.Connect, and
keep the existing error handling and conn.Close(ctx) flow in the kb connector
logic.
---
Nitpick comments:
In `@packages/engine/internal/connector/kb_test.go`:
- Around line 64-85: Add test coverage in TestPrepareUpsert_Errors for the
kbMetadata mutual-exclusivity branch by including a case where both metadata and
metadatas are set, and assert prepareUpsert returns an error. Use the existing
table-driven pattern alongside the content/contents conflict so the new case is
easy to locate, and optionally add a similar case for vector and vectors if you
want to cover that ambiguity path too.
- Around line 137-156: Add a boundary test for `prepareQuery` to cover `top_k:
0`, since `TestPrepareQuery_Errors` already checks negative values and
`TestPrepareQuery_TopKCap` covers the upper cap. Extend the existing
`TestPrepareQuery_Errors` table or add a dedicated case so it asserts the
intended behavior for `top_k` zero in `prepareQuery`, confirming whether it is
rejected or treated as the default.
In `@packages/engine/internal/connector/kb.go`:
- Around line 24-48: The identifier validation in kbIdentParam is too strict for
conflict_target and blocks valid composite unique keys. Update the
conflict_target handling in the kb.go path that reads params so it accepts a
comma-separated list of identifiers, validating each piece with kbIdentRe
instead of treating the whole string as one identifier. Keep the existing
single-identifier behavior for other params, and mirror the list-splitting
approach already used for columns.
- Around line 186-198: The kbDistanceOperator error message is missing the
accepted alias names, so update the default branch in kbDistanceOperator to
mention all documented valid metrics and aliases, including euclidean and ip
alongside cosine, l2, and inner_product. Keep the existing switch behavior
unchanged; just adjust the fmt.Errorf message so callers who pass an invalid
metric can self-correct using the full set of accepted values.
In `@packages/site/src/content/examples/rag-kb-schema.sql`:
- Around line 32-33: The kb_documents HNSW index only covers vector_cosine_ops,
so kb/query requests using metric values like l2 or inner_product can’t use an
index. Update the rag-kb-schema.sql example to either add matching HNSW indexes
for the other supported metrics or clearly note in the schema/comment that only
cosine is indexed; use the kb/query metric parameter and the existing CREATE
INDEX on kb_documents as the reference points.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 764978be-aca8-49b5-abeb-99f38f47b83f
📒 Files selected for processing (9)
.changeset/kb-connectors.mdpackages/engine/internal/connector/connector.gopackages/engine/internal/connector/kb.gopackages/engine/internal/connector/kb_test.gopackages/site/src/content/docs/rag-guide.mdpackages/site/src/content/docs/workflow-reference/connectors.mdpackages/site/src/content/examples/rag-ask.yamlpackages/site/src/content/examples/rag-ingest.yamlpackages/site/src/content/examples/rag-kb-schema.sql
…ests) CodeRabbit review follow-ups on the kb/* connectors: - Bound pgx.Connect with a 15s connect timeout (Major) so a stalled Postgres handshake can't hang a step that has no timeout of its own; it only shortens an existing deadline. - conflict_target accepts a comma-separated composite key (each part validated as an identifier), not just a single column. - kbDistanceOperator error message lists the l2/euclidean and inner_product/ip aliases it accepts. - Tests: both-metadata-forms, both-vector-forms, and top_k: 0 error cases. - Schema example: note that only cosine is indexed; l2/inner_product need a matching index or they seq-scan. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79
|
Addressed the review in
Build, vet, gofmt, and golangci-lint clean. Generated by Claude Code |
michaelmcnees
deleted the
claude/office-assistant-mantle-feasibility-n8f4bk
branch
2 participants
Summary
Continues RAG (#153): native
kb/upsertandkb/queryconnectors so users don't hand-write pgvector SQL. Per the agreed design, they're thin sugar over a bring-your-own Postgres database (the step'scredential) and compose with the existingai/embed— no engine-managed store, single credential per step, fitting Mantle's BYO-Postgres ethos.Changes
kb/upsert— stores document text + embedding (+ optional JSONB metadata) into a pgvector table. Takesai/embed'soutput.vectorliteral directly, supports single or batch (chunk) inserts viaunnest(solving multi-row ingest without a loop construct), and idempotentON CONFLICT (conflict_target) DO NOTHING.kb/query— distance-ordered nearest-neighbour search (cosinedefault,l2,inner_product), returning the closest rows plus adistancefield.prepareUpsert/prepareQueryfunctions with thorough unit tests (happy paths, arg alignment, and injection-rejection cases) — no DB required.postgres.go's connection + exec/select helpers.rag-ingest.yaml,rag-ask.yaml,rag-kb-schema.sql) and guide are rewritten to usekb/*; the connector reference documents both.Scope
kb/*don't manage schema (you create the pgvector table). Native chunking and metadata filtering onkb/queryremain follow-ups on #153, as do Cohere Bedrock embeddings.Testing
go test ./internal/connector -run TestPrepare— pure builder tests incl. SQL-injection rejection for table/column/conflict identifiers. Build, vet, gofmt, golangci-lint clean.rag-ingest.yamlandrag-ask.yamlpassmantle validate.Related Issues
Part of #153 (RAG subsystem).
Generated by Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Documentation