first commit

qdrant.go

@@ -0,0 +1,149 @@
+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
+}