Customize embedding model

example/webui/main.go

@@ -23,6 +23,7 @@ var apiKey = os.Getenv("API_KEY")
 var vectorStore = os.Getenv("VECTOR_STORE")
 var kbdisableIndexing = os.Getenv("KBDISABLEINDEX")
 var timeout = os.Getenv("TIMEOUT")
+var embeddingModel = os.Getenv("EMBEDDING_MODEL")
 
 const defaultChunkSize = 4098
 
@@ -63,7 +64,7 @@ func main() {
 		dbStore = rag.NewLocalAIRAGDB(laiStore, lai)
 	default:
 		var err error
-		dbStore, err = rag.NewChromemDB("local-agent-framework", stateDir, lai)
+		dbStore, err = rag.NewChromemDB("local-agent-framework", stateDir, lai, embeddingModel)
 		if err != nil {
 			panic(err)
 		}

llm/rag/chromem.go

@@ -15,9 +15,10 @@ type ChromemDB struct {
 	index           int
 	client          *openai.Client
 	db              *chromem.DB
+	embeddingsModel string
 }
 
-func NewChromemDB(collection, path string, openaiClient *openai.Client) (*ChromemDB, error) {
+func NewChromemDB(collection, path string, openaiClient *openai.Client, embeddingsModel string) (*ChromemDB, error) {
 	// db, err := chromem.NewPersistentDB(path, true)
 	// if err != nil {
 	// 	return nil, err
@@ -29,6 +30,7 @@ func NewChromemDB(collection, path string, openaiClient *openai.Client) (*Chrome
 		index:           1,
 		db:              db,
 		client:          openaiClient,
+		embeddingsModel: embeddingsModel,
 	}
 
 	c, err := db.GetOrCreateCollection(collection, nil, chromem.embedding())
@@ -59,7 +61,7 @@ func (c *ChromemDB) embedding() chromem.EmbeddingFunc {
 			resp, err := c.client.CreateEmbeddings(ctx,
 				openai.EmbeddingRequestStrings{
 					Input: []string{text},
-					Model: openai.AdaEmbeddingV2,
+					Model: openai.EmbeddingModel(c.embeddingsModel),
 				},
 			)
 			if err != nil {