Coverage for src / dataknobs_bots / memory / vector.py: 24%
50 statements
« prev ^ index » next coverage.py v7.13.0, created at 2025-12-16 10:13 -0700
1"""Vector-based semantic memory implementation."""
3from datetime import datetime
4from typing import Any
5from uuid import uuid4
7import numpy as np
9from .base import Memory
12class VectorMemory(Memory):
13 """Vector-based semantic memory using dataknobs-data vector stores.
15 This implementation stores messages with vector embeddings and retrieves
16 relevant messages based on semantic similarity.
18 Attributes:
19 vector_store: Vector store backend from dataknobs_data.vector.stores
20 embedding_provider: LLM provider for generating embeddings
21 max_results: Maximum number of results to return
22 similarity_threshold: Minimum similarity score for results
23 """
25 def __init__(
26 self,
27 vector_store: Any,
28 embedding_provider: Any,
29 max_results: int = 5,
30 similarity_threshold: float = 0.7,
31 ):
32 """Initialize vector memory.
34 Args:
35 vector_store: Vector store backend instance
36 embedding_provider: LLM provider with embed() method
37 max_results: Maximum number of similar messages to return
38 similarity_threshold: Minimum similarity score (0-1)
39 """
40 self.vector_store = vector_store
41 self.embedding_provider = embedding_provider
42 self.max_results = max_results
43 self.similarity_threshold = similarity_threshold
45 @classmethod
46 async def from_config(cls, config: dict[str, Any]) -> "VectorMemory":
47 """Create VectorMemory from configuration.
49 Args:
50 config: Configuration dictionary with:
51 - backend: Vector store backend type
52 - dimension: Vector dimension (optional, depends on backend)
53 - collection: Collection/index name (optional)
54 - embedding_provider: LLM provider name for embeddings
55 - embedding_model: Model to use for embeddings
56 - max_results: Max results to return (default 5)
57 - similarity_threshold: Min similarity score (default 0.7)
59 Returns:
60 Configured VectorMemory instance
61 """
62 from dataknobs_data.vector.stores import VectorStoreFactory
63 from dataknobs_llm.llm import LLMProviderFactory
65 # Create vector store
66 store_config = {
67 "backend": config.get("backend", "memory"),
68 "dimensions": config.get("dimension", 1536),
69 }
71 # Add optional store parameters
72 if "collection" in config:
73 store_config["collection_name"] = config["collection"]
74 if "persist_path" in config:
75 store_config["persist_path"] = config["persist_path"]
77 # Merge any additional store_params
78 if "store_params" in config:
79 store_config.update(config["store_params"])
81 factory = VectorStoreFactory()
82 vector_store = factory.create(**store_config)
83 await vector_store.initialize()
85 # Create embedding provider
86 llm_factory = LLMProviderFactory(is_async=True)
87 embedding_provider = llm_factory.create({
88 "provider": config.get("embedding_provider", "openai"),
89 "model": config.get("embedding_model", "text-embedding-ada-002"),
90 })
91 await embedding_provider.initialize()
93 return cls(
94 vector_store=vector_store,
95 embedding_provider=embedding_provider,
96 max_results=config.get("max_results", 5),
97 similarity_threshold=config.get("similarity_threshold", 0.7),
98 )
100 async def add_message(
101 self, content: str, role: str, metadata: dict[str, Any] | None = None
102 ) -> None:
103 """Add message with vector embedding.
105 Args:
106 content: Message content
107 role: Message role
108 metadata: Optional metadata
109 """
110 # Generate embedding
111 embedding = await self.embedding_provider.embed(content)
113 # Convert to numpy array if needed
114 if not isinstance(embedding, np.ndarray):
115 embedding = np.array(embedding, dtype=np.float32)
117 # Prepare metadata
118 msg_metadata = {
119 "content": content,
120 "role": role,
121 "timestamp": datetime.now().isoformat(),
122 "id": str(uuid4()),
123 }
124 if metadata:
125 msg_metadata.update(metadata)
127 # Store in vector store
128 await self.vector_store.add_vectors(
129 vectors=[embedding], ids=[msg_metadata["id"]], metadata=[msg_metadata]
130 )
132 async def get_context(self, current_message: str) -> list[dict[str, Any]]:
133 """Get semantically relevant messages.
135 Args:
136 current_message: Current message to find context for
138 Returns:
139 List of relevant message dictionaries sorted by similarity
140 """
141 # Generate query embedding
142 query_embedding = await self.embedding_provider.embed(current_message)
144 # Convert to numpy array if needed
145 if not isinstance(query_embedding, np.ndarray):
146 query_embedding = np.array(query_embedding, dtype=np.float32)
148 # Search for similar vectors
149 results = await self.vector_store.search(
150 query_vector=query_embedding,
151 k=self.max_results,
152 include_metadata=True,
153 )
155 # Format results
156 context = []
157 for _vector_id, similarity, msg_metadata in results:
158 if msg_metadata and similarity >= self.similarity_threshold:
159 context.append(
160 {
161 "content": msg_metadata.get("content", ""),
162 "role": msg_metadata.get("role", ""),
163 "similarity": similarity,
164 "metadata": msg_metadata,
165 }
166 )
168 return context
170 async def clear(self) -> None:
171 """Clear all vectors from memory.
173 Note: This deletes all vectors in the store. Use with caution
174 if the store is shared across multiple memory instances.
175 """
176 # Get all vector IDs and delete them
177 # Note: This is a simplified implementation
178 # In production, you might want to track IDs separately
179 # or use collection-level clearing if supported
180 if hasattr(self.vector_store, "clear"):
181 await self.vector_store.clear()
182 else:
183 # Fallback: delete individual vectors if we track them
184 # For now, we'll raise an error suggesting to use a new instance
185 raise NotImplementedError(
186 "Vector store does not support clearing. "
187 "Consider creating a new VectorMemory instance with a fresh collection."
188 )