Merge pull request #7 from redis-developer/extract-memories

Extract memories, better summarization, etc.

Author: abrookins

.gitignore

@@ -221,3 +221,4 @@ libs/redis/docs/.Trash*
 .python-version
 .idea/*
 .vscode/settings.json
+.cursor

Dockerfile

@@ -28,4 +28,4 @@ ENTRYPOINT []
 
 
 # Run the API server
-CMD ["python", "-m", "agent_memory_server.main"]
+CMD ["uv", "run", "agent-memory", "api"]

README.md

@@ -197,6 +197,18 @@ Example:
 agent-memory task-worker --concurrency 5 --redelivery-timeout 60
 ```
 
+#### `rebuild_index`
+Rebuilds the search index for Redis Memory Server.
+```bash
+agent-memory rebuild_index
+```
+
+#### `migrate_memories`
+Runs data migrations. Migrations are reentrant.
+```bash
+agent-memory migrate_memories
+```
+
 To see help for any command, you can use `--help`:
 ```bash
 agent-memory --help
@@ -207,25 +219,39 @@ agent-memory mcp --help
 
 ## Getting Started
 
-### Local Install
+### Installation
 
 First, you'll need to download this repository. After you've downloaded it, you can install and run the servers.
 
-1. Install the package and required dependencies with pip, ideally into a virtual environment:
-   ```bash
-   pip install -e .
-   ```
-
-**NOTE:** This project uses `uv` for dependency management, so if you have uv installed, you can run `uv sync` instead of `pip install ...` to install the project's dependencies.
+This project uses [uv](https://github.com/astral-sh/uv) for dependency management.
 
-2 (a). The easiest way to start the REST API server and MCP server in SSE mode is to use Docker Compose. See the Docker Compose section of this file for more details.
+1. Install uv:
+  ```bash
+  pip install uv
+  ```
 
-2 (b). You can also run the REST API and MCP servers directly, e.g.:
-#### REST API (direct, without CLI)
+2. Install the package and required dependencies:
   ```bash
-  python -m agent_memory_server.main
+  uv sync
   ```
 
+2. Set up environment variables (see Configuration section)
+
+### Running
+
+The easiest way to start the REST and MCP servers is to use Docker Compose. See the Docker Compose section of this file for more details.
+
+But you can also run these servers via the CLI commands. Here's how you
+run the REST API server:
+```bash
+uv run agent-memory api
+```
+
+And the MCP server:
+```
+uv run agent-memory mcp --mode <stdio|sse>
+```
+
 **NOTE:** With uv, prefix the command with `uv`, e.g.: `uv run agent-memory --mode sse`. If you installed from source, you'll probably need to add `--directory` to tell uv where to find the code: `uv run --directory <path/to/checkout> run agent-memory --mode stdio`.
 
 ### Docker Compose
@@ -293,52 +319,12 @@ Cursor's MCP config is similar to Claude's, but it also supports SSE servers, so
 
 ## Configuration
 
-You can configure the service using environment variables:
-
-| Variable | Description | Default |
-|----------|-------------|---------|
-| `REDIS_URL` | URL for Redis connection | `redis://localhost:6379` |
-| `LONG_TERM_MEMORY` | Enable/disable long-term memory | `True` |
-| `WINDOW_SIZE` | Maximum messages in short-term memory | `20` |
-| `OPENAI_API_KEY` | API key for OpenAI | - |
-| `ANTHROPIC_API_KEY` | API key for Anthropic | - |
-| `GENERATION_MODEL` | Model for text generation | `gpt-4o-mini` |
-| `EMBEDDING_MODEL` | Model for text embeddings | `text-embedding-3-small` |
-| `PORT` | REST API server port | `8000` |
-| `TOPIC_MODEL` | BERTopic model for topic extraction | `MaartenGr/BERTopic_Wikipedia` |
-| `NER_MODEL` | BERT model for NER | `dbmdz/bert-large-cased-finetuned-conll03-english` |
-| `ENABLE_TOPIC_EXTRACTION` | Enable/disable topic extraction | `True` |
-| `ENABLE_NER` | Enable/disable named entity recognition | `True` |
-| `MCP_PORT` | MCP server port |9000|
-
+You can configure the MCP and REST servers and task worker using environment
+variables. See the file `config.py` for all the available settings.
 
-## Development
-
-### Installation
-
-This project uses [uv](https://github.com/astral-sh/uv) for dependency management.
-
-1. Install dependencies:
-```bash
-uv sync --all-extras
-```
-
-2. Set up environment variables (see Configuration section)
-
-3. Run the API server:
-```bash
-agent-memory api
-```
-
-4. In a separate terminal, run the MCP server (use either the "stdio" or "sse" options to set the running mode) if you want to test with tools like Cursor or Claude:
-```bash
-agent-memory mcp --mode <stdio|sse>
-```
-
-### Running Tests
-```bash
-python -m pytest
-```
+The names of the settings map directly to an environment variable, so for
+example, you can set the `openai_api_key` setting with the `OPENAI_API_KEY`
+environment variable.
 
 ## Running the Background Task Worker
 
@@ -382,16 +368,24 @@ agent-memory schedule-task "agent_memory_server.long_term_memory.compact_long_te
 - **Semantic Deduplication**: Finds and merges memories with similar meaning using vector search
 - **LLM-powered Merging**: Uses language models to intelligently combine memories
 
+## Running Migrations
+When the data model changes, we add a migration in `migrations.py`. You can run
+these to make sure your schema is up to date, like so:
+
+```bash
+uv run agent-memory migrate-memories
+```
+
+## Development
+
+### Running Tests
+```bash
+uv run pytest
+```
+
 ## Contributing
 1. Fork the repository
 2. Create a feature branch
 3. Commit your changes
 4. Push to the branch
 5. Create a Pull Request
-
-### Running Tests
-
-```bash
-# Run all tests
-pytest tests
-```

agent_memory_server/api.py

@@ -223,12 +223,16 @@ async def search_long_term_memory(payload: SearchPayload):
     # Extract filter objects from the payload
     filters = payload.get_filters()
 
-    # Pass text, redis, and filter objects to the search function
-    return await long_term_memory.search_long_term_memories(
-        redis=redis,
-        text=payload.text,
-        distance_threshold=payload.distance_threshold,
-        limit=payload.limit,
-        offset=payload.offset,
+    kwargs = {
+        "redis": redis,
+        "distance_threshold": payload.distance_threshold,
+        "limit": payload.limit,
+        "offset": payload.offset,
         **filters,
-    )
+    }
+
+    if payload.text:
+        kwargs["text"] = payload.text
+
+    # Pass text, redis, and filter objects to the search function
+    return await long_term_memory.search_long_term_memories(**kwargs)

agent_memory_server/cli.py

@@ -13,6 +13,11 @@
 
 from agent_memory_server.config import settings
 from agent_memory_server.logging import configure_logging, get_logger
+from agent_memory_server.migrations import (
+    migrate_add_discrete_memory_extracted_2,
+    migrate_add_memory_hashes_1,
+    migrate_add_memory_type_3,
+)
 from agent_memory_server.utils.redis import ensure_search_index_exists, get_redis_conn
 
 
@@ -34,17 +39,51 @@ def version():
     click.echo(f"agent-memory-server version {VERSION}")
 
 
+@cli.command()
+def rebuild_index():
+    """Rebuild the search index."""
+    import asyncio
+
+    async def setup_and_run():
+        redis = await get_redis_conn()
+        await ensure_search_index_exists(redis, overwrite=True)
+
+    asyncio.run(setup_and_run())
+
+
+@cli.command()
+def migrate_memories():
+    """Migrate memories from the old format to the new format."""
+    import asyncio
+
+    click.echo("Starting memory migrations...")
+
+    async def run_migrations():
+        redis = await get_redis_conn()
+        migrations = [
+            migrate_add_memory_hashes_1,
+            migrate_add_discrete_memory_extracted_2,
+            migrate_add_memory_type_3,
+        ]
+        for migration in migrations:
+            await migration(redis=redis)
+
+    asyncio.run(run_migrations())
+
+    click.echo("Memory migrations completed successfully.")
+
+
 @cli.command()
 @click.option("--port", default=settings.port, help="Port to run the server on")
 @click.option("--host", default="0.0.0.0", help="Host to run the server on")
 @click.option("--reload", is_flag=True, help="Enable auto-reload")
 def api(port: int, host: str, reload: bool):
     """Run the REST API server."""
-    from agent_memory_server.main import app, on_start_logger
+    from agent_memory_server.main import on_start_logger
 
     on_start_logger(port)
     uvicorn.run(
-        app,
+        "agent_memory_server.main:app",
         host=host,
         port=port,
         reload=reload,

agent_memory_server/client/api.py

@@ -13,6 +13,7 @@
     CreatedAt,
     Entities,
     LastAccessed,
+    MemoryType,
     Namespace,
     SessionId,
     Topics,
@@ -273,6 +274,7 @@ async def search_long_term_memory(
         last_accessed: LastAccessed | dict[str, Any] | None = None,
         user_id: UserId | dict[str, Any] | None = None,
         distance_threshold: float | None = None,
+        memory_type: MemoryType | dict[str, Any] | None = None,
         limit: int = 10,
         offset: int = 0,
     ) -> LongTermMemoryResults:
@@ -313,6 +315,8 @@ async def search_long_term_memory(
             last_accessed = LastAccessed(**last_accessed)
         if isinstance(user_id, dict):
             user_id = UserId(**user_id)
+        if isinstance(memory_type, dict):
+            memory_type = MemoryType(**memory_type)
 
         # Apply default namespace if needed and no namespace filter specified
         if namespace is None and self.config.default_namespace is not None:
@@ -328,6 +332,7 @@ async def search_long_term_memory(
             last_accessed=last_accessed,
             user_id=user_id,
             distance_threshold=distance_threshold,
+            memory_type=memory_type,
             limit=limit,
             offset=offset,
         )

agent_memory_server/config.py

@@ -20,10 +20,12 @@ class Settings(BaseSettings):
     mcp_port: int = 9000
 
     # Topic and NER model settings
-    topic_model: str = "MaartenGr/BERTopic_Wikipedia"
+    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"
     enable_topic_extraction: bool = True
     enable_ner: bool = True
+    top_k_topics: int = 3
 
     # RedisVL Settings
     redisvl_distance_metric: str = "COSINE"

agent_memory_server/docket_tasks.py

@@ -7,6 +7,7 @@
 from docket import Docket
 
 from agent_memory_server.config import settings
+from agent_memory_server.extraction import extract_discrete_memories
 from agent_memory_server.long_term_memory import (
     compact_long_term_memories,
     extract_memory_structure,
@@ -24,6 +25,7 @@
     summarize_session,
     index_long_term_memories,
     compact_long_term_memories,
+    extract_discrete_memories,
 ]