Configurable model selection

Author: abrookins

.github/workflows/python-tests.yml

@@ -15,7 +15,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v4
       with:
-        python-version: '3.11'
+        python-version: '3.12'
         cache: 'pip'
 
     - name: Install dependencies

README.md

@@ -1,55 +1,106 @@
 # Redis Memory Server
 
-A Python memory server for agents and LLM applications. This application
-provides memory features for LLM conversations, including short-term memory
-(message history) and long-term memory (vector embeddings for semantic search).
+A service that provides memory management for AI applications using Redis.
 
 ## Features
 
-- Short-term memory storage for conversation history
-- Optional long-term memory with semantic search capabilities
-- Automatic context summarization to handle long conversations
-- Integration with OpenAI API (more coming soon)
-- Redis-based storage with vector search
+- Short-term memory management with configurable window size
+- Long-term memory with semantic search capabilities
+- Automatic context summarization using LLMs
+- Support for multiple model providers (OpenAI and Anthropic)
+- Configurable token limits based on selected model
+
+## Configuration
+
+The service can be configured using environment variables:
+
+- `REDIS_URL`: URL for Redis connection (default: `redis://localhost:6379`)
+- `LONG_TERM_MEMORY`: Enable/disable long-term memory (default: `True`)
+- `WINDOW_SIZE`: Maximum number of messages to keep in short-term memory (default: `20`)
+- `OPENAI_API_KEY`: API key for OpenAI
+- `ANTHROPIC_API_KEY`: API key for Anthropic
+- `GENERATION_MODEL`: Model to use for text generation (default: `gpt-4o-mini`)
+- `EMBEDDING_MODEL`: Model to use for text embeddings (default: `text-embedding-3-small`)
+- `PORT`: Port to run the server on (default: `8000`)
+
+## Supported Models
+
+### OpenAI Models
+
+- `gpt-3.5-turbo`: 4K context window
+- `gpt-3.5-turbo-16k`: 16K context window
+- `gpt-4`: 8K context window
+- `gpt-4-32k`: 32K context window
+- `gpt-4o`: 128K context window
+- `gpt-4o-mini`: 128K context window
+
+### Anthropic Models
+
+- `claude-3-opus-20240229`: 200K context window
+- `claude-3-sonnet-20240229`: 200K context window
+- `claude-3-haiku-20240307`: 200K context window
+- `claude-3-5-sonnet-20240620`: 200K context window
+
+**Note**: Embedding operations always use OpenAI models, as Anthropic does not provide embedding API.
 
 ## Installation
 
 1. Clone the repository
-2. Install dependencies:
+2. Install dependencies: `pip install -r requirements.txt`
+3. Set up environment variables (see Configuration section)
+4. Run the server: `python main.py`
+
+## Usage
+
+### Add Messages to Memory
+
 ```
-pip install -r requirements.txt
+POST /sessions/{session_id}/memory
 ```
-3. Set up environment variables:
+
+Request body:
+```json
+{
+  "messages": [
+    {
+      "role": "user",
+      "content": "Hello, how are you?"
+    },
+    {
+      "role": "assistant",
+      "content": "I'm doing well, thank you for asking!"
+    }
+  ],
+  "context": "Optional previous summary"
+}
 ```
-# Required
-REDIS_URL=redis://localhost:6379
 
-# Optional
-PORT=8000
-LONG_TERM_MEMORY=true
-MAX_WINDOW_SIZE=12
-MODEL=gpt-3.5-turbo
+### Get Memory
 
-# For OpenAI
-OPENAI_API_KEY=your_openai_api_key
+```
+GET /sessions/{session_id}/memory
 ```
 
-## Usage
-
-Start the server:
+### Search Memory
 
 ```
-python main.py
+POST /sessions/{session_id}/retrieval
 ```
 
-## API Endpoints
+Request body:
+```json
+{
+  "text": "What was the conversation about?"
+}
+```
+
+## Development
+
+To run tests:
 
-- `GET /health`: Health check endpoint
-- `GET /sessions`: Get a list of session IDs
-- `GET /sessions/{session_id}/memory`: Get memory for a session
-- `POST /sessions/{session_id}/memory`: Add messages to a session
-- `DELETE /sessions/{session_id}/memory`: Delete a session's memory
-- `POST /sessions/{session_id}/retrieval`: Perform semantic search on session memory
+```
+python -m pytest
+```
 
 ## License
 TBD

config.py

@@ -7,6 +7,7 @@ class Settings(BaseSettings):
     long_term_memory: bool = True
     window_size: int = 20
     openai_api_key: str = os.getenv("OPENAI_API_KEY", "")
+    anthropic_api_key: str = os.getenv("ANTHROPIC_API_KEY", "")
     generation_model: str = "gpt-4o-mini"
     embedding_model: str = "text-embedding-3-small"
     port: int = 8000

long_term_memory.py

@@ -1,13 +1,15 @@
-from typing import List, Type
+from typing import List, Type, Union, Any
 import nanoid
-import numpy as np
 from redis.asyncio import Redis
 from redis.commands.search.query import Query
 from models import (
     MemoryMessage,
     OpenAIClientWrapper,
+    AnthropicClientWrapper,
     RedisearchResult,
     SearchResults,
+    ModelProvider,
+    get_model_config,
 )
 import logging
 
@@ -19,7 +21,7 @@
 async def index_messages(
     messages: List[MemoryMessage],
     session_id: str,
-    openai_client: OpenAIClientWrapper,
+    client: OpenAIClientWrapper,  # Only OpenAI supports embeddings currently
     redis_conn: Redis,
 ) -> None:
     """Index messages in Redis for vector search"""
@@ -28,7 +30,7 @@ async def index_messages(
         contents = [msg.content for msg in messages]
 
         # Get embeddings from OpenAI
-        embeddings = await openai_client.create_embedding(contents)
+        embeddings = await client.create_embedding(contents)
 
         # Index each message with its embedding
         for index, embedding in enumerate(embeddings):
@@ -64,16 +66,18 @@ class Unset:
 async def search_messages(
     query: str,
     session_id: str,
-    openai_client: OpenAIClientWrapper,
+    client: OpenAIClientWrapper,  # Only OpenAI supports embeddings currently
     redis_conn: Redis,
     distance_threshold: float | Type[Unset] = Unset,
     limit: int = 10,
 ) -> SearchResults:
     """Search for messages using vector similarity"""
     try:
         # Get embedding for query
-        query_embedding = await openai_client.create_embedding([query])
+        query_embedding = await client.create_embedding([query])
         vector = query_embedding.tobytes()
+
+        # Set up query parameters
         params = {"vec": vector}
 
         if distance_threshold and distance_threshold is not Unset:
@@ -85,26 +89,46 @@ async def search_messages(
             base_query = Query(
                 f"@session:{{{session_id}}}=>[KNN {limit} @vector $vec AS dist]"
             )
+
         q = (
             base_query.return_fields("role", "content", "dist")
             .sort_by("dist", asc=True)
             .paging(0, limit)
             .dialect(2)
         )
 
+        # Execute search
         raw_results = await redis_conn.ft(REDIS_INDEX_NAME).search(
             q,
             query_params=params,  # type: ignore
         )
 
-        # Parse results
-        results = [
-            RedisearchResult(role=doc.role, content=doc.content, dist=doc.dist)
-            for doc in raw_results.docs
-        ]
+        # Parse results safely
+        results = []
+        total_results = 0
+
+        # Check if raw_results has the expected attributes
+        if hasattr(raw_results, "docs") and isinstance(raw_results.docs, list):
+            for doc in raw_results.docs:
+                if (
+                    hasattr(doc, "role")
+                    and hasattr(doc, "content")
+                    and hasattr(doc, "dist")
+                ):
+                    results.append(
+                        RedisearchResult(
+                            role=doc.role, content=doc.content, dist=float(doc.dist)
+                        )
+                    )
+
+            total_results = getattr(raw_results, "total", len(results))
+        else:
+            # Handle the case where raw_results doesn't have the expected structure
+            logger.warning("Unexpected search result format")
+            total_results = 0
 
         logger.info(f"Found {len(results)} results for query in session {session_id}")
-        return SearchResults(total=raw_results.total, docs=results)
+        return SearchResults(total=total_results, docs=results)
     except Exception as e:
         logger.error(f"Error searching messages: {e}")
         raise

main.py

@@ -6,6 +6,7 @@
 from fastapi import FastAPI
 
 import utils
+from models import ModelProvider, MODEL_CONFIGS
 
 load_dotenv()
 
@@ -36,12 +37,58 @@ async def startup_event():
     """Initialize the application on startup"""
     logger.info("Starting Redis Memory Server 🤘")
 
+    # Check for required API keys
+    available_providers = []
+
+    if settings.openai_api_key:
+        available_providers.append(ModelProvider.OPENAI)
+    else:
+        logger.warning("OpenAI API key not set, OpenAI models will not be available")
+
+    if settings.anthropic_api_key:
+        available_providers.append(ModelProvider.ANTHROPIC)
+    else:
+        logger.warning(
+            "Anthropic API key not set, Anthropic models will not be available"
+        )
+
+    # Check if the configured models are available
+    generation_model_config = MODEL_CONFIGS.get(settings.generation_model)
+    embedding_model_config = MODEL_CONFIGS.get(settings.embedding_model)
+
+    if (
+        generation_model_config
+        and generation_model_config.provider not in available_providers
+    ):
+        logger.warning(
+            f"Selected generation model {settings.generation_model} requires {generation_model_config.provider} API key"
+        )
+
+    if (
+        embedding_model_config
+        and embedding_model_config.provider not in available_providers
+    ):
+        logger.warning(
+            f"Selected embedding model {settings.embedding_model} requires {embedding_model_config.provider} API key"
+        )
+
+    # If long-term memory is enabled but OpenAI isn't available, warn user
+    if settings.long_term_memory and ModelProvider.OPENAI not in available_providers:
+        logger.warning(
+            "Long-term memory requires OpenAI for embeddings, but OpenAI API key is not set"
+        )
+
     # Set up RediSearch index if long-term memory is enabled
     if settings.long_term_memory:
         redis = get_redis_conn()
 
-        # For now, just ada support
-        vector_dimensions = 1536
+        # Get embedding dimensions from model config
+        embedding_model_config = MODEL_CONFIGS.get(settings.embedding_model)
+        vector_dimensions = (
+            embedding_model_config.embedding_dimensions
+            if embedding_model_config
+            else 1536
+        )
         distance_metric = "COSINE"
 
         try:
@@ -50,6 +97,25 @@ async def startup_event():
             logger.error(f"Failed to ensure RediSearch index: {e}")
             raise
 
+    # Show available models
+    openai_models = [
+        model
+        for model, config in MODEL_CONFIGS.items()
+        if config.provider == ModelProvider.OPENAI
+        and ModelProvider.OPENAI in available_providers
+    ]
+    anthropic_models = [
+        model
+        for model, config in MODEL_CONFIGS.items()
+        if config.provider == ModelProvider.ANTHROPIC
+        and ModelProvider.ANTHROPIC in available_providers
+    ]
+
+    if openai_models:
+        logger.info(f"Available OpenAI models: {', '.join(openai_models)}")
+    if anthropic_models:
+        logger.info(f"Available Anthropic models: {', '.join(anthropic_models)}")
+
     logger.info(
         "Redis Memory Server initialized",
         window_size=settings.window_size,