fix(pg): align SearchByEmbedding SQL placeholders with scope args (#623)

This commit is contained in:
MK authored and GitHub committed 2026-04-01 18:28:56 +07:00
1 parent d819e08071
commit d334a2082f
2 files changed
+45 -10

No files matched your search

+17 -10
View File
@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"log/slog"
"strings"
"github.com/google/uuid"
@@ -22,16 +23,13 @@ func (s *PGSkillStore) SearchByEmbedding(ctx context.Context, embedding []float3
}
vecStr := vectorToString(embedding)
// $1=vec, $2=vec → tenant at $3 (if needed), ORDER vec at $3+len(tcArgs), LIMIT after
tc, tcArgs, _, err := scopeClause(ctx, 3)
// $1=vec, scope starts at $2 (if present), ORDER vec uses next available param, LIMIT after.
tc, tcArgs, nextParam, err := scopeClause(ctx, 2)
if err != nil {
return nil, err
}
tenantCond := ""
if tc != "" {
tenantCond = fmt.Sprintf(" AND (is_system = true OR tenant_id = $%d)", 3)
}
orderN := 3 + len(tcArgs)
tenantCond := buildSkillEmbeddingTenantCond(tc)
orderN := nextParam
limitN := orderN + 1
q := fmt.Sprintf(`SELECT name, slug, COALESCE(description, ''), version, file_path,
1 - (embedding <=> $1::vector) AS score
@@ -41,9 +39,9 @@ func (s *PGSkillStore) SearchByEmbedding(ctx context.Context, embedding []float3
ORDER BY embedding <=> $%d::vector
LIMIT $%d`, tenantCond, orderN, limitN)
rows, err := s.db.QueryContext(ctx, q,
append(append([]any{vecStr}, tcArgs...), vecStr, limit)...,
)
args := append([]any{vecStr}, tcArgs...)
args = append(args, vecStr, limit)
rows, err := s.db.QueryContext(ctx, q, args...)
if err != nil {
return nil, fmt.Errorf("embedding skill search: %w", err)
}
@@ -68,6 +66,15 @@ func (s *PGSkillStore) SearchByEmbedding(ctx context.Context, embedding []float3
return results, nil
}
func buildSkillEmbeddingTenantCond(scope string) string {
if scope == "" {
return ""
}
tenantExpr := strings.TrimPrefix(scope, " AND ")
return fmt.Sprintf(" AND (is_system = true OR (%s))", tenantExpr)
}
// BackfillSkillEmbeddings generates embeddings for all active skills that don't have one yet.
func (s *PGSkillStore) BackfillSkillEmbeddings(ctx context.Context) (int, error) {
if s.embProvider == nil {
@@ -0,0 +1,28 @@
package pg
import (
"testing"
)
func TestBuildSkillEmbeddingTenantCond_Empty(t *testing.T) {
got := buildSkillEmbeddingTenantCond("")
if got != "" {
t.Fatalf("expected empty tenant condition, got %q", got)
}
}
func TestBuildSkillEmbeddingTenantCond_WithTenant(t *testing.T) {
got := buildSkillEmbeddingTenantCond(" AND tenant_id = $2")
want := " AND (is_system = true OR (tenant_id = $2))"
if got != want {
t.Fatalf("unexpected tenant condition\nwant: %q\n got: %q", want, got)
}
}
func TestBuildSkillEmbeddingTenantCond_WithTenantAndProject(t *testing.T) {
got := buildSkillEmbeddingTenantCond(" AND tenant_id = $2 AND project_id = $3")
want := " AND (is_system = true OR (tenant_id = $2 AND project_id = $3))"
if got != want {
t.Fatalf("unexpected tenant/project condition\nwant: %q\n got: %q", want, got)
}
}