Add rag commands

llm/rag.go

@@ -0,0 +1,71 @@
+package llm
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/sashabaranov/go-openai"
+)
+
+func StoreStringEmbeddingInVectorDB(apiHost string, openaiClient *openai.Client, s string) error {
+	// Example usage
+	client := NewStoreClient(apiHost)
+
+	resp, err := openaiClient.CreateEmbeddings(context.TODO(),
+		openai.EmbeddingRequestStrings{
+			Input: []string{s},
+			Model: openai.AdaEmbeddingV2,
+		},
+	)
+	if err != nil {
+		return fmt.Errorf("error getting keys: %v", err)
+	}
+
+	if len(resp.Data) == 0 {
+		return fmt.Errorf("no response from OpenAI API")
+	}
+
+	embedding := resp.Data[0].Embedding
+
+	setReq := SetRequest{
+		Keys:   [][]float32{embedding},
+		Values: []string{s},
+	}
+	err = client.Set(setReq)
+	if err != nil {
+		return fmt.Errorf("error setting keys: %v", err)
+	}
+
+	return nil
+}
+
+func FindSimilarStrings(apiHost string, openaiClient *openai.Client, s string, similarEntries int) ([]string, error) {
+	client := NewStoreClient(apiHost)
+
+	resp, err := openaiClient.CreateEmbeddings(context.TODO(),
+		openai.EmbeddingRequestStrings{
+			Input: []string{s},
+			Model: openai.AdaEmbeddingV2,
+		},
+	)
+	if err != nil {
+		return []string{}, fmt.Errorf("error getting keys: %v", err)
+	}
+
+	if len(resp.Data) == 0 {
+		return []string{}, fmt.Errorf("no response from OpenAI API")
+	}
+	embedding := resp.Data[0].Embedding
+
+	// Find example
+	findReq := FindRequest{
+		TopK: similarEntries, // Number of similar entries you want to find
+		Key:  embedding,      // The key you're looking for similarities to
+	}
+	findResp, err := client.Find(findReq)
+	if err != nil {
+		return []string{}, fmt.Errorf("error finding keys: %v", err)
+	}
+
+	return findResp.Values, nil
+}