150 lines
3.4 KiB
Go
150 lines
3.4 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/json"
|
|
"fmt"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/qdrant/go-client/qdrant"
|
|
)
|
|
|
|
const collectionName = "knowledge"
|
|
|
|
type qdrantStore struct {
|
|
client *qdrant.Client
|
|
embedder *embedder
|
|
}
|
|
|
|
func newQdrantStore(client *qdrant.Client, embedder *embedder) *qdrantStore {
|
|
return &qdrantStore{client: client, embedder: embedder}
|
|
}
|
|
|
|
func (s *qdrantStore) ensureCollection(ctx context.Context) error {
|
|
exists, err := s.client.CollectionExists(ctx, collectionName)
|
|
if err != nil {
|
|
return fmt.Errorf("check collection exists: %w", err)
|
|
}
|
|
if exists {
|
|
return nil
|
|
}
|
|
|
|
err = s.client.CreateCollection(ctx, &qdrant.CreateCollection{
|
|
CollectionName: collectionName,
|
|
VectorsConfig: qdrant.NewVectorsConfig(&qdrant.VectorParams{
|
|
Size: 1536, // text-embedding-3-small dimension
|
|
Distance: qdrant.Distance_Cosine,
|
|
}),
|
|
})
|
|
if err != nil {
|
|
return fmt.Errorf("create collection: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
type storedKnowledge struct {
|
|
ID string `json:"id"`
|
|
Content string `json:"content"`
|
|
Metadata map[string]any `json:"metadata,omitempty"`
|
|
}
|
|
|
|
type searchResult struct {
|
|
ID string `json:"id"`
|
|
Content string `json:"content"`
|
|
Score float32 `json:"score"`
|
|
Metadata map[string]any `json:"metadata,omitempty"`
|
|
}
|
|
|
|
func (s *qdrantStore) store(ctx context.Context, content string, metadata map[string]any) (*storedKnowledge, error) {
|
|
vec, err := s.embedder.embedText(ctx, content)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("embed text: %w", err)
|
|
}
|
|
|
|
id := uuid.New().String()
|
|
payload := map[string]*qdrant.Value{
|
|
"content": qdrant.NewValueString(content),
|
|
}
|
|
|
|
if metadata != nil {
|
|
metaBytes, err := json.Marshal(metadata)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("marshal metadata: %w", err)
|
|
}
|
|
payload["_metadata"] = qdrant.NewValueString(string(metaBytes))
|
|
}
|
|
|
|
idHash := sha256.Sum256([]byte(id))
|
|
idNum := uint64(0)
|
|
for i := 0; i < 8; i++ {
|
|
idNum = (idNum << 8) | uint64(idHash[i])
|
|
}
|
|
|
|
_, err = s.client.Upsert(ctx, &qdrant.UpsertPoints{
|
|
CollectionName: collectionName,
|
|
Points: []*qdrant.PointStruct{
|
|
{
|
|
Id: qdrant.NewIDNum(idNum),
|
|
Vectors: qdrant.NewVectors(vec...),
|
|
Payload: payload,
|
|
},
|
|
},
|
|
})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("upsert point: %w", err)
|
|
}
|
|
|
|
return &storedKnowledge{
|
|
ID: id,
|
|
Content: content,
|
|
Metadata: metadata,
|
|
}, nil
|
|
}
|
|
|
|
func (s *qdrantStore) search(ctx context.Context, query string, limit uint64) ([]searchResult, error) {
|
|
vec, err := s.embedder.embedText(ctx, query)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("embed query: %w", err)
|
|
}
|
|
|
|
if limit == 0 {
|
|
limit = 5
|
|
}
|
|
|
|
results, err := s.client.Query(ctx, &qdrant.QueryPoints{
|
|
CollectionName: collectionName,
|
|
Query: qdrant.NewQueryDense(vec),
|
|
Limit: &limit,
|
|
WithPayload: qdrant.NewWithPayload(true),
|
|
})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("query points: %w", err)
|
|
}
|
|
|
|
var searchResults []searchResult
|
|
for _, pt := range results {
|
|
r := searchResult{
|
|
Score: pt.Score,
|
|
}
|
|
|
|
if payload := pt.GetPayload(); payload != nil {
|
|
if v, ok := payload["content"]; ok {
|
|
r.ID = fmt.Sprintf("%d", pt.GetId().GetNum())
|
|
r.Content = v.GetStringValue()
|
|
}
|
|
if v, ok := payload["_metadata"]; ok {
|
|
var meta map[string]any
|
|
if err := json.Unmarshal([]byte(v.GetStringValue()), &meta); err == nil {
|
|
r.Metadata = meta
|
|
}
|
|
}
|
|
}
|
|
|
|
searchResults = append(searchResults, r)
|
|
}
|
|
|
|
return searchResults, nil
|
|
}
|