114 lines
2.4 KiB
Go
114 lines
2.4 KiB
Go
package rag
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"runtime"
|
|
|
|
"github.com/philippgille/chromem-go"
|
|
"github.com/sashabaranov/go-openai"
|
|
)
|
|
|
|
type ChromemDB struct {
|
|
collectionName string
|
|
collection *chromem.Collection
|
|
index int
|
|
client *openai.Client
|
|
db *chromem.DB
|
|
embeddingsModel string
|
|
}
|
|
|
|
func NewChromemDB(collection, path string, openaiClient *openai.Client, embeddingsModel string) (*ChromemDB, error) {
|
|
// db, err := chromem.NewPersistentDB(path, true)
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
db := chromem.NewDB()
|
|
|
|
chromem := &ChromemDB{
|
|
collectionName: collection,
|
|
index: 1,
|
|
db: db,
|
|
client: openaiClient,
|
|
embeddingsModel: embeddingsModel,
|
|
}
|
|
|
|
c, err := db.GetOrCreateCollection(collection, nil, chromem.embedding())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
chromem.collection = c
|
|
|
|
return chromem, nil
|
|
}
|
|
|
|
func (c *ChromemDB) Count() int {
|
|
return c.collection.Count()
|
|
}
|
|
|
|
func (c *ChromemDB) Reset() error {
|
|
if err := c.db.DeleteCollection(c.collectionName); err != nil {
|
|
return err
|
|
}
|
|
collection, err := c.db.GetOrCreateCollection(c.collectionName, nil, c.embedding())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
c.collection = collection
|
|
|
|
return nil
|
|
}
|
|
|
|
func (c *ChromemDB) embedding() chromem.EmbeddingFunc {
|
|
return chromem.EmbeddingFunc(
|
|
func(ctx context.Context, text string) ([]float32, error) {
|
|
resp, err := c.client.CreateEmbeddings(ctx,
|
|
openai.EmbeddingRequestStrings{
|
|
Input: []string{text},
|
|
Model: openai.EmbeddingModel(c.embeddingsModel),
|
|
},
|
|
)
|
|
if err != nil {
|
|
return []float32{}, fmt.Errorf("error getting keys: %v", err)
|
|
}
|
|
|
|
if len(resp.Data) == 0 {
|
|
return []float32{}, fmt.Errorf("no response from OpenAI API")
|
|
}
|
|
|
|
embedding := resp.Data[0].Embedding
|
|
|
|
return embedding, nil
|
|
},
|
|
)
|
|
}
|
|
|
|
func (c *ChromemDB) Store(s string) error {
|
|
defer func() {
|
|
c.index++
|
|
}()
|
|
if s == "" {
|
|
return fmt.Errorf("empty string")
|
|
}
|
|
return c.collection.AddDocuments(context.Background(), []chromem.Document{
|
|
{
|
|
Content: s,
|
|
ID: fmt.Sprint(c.index),
|
|
},
|
|
}, runtime.NumCPU())
|
|
}
|
|
|
|
func (c *ChromemDB) Search(s string, similarEntries int) ([]string, error) {
|
|
res, err := c.collection.Query(context.Background(), s, similarEntries, nil, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var results []string
|
|
for _, r := range res {
|
|
results = append(results, r.Content)
|
|
}
|
|
|
|
return results, nil
|
|
}
|