From 1c4590916caf946068f38f838c7191fd6409b203 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Te=C3=AFlo=20M?= <teilomillet@gmail.com>
Date: Thu, 21 Nov 2024 18:14:19 +0100
Subject: [PATCH] feat: add configurable vector store settings to MemoryContext

Adds new options to configure vector database type and address:
- Adds DBType and DBAddress fields to MemoryContextOptions
- Adds MemoryVectorDB functional option for configuration
- Maintains backward compatibility with default Milvus settings
---
 memory_context.go | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/memory_context.go b/memory_context.go
index f4df51d..2deebf5 100644
--- a/memory_context.go
+++ b/memory_context.go
@@ -26,6 +26,12 @@ type MemoryContextOptions struct {
 	// Collection specifies the vector database collection for storing memories
 	Collection string
 
+	// DBType specifies the vector database type (e.g., "chromem", "milvus")
+	DBType string
+
+	// DBAddress specifies the vector database address (e.g., "./.chromem/name" or "localhost:19530")
+	DBAddress string
+
 	// IncludeScore determines whether to include relevance scores in results
 	IncludeScore bool
 
@@ -94,6 +100,21 @@ func MemoryCollection(collection string) func(*MemoryContextOptions) {
 	}
 }
 
+// MemoryVectorDB configures the vector database type and address.
+// This allows choosing between different vector store implementations and their locations.
+//
+// Example:
+//
+//	ctx, err := raggo.NewMemoryContext(apiKey,
+//	    raggo.MemoryVectorDB("chromem", "./.chromem/chat"),  // Use Chromem store
+//	)
+func MemoryVectorDB(dbType, address string) func(*MemoryContextOptions) {
+	return func(o *MemoryContextOptions) {
+		o.DBType = dbType
+		o.DBAddress = address
+	}
+}
+
 // MemoryScoreInclusion controls whether similarity scores are included in results.
 // Useful for debugging or implementing custom relevance filtering.
 //
@@ -165,6 +186,8 @@ func NewMemoryContext(apiKey string, opts ...func(*MemoryContextOptions)) (*Memo
 		TopK:         3,
 		MinScore:     0.7,
 		Collection:   "memory_store",
+		DBType:       "milvus",           // Default to Milvus
+		DBAddress:    "localhost:19530",   // Default Milvus address
 		IncludeScore: false,
 		StoreLastN:   0,
 		StoreRAGInfo: false,
@@ -181,7 +204,7 @@ func NewMemoryContext(apiKey string, opts ...func(*MemoryContextOptions)) (*Memo
 		WithTopK(options.TopK),
 		WithMinScore(options.MinScore),
 		WithRetrieveEmbedding("openai", "text-embedding-3-small", apiKey),
-		WithRetrieveDB("milvus", "localhost:19530"), // Add default DB config
+		WithRetrieveDB(options.DBType, options.DBAddress),
 	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to initialize retriever: %w", err)