Make memory extraction configurable

Author: abrookins

agent_memory_server/config.py

@@ -28,14 +28,23 @@ class Settings(BaseSettings):
     port: int = 8000
     mcp_port: int = 9000
 
-    # Topic and NER model settings
-    topic_model_source: Literal["NER", "LLM"] = "LLM"
-    topic_model: str = "MaartenGr/BERTopic_Wikipedia"  # LLM model here if using LLM
-    ner_model: str = "dbmdz/bert-large-cased-finetuned-conll03-english"
+    # The server indexes messages in long-term memory by default. If this
+    # setting is enabled, we also extract discrete memories from message text
+    # and save them as separate long-term memory records.
+    enable_discrete_memory_extraction: bool = True
+
+    # Topic modeling
+    topic_model_source: Literal["BERTopic", "LLM"] = "LLM"
+    topic_model: str = (
+        "MaartenGr/BERTopic_Wikipedia"  # Use an LLM model name here if using LLM
+    )
     enable_topic_extraction: bool = True
-    enable_ner: bool = True
     top_k_topics: int = 3
 
+    # Used for extracting entities from text
+    ner_model: str = "dbmdz/bert-large-cased-finetuned-conll03-english"
+    enable_ner: bool = True
+
     # RedisVL Settings
     redisvl_distance_metric: str = "COSINE"
     redisvl_vector_dimensions: str = "1536"

agent_memory_server/extraction.py

@@ -146,7 +146,7 @@ async def extract_topics_llm(
     return topics
 
 
-def extract_topics_ner(text: str, num_topics: int | None = None) -> list[str]:
+def extract_topics_bertopic(text: str, num_topics: int | None = None) -> list[str]:
     """
     Extract topics from text using the BERTopic model.
 
@@ -193,12 +193,8 @@ async def handle_extraction(text: str) -> tuple[list[str], list[str]]:
     # Extract topics if enabled
     topics = []
     if settings.enable_topic_extraction:
-        # Check if the topic_model_source setting exists and use appropriate function
-        if (
-            hasattr(settings, "topic_model_source")
-            and settings.topic_model_source == "NER"
-        ):
-            topics = extract_topics_ner(text)
+        if settings.topic_model_source == "BERTopic":
+            topics = extract_topics_bertopic(text)
         else:
             topics = await extract_topics_llm(text)
 
@@ -263,7 +259,10 @@ async def handle_extraction(text: str) -> tuple[list[str], list[str]]:
     """
 
 
-async def extract_discrete_memories(redis: Redis | None = None):
+async def extract_discrete_memories(
+    redis: Redis | None = None,
+    deduplicate: bool = True,
+):
     """
     Extract episodic and semantic memories from text using an LLM.
     """
@@ -345,5 +344,5 @@ async def extract_discrete_memories(redis: Redis | None = None):
 
         await index_long_term_memories(
             long_term_memories,
-            deduplicate=True,
+            deduplicate=deduplicate,
         )

agent_memory_server/long_term_memory.py

@@ -500,8 +500,6 @@ async def index_long_term_memories(
     memories: list[LongTermMemory],
     redis_client: Redis | None = None,
     deduplicate: bool = False,
-    deduplicate_hash: bool = True,
-    deduplicate_semantic: bool = True,
     vector_distance_threshold: float = 0.12,
     llm_client: Any = None,
 ) -> None:
@@ -612,7 +610,14 @@ async def index_long_term_memories(
         await pipe.execute()
 
     logger.info(f"Indexed {len(processed_memories)} memories")
-    await background_tasks.add_task(extract_discrete_memories)
+    if settings.enable_discrete_memory_extraction:
+        # Extract discrete memories from the indexed messages and persist
+        # them as separate long-term memory records. This process also
+        # runs deduplication if requested.
+        await background_tasks.add_task(
+            extract_discrete_memories,
+            deduplicate=deduplicate,
+        )
 
 
 async def search_long_term_memories(

tests/test_extraction.py

@@ -6,7 +6,7 @@
 from agent_memory_server.config import settings
 from agent_memory_server.extraction import (
     extract_entities,
-    extract_topics_ner,
+    extract_topics_bertopic,
     handle_extraction,
 )
 
@@ -45,7 +45,7 @@ async def test_extract_topics_success(self, mock_get_topic_model, mock_bertopic)
         mock_get_topic_model.return_value = mock_bertopic
         text = "Discussion about AI technology and business"
 
-        topics = extract_topics_ner(text)
+        topics = extract_topics_bertopic(text)
 
         assert set(topics) == {"technology", "business"}
         mock_bertopic.transform.assert_called_once_with([text])
@@ -58,7 +58,7 @@ async def test_extract_topics_no_valid_topics(
         mock_bertopic.transform.return_value = (np.array([-1]), np.array([0.0]))
         mock_get_topic_model.return_value = mock_bertopic
 
-        topics = extract_topics_ner("Test message")
+        topics = extract_topics_bertopic("Test message")
 
         assert topics == []
         mock_bertopic.transform.assert_called_once()