Add reset, some minor UX tweaks
This commit is contained in:
@@ -13,6 +13,8 @@ type ChromemDB struct {
|
||||
collectionName string
|
||||
collection *chromem.Collection
|
||||
index int
|
||||
client *openai.Client
|
||||
db *chromem.DB
|
||||
}
|
||||
|
||||
func NewChromemDB(collection, path string, openaiClient *openai.Client) (*ChromemDB, error) {
|
||||
@@ -22,9 +24,39 @@ func NewChromemDB(collection, path string, openaiClient *openai.Client) (*Chrome
|
||||
// }
|
||||
db := chromem.NewDB()
|
||||
|
||||
embeddingFunc := chromem.EmbeddingFunc(
|
||||
chromem := &ChromemDB{
|
||||
collectionName: collection,
|
||||
index: 1,
|
||||
db: db,
|
||||
client: openaiClient,
|
||||
}
|
||||
|
||||
c, err := db.GetOrCreateCollection(collection, nil, chromem.embedding())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
chromem.collection = c
|
||||
|
||||
return chromem, nil
|
||||
}
|
||||
|
||||
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 := openaiClient.CreateEmbeddings(ctx,
|
||||
resp, err := c.client.CreateEmbeddings(ctx,
|
||||
openai.EmbeddingRequestStrings{
|
||||
Input: []string{text},
|
||||
Model: openai.AdaEmbeddingV2,
|
||||
@@ -43,17 +75,6 @@ func NewChromemDB(collection, path string, openaiClient *openai.Client) (*Chrome
|
||||
return embedding, nil
|
||||
},
|
||||
)
|
||||
|
||||
c, err := db.GetOrCreateCollection(collection, nil, embeddingFunc)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &ChromemDB{
|
||||
collectionName: collection,
|
||||
collection: c,
|
||||
index: 1,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *ChromemDB) Store(s string) error {
|
||||
|
||||
@@ -19,6 +19,10 @@ func NewLocalAIRAGDB(storeClient *StoreClient, openaiClient *openai.Client) *Loc
|
||||
}
|
||||
}
|
||||
|
||||
func (db *LocalAIRAGDB) Reset() error {
|
||||
return fmt.Errorf("not implemented")
|
||||
}
|
||||
|
||||
func (db *LocalAIRAGDB) Store(s string) error {
|
||||
resp, err := db.openaiClient.CreateEmbeddings(context.TODO(),
|
||||
openai.EmbeddingRequestStrings{
|
||||
|
||||
Reference in New Issue
Block a user