diff --git a/internal/cli/search.go b/internal/cli/search.go index 45a97f2..154dc1d 100644 --- a/internal/cli/search.go +++ b/internal/cli/search.go @@ -8,6 +8,7 @@ import ( "github.com/tae2089/code-context-graph/internal/ctxns" "github.com/tae2089/code-context-graph/internal/pathutil" + "github.com/tae2089/code-context-graph/internal/searchrank" ) // newSearchCmd creates the full-text search command. @@ -31,12 +32,10 @@ func newSearchCmd(deps *Deps) *cobra.Command { ns, _ := cmd.Flags().GetString("namespace") ctx = ctxns.WithNamespace(ctx, ns) - fetchLimit := limit - if pathPrefix != "" { - fetchLimit = max(limit*5, 50) - } - - nodes, err := deps.SearchBackend.Query(ctx, deps.DB, query, fetchLimit) + // Over-fetch a wider candidate pool so structural reranking can + // promote good matches FTS ranked below the limit, and so path + // filtering still leaves up to 'limit' results. + nodes, err := deps.SearchBackend.Query(ctx, deps.DB, query, searchrank.FetchLimit(limit)) if err != nil { return trace.Wrap(err, "search") } @@ -51,11 +50,11 @@ func newSearchCmd(deps *Deps) *cobra.Command { } } nodes = filtered - if len(nodes) > limit { - nodes = nodes[:limit] - } } + // Rerank FTS candidates with structural signals, then bound to limit. + nodes = searchrank.Rerank(query, nodes, limit) + if len(nodes) == 0 { fmt.Fprintln(out, "No results") return nil diff --git a/internal/mcp/handler_query.go b/internal/mcp/handler_query.go index be3c334..9ee3a7f 100644 --- a/internal/mcp/handler_query.go +++ b/internal/mcp/handler_query.go @@ -15,6 +15,7 @@ import ( "github.com/tae2089/code-context-graph/internal/ctxns" "github.com/tae2089/code-context-graph/internal/model" "github.com/tae2089/code-context-graph/internal/pathutil" + "github.com/tae2089/code-context-graph/internal/searchrank" ) var strictFalse = false @@ -22,9 +23,6 @@ var strictFalse = false const ( defaultQueryGraphLimit = 50 maxQueryGraphLimit = 500 - searchPathFetchFactor = 5 - searchPathFetchFloor = 50 - searchPathFetchCap = 500 ) // annotationTagItem serializes one stored annotation tag. @@ -197,14 +195,10 @@ func (h *handlers) search(ctx context.Context, request mcp.CallToolRequest) (*mc } return finalizeToolResult(h.cachedExecute(ctx, "search:", map[string]any{"query": query, "limit": limit, "path": pathPrefix, "namespace": requestNamespace(request)}, func() (string, error) { - // When path filtering is active, fetch more results from FTS so - // that after filtering we still have up to 'limit' results. - fetchLimit := limit - if pathPrefix != "" { - fetchLimit = min(max(limit*searchPathFetchFactor, searchPathFetchFloor), searchPathFetchCap) - } - - nodes, err := h.deps.SearchBackend.Query(ctx, h.deps.DB, query, fetchLimit) + // Over-fetch a wider candidate pool so structural reranking can promote + // good matches that FTS ranked below the caller's limit, and so path + // filtering still leaves up to 'limit' results. + nodes, err := h.deps.SearchBackend.Query(ctx, h.deps.DB, query, searchrank.FetchLimit(limit)) if err != nil { log.Error("search error", "query", query, trace.SlogError(err)) return "", trace.Wrap(err, "search error") @@ -218,11 +212,11 @@ func (h *handlers) search(ctx context.Context, request mcp.CallToolRequest) (*mc } } nodes = filtered - if len(nodes) > limit { - nodes = nodes[:limit] - } } + // Rerank FTS candidates with structural signals, then bound to limit. + nodes = searchrank.Rerank(query, nodes, limit) + log.Info("search completed", "query", query, "result_count", len(nodes)) searchResult := make([]searchResultItem, len(nodes)) diff --git a/internal/searchrank/searchrank.go b/internal/searchrank/searchrank.go new file mode 100644 index 0000000..d8780f1 --- /dev/null +++ b/internal/searchrank/searchrank.go @@ -0,0 +1,261 @@ +// Package searchrank reranks FTS-ranked code-search candidates using +// dependency-free structural signals (name fuzzy similarity, path proximity), +// so both the CLI `search` command and the MCP `search` tool share one ranking. +package searchrank + +import ( + "sort" + "strings" + "unicode" + + "github.com/tae2089/code-context-graph/internal/model" +) + +// Reciprocal Rank Fusion constants. +// rrfK dampens how sharply rank position affects the fused score (60 is the +// conventional value). rrfStructWeight (>1) lets the structural signal override +// a small FTS-rank gap; at weight 1 RRF is symmetric and a two-item rank swap +// always ties, leaving FTS order untouched. +const ( + rrfK = 60.0 + rrfStructWeight = 2.0 +) + +// Structural signal weights. Name match is more discriminating than path +// proximity: every node in one file shares the path, so an unweighted max lets +// a path hit saturate the score and bury the exact-name match. Weighting name +// above path keeps name the dominant signal with path as a nudge. +const ( + nameSignalWeight = 1.0 + pathSignalWeight = 0.25 +) + +// Candidate-pool sizing. Callers over-fetch a wider pool than the requested +// limit so reranking has more than `limit` rows to reorder before bounding. +const ( + fetchFactor = 5 + fetchFloor = 50 + fetchCap = 500 +) + +// FetchLimit widens the candidate pool pulled from FTS so structural reranking +// (and any path filtering) has more than the caller's `limit` rows to reorder; +// the final slice is bounded back to `limit` after reranking. +func FetchLimit(limit int) int { + return min(max(limit*fetchFactor, fetchFloor), fetchCap) +} + +// Rerank reorders FTS-ranked search candidates using structural signals fused +// with the backend rank via Reciprocal Rank Fusion. +// +// @requires nodes is the backend's rank-ordered candidate slice (index == FTS rank). +// @ensures deterministic output; empty query or empty nodes returns the input +// bounded by limit, preserving FTS order. +func Rerank(query string, nodes []model.Node, limit int) []model.Node { + if strings.TrimSpace(query) == "" || len(nodes) == 0 { + return applyLimit(nodes, limit) + } + qTokens := tokenize(query) + if len(qTokens) == 0 { + return applyLimit(nodes, limit) + } + + structScores := make([]float64, len(nodes)) + for i := range nodes { + structScores[i] = structScore(qTokens, nodes[i]) + } + structRank := rankDesc(structScores) + + final := make([]float64, len(nodes)) + for i := range nodes { + final[i] = 1.0/(rrfK+float64(i)) + rrfStructWeight/(rrfK+float64(structRank[i])) + } + + order := make([]int, len(nodes)) + for i := range order { + order[i] = i + } + // Stable sort by fused score descending; equal scores keep the original + // FTS order, so ranking stays deterministic. + sort.SliceStable(order, func(a, b int) bool { + return final[order[a]] > final[order[b]] + }) + + out := make([]model.Node, len(nodes)) + for pos, idx := range order { + out[pos] = nodes[idx] + } + return applyLimit(out, limit) +} + +// applyLimit bounds the result slice, treating a non-positive limit as unbounded. +func applyLimit(nodes []model.Node, limit int) []model.Node { + if limit > 0 && len(nodes) > limit { + return nodes[:limit] + } + return nodes +} + +// structScore is a weighted sum of the name and path signals, with name +// dominant so an exact-name match outranks a mere same-file (path) match. +func structScore(qTokens []string, node model.Node) float64 { + return nameSignalWeight*nameSim(qTokens, node) + pathSignalWeight*pathScore(qTokens, node) +} + +// nameSim scores fuzzy similarity of the query against the node name and the +// last segment of its qualified name. For each target it takes the stronger of: +// - token-level: every query token's best match against the whole name or any +// of its identifier sub-tokens (so "user" or "id" matches getUserById), and +// - joined-whole: the run-together query vs the whole name (so a typo like +// "getUsrById" still matches getUserById). +func nameSim(qTokens []string, node model.Node) float64 { + joined := strings.Join(qTokens, "") + rawTargets := []string{node.Name, lastSegment(node.QualifiedName, '.')} + best := 0.0 + for _, raw := range rawTargets { + if raw == "" { + continue + } + lower := strings.ToLower(raw) + subs := splitIdentifier(raw) // original case: camelCase boundaries matter + sum := 0.0 + for _, tok := range qTokens { + b := normLevSim(tok, lower) + for _, st := range subs { + b = max(b, normLevSim(tok, st)) + } + sum += b + } + cand := max(sum/float64(len(qTokens)), normLevSim(joined, lower)) + best = max(best, cand) + } + return best +} + +// splitIdentifier breaks an identifier into lowercased sub-tokens on separators, +// camelCase boundaries, and letter/digit transitions ("getUserById" -> get, +// user, by, id; "HTTPServer" -> http, server; "parseHTML5" -> parse, html, 5). +func splitIdentifier(s string) []string { + runes := []rune(s) + var tokens []string + var cur []rune + flush := func() { + if len(cur) > 0 { + tokens = append(tokens, strings.ToLower(string(cur))) + cur = cur[:0] + } + } + for i, r := range runes { + if !isAlnum(r) { + flush() + continue + } + if len(cur) > 0 { + prev := cur[len(cur)-1] + switch { + case unicode.IsUpper(r) && (unicode.IsLower(prev) || unicode.IsDigit(prev)): + flush() // lower/digit -> Upper: new word + case unicode.IsUpper(r) && unicode.IsUpper(prev) && i+1 < len(runes) && unicode.IsLower(runes[i+1]): + flush() // acronym tail begins a new word (HTTPServer -> http, server) + case unicode.IsDigit(r) != unicode.IsDigit(prev): + flush() // letter/digit transition + } + } + cur = append(cur, r) + } + flush() + return tokens +} + +func isAlnum(r rune) bool { + return unicode.IsLetter(r) || unicode.IsDigit(r) +} + +// pathScore is the fraction of query tokens that appear as file-path segments. +func pathScore(qTokens []string, node model.Node) float64 { + segs := map[string]struct{}{} + for _, seg := range strings.FieldsFunc(strings.ToLower(node.FilePath), isPathSep) { + segs[seg] = struct{}{} + } + if len(segs) == 0 { + return 0 + } + hits := 0 + for _, tok := range qTokens { + if _, ok := segs[tok]; ok { + hits++ + } + } + return float64(hits) / float64(len(qTokens)) +} + +// tokenize lowercases and splits input into alphanumeric tokens. +func tokenize(s string) []string { + return strings.FieldsFunc(strings.ToLower(s), func(r rune) bool { + return !unicode.IsLetter(r) && !unicode.IsDigit(r) + }) +} + +func isPathSep(r rune) bool { + switch r { + case '/', '.', '_', '-': + return true + default: + return false + } +} + +func lastSegment(s string, sep rune) string { + if i := strings.LastIndexByte(s, byte(sep)); i >= 0 { + return s[i+1:] + } + return s +} + +// rankDesc returns each index's position when scores are ordered descending; +// ties resolve by original index so the ranking is deterministic. +func rankDesc(scores []float64) []int { + order := make([]int, len(scores)) + for i := range order { + order[i] = i + } + sort.SliceStable(order, func(a, b int) bool { + return scores[order[a]] > scores[order[b]] + }) + rank := make([]int, len(scores)) + for pos, idx := range order { + rank[idx] = pos + } + return rank +} + +// normLevSim is the Levenshtein distance normalized to a [0,1] similarity. +func normLevSim(a, b string) float64 { + ra, rb := []rune(a), []rune(b) + maxLen := max(len(ra), len(rb)) + if maxLen == 0 { + return 0 + } + return 1.0 - float64(levenshtein(ra, rb))/float64(maxLen) +} + +// levenshtein computes edit distance with a rolling single-row DP. +func levenshtein(a, b []rune) int { + prev := make([]int, len(b)+1) + for j := range prev { + prev[j] = j + } + for i := 1; i <= len(a); i++ { + curr := make([]int, len(b)+1) + curr[0] = i + for j := 1; j <= len(b); j++ { + cost := 1 + if a[i-1] == b[j-1] { + cost = 0 + } + curr[j] = min(prev[j]+1, curr[j-1]+1, prev[j-1]+cost) + } + prev = curr + } + return prev[len(b)] +} diff --git a/internal/searchrank/searchrank_test.go b/internal/searchrank/searchrank_test.go new file mode 100644 index 0000000..1744f84 --- /dev/null +++ b/internal/searchrank/searchrank_test.go @@ -0,0 +1,192 @@ +package searchrank + +import ( + "testing" + + "github.com/tae2089/code-context-graph/internal/model" +) + +// 경계: 빈 쿼리 또는 빈 노드 슬라이스는 입력을 그대로 돌려준다. +func TestRerank_EmptyInputsReturnedUnchanged(t *testing.T) { + nodes := []model.Node{ + {ID: 1, Name: "getUserById", QualifiedName: "svc.getUserById", FilePath: "svc/user.go"}, + {ID: 2, Name: "deleteUser", QualifiedName: "svc.deleteUser", FilePath: "svc/user.go"}, + } + + t.Run("empty query preserves order", func(t *testing.T) { + got := Rerank("", nodes, 10) + assertNodeIDOrder(t, got, []uint{1, 2}) + }) + + t.Run("nil nodes returns empty", func(t *testing.T) { + got := Rerank("user", nil, 10) + if len(got) != 0 { + t.Fatalf("expected empty result, got %d nodes", len(got)) + } + }) +} + +// 오타 쿼리라도 이름 fuzzy 신호가 정확한 심볼을 FTS 하위에서 상위로 끌어올린다. +func TestRerank_TypoPromotesFuzzyNameMatch(t *testing.T) { + nodes := []model.Node{ + {ID: 1, Name: "handleRequest", QualifiedName: "http.handleRequest", FilePath: "http/handler.go"}, + {ID: 2, Name: "getUserById", QualifiedName: "svc.getUserById", FilePath: "svc/user.go"}, + } + got := Rerank("getUsrById", nodes, 10) + if got[0].ID != 2 { + t.Fatalf("expected fuzzy-name match (id=2) promoted to top, got id=%d", got[0].ID) + } +} + +// 경로 세그먼트와 겹치는 쿼리 토큰이 해당 노드를 부스트한다. +func TestRerank_PathProximityBoost(t *testing.T) { + nodes := []model.Node{ + {ID: 1, Name: "helper", QualifiedName: "util.helper", FilePath: "util/helper.go"}, + {ID: 2, Name: "handler", QualifiedName: "svc.handler", FilePath: "svc/auth/login.go"}, + } + got := Rerank("auth login", nodes, 10) + if got[0].ID != 2 { + t.Fatalf("expected path-matching node (id=2) promoted, got id=%d", got[0].ID) + } +} + +// FTS 1위가 구조 신호 0이어도 RRF의 FTS 항 덕분에 완전히 탈락하지 않는다. +func TestRerank_TopFTSHitSurvivesZeroStructSignal(t *testing.T) { + nodes := []model.Node{ + {ID: 1, Name: "alpha", QualifiedName: "pkg.alpha", FilePath: "pkg/alpha.go"}, + {ID: 2, Name: "getUserById", QualifiedName: "svc.getUserById", FilePath: "svc/user.go"}, + {ID: 3, Name: "beta", QualifiedName: "pkg.beta", FilePath: "pkg/beta.go"}, + } + got := Rerank("getUserById", nodes, 2) + if len(got) != 2 { + t.Fatalf("expected 2 results after limit, got %d", len(got)) + } + assertNodeIDOrder(t, got, []uint{2, 1}) +} + +// 구조 신호가 전부 동점이면 원래 FTS 순서를 보존한다(stable). +func TestRerank_TiePreservesFTSOrder(t *testing.T) { + nodes := []model.Node{ + {ID: 1, Name: "same", QualifiedName: "pkg.same", FilePath: "pkg/same.go"}, + {ID: 2, Name: "same", QualifiedName: "pkg.same", FilePath: "pkg/same.go"}, + {ID: 3, Name: "same", QualifiedName: "pkg.same", FilePath: "pkg/same.go"}, + } + got := Rerank("anything", nodes, 10) + assertNodeIDOrder(t, got, []uint{1, 2, 3}) +} + +// limit은 리랭크 후에 적용된다: 승격된 노드가 잘려나가면 안 된다. +func TestRerank_LimitAppliedAfterRerank(t *testing.T) { + nodes := []model.Node{ + {ID: 1, Name: "handleRequest", QualifiedName: "http.handleRequest", FilePath: "http/handler.go"}, + {ID: 2, Name: "getUserById", QualifiedName: "svc.getUserById", FilePath: "svc/user.go"}, + } + got := Rerank("getUsrById", nodes, 1) + assertNodeIDOrder(t, got, []uint{2}) +} + +// 공백으로 나뉜 멀티토큰 쿼리도 심볼 서브토큰과 매칭돼 승격된다. +func TestRerank_MultiTokenNameMatch(t *testing.T) { + nodes := []model.Node{ + {ID: 1, Name: "handleRequest", QualifiedName: "http.handleRequest", FilePath: "http/handler.go"}, + {ID: 2, Name: "getUserById", QualifiedName: "svc.getUserById", FilePath: "svc/user.go"}, + } + got := Rerank("get user by id", nodes, 10) + if got[0].ID != 2 { + t.Fatalf("expected multi-token match (id=2) promoted, got id=%d", got[0].ID) + } +} + +// 이름 안의 서브토큰(camelCase 조각)과 일치하는 짧은 토큰도 강한 신호가 된다. +func TestRerank_SubtokenMatch(t *testing.T) { + nodes := []model.Node{ + {ID: 1, Name: "handleRequest", QualifiedName: "http.handleRequest", FilePath: "http/handler.go"}, + {ID: 2, Name: "getUserById", QualifiedName: "svc.getUserById", FilePath: "svc/user.go"}, + } + got := Rerank("user", nodes, 10) + if got[0].ID != 2 { + t.Fatalf("expected subtoken 'user' match (id=2) promoted, got id=%d", got[0].ID) + } +} + +// 쿼리 토큰이 경로에도 있어 path 신호가 포화될 때, 이름이 일치하는 노드가 +// 단지 같은 파일에 있는 노드보다 위로 올라와야 한다(name > path). +func TestRerank_NameOutranksSaturatedPath(t *testing.T) { + nodes := []model.Node{ + {ID: 1, Name: "isAlnum", QualifiedName: "mcp.isAlnum", FilePath: "mcp/search_rerank.go"}, + {ID: 2, Name: "rerankSearch", QualifiedName: "mcp.rerankSearch", FilePath: "mcp/search_rerank.go"}, + } + got := Rerank("rerank", nodes, 10) + if got[0].ID != 2 { + t.Fatalf("expected name match (id=2) above same-file node, got id=%d", got[0].ID) + } +} + +func TestFetchLimit(t *testing.T) { + if got := FetchLimit(10); got <= 10 { + t.Fatalf("expected fetch limit wider than 10, got %d", got) + } + if got := FetchLimit(200); got != fetchCap { + t.Fatalf("expected fetch limit capped at %d, got %d", fetchCap, got) + } +} + +func TestSplitIdentifier(t *testing.T) { + cases := []struct { + in string + want []string + }{ + {"getUserById", []string{"get", "user", "by", "id"}}, + {"HTTPServer", []string{"http", "server"}}, + {"user_id", []string{"user", "id"}}, + {"parseHTML5", []string{"parse", "html", "5"}}, + } + for _, c := range cases { + got := splitIdentifier(c.in) + if len(got) != len(c.want) { + t.Errorf("splitIdentifier(%q)=%v, want %v", c.in, got, c.want) + continue + } + for i := range c.want { + if got[i] != c.want[i] { + t.Errorf("splitIdentifier(%q)=%v, want %v", c.in, got, c.want) + break + } + } + } +} + +func TestLevenshtein(t *testing.T) { + cases := []struct { + a, b string + want int + }{ + {"", "", 0}, + {"abc", "abc", 0}, + {"abc", "abd", 1}, + {"abc", "abxc", 1}, + {"abc", "ac", 1}, + {"kitten", "sitting", 3}, + } + for _, c := range cases { + if got := levenshtein([]rune(c.a), []rune(c.b)); got != c.want { + t.Errorf("levenshtein(%q,%q)=%d, want %d", c.a, c.b, got, c.want) + } + } +} + +func assertNodeIDOrder(t *testing.T, got []model.Node, wantIDs []uint) { + t.Helper() + if len(got) != len(wantIDs) { + t.Fatalf("length mismatch: got %d, want %d", len(got), len(wantIDs)) + } + for i, id := range wantIDs { + if got[i].ID != id { + gotIDs := make([]uint, len(got)) + for j, n := range got { + gotIDs[j] = n.ID + } + t.Fatalf("order mismatch: got %v, want %v", gotIDs, wantIDs) + } + } +}