diff --git a/athena/app/dependencies.py b/athena/app/dependencies.py index 52f0f6a..92f94c6 100644 --- a/athena/app/dependencies.py +++ b/athena/app/dependencies.py @@ -5,6 +5,7 @@ from athena.app.services.base_service import BaseService from athena.app.services.database_service import DatabaseService from athena.app.services.llm_service import LLMService +from athena.app.services.memory_service import MemoryService from athena.configuration.config import settings @@ -40,7 +41,10 @@ def initialize_services() -> Dict[str, BaseService]: settings.GEMINI_API_KEY, ) + memory_service = MemoryService(storage_backend=settings.MEMORY_STORAGE_BACKEND) + return { "llm_service": llm_service, "database_service": database_service, + "memory_service": memory_service, } diff --git a/athena/app/services/memory_service.py b/athena/app/services/memory_service.py new file mode 100644 index 0000000..0edde6e --- /dev/null +++ b/athena/app/services/memory_service.py @@ -0,0 +1,136 @@ +from typing import List, Optional + +from athena.app.services.base_service import BaseService +from init_memory_base.trajectory_memory_extractor_deepseek import MemoryUnit + + +class MemoryService(BaseService): + """ + Memory Service for managing memory information for software engineering agents. + + This service provides a unified interface for storing and retrieving memory units + that capture the execution trajectories of coding agents. Memory units contain + information about tasks, states, actions, and results that can be used to + improve agent performance through experience and contextual memory. + + Design Principles: + - Store complete execution trajectories as memory units + - Support semantic search for relevant past experiences + - Enable agents to learn from previous successes and failures + - Provide scalable storage backend options (in-memory, database, vector store) + - Support deduplication and relevance ranking of memories + + The service follows the MemoryBaseProtocol interface and can be implemented + with various storage backends (in-memory, SQL database, vector database). + """ + + def __init__(self, storage_backend: str = "in_memory"): + """ + Initialize the Memory Service with a specified storage backend. + + Args: + storage_backend: Type of storage backend to use. Options: + - "in_memory": Simple in-memory storage (default) + - "database": Persistent database storage + - "vector": Vector database for semantic search + """ + self.storage_backend = storage_backend + self._store = None + + def start(self): + """ + Start the memory service and initialize the storage backend. + + This method should be overridden to implement specific initialization + logic for the chosen storage backend (e.g., database connection, + vector index creation). + """ + # TODO: Implement storage backend initialization based on storage_backend + # For in_memory: create empty dictionary + # For database: establish connection and create tables if needed + # For vector: initialize vector index and embeddings + pass + + def close(self): + """ + Close the memory service and release any resources. + + This method should be overridden to implement specific cleanup + logic for the storage backend (e.g., close database connections, + persist in-memory data). + """ + # TODO: Implement resource cleanup based on storage_backend + pass + + def store_memory(self, memory_unit: MemoryUnit) -> None: + """ + Store a memory unit in the memory service. + + Args: + memory_unit: The MemoryUnit object containing task, state, + action, and result information to be stored. + + This method should handle deduplication using the memory unit's + canonical key and support bulk operations for efficiency. + """ + # TODO: Implement memory storage logic + # - Use memory_unit.key() for deduplication + # - Support different storage backends + # - Handle bulk operations efficiently + pass + + def search_memory(self, query: str, limit: int = 10) -> List[MemoryUnit]: + """ + Search for relevant memory units based on a query. + + Args: + query: Search query string to find relevant memories + limit: Maximum number of results to return (default: 10) + + Returns: + List of MemoryUnit objects matching the search query, + ordered by relevance. + + This method should support semantic search across multiple + fields (task description, state information, action details, + result summaries) and return the most relevant memories. + """ + # TODO: Implement memory search logic + # - Support semantic search across multiple fields + # - Implement relevance scoring and ranking + # - Handle different storage backend search capabilities + # - Return results ordered by relevance + pass + + def get_memory_by_key(self, key: str) -> Optional[MemoryUnit]: + """ + Retrieve a specific memory unit by its canonical key. + + Args: + key: The canonical key of the memory unit (generated by MemoryUnit.key()) + + Returns: + The MemoryUnit object if found, None otherwise. + """ + # TODO: Implement direct key-based retrieval + pass + + def get_all_memories(self) -> List[MemoryUnit]: + """ + Retrieve all memory units stored in the service. + + Returns: + List of all MemoryUnit objects in the service. + """ + # TODO: Implement retrieval of all memories + pass + + def clear_memories(self) -> None: + """ + Clear all memory units from the service. + + This method is primarily for testing and should be used with caution + in production environments. + """ + # TODO: Implement memory clearing functionality + pass \ No newline at end of file diff --git a/athena/configuration/config.py b/athena/configuration/config.py index ac2f7a9..62366ed 100644 --- a/athena/configuration/config.py +++ b/athena/configuration/config.py @@ -35,5 +35,10 @@ class Settings(BaseSettings): # Database DATABASE_URL: str + # Memory Service + MEMORY_STORAGE_BACKEND: str = "in_memory" # Options: in_memory, database, vector + MEMORY_MAX_STORED_UNITS: int = 1000 # Maximum number of memory units to store + MEMORY_SEARCH_LIMIT: int = 10 # Default search result limit + settings = Settings()