Clean up the vectore store init options

Author: abrookins

agent_memory_server/config.py

@@ -78,9 +78,7 @@ class Settings(BaseSettings):
 
     # Topic modeling
     topic_model_source: Literal["BERTopic", "LLM"] = "LLM"
-    topic_model: str = (
-        "MaartenGr/BERTopic_Wikipedia"  # Use an LLM model name here if using LLM
-    )
+    topic_model: str = "gpt-4o-mini"
     enable_topic_extraction: bool = True
     top_k_topics: int = 3
 
@@ -89,9 +87,11 @@ class Settings(BaseSettings):
     enable_ner: bool = True
 
     # RedisVL Settings
+    # TODO: Adapt to vector store settings
     redisvl_distance_metric: str = "COSINE"
     redisvl_vector_dimensions: str = "1536"
     redisvl_index_prefix: str = "memory_idx"
+    redisvl_indexing_algorithm: str = "HNSW"
 
     # Docket settings
     docket_name: str = "memory-server"

agent_memory_server/extraction.py

@@ -1,9 +1,8 @@
 import json
 import os
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import ulid
-from bertopic import BERTopic
 from redis.asyncio.client import Redis
 from tenacity.asyncio import AsyncRetrying
 from tenacity.stop import stop_after_attempt
@@ -22,24 +21,30 @@
 from agent_memory_server.utils.redis import get_redis_conn
 
 
+if TYPE_CHECKING:
+    from bertopic import BERTopic
+
+
 logger = get_logger(__name__)
 
 # Set tokenizer parallelism environment variable
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 # Global model instances
-_topic_model: BERTopic | None = None
+_topic_model: "BERTopic | None" = None
 _ner_model: Any | None = None
 _ner_tokenizer: Any | None = None
 
 
-def get_topic_model() -> BERTopic:
+def get_topic_model() -> "BERTopic":
     """
     Get or initialize the BERTopic model.
 
     Returns:
         The BERTopic model instance
     """
+    from bertopic import BERTopic
+
     global _topic_model
     if _topic_model is None:
         # TODO: Expose this as a config option
@@ -112,7 +117,7 @@ async def extract_topics_llm(
     """
     Extract topics from text using the LLM model.
     """
-    _client = client or await get_model_client(settings.generation_model)
+    _client = client or await get_model_client(settings.topic_model)
     _num_topics = num_topics if num_topics is not None else settings.top_k_topics
 
     prompt = f"""

agent_memory_server/long_term_memory.py

@@ -208,11 +208,19 @@ async def merge_memories_with_llm(memories: list[dict], llm_client: Any = None)
             # Fallback if the structure is different
             merged_text = str(response.choices[0])
 
+    def float_or_datetime(m: dict, key: str) -> float:
+        val = m.get(key, time.time())
+        if val is None:
+            return time.time()
+        if isinstance(val, datetime):
+            return int(val.timestamp())
+        return float(val)
+
     # Use the earliest creation timestamp
-    created_at = min(int(m.get("created_at", int(time.time()))) for m in memories)
+    created_at = min(float_or_datetime(m, "created_at") for m in memories)
 
     # Use the most recent last_accessed timestamp
-    last_accessed = max(int(m.get("last_accessed", int(time.time()))) for m in memories)
+    last_accessed = max(float_or_datetime(m, "last_accessed") for m in memories)
 
     # Prefer non-empty namespace, user_id, session_id from memories
     namespace = next((m["namespace"] for m in memories if m.get("namespace")), None)
@@ -616,6 +624,7 @@ async def index_long_term_memories(
 
             # Add the memory to be indexed if not a pure duplicate
             if not was_deduplicated:
+                current_memory.discrete_memory_extracted = "t"
                 processed_memories.append(current_memory)
     else:
         processed_memories = memories

agent_memory_server/vectorstore_adapter.py

@@ -123,10 +123,12 @@ def convert_filters_to_backend_format(
         last_accessed: LastAccessed | None = None,
         event_date: EventDate | None = None,
         memory_hash: MemoryHash | None = None,
+        discrete_memory_extracted: DiscreteMemoryExtracted | None = None,
     ) -> dict[str, Any] | None:
         """Convert filter objects to backend format for LangChain vectorstores."""
         filter_dict: dict[str, Any] = {}
 
+        # TODO: Seems like we could take *args filters and decide what to do based on type.
         # Apply tag/string filters using the helper function
         self.process_tag_filter(session_id, "session_id", filter_dict)
         self.process_tag_filter(user_id, "user_id", filter_dict)
@@ -135,6 +137,9 @@ def convert_filters_to_backend_format(
         self.process_tag_filter(topics, "topics", filter_dict)
         self.process_tag_filter(entities, "entities", filter_dict)
         self.process_tag_filter(memory_hash, "memory_hash", filter_dict)
+        self.process_tag_filter(
+            discrete_memory_extracted, "discrete_memory_extracted", filter_dict
+        )
 
         # Apply datetime filters using the helper function (uses instance method for backend-specific formatting)
         self.process_datetime_filter(created_at, "created_at", filter_dict)
@@ -374,6 +379,7 @@ def _convert_filters_to_backend_format(
         last_accessed: LastAccessed | None = None,
         event_date: EventDate | None = None,
         memory_hash: MemoryHash | None = None,
+        discrete_memory_extracted: DiscreteMemoryExtracted | None = None,
     ) -> dict[str, Any] | None:
         """Convert filter objects to standard LangChain dictionary format.
 
@@ -391,6 +397,7 @@ def _convert_filters_to_backend_format(
             Dictionary filter in format: {"field": {"$eq": "value"}} or None
         """
         processor = LangChainFilterProcessor(self.vectorstore)
+        # TODO: Seems like we could take *args and pass them to the processor
         filter_dict = processor.convert_filters_to_backend_format(
             session_id=session_id,
             user_id=user_id,
@@ -489,6 +496,7 @@ async def search_memories(
                 last_accessed=last_accessed,
                 event_date=event_date,
                 memory_hash=memory_hash,
+                discrete_memory_extracted=discrete_memory_extracted,
             )
 
             # Use LangChain's similarity search with filters
@@ -497,6 +505,8 @@ async def search_memories(
                 search_kwargs["filter"] = filter_dict
 
             # Perform similarity search
+            logger.info(f"Searching for memories with filters: {search_kwargs}")
+
             docs_with_scores = (
                 await self.vectorstore.asimilarity_search_with_relevance_scores(
                     query, **search_kwargs

agent_memory_server/vectorstore_factory.py

@@ -23,6 +23,7 @@
 
 from langchain_core.embeddings import Embeddings
 from langchain_core.vectorstores import VectorStore
+from langchain_redis.config import RedisConfig
 from pydantic.types import SecretStr
 
 
@@ -207,9 +208,15 @@ def create_redis_vectorstore(embeddings: Embeddings) -> VectorStore:
         # Always use MemoryRedisVectorStore for consistency and to fix relevance score issues
         return MemoryRedisVectorStore(
             embeddings=embeddings,
-            redis_url=settings.redis_url,
-            index_name=settings.redisvl_index_name,
-            metadata_schema=metadata_schema,
+            config=RedisConfig(
+                redis_url=settings.redis_url,
+                key_prefix=settings.redisvl_index_prefix,
+                indexing_algorithm=settings.redisvl_indexing_algorithm,
+                index_name=settings.redisvl_index_name,
+                metadata_schema=metadata_schema,
+                distance_metric=settings.redisvl_distance_metric,
+                embedding_dimensions=int(settings.redisvl_vector_dimensions),
+            ),
         )
     except ImportError:
         logger.error(