More tweaks

Author: abrookins

README.md

@@ -109,7 +109,7 @@ To start the API using Docker Compose, follow these steps:
 4. Build and start the containers by running:
    docker-compose up --build
 
-5. Once the containers are up, the API will be available at http://localhost:8000. You can also access the interactive API documentation at http://localhost:8000/docs.
+5. Once the containers are up, the REST API will be available at http://localhost:8000. You can also access the interactive API documentation at http://localhost:8000/docs. The MCP server will be available at http://localhost:9000/sse.
 
 6. To stop the containers, press Ctrl+C in the terminal and then run:
    docker-compose down
@@ -127,11 +127,12 @@ You can configure the service using environment variables:
 | `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` | Server port | `8000` |
+| `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|
 
 
 ## Development
@@ -145,16 +146,25 @@ pip install -e ".[dev]"
 
 2. Set up environment variables (see Configuration section)
 
-3. Run the server:
+3. Run the API server:
 ```bash
 python -m redis_memory_server.main
 ```
 
+4. In a separate terminal, run the MCP server (use either the "stdio" or "sse" options to set the running mode):
+```bash
+python -m redis_memory_server.mcp [stdio|sse]
+```
+
 ### Running Tests
 ```bash
 python -m pytest
 ```
 
+## Known Issues
+- The MCP server from the Python MCP SDK often refuses to shut down with Control-C if it's connected to a client
+- All background tasks run as async coroutines in the same process as the REST API server, using Starlette's `BackgroundTask`
+
 ### Contributing
 1. Fork the repository
 2. Create a feature branch

docker-compose.yml

@@ -28,6 +28,29 @@ services:
       timeout: 10s
       retries: 3
 
+  mcp:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    environment:
+      - REDIS_URL=redis://redis:6379
+      - PORT=9000
+      # Add your API keys here or use a .env file
+      - OPENAI_API_KEY=${OPENAI_API_KEY}
+      - ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY}
+      # Optional configurations with defaults
+      - LONG_TERM_MEMORY=True
+      - WINDOW_SIZE=20
+      - GENERATION_MODEL=gpt-4o-mini
+      - EMBEDDING_MODEL=text-embedding-3-small
+      - ENABLE_TOPIC_EXTRACTION=True
+      - ENABLE_NER=True
+    ports:
+      - "9000:9000"
+    depends_on:
+      - redis
+    command: ["python", "-m", "redis_memory_server.mcp", "sse"]
+
   redis:
     image: redis/redis-stack:latest
     ports:

redis_memory_server/api.py

@@ -134,7 +134,9 @@ async def delete_session_memory(
 
 
 @router.post("/long-term-memory", response_model=AckResponse)
-async def create_long_term_memory(payload: CreateLongTermMemoryPayload):
+async def create_long_term_memory(
+    payload: CreateLongTermMemoryPayload, background_tasks: BackgroundTasks
+):
     """
     Create a long-term memory
 
@@ -152,6 +154,7 @@ async def create_long_term_memory(payload: CreateLongTermMemoryPayload):
     await long_term_memory.index_long_term_memories(
         redis=redis,
         memories=payload.memories,
+        background_tasks=background_tasks,
     )
     return AckResponse(status="ok")
 

redis_memory_server/config.py

@@ -16,6 +16,7 @@ class Settings(BaseSettings):
     generation_model: str = "gpt-4o-mini"
     embedding_model: str = "text-embedding-3-small"
     port: int = 8000
+    mcp_port: int = 9000
 
     # Topic and NER model settings
     topic_model: str = "MaartenGr/BERTopic_Wikipedia"

redis_memory_server/long_term_memory.py

@@ -36,7 +36,7 @@ async def extract_memory_structure(
     entities_joined = ",".join(entities) if entities else ""
 
     await redis.hset(
-        Keys.memory_key(_id, ""),
+        Keys.memory_key(_id, namespace),
         mapping={
             "topics": topics_joined,
             "entities": entities_joined,
@@ -62,7 +62,6 @@ async def index_long_term_memories(
             id_ = memory.id_ if memory.id_ else nanoid.generate()
             key = Keys.memory_key(id_, memory.namespace)
             vector = embedding.tobytes()
-            id_ = memory.id_ if memory.id_ else nanoid.generate()
 
             await pipe.hset(  # type: ignore
                 key,
@@ -190,8 +189,12 @@ async def search_long_term_memories(
                             user_id=doc.user_id,
                             session_id=doc.session_id,
                             namespace=doc.namespace,
-                            topics=doc.topics.split(",") if doc.topics else [],
-                            entities=doc.entities.split(",") if doc.entities else [],
+                            topics=doc.get("topics", "").split(",")
+                            if doc.get("topics")
+                            else [],
+                            entities=doc.get("entities", "").split(",")
+                            if doc.get("entities")
+                            else [],
                         )
                     )
 

redis_memory_server/main.py

@@ -10,7 +10,6 @@
 from redis_memory_server.healthcheck import router as health_router
 from redis_memory_server.llms import MODEL_CONFIGS, ModelProvider
 from redis_memory_server.logging import configure_logging, get_logger
-from redis_memory_server.mcp import mcp_app
 from redis_memory_server.utils import ensure_redisearch_index, get_redis_conn
 
 
@@ -125,10 +124,6 @@ async def lifespan(app: FastAPI):
 app.include_router(memory_router)
 
 
-# Mount MCP server
-app.mount("/", mcp_app.sse_app())
-
-
 def on_start_logger(port: int):
     """Log startup information"""
     print("\n-----------------------------------")

redis_memory_server/mcp.py

@@ -1,4 +1,5 @@
 import logging
+import sys
 
 from fastapi import HTTPException
 from mcp.server.fastmcp import FastMCP
@@ -10,6 +11,7 @@
     get_session_memory as core_get_session_memory,
     search_long_term_memory as core_search_long_term_memory,
 )
+from redis_memory_server.config import settings
 from redis_memory_server.models import (
     AckResponse,
     CreateLongTermMemoryPayload,
@@ -20,7 +22,7 @@
 
 
 logger = logging.getLogger(__name__)
-mcp_app = FastMCP("Redis Agent Memory Server")
+mcp_app = FastMCP("Redis Agent Memory Server", port=settings.mcp_port)
 
 
 @mcp_app.tool()
@@ -146,3 +148,10 @@ async def memory_prompt(
     )
 
     return messages
+
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1 and sys.argv[1] == "sse":
+        mcp_app.run(transport="sse")
+    else:
+        mcp_app.run(transport="stdio")

redis_memory_server/messages.py

@@ -3,6 +3,7 @@
 import time
 
 from fastapi import BackgroundTasks
+from redis import WatchError
 from redis.asyncio import Redis
 
 from redis_memory_server.config import settings
@@ -116,16 +117,25 @@ async def set_session_memory(
     messages_key = Keys.messages_key(session_id, namespace=memory.namespace)
     metadata_key = Keys.metadata_key(session_id, namespace=memory.namespace)
     messages_json = [json.dumps(msg.model_dump()) for msg in memory.messages]
-
     metadata = memory.model_dump(
         exclude_none=True,
         exclude={"messages"},
     )
 
-    current_time = int(time.time())
-    await redis.zadd(sessions_key, {session_id: current_time})
-    await redis.rpush(messages_key, *messages_json)  # type: ignore
-    await redis.hset(metadata_key, mapping=metadata)  # type: ignore
+    async with redis.pipeline(transaction=True) as pipe:
+        await pipe.watch(messages_key, metadata_key)
+        pipe.multi()
+
+        while True:
+            try:
+                current_time = int(time.time())
+                pipe.zadd(sessions_key, {session_id: current_time})
+                pipe.rpush(messages_key, *messages_json)  # type: ignore
+                pipe.hset(metadata_key, mapping=metadata)  # type: ignore
+                await pipe.execute()
+            except WatchError:
+                continue
+            break
 
     # Check if window size is exceeded
     current_size = await redis.llen(messages_key)  # type: ignore
@@ -155,6 +165,7 @@ async def set_session_memory(
                 )
                 for msg in memory.messages
             ],
+            background_tasks,
         )
 
 

tests/conftest.py

@@ -118,6 +118,7 @@ async def session(use_test_redis_connection, async_redis_client):
                 namespace="test-namespace",
             ),
         ],
+        background_tasks=BackgroundTasks(),
     )
 
     # Add messages to session memory

tests/test_long_term_memory.py

@@ -5,6 +5,7 @@
 import nanoid
 import numpy as np
 import pytest
+from fastapi import BackgroundTasks
 from redis.commands.search.document import Document
 
 from redis_memory_server.long_term_memory import (
@@ -40,6 +41,7 @@ async def test_index_messages(
             await index_long_term_memories(
                 mock_async_redis_client,
                 long_term_memories,
+                background_tasks=BackgroundTasks(),
             )
 
         # Check that create_embedding was called with the right arguments
@@ -169,6 +171,7 @@ async def test_search_messages(self, async_redis_client):
         await index_long_term_memories(
             async_redis_client,
             long_term_memories,
+            background_tasks=BackgroundTasks(),
         )
 
         results = await search_long_term_memories(
@@ -194,6 +197,7 @@ async def test_search_messages_with_distance_threshold(self, async_redis_client)
         await index_long_term_memories(
             async_redis_client,
             long_term_memories,
+            background_tasks=BackgroundTasks(),
         )
 
         results = await search_long_term_memories(