Move into package

Author: abrookins

pyproject.toml

@@ -1,3 +1,54 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "redis-memory-server"
+version = "0.1.0"
+description = "A memory system for conversational AI using Redis"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+authors = [
+    { name = "Andrew Brookins", email = "[email protected]" },
+]
+dependencies = [
+    "fastapi>=0.109.0",
+    "uvicorn>=0.27.0",
+    "redis>=5.0.1",
+    "openai>=1.0.0",
+    "anthropic>=0.18.1",
+    "python-dotenv>=1.0.0",
+    "structlog>=24.1.0",
+    "tiktoken>=0.5.2",
+    "numpy>=1.26.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.4.0",
+    "pytest-asyncio>=0.23.0",
+    "pytest-xdist>=3.5.0",
+    "black>=24.1.0",
+    "ruff>=0.2.0",
+    "testcontainers>=3.7.0",
+    "pre-commit>=3.6.0",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["redis_memory_server"]
+
+[tool.pytest.ini_options]
+addopts = "-v"
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+asyncio_mode = "auto"
+
+[tool.black]
+line-length = 88
+target-version = ['py310']
+include = '\.pyi?$'
+
 [tool.ruff]
 # Exclude a variety of commonly ignored directories
 exclude = [
@@ -49,3 +100,6 @@ ban-relative-imports = "all"
 quote-style = "double"
 # Use spaces for indentation
 indent-style = "space"
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]

redis_memory_server/__init__.py

@@ -0,0 +1,3 @@
+"""Redis Memory Server - A memory system for conversational AI."""
+
+__version__ = "0.1.0"

config.py redis_memory_server/config.pyconfig.py renamed to redis_memory_server/config.py

@@ -1,8 +1,12 @@
 import os
 
+from dotenv import load_dotenv
 from pydantic_settings import BaseSettings
 
 
+load_dotenv()
+
+
 class Settings(BaseSettings):
     redis_url: str = "redis://localhost:6379"
     long_term_memory: bool = True

redis_memory_server/extraction.py

@@ -0,0 +1,104 @@
+import json
+import logging
+
+from redis_memory_server.models import (
+    AnthropicClientWrapper,
+    MemoryMessage,
+    OpenAIClientWrapper,
+)
+
+
+logger = logging.getLogger(__name__)
+
+EXTRACTION_PROMPT = """Analyze the following message and extract:
+1. Key topics (as single words or short phrases)
+2. Named entities (people, places, organizations, etc.)
+
+Message: {message}
+
+Respond in JSON format:
+{{
+    "topics": ["topic1", "topic2", ...],
+    "entities": ["entity1", "entity2", ...]
+}}
+
+Keep topics and entities concise and relevant."""
+
+
+async def extract_topics_and_entities(
+    message: str,
+    model_client: OpenAIClientWrapper | AnthropicClientWrapper,
+) -> tuple[list[str], list[str]]:
+    """
+    Extract topics and entities from a message using the LLM.
+
+    Args:
+        message: The message to analyze
+        model_client: The LLM client to use
+
+    Returns:
+        Tuple of (topics, entities) lists
+    """
+    try:
+        # Get LLM response
+        response = await model_client.create_chat_completion(
+            "gpt-4o-mini",  # TODO: Make configurable
+            EXTRACTION_PROMPT.format(message=message),
+        )
+
+        # Parse JSON response from content field
+        content = response.choices[0]["message"]["content"].strip()
+        result = json.loads(content)
+
+        # Extract and validate topics and entities
+        topics = result.get("topics", [])
+        entities = result.get("entities", [])
+
+        # Ensure we have lists
+        if not isinstance(topics, list) or not isinstance(entities, list):
+            logger.error("Invalid extraction response format")
+            return [], []
+
+        return topics, entities
+
+    except Exception as e:
+        logger.error(f"Error in topic/entity extraction: {e}")
+        return [], []
+
+
+async def handle_extraction(
+    message: MemoryMessage,
+    model_client: OpenAIClientWrapper | AnthropicClientWrapper,
+) -> MemoryMessage:
+    """
+    Handle topic and entity extraction for a message.
+
+    Args:
+        message: The message to process
+        model_client: The LLM client to use
+
+    Returns:
+        Updated message with extracted topics and entities
+    """
+    # Skip if message already has both topics and entities
+    if message.topics and message.entities:
+        return message
+
+    # Extract topics and entities
+    extracted_topics, extracted_entities = await extract_topics_and_entities(
+        message.content, model_client
+    )
+
+    # Merge with existing topics and entities
+    message.topics = (
+        list(set(message.topics + extracted_topics))
+        if message.topics
+        else extracted_topics
+    )
+    message.entities = (
+        list(set(message.entities + extracted_entities))
+        if message.entities
+        else extracted_entities
+    )
+
+    return message

healthcheck.py redis_memory_server/healthcheck.pyhealthcheck.py renamed to redis_memory_server/healthcheck.py

@@ -2,7 +2,7 @@
 
 from fastapi import APIRouter
 
-from models import HealthCheckResponse
+from redis_memory_server.models import HealthCheckResponse
 
 
 router = APIRouter()

long_term_memory.py redis_memory_server/long_term_memory.pylong_term_memory.py renamed to redis_memory_server/long_term_memory.py

@@ -4,13 +4,13 @@
 from redis.asyncio import Redis
 from redis.commands.search.query import Query
 
-from models import (
+from redis_memory_server.models import (
     MemoryMessage,
     OpenAIClientWrapper,
     RedisearchResult,
     SearchResults,
 )
-from utils import REDIS_INDEX_NAME, Keys
+from redis_memory_server.utils import REDIS_INDEX_NAME, Keys
 
 
 logger = logging.getLogger(__name__)

main.py redis_memory_server/main.pymain.py renamed to redis_memory_server/main.py

@@ -2,20 +2,15 @@
 
 import structlog
 import uvicorn
-from dotenv import load_dotenv
 from fastapi import FastAPI
 
-import utils
-from models import MODEL_CONFIGS, ModelProvider
-
-
-load_dotenv()
-
-from config import settings
-from healthcheck import router as health_router
-from memory import router as memory_router
-from retrieval import router as retrieval_router
-from utils import ensure_redisearch_index, get_redis_conn
+from redis_memory_server import utils
+from redis_memory_server.config import settings
+from redis_memory_server.healthcheck import router as health_router
+from redis_memory_server.memory import router as memory_router
+from redis_memory_server.models import MODEL_CONFIGS, ModelProvider
+from redis_memory_server.retrieval import router as retrieval_router
+from redis_memory_server.utils import ensure_redisearch_index, get_redis_conn
 
 
 # Configure logging

memory.py redis_memory_server/memory.pymemory.py renamed to redis_memory_server/memory.py

@@ -4,17 +4,23 @@
 
 from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
 
-from config import settings
-from long_term_memory import index_messages
-from models import (
+from redis_memory_server.config import settings
+from redis_memory_server.extraction import handle_extraction
+from redis_memory_server.long_term_memory import index_messages
+from redis_memory_server.models import (
     AckResponse,
     GetSessionsQuery,
     MemoryMessage,
     MemoryMessagesAndContext,
     MemoryResponse,
 )
-from summarization import handle_compaction
-from utils import Keys, get_model_client, get_openai_client, get_redis_conn
+from redis_memory_server.summarization import handle_compaction
+from redis_memory_server.utils import (
+    Keys,
+    get_model_client,
+    get_openai_client,
+    get_redis_conn,
+)
 
 
 logger = logging.getLogger(__name__)
@@ -180,9 +186,23 @@ async def post_memory(
         current_time = int(time.time())
         await redis.zadd(sessions_key, {session_id: current_time})
 
+        # Get model client for extraction
+        model_client = await get_model_client(settings.generation_model)
+
+        # Process messages for topic/entity extraction
+        processed_messages = []
+        for msg in memory_messages.messages:
+            # Handle extraction in background for each message
+            background_tasks.add_task(
+                handle_extraction,
+                msg,
+                model_client,
+            )
+            processed_messages.append(msg)
+
         # Convert messages to JSON, handling topics and entities
         messages_json = []
-        for msg in memory_messages.messages:
+        for msg in processed_messages:
             msg_dict = msg.model_dump()
             # Convert lists to comma-separated strings for TAG fields
             msg_dict["topics"] = ",".join(msg.topics) if msg.topics else ""
@@ -193,10 +213,9 @@ async def post_memory(
         await redis.lpush(messages_key, *messages_json)  # type: ignore
 
         # Check if window size is exceeded
-        current_size = await redis.llen(messages_key)
+        current_size = await redis.llen(messages_key)  # type: ignore
         if current_size > settings.window_size:
             # Handle compaction in background
-            model_client = await get_model_client(settings.generation_model)
             background_tasks.add_task(
                 handle_compaction,
                 session_id,

models.py redis_memory_server/models.pymodels.py renamed to redis_memory_server/models.py

@@ -2,10 +2,10 @@
 
 from fastapi import APIRouter, HTTPException
 
-from config import settings
-from long_term_memory import search_messages
-from models import SearchPayload, SearchResults
-from utils import get_openai_client, get_redis_conn
+from redis_memory_server.config import settings
+from redis_memory_server.long_term_memory import search_messages
+from redis_memory_server.models import SearchPayload, SearchResults
+from redis_memory_server.utils import get_openai_client, get_redis_conn
 
 
 logger = logging.getLogger(__name__)
@@ -37,8 +37,7 @@ async def run_retrieval(
     client = await get_openai_client()
 
     try:
-        results = await search_messages(payload.text, session_id, client, redis)
-        return results
+        return await search_messages(payload.text, session_id, client, redis)
     except Exception as e:
         logger.error(f"Error in retrieval API: {e}")
-        raise HTTPException(status_code=500, detail="Internal server error")
+        raise HTTPException(status_code=500, detail="Internal server error") from e