Repository navigation
Expand file tree
/
Copy pathvector_store.go
More file actions
241 lines (204 loc) · 5.74 KB
/
Copy pathvector_store.go
File metadata and controls
241 lines (204 loc) · 5.74 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
package codegraph
import (
"context"
"fmt"
)
// VectorStore provides zero-dependency vector search using chromem-go.
// This is a pure Go implementation that requires no external services.
// Based on research: chromem-go provides ChromaDB-compatible vector search
// with zero CGO dependencies, making it ideal for cross-platform builds.
type VectorStore struct {
collection *Collection
dim int
}
// Collection represents a vector collection (simplified for non-CGO builds).
type Collection struct {
name string
vectors map[string]Vector
}
// Vector represents a stored vector.
type Vector struct {
ID string
Embedding []float32
Metadata map[string]string
}
// NewVectorStore creates a new vector store.
func NewVectorStore(dim int) *VectorStore {
return &VectorStore{
collection: &Collection{
name: "codegraph",
vectors: make(map[string]Vector),
},
dim: dim,
}
}
// Add inserts a vector into the store.
func (vs *VectorStore) Add(ctx context.Context, id string, embedding []float32, metadata map[string]string) error {
if len(embedding) != vs.dim {
return fmt.Errorf("embedding dimension mismatch: expected %d, got %d", vs.dim, len(embedding))
}
vs.collection.vectors[id] = Vector{
ID: id,
Embedding: embedding,
Metadata: metadata,
}
return nil
}
// Search finds the k nearest neighbors to the query vector.
func (vs *VectorStore) Search(ctx context.Context, query []float32, k int) ([]SearchResult, error) {
if len(query) != vs.dim {
return nil, fmt.Errorf("query dimension mismatch: expected %d, got %d", vs.dim, len(query))
}
type scored struct {
id string
score float32
meta map[string]string
}
var results []scored
for id, vec := range vs.collection.vectors {
sim := CosineSimilarity(query, vec.Embedding)
if sim > 0.1 { // threshold
results = append(results, scored{id: id, score: sim, meta: vec.Metadata})
}
}
// Sort by score descending
for i := 0; i < len(results); i++ {
for j := i + 1; j < len(results); j++ {
if results[j].score > results[i].score {
results[i], results[j] = results[j], results[i]
}
}
}
// Return top k
if k > len(results) {
k = len(results)
}
searchResults := make([]SearchResult, k)
for i := 0; i < k; i++ {
searchResults[i] = SearchResult{
ID: results[i].id,
Score: results[i].score,
Metadata: results[i].meta,
}
}
return searchResults, nil
}
// Delete removes a vector from the store.
func (vs *VectorStore) Delete(ctx context.Context, id string) error {
delete(vs.collection.vectors, id)
return nil
}
// Count returns the number of vectors in the store.
func (vs *VectorStore) Count(ctx context.Context) int {
return len(vs.collection.vectors)
}
// SearchResult represents a search result.
type SearchResult struct {
ID string
Score float32
Metadata map[string]string
}
// CodeVectorStore extends VectorStore with code-specific functionality.
type CodeVectorStore struct {
store *VectorStore
nodes map[string]Node
}
// NewCodeVectorStore creates a vector store for code symbols.
func NewCodeVectorStore() *CodeVectorStore {
return &CodeVectorStore{
store: NewVectorStore(EmbeddingDimension),
nodes: make(map[string]Node),
}
}
// IndexNode adds a code symbol to the vector store.
func (cvs *CodeVectorStore) IndexNode(ctx context.Context, node Node) error {
embedding := GenerateEmbedding(node)
metadata := map[string]string{
"name": node.Name,
"kind": node.Kind,
"file": node.FilePath,
"language": node.Language,
}
if err := cvs.store.Add(ctx, node.ID, embedding, metadata); err != nil {
return err
}
cvs.nodes[node.ID] = node
return nil
}
// SearchCode finds code symbols similar to the query.
func (cvs *CodeVectorStore) SearchCode(ctx context.Context, query string, limit int) ([]Node, error) {
queryNode := Node{
Name: query,
QualifiedName: query,
Kind: "query",
Language: "query",
}
queryEmbedding := GenerateEmbedding(queryNode)
results, err := cvs.store.Search(ctx, queryEmbedding, limit)
if err != nil {
return nil, err
}
nodes := make([]Node, 0, len(results))
for _, r := range results {
if node, ok := cvs.nodes[r.ID]; ok {
nodes = append(nodes, node)
}
}
return nodes, nil
}
// HybridSearch combines vector search with keyword search.
func (cvs *CodeVectorStore) HybridSearch(ctx context.Context, query string, limit int, keywordResults []Node) ([]Node, error) {
// Get vector results
vectorResults, err := cvs.SearchCode(ctx, query, limit*2)
if err != nil {
return nil, err
}
// Merge with RRF (Reciprocal Rank Fusion)
k := 60.0
scores := make(map[string]float64)
nodeMap := make(map[string]Node)
// Add keyword results
for rank, node := range keywordResults {
scores[node.ID] += 1.0 / (k + float64(rank+1))
nodeMap[node.ID] = node
}
// Add vector results
for rank, node := range vectorResults {
scores[node.ID] += 1.0 / (k + float64(rank+1))
nodeMap[node.ID] = node
}
// Sort by RRF score
type scored struct {
id string
score float64
}
var sorted []scored
for id, score := range scores {
sorted = append(sorted, scored{id, score})
}
for i := 0; i < len(sorted); i++ {
for j := i + 1; j < len(sorted); j++ {
if sorted[j].score > sorted[i].score {
sorted[i], sorted[j] = sorted[j], sorted[i]
}
}
}
// Return top results
if limit > len(sorted) {
limit = len(sorted)
}
nodes := make([]Node, limit)
for i := 0; i < limit; i++ {
nodes[i] = nodeMap[sorted[i].id]
}
return nodes, nil
}
// Stats returns vector store statistics.
func (cvs *CodeVectorStore) Stats() map[string]interface{} {
return map[string]interface{}{
"vectors": cvs.store.Count(context.Background()),
"nodes": len(cvs.nodes),
"dim": EmbeddingDimension,
}
}
// Helper to check if string contains substring