diff --git a/internal/store/pg/skills_embedding.go b/internal/store/pg/skills_embedding.go index 97737553..2195fdf0 100644 --- a/internal/store/pg/skills_embedding.go +++ b/internal/store/pg/skills_embedding.go @@ -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 { diff --git a/internal/store/pg/skills_embedding_test.go b/internal/store/pg/skills_embedding_test.go new file mode 100644 index 00000000..94cdb1fa --- /dev/null +++ b/internal/store/pg/skills_embedding_test.go @@ -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) + } +}