mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-13 16:14:57 +00:00
fix(pg): align SearchByEmbedding SQL placeholders with scope args (#623)
This commit is contained in:
2 files changed
+45
-10
No files matched your search
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user