Skip to content

Commit

Permalink
feat: add configurable vector store settings to MemoryContext
Browse files Browse the repository at this point in the history
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
  • Loading branch information
teilomillet committed Nov 21, 2024
1 parent e903f43 commit 1c45909
Showing 1 changed file with 24 additions and 1 deletion.
25 changes: 24 additions & 1 deletion memory_context.go
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.
//
Expand Down Expand Up @@ -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,
Expand All @@ -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)
Expand Down

0 comments on commit 1c45909

Please sign in to comment.