From 37d15d09ea4562f9dd4e845a2698c9eeaba4a81f Mon Sep 17 00:00:00 2001 From: Zhaoyang-Chu Date: Wed, 1 Oct 2025 04:26:18 +0800 Subject: [PATCH 01/30] feat: refine memory storage --- athena/app/services/database_service.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/athena/app/services/database_service.py b/athena/app/services/database_service.py index 57c903e..07331e3 100644 --- a/athena/app/services/database_service.py +++ b/athena/app/services/database_service.py @@ -20,7 +20,7 @@ def __init__(self, DATABASE_URL: str, max_retries: int = 5, initial_backoff: flo self._initial_backoff = initial_backoff # Create the database and tables - async def create_db_and_tables(self): + async def create_db_and_tables(self): # TODO: create dbs with custom names, "memories_xxx" async with self.engine.begin() as conn: # Ensure pgvector extension exists (safe to ignore if unavailable) try: From 86e982f8f3b3c450ea79ac8ef54a08823ad30257 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Thu, 2 Oct 2025 15:28:49 +0100 Subject: [PATCH 02/30] feat: enhance CustomChatOpenAI with logging and retry mechanism --- athena/chat_models/custom_chat_openai.py | 28 ++++------- athena/utils/llm_util.py | 62 ------------------------ 2 files changed, 9 insertions(+), 81 deletions(-) delete mode 100644 athena/utils/llm_util.py diff --git a/athena/chat_models/custom_chat_openai.py b/athena/chat_models/custom_chat_openai.py index a459ec7..f1ac826 100644 --- a/athena/chat_models/custom_chat_openai.py +++ b/athena/chat_models/custom_chat_openai.py @@ -1,20 +1,18 @@ -from typing import Any, List, Optional +import logging +import threading +from typing import Any, Optional from langchain_core.language_models import LanguageModelInput -from langchain_core.messages import BaseMessage, trim_messages +from langchain_core.messages import BaseMessage from langchain_core.runnables import RunnableConfig from langchain_openai import ChatOpenAI -from pydantic import PrivateAttr - -from athena.utils.llm_util import tiktoken_counter class CustomChatOpenAI(ChatOpenAI): - _max_input_tokens: int = PrivateAttr() - - def __init__(self, max_input_tokens: int, *args: Any, **kwargs: Any): + def __init__(self, *args: Any, **kwargs: Any): super().__init__(*args, **kwargs) - self._max_input_tokens = max_input_tokens + self.max_retries = 3 # Set the maximum number of retries + self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}") def bind_tools(self, tools, tool_choice=None, **kwargs): kwargs["parallel_tool_calls"] = False @@ -25,19 +23,11 @@ def invoke( input: LanguageModelInput, config: Optional[RunnableConfig] = None, *, - stop: Optional[List[str]] = None, + stop: Optional[list[str]] = None, **kwargs: Any, ) -> BaseMessage: return super().invoke( - input=trim_messages( - input, - token_counter=tiktoken_counter, - strategy="last", - max_tokens=self._max_input_tokens, - start_on="human", - end_on=("human", "tool"), - include_system=True, - ), + input=input, config=config, stop=stop, **kwargs, diff --git a/athena/utils/llm_util.py b/athena/utils/llm_util.py deleted file mode 100644 index db02477..0000000 --- a/athena/utils/llm_util.py +++ /dev/null @@ -1,62 +0,0 @@ -from typing import Sequence - -import tiktoken -from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage -from langchain_core.output_parsers import StrOutputParser - - -def str_token_counter(text: str) -> int: - """Counts the number of tokens in a string using tiktoken's o200k_base encoding. - - Args: - text: The input string to count tokens for. - - Returns: - The number of tokens in the input string. - """ - enc = tiktoken.get_encoding("o200k_base") - return len(enc.encode(text)) - - -def tiktoken_counter(messages: Sequence[BaseMessage]) -> int: - """Counts tokens across multiple message types using tiktoken tokenization. - - Approximately reproduces the token counting methodology from OpenAI's cookbook: - https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb - - Args: - messages: A sequence of BaseMessage objects (HumanMessage, AIMessage, - ToolMessage, or SystemMessage) to count tokens for. - - Returns: - The total number of tokens across all messages, including overhead tokens. - - Raises: - ValueError: If an unsupported message type is encountered. - - Notes: - - Uses a fixed overhead of 3 tokens for reply priming - - Adds 3 tokens per message for message formatting - - Adds 1 token per message name if present - - For simplicity, only supports string message contents - """ - output_parser = StrOutputParser() - num_tokens = 3 # every reply is primed with <|start|>assistant<|message|> - tokens_per_message = 3 - tokens_per_name = 1 - for msg in messages: - if isinstance(msg, HumanMessage): - role = "user" - elif isinstance(msg, AIMessage): - role = "assistant" - elif isinstance(msg, ToolMessage): - role = "tool" - elif isinstance(msg, SystemMessage): - role = "system" - else: - raise ValueError(f"Unsupported messages type {msg.__class__}") - msg_content = output_parser.invoke(msg) - num_tokens += tokens_per_message + str_token_counter(role) + str_token_counter(msg_content) - if msg.name: - num_tokens += tokens_per_name + str_token_counter(msg.name) - return num_tokens From 60758f91ccd292749ace3d6363ee0e3c5a73f838 Mon Sep 17 00:00:00 2001 From: Zhaoyang-Chu Date: Sat, 4 Oct 2025 15:40:02 +0000 Subject: [PATCH 03/30] Adjust service names --- athena/app/api/main.py | 26 +++++++++---------- athena/app/dependencies.py | 12 ++++----- ... => episodic_memory_extraction_service.py} | 6 ++--- ..._service.py => episodic_memory_service.py} | 4 +-- ....py => episodic_memory_storage_service.py} | 6 ++--- athena/scripts/offline_ingest_hf.py | 8 +++--- 6 files changed, 31 insertions(+), 31 deletions(-) rename athena/app/services/{memory_extraction_service.py => episodic_memory_extraction_service.py} (99%) rename athena/app/services/{memory_service.py => episodic_memory_service.py} (97%) rename athena/app/services/{memory_storage_service.py => episodic_memory_storage_service.py} (97%) diff --git a/athena/app/api/main.py b/athena/app/api/main.py index 9b884b2..f1e5356 100644 --- a/athena/app/api/main.py +++ b/athena/app/api/main.py @@ -5,30 +5,30 @@ api_router = APIRouter() -@api_router.get("/memory/search", tags=["memory"]) -async def search_memory( +@api_router.get("/episodic_memory/search", tags=["episodic_memory"]) +async def search_episodic_memory( request: Request, q: str, field: str = "task_state", limit: int = Query(default=settings.MEMORY_SEARCH_LIMIT, ge=1, le=100), ): services = getattr(request.app.state, "service", {}) - memory_service = services.get("memory_service") - if memory_service is None: - raise HTTPException(status_code=503, detail="Memory service not initialized") + episodic_memory_service = services.get("episodic_memory_service") + if episodic_memory_service is None: + raise HTTPException(status_code=503, detail="Episodic memory service not initialized") if field not in {"task_state", "task", "state"}: raise HTTPException(status_code=400, detail=f"Invalid field: {field}") - results = await memory_service.search_memory(q, limit=limit, field=field) + results = await episodic_memory_service.search_memory(q, limit=limit, field=field) return [m.model_dump() for m in results] -@api_router.get("/memory/{memory_id}", tags=["memory"]) -async def get_memory(request: Request, memory_id: str): +@api_router.get("/episodic_memory/{memory_id}", tags=["episodic_memory"]) +async def get_episodic_memory(request: Request, memory_id: str): services = getattr(request.app.state, "service", {}) - memory_service = services.get("memory_service") - if memory_service is None: - raise HTTPException(status_code=503, detail="Memory service not initialized") - result = await memory_service.get_memory_by_key(memory_id) + episodic_memory_service = services.get("episodic_memory_service") + if episodic_memory_service is None: + raise HTTPException(status_code=503, detail="Episodic memory service not initialized") + result = await episodic_memory_service.get_memory_by_key(memory_id) if result is None: - raise HTTPException(status_code=404, detail="Memory not found") + raise HTTPException(status_code=404, detail="Episodic memory not found") return result.model_dump() diff --git a/athena/app/dependencies.py b/athena/app/dependencies.py index d292359..8a88f50 100644 --- a/athena/app/dependencies.py +++ b/athena/app/dependencies.py @@ -6,8 +6,8 @@ from athena.app.services.database_service import DatabaseService from athena.app.services.embedding_service import EmbeddingService from athena.app.services.llm_service import LLMService -from athena.app.services.memory_service import MemoryService -from athena.app.services.memory_storage_service import MemoryStorageService +from athena.app.services.episodic_memory_service import EpisodicMemoryService +from athena.app.services.episodic_memory_storage_service import EpisodicMemoryStorageService from athena.configuration.config import settings @@ -52,14 +52,14 @@ def initialize_services() -> Dict[str, BaseService]: embed_dim=settings.EMBEDDING_DIM or 1024, ) - memory_store = MemoryStorageService(database_service.get_sessionmaker(), embedding_service) + episodic_memory_store = EpisodicMemoryStorageService(database_service.get_sessionmaker(), embedding_service) - memory_service = MemoryService( - storage_backend=settings.MEMORY_STORAGE_BACKEND, store=memory_store + episodic_memory_service = EpisodicMemoryService( + storage_backend=settings.MEMORY_STORAGE_BACKEND, store=episodic_memory_store ) return { "llm_service": llm_service, "database_service": database_service, - "memory_service": memory_service, + "episodic_memory_service": episodic_memory_service, } diff --git a/athena/app/services/memory_extraction_service.py b/athena/app/services/episodic_memory_extraction_service.py similarity index 99% rename from athena/app/services/memory_extraction_service.py rename to athena/app/services/episodic_memory_extraction_service.py index b6a5071..36b0635 100644 --- a/athena/app/services/memory_extraction_service.py +++ b/athena/app/services/episodic_memory_extraction_service.py @@ -43,7 +43,7 @@ def __init__(self, source_name: str, run_id: str): super().__init__(f"Failed to extract memory units from {run_id} in {source_name}") -class MemoryExtractionService(BaseService): +class EpisodicMemoryExtractionService(BaseService): """ Service for extracting structured memory units from interaction trajectories. @@ -69,7 +69,7 @@ class MemoryExtractionService(BaseService): 5. Storage -> Persist to memory system Usage: - service = MemoryExtractionService(llm_service) + service = EpisodicMemoryExtractionService(llm_service) memory_units = service.extract_from_trajectories(trajectory_data) """ @@ -868,7 +868,7 @@ def _repair_json_with_llm(self, raw: str, expect_key: Optional[str]) -> str: if __name__ == "__main__": - service = MemoryExtractionService( + service = EpisodicMemoryExtractionService( llm_service=LLMService( model_name="vertex:gemini-2.5-flash", model_temperature=0.0, diff --git a/athena/app/services/memory_service.py b/athena/app/services/episodic_memory_service.py similarity index 97% rename from athena/app/services/memory_service.py rename to athena/app/services/episodic_memory_service.py index af740ac..e2247a4 100644 --- a/athena/app/services/memory_service.py +++ b/athena/app/services/episodic_memory_service.py @@ -6,7 +6,7 @@ from athena.models.memory import MemoryUnit -class MemoryService(BaseService): +class EpisodicMemoryService(BaseService): """ Memory Service for managing memory information for software engineering agents. @@ -45,7 +45,7 @@ async def start(self): """Initialize the storage backend if needed and validate configuration.""" if self.storage_backend in {"database", "vector"} and self._store is None: raise RuntimeError( - "MemoryService requires a storage store for database/vector backends" + "EpisodicMemoryService requires a storage store for database/vector backends" ) if self._store is not None and hasattr(self._store, "start"): if inspect.iscoroutinefunction(self._store.start): diff --git a/athena/app/services/memory_storage_service.py b/athena/app/services/episodic_memory_storage_service.py similarity index 97% rename from athena/app/services/memory_storage_service.py rename to athena/app/services/episodic_memory_storage_service.py index 1d2e308..4f85da7 100644 --- a/athena/app/services/memory_storage_service.py +++ b/athena/app/services/episodic_memory_storage_service.py @@ -21,7 +21,7 @@ def _ensure_dim(vec: List[float]) -> List[float]: return vec + [0.0] * (dim - len(vec)) -class MemoryStorageService(BaseService): +class EpisodicMemoryStorageService(BaseService): """ Postgres-backed memory store using SQLModel and optional embeddings for semantic search. @@ -179,7 +179,7 @@ def _serialize_state(u: MemoryUnit) -> str: @staticmethod def _serialize_task_state(u: MemoryUnit) -> str: # TODO: text templates for static embeddings return ( - MemoryStorageService._serialize_task(u) + EpisodicMemoryStorageService._serialize_task(u) + "\n\n" - + MemoryStorageService._serialize_state(u) + + EpisodicMemoryStorageService._serialize_state(u) ) diff --git a/athena/scripts/offline_ingest_hf.py b/athena/scripts/offline_ingest_hf.py index 2baf223..d38ad7e 100644 --- a/athena/scripts/offline_ingest_hf.py +++ b/athena/scripts/offline_ingest_hf.py @@ -5,8 +5,8 @@ from athena.app.services.database_service import DatabaseService from athena.app.services.embedding_service import EmbeddingService from athena.app.services.llm_service import LLMService -from athena.app.services.memory_extraction_service import MemoryExtractionService -from athena.app.services.memory_storage_service import MemoryStorageService +from athena.app.services.episodic_memory_extraction_service import EpisodicMemoryExtractionService +from athena.app.services.episodic_memory_storage_service import EpisodicMemoryStorageService from athena.configuration.config import settings @@ -29,7 +29,7 @@ async def main(): embed_dim=settings.EMBEDDING_DIM or 1024, ) - store = MemoryStorageService(db.get_sessionmaker(), embedding_service) + store = EpisodicMemoryStorageService(db.get_sessionmaker(), embedding_service) llm = LLMService( model_name=settings.MODEL_NAME, @@ -42,7 +42,7 @@ async def main(): gemini_api_key=settings.GEMINI_API_KEY, ) - extractor = MemoryExtractionService(llm_service=llm, memory_store=store) + extractor = EpisodicMemoryExtractionService(llm_service=llm, memory_store=store) extractor.extract_from_huggingface_trajectory_repository(args.repo, args.split) From 544f4e1863ec83a7683e81c36c4280d871c42ee0 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Sat, 4 Oct 2025 16:59:38 +0100 Subject: [PATCH 04/30] feat: add MemoryContext and Query models, and implement SemanticMemoryService --- athena/app/dependencies.py | 6 ++++-- athena/app/services/semantic_memory_service.py | 11 +++++++++++ athena/models/memory_context.py | 6 ++++++ athena/models/query.py | 7 +++++++ athena/models/semantic_memory.py | 9 +++++++++ athena/scripts/offline_ingest_hf.py | 2 +- 6 files changed, 38 insertions(+), 3 deletions(-) create mode 100644 athena/app/services/semantic_memory_service.py create mode 100644 athena/models/memory_context.py create mode 100644 athena/models/query.py create mode 100644 athena/models/semantic_memory.py diff --git a/athena/app/dependencies.py b/athena/app/dependencies.py index 8a88f50..49511ab 100644 --- a/athena/app/dependencies.py +++ b/athena/app/dependencies.py @@ -5,9 +5,9 @@ from athena.app.services.base_service import BaseService from athena.app.services.database_service import DatabaseService from athena.app.services.embedding_service import EmbeddingService -from athena.app.services.llm_service import LLMService from athena.app.services.episodic_memory_service import EpisodicMemoryService from athena.app.services.episodic_memory_storage_service import EpisodicMemoryStorageService +from athena.app.services.llm_service import LLMService from athena.configuration.config import settings @@ -52,7 +52,9 @@ def initialize_services() -> Dict[str, BaseService]: embed_dim=settings.EMBEDDING_DIM or 1024, ) - episodic_memory_store = EpisodicMemoryStorageService(database_service.get_sessionmaker(), embedding_service) + episodic_memory_store = EpisodicMemoryStorageService( + database_service.get_sessionmaker(), embedding_service + ) episodic_memory_service = EpisodicMemoryService( storage_backend=settings.MEMORY_STORAGE_BACKEND, store=episodic_memory_store diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py new file mode 100644 index 0000000..8c519b2 --- /dev/null +++ b/athena/app/services/semantic_memory_service.py @@ -0,0 +1,11 @@ +class SemanticMemoryService: + def __init__(self): + pass + + def store_memory(self, memory): + # Logic to store memory + pass + + def retrieve_memory(self, query): + # Logic to retrieve memory based on a query + pass diff --git a/athena/models/memory_context.py b/athena/models/memory_context.py new file mode 100644 index 0000000..ce4afee --- /dev/null +++ b/athena/models/memory_context.py @@ -0,0 +1,6 @@ +from pydantic import BaseModel + + +class MemoryContext(BaseModel): + overview: str + content: str diff --git a/athena/models/query.py b/athena/models/query.py new file mode 100644 index 0000000..b9c18b6 --- /dev/null +++ b/athena/models/query.py @@ -0,0 +1,7 @@ +from pydantic import BaseModel + + +class Query(BaseModel): + essential_query: str + extra_requirements: str = "" + purpose: str = "" diff --git a/athena/models/semantic_memory.py b/athena/models/semantic_memory.py new file mode 100644 index 0000000..62db680 --- /dev/null +++ b/athena/models/semantic_memory.py @@ -0,0 +1,9 @@ +from pydantic import BaseModel + +from athena.models.memory_context import MemoryContext +from athena.models.query import Query + + +class SemanticMemoryUnit(BaseModel): + query: Query + memory_context: MemoryContext diff --git a/athena/scripts/offline_ingest_hf.py b/athena/scripts/offline_ingest_hf.py index d38ad7e..9da2e1a 100644 --- a/athena/scripts/offline_ingest_hf.py +++ b/athena/scripts/offline_ingest_hf.py @@ -4,9 +4,9 @@ from athena.app.services.database_service import DatabaseService from athena.app.services.embedding_service import EmbeddingService -from athena.app.services.llm_service import LLMService from athena.app.services.episodic_memory_extraction_service import EpisodicMemoryExtractionService from athena.app.services.episodic_memory_storage_service import EpisodicMemoryStorageService +from athena.app.services.llm_service import LLMService from athena.configuration.config import settings From 6ae8cde3012ac9a1db02658b6bf9c2a1953123d0 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Sat, 4 Oct 2025 17:07:54 +0100 Subject: [PATCH 05/30] feat: update SemanticMemoryService to support async memory storage and retrieval with URL and commit ID --- athena/app/services/semantic_memory_service.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index 8c519b2..0f025a8 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -1,11 +1,20 @@ -class SemanticMemoryService: +from typing import List + +from athena.app.services.base_service import BaseService +from athena.models.query import Query +from athena.models.semantic_memory import SemanticMemoryUnit + + +class SemanticMemoryService(BaseService): def __init__(self): pass - def store_memory(self, memory): + async def store_memory(self, url: str, commit_id: str, memory: SemanticMemoryUnit): # Logic to store memory pass - def retrieve_memory(self, query): + async def retrieve_memory( + self, url: str, commit_id: str, query: Query + ) -> List[SemanticMemoryUnit]: # Logic to retrieve memory based on a query pass From d17e9290f04f730ccc2c7e332fb79efe6100c616 Mon Sep 17 00:00:00 2001 From: Zhaoyang-Chu Date: Sat, 4 Oct 2025 16:30:32 +0000 Subject: [PATCH 06/30] Adjust service name --- .../episodic_memory_extraction_service.py | 42 ++++++------ .../app/services/episodic_memory_service.py | 4 +- .../episodic_memory_storage_service.py | 44 ++++++------ .../semantic_memory_extraction_service.py | 13 ++++ .../semantic_memory_storage_service.py | 11 +++ athena/entity/__init__.py | 4 +- athena/entity/memory.py | 68 +++++++++---------- athena/models/__init__.py | 14 ++-- .../models/{memory.py => episodic_memory.py} | 32 ++++----- ...ction.py => episodic_memory_extraction.py} | 0 10 files changed, 128 insertions(+), 104 deletions(-) create mode 100644 athena/app/services/semantic_memory_extraction_service.py create mode 100644 athena/app/services/semantic_memory_storage_service.py rename athena/models/{memory.py => episodic_memory.py} (78%) rename athena/prompts/{memory_extraction.py => episodic_memory_extraction.py} (100%) diff --git a/athena/app/services/episodic_memory_extraction_service.py b/athena/app/services/episodic_memory_extraction_service.py index 36b0635..836f507 100644 --- a/athena/app/services/episodic_memory_extraction_service.py +++ b/athena/app/services/episodic_memory_extraction_service.py @@ -11,12 +11,12 @@ from athena.app.services.base_service import BaseService from athena.app.services.llm_service import LLMService -from athena.app.services.memory_storage_service import MemoryStorageService +from athena.app.services.episodic_memory_storage_service import EpisodicMemoryStorageService from athena.models import ( Action, - MemorySource, - MemoryTimestamp, - MemoryUnit, + EpisodicMemorySource, + EpisodicMemoryTimestamp, + EpisodicMemoryUnit, Message, Result, State, @@ -64,7 +64,7 @@ class EpisodicMemoryExtractionService(BaseService): The service follows a pipeline pattern: 1. Data Source -> Load trajectories 2. Extraction Strategy -> Extract components (task, action, result, state) - 3. Memory Unit Assembly -> Combine components into MemoryUnit + 3. Memory Unit Assembly -> Combine components into EpisodicMemoryUnit 4. Validation -> Ensure data quality and consistency 5. Storage -> Persist to memory system @@ -78,7 +78,7 @@ def __init__( llm_service: LLMService, batch_size: int = 100, max_retries: int = 3, - memory_store: Optional[MemoryStorageService] = None, + memory_store: Optional[EpisodicMemoryStorageService] = None, ): """ Initialize the Memory Extraction Service. @@ -91,7 +91,7 @@ def __init__( self.llm_service = llm_service self.batch_size = batch_size self.max_retries = max_retries - self._extraction_cache: Dict[str, MemoryUnit] = {} + self._extraction_cache: Dict[str, EpisodicEpisodicMemoryUnit] = {} self.memory_store = memory_store # self._logger = get_logger(__name__) @@ -106,7 +106,7 @@ def close(self): def extract_from_huggingface_trajectory_repository( # TODO: batch extraction self, repo_name: str, split: str - ) -> List[MemoryUnit]: + ) -> List[EpisodicMemoryUnit]: """ Extract memory units from a HuggingFace trajectory repository. @@ -185,7 +185,7 @@ def _pick(d: Dict[str, Any], keys) -> Optional[Any]: def _extract_memory_source( self, source: str, run_id: str, metadata: Optional[Dict[str, Any]] = None - ) -> MemorySource: + ) -> EpisodicMemorySource: """ Extract memory source for a trajectory. @@ -195,9 +195,9 @@ def _extract_memory_source( metadata: Optional additional metadata Returns: - MemorySource object + EpisodicMemorySource object """ - return MemorySource( + return EpisodicMemorySource( source_name=source, run_id=run_id, metadata=metadata or {}, @@ -206,9 +206,9 @@ def _extract_memory_source( def _extract_memory_units_by_action_windows( self, messages: List[Message], - memory_source: MemorySource, - ) -> List[MemoryUnit]: - ordered_memory_units: List[MemoryUnit] = [] + memory_source: EpisodicMemorySource, + ) -> List[EpisodicMemoryUnit]: + ordered_memory_units: List[EpisodicMemoryUnit] = [] window_msgs: List[Message] = [] window_first_action: Optional[Message] = None @@ -285,11 +285,11 @@ def _is_action_message(self, message: Message) -> bool: def _create_memory_unit( self, - source: MemorySource, + source: EpisodicMemorySource, task: Task, window_msgs: List[Message], - prior_units: List[MemoryUnit], - ) -> MemoryUnit: + prior_units: List[EpisodicMemoryUnit], + ) -> EpisodicMemoryUnit: """ Extract a single memory unit from a window of messages. - Synthesize state.done from prior actions @@ -301,9 +301,9 @@ def _create_memory_unit( state_done = self._synthesize_state_done_from_context(prior_units, task) state_todo = self._synthesize_state_todo_from_window(window_msgs, task, state_done, action) result = self._extract_result_from_window(window_msgs, action) - return MemoryUnit( + return EpisodicMemoryUnit( memory_id=str(uuid.uuid4()), - timestamp=MemoryTimestamp( + timestamp=EpisodicMemoryTimestamp( created_at=datetime.now(timezone.utc), updated_at=None, invalid_at=None, @@ -384,7 +384,7 @@ def s(x): def _synthesize_state_from_context( self, - prior_units: List[MemoryUnit], + prior_units: List[EpisodicMemoryUnit], task: Task, window_msgs: List[Message], current_action: Optional[Action] = None, @@ -398,7 +398,7 @@ def _synthesize_state_from_context( ) return State(done=state_done, todo=state_todo) # TODO: open_file, working_dir - def _synthesize_state_done_from_context(self, prior_units: List[MemoryUnit], task: Task) -> str: + def _synthesize_state_done_from_context(self, prior_units: List[EpisodicMemoryUnit], task: Task) -> str: """ Summarize previous context into a concise `state.done` string. Summary of what has ALREADY BEEN COMPLETED (no plans). diff --git a/athena/app/services/episodic_memory_service.py b/athena/app/services/episodic_memory_service.py index e2247a4..fd64356 100644 --- a/athena/app/services/episodic_memory_service.py +++ b/athena/app/services/episodic_memory_service.py @@ -2,8 +2,8 @@ from typing import List, Optional, Sequence from athena.app.services.base_service import BaseService -from athena.app.services.memory_storage_service import MemoryStorageService -from athena.models.memory import MemoryUnit +from athena.app.services.episodic_memory_storage_service import EpisodicMemoryStorageService +from athena.models.memory import EpisodicMemoryUnit class EpisodicMemoryService(BaseService): diff --git a/athena/app/services/episodic_memory_storage_service.py b/athena/app/services/episodic_memory_storage_service.py index 4f85da7..a751bb4 100644 --- a/athena/app/services/episodic_memory_storage_service.py +++ b/athena/app/services/episodic_memory_storage_service.py @@ -7,8 +7,8 @@ from athena.app.services.base_service import BaseService from athena.app.services.embedding_service import EmbeddingService from athena.configuration.config import settings -from athena.entity.memory import MemoryUnitDB -from athena.models.memory import MemoryUnit +from athena.entity.memory import EpisodicMemoryUnitDB +from athena.models.memory import EpisodicMemoryUnit def _ensure_dim(vec: List[float]) -> List[float]: @@ -39,7 +39,7 @@ def __init__( self._embeddings = embedding_service self._max = max_stored_units - async def upsert(self, units: List[MemoryUnit]) -> None: + async def upsert(self, units: List[EpisodicMemoryUnit]) -> None: if not units: return async with self._sessionmaker() as session: @@ -47,7 +47,7 @@ async def upsert(self, units: List[MemoryUnit]) -> None: await self._upsert_one(session, u) await session.commit() - async def _upsert_one(self, session: AsyncSession, unit: MemoryUnit) -> None: + async def _upsert_one(self, session: AsyncSession, unit: EpisodicMemoryUnit) -> None: # Prepare embeddings if configured task_text = self._serialize_task(unit) state_text = self._serialize_state(unit) @@ -66,18 +66,18 @@ async def _upsert_one(self, session: AsyncSession, unit: MemoryUnit) -> None: # Check existing existing = await session.scalar( - select(MemoryUnitDB).where(col(MemoryUnitDB.memory_id) == unit.memory_id) + select(EpisodicMemoryUnitDB).where(col(EpisodicMemoryUnitDB.memory_id) == unit.memory_id) ) if existing is None: - row = MemoryUnitDB.from_memory_unit(unit) + row = EpisodicMemoryUnitDB.from_memory_unit(unit) row.task_embedding = task_vec row.state_embedding = state_vec row.task_state_embedding = task_state_vec session.add(row) else: # Update fields - fresh = MemoryUnitDB.from_memory_unit(unit) + fresh = EpisodicMemoryUnitDB.from_memory_unit(unit) for attr in ( "memory_created_at", "memory_updated_at", @@ -114,42 +114,42 @@ async def _upsert_one(self, session: AsyncSession, unit: MemoryUnit) -> None: async def search_by_text( self, text: str, field: str = "task_state", limit: int = 10 - ) -> List[MemoryUnit]: + ) -> List[EpisodicMemoryUnit]: if self._embeddings is None: return [] q_vec = _ensure_dim(self._embeddings.embed([text])[0]) async with self._sessionmaker() as session: # Choose column by field col_expr = { - "task": MemoryUnitDB.task_embedding, - "state": MemoryUnitDB.state_embedding, - "task_state": MemoryUnitDB.task_state_embedding, - }.get(field, MemoryUnitDB.task_state_embedding) + "task": EpisodicMemoryUnitDB.task_embedding, + "state": EpisodicMemoryUnitDB.state_embedding, + "task_state": EpisodicMemoryUnitDB.task_state_embedding, + }.get(field, EpisodicMemoryUnitDB.task_state_embedding) # Order by cosine distance using pgvector `<=>` res = await session.execute( - select(MemoryUnitDB) + select(EpisodicMemoryUnitDB) .where(col_expr.is_not(None)) .order_by(col_expr.cosine_distance(q_vec)) .limit(limit) ) - rows: List[MemoryUnitDB] = list(res.scalars()) + rows: List[EpisodicMemoryUnitDB] = list(res.scalars()) return [r.to_memory_unit() for r in rows] - async def get_by_memory_id(self, memory_id: str) -> Optional[MemoryUnit]: + async def get_by_memory_id(self, memory_id: str) -> Optional[EpisodicMemoryUnit]: async with self._sessionmaker() as session: row = await session.scalar( - select(MemoryUnitDB).where(col(MemoryUnitDB.memory_id) == memory_id) + select(EpisodicMemoryUnitDB).where(col(EpisodicMemoryUnitDB.memory_id) == memory_id) ) return None if row is None else row.to_memory_unit() - async def list_all(self, limit: Optional[int] = None) -> List[MemoryUnit]: + async def list_all(self, limit: Optional[int] = None) -> List[EpisodicMemoryUnit]: async with self._sessionmaker() as session: - stmt = select(MemoryUnitDB).order_by(MemoryUnitDB.id.desc()) + stmt = select(EpisodicMemoryUnitDB).order_by(EpisodicMemoryUnitDB.id.desc()) if limit is not None and limit > 0: stmt = stmt.limit(limit) res = await session.execute(stmt) - rows: List[MemoryUnitDB] = list(res.scalars()) + rows: List[EpisodicMemoryUnitDB] = list(res.scalars()) return [r.to_memory_unit() for r in rows] async def clear_all(self) -> None: @@ -159,13 +159,13 @@ async def clear_all(self) -> None: await session.commit() @staticmethod - def _serialize_task(u: MemoryUnit) -> str: + def _serialize_task(u: EpisodicMemoryUnit) -> str: t = u.task parts = [t.issue_title, t.issue_type, t.repository, t.issue_body, t.issue_comments] return "\n".join([p for p in parts if p]) @staticmethod - def _serialize_state(u: MemoryUnit) -> str: + def _serialize_state(u: EpisodicMemoryUnit) -> str: s = u.state parts = [s.done, s.todo] if s.open_file: @@ -177,7 +177,7 @@ def _serialize_state(u: MemoryUnit) -> str: return "\n".join([str(p) for p in parts if p]) @staticmethod - def _serialize_task_state(u: MemoryUnit) -> str: # TODO: text templates for static embeddings + def _serialize_task_state(u: EpisodicMemoryUnit) -> str: # TODO: text templates for static embeddings return ( EpisodicMemoryStorageService._serialize_task(u) + "\n\n" diff --git a/athena/app/services/semantic_memory_extraction_service.py b/athena/app/services/semantic_memory_extraction_service.py new file mode 100644 index 0000000..009390a --- /dev/null +++ b/athena/app/services/semantic_memory_extraction_service.py @@ -0,0 +1,13 @@ +import asyncio +import html +import json +import re +import uuid +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional + +from datasets import load_dataset +from tqdm import tqdm + +from athena.app.services.base_service import BaseService +from athena.app.services.llm_service import LLMService diff --git a/athena/app/services/semantic_memory_storage_service.py b/athena/app/services/semantic_memory_storage_service.py new file mode 100644 index 0000000..3815660 --- /dev/null +++ b/athena/app/services/semantic_memory_storage_service.py @@ -0,0 +1,11 @@ +from typing import List, Optional + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker +from sqlmodel import col + +from athena.app.services.base_service import BaseService +from athena.app.services.embedding_service import EmbeddingService +from athena.configuration.config import settings +from athena.entity.memory import SemanticMemoryUnitDB +from athena.models.memory import SemanticMemoryUnit diff --git a/athena/entity/__init__.py b/athena/entity/__init__.py index f5b71d1..51367cc 100644 --- a/athena/entity/__init__.py +++ b/athena/entity/__init__.py @@ -1,3 +1,3 @@ -from .memory import MemoryUnitDB +from .memory import EpisodicMemoryUnitDB -__all__ = ["MemoryUnitDB"] +__all__ = ["EpisodicMemoryUnitDB"] diff --git a/athena/entity/memory.py b/athena/entity/memory.py index eb2730c..f22aa2d 100644 --- a/athena/entity/memory.py +++ b/athena/entity/memory.py @@ -7,30 +7,30 @@ from sqlmodel import Field, SQLModel, UniqueConstraint from athena.configuration.config import settings -from athena.models import Action, MemorySource, MemoryTimestamp, MemoryUnit, Result, State, Task +from athena.models import Action, EpisodicMemorySource, EpisodicMemoryTimestamp, EpisodicMemoryUnit, Result, State, Task # Database models for persistent storage -class MemoryUnitDB(SQLModel, table=True): - """Database model for persistent storage of memory units.""" +class EpisodicMemoryUnitDB(SQLModel, table=True): + """Database model for persistent storage of episodic memory units.""" - __tablename__ = "memory_units" - __table_args__ = (UniqueConstraint("memory_id", name="uq_memory_id"),) + __tablename__ = "episodic_memories" + __table_args__ = (UniqueConstraint("episodic_memory_id", name="uq_episodic_memory_id"),) id: Optional[int] = Field(default=None, primary_key=True) # Source information - memory_id: str - memory_source_name: str - memory_run_id: str - memory_created_at: datetime = Field(sa_column=Column(DateTime(timezone=True))) - memory_updated_at: Optional[datetime] = Field( + episodic_memory_id: str + episodic_memory_source_name: str + episodic_memory_run_id: str + episodic_memory_created_at: datetime = Field(sa_column=Column(DateTime(timezone=True))) + episodic_memory_updated_at: Optional[datetime] = Field( default=None, sa_column=Column(DateTime(timezone=True)) ) - memory_invalid_at: Optional[datetime] = Field( + episodic_memory_invalid_at: Optional[datetime] = Field( default=None, sa_column=Column(DateTime(timezone=True)) ) - memory_metadata: str = Field( + episodic_memory_metadata: str = Field( default="{}", sa_column=Column(Text) ) # JSON string for Dict[str, Any] @@ -70,16 +70,16 @@ class MemoryUnitDB(SQLModel, table=True): ) @classmethod - def from_memory_unit(cls, memory_unit: MemoryUnit) -> "MemoryUnitDB": - """Create a database model from a MemoryUnit.""" + def from_episodic_memory_unit(cls, memory_unit: EpisodicMemoryUnit) -> "EpisodicMemoryUnitDB": + """Create a database model from a EpisodicMemoryUnit.""" return cls( - memory_id=memory_unit.memory_id, - memory_source_name=memory_unit.source.source_name, - memory_run_id=memory_unit.source.run_id, - memory_created_at=memory_unit.timestamp.created_at, - memory_updated_at=memory_unit.timestamp.updated_at, - memory_invalid_at=memory_unit.timestamp.invalid_at, - memory_metadata=json.dumps(memory_unit.source.metadata) + episodic_memory_id=memory_unit.episodic_memory_id, + episodic_memory_source_name=memory_unit.source.source_name, + episodic_memory_run_id=memory_unit.source.run_id, + episodic_memory_created_at=memory_unit.timestamp.created_at, + episodic_memory_updated_at=memory_unit.timestamp.updated_at, + episodic_memory_invalid_at=memory_unit.timestamp.invalid_at, + episodic_memory_metadata=json.dumps(memory_unit.source.metadata) if memory_unit.source.metadata else "{}", task_issue_title=memory_unit.task.issue_title, @@ -103,20 +103,20 @@ def from_memory_unit(cls, memory_unit: MemoryUnit) -> "MemoryUnitDB": result_exit_code=memory_unit.result.exit_code, ) - def to_memory_unit(self) -> MemoryUnit: - """Convert database model back to MemoryUnit.""" - return MemoryUnit( - memory_id=self.memory_id, - timestamp=MemoryTimestamp( - created_at=self.memory_created_at, - updated_at=self.memory_updated_at, - invalid_at=self.memory_invalid_at, + def to_episodic_memory_unit(self) -> EpisodicMemoryUnit: + """Convert database model back to EpisodicMemoryUnit.""" + return EpisodicMemoryUnit( + episodic_memory_id=self.episodic_memory_id, + timestamp=EpisodicMemoryTimestamp( + created_at=self.episodic_memory_created_at, + updated_at=self.episodic_memory_updated_at, + invalid_at=self.episodic_memory_invalid_at, ), - source=MemorySource( - source_name=self.memory_source_name, - run_id=self.memory_run_id, - metadata=json.loads(self.memory_metadata) - if self.memory_metadata not in (None, "", "null") + source=EpisodicMemorySource( + source_name=self.episodic_memory_source_name, + run_id=self.episodic_memory_run_id, + metadata=json.loads(self.episodic_memory_metadata) + if self.episodic_memory_metadata not in (None, "", "null") else {}, ), task=Task( diff --git a/athena/models/__init__.py b/athena/models/__init__.py index 5228135..408202e 100644 --- a/athena/models/__init__.py +++ b/athena/models/__init__.py @@ -1,8 +1,8 @@ -from .memory import ( +from .episodic_memory import ( Action, - MemorySource, - MemoryTimestamp, - MemoryUnit, + EpisodicMemorySource, + EpisodicMemoryTimestamp, + EpisodicMemoryUnit, Result, State, Task, @@ -11,9 +11,9 @@ __all__ = [ "Message", - "MemoryUnit", - "MemorySource", - "MemoryTimestamp", + "EpisodicMemoryUnit", + "EpisodicMemorySource", + "EpisodicMemoryTimestamp", "Task", "State", "Action", diff --git a/athena/models/memory.py b/athena/models/episodic_memory.py similarity index 78% rename from athena/models/memory.py rename to athena/models/episodic_memory.py index 4776665..623fc98 100644 --- a/athena/models/memory.py +++ b/athena/models/episodic_memory.py @@ -5,23 +5,23 @@ from pydantic import BaseModel, Field -class MemoryTimestamp(BaseModel): - """Lifecycle timestamps for a memory unit.""" +class EpisodicMemoryTimestamp(BaseModel): + """Lifecycle timestamps for a episodic memory unit.""" created_at: datetime = Field( default_factory=lambda: datetime.now(timezone.utc), - description="When the memory unit was first created", + description="When the episodic memory unit was first created", ) updated_at: Optional[datetime] = Field( - None, description="When the memory was last updated/refreshed" + None, description="When the episodic memory was last updated/refreshed" ) invalid_at: Optional[datetime] = Field( - None, description="When the memory was invalidated or expired" + None, description="When the episodic memory was invalidated or expired" ) -class MemorySource(BaseModel): - """Source information for a memory unit.""" +class EpisodicMemorySource(BaseModel): + """Source information for a episodic memory unit.""" source_name: str = Field( ..., description="Memory source, e.g., agent name, model name, dataset name, or file path" @@ -93,26 +93,26 @@ class Result(BaseModel): ) -class MemoryUnit(BaseModel): +class EpisodicMemoryUnit(BaseModel): """ - Core memory unit capturing one action of agent execution. + Core episodic memory unit capturing one action of agent execution. This includes: - - The memory id (memory_id) - - The memory timestamp (timestamp) - - The memory source (source_name, run_id, metadata) + - The episodic memory id (episodic_memory_id) + - The episodic memory timestamp (timestamp) + - The episodic memory source (source_name, run_id, metadata) - The task being worked on (issue and repository details) - The current state (what's done, what's todo) - The action taken by the agent - The result of the action """ - memory_id: str = Field( + episodic_memory_id: str = Field( default_factory=lambda: str(uuid.uuid4()), - description="Unique identifier for this memory unit", + description="Unique identifier for this episodic memory unit", ) - timestamp: MemoryTimestamp = Field(..., description="Timestamp of the memory") - source: MemorySource + timestamp: EpisodicMemoryTimestamp = Field(..., description="Timestamp of the episodic memory") + source: EpisodicMemorySource task: Task state: State action: Action diff --git a/athena/prompts/memory_extraction.py b/athena/prompts/episodic_memory_extraction.py similarity index 100% rename from athena/prompts/memory_extraction.py rename to athena/prompts/episodic_memory_extraction.py From 257b3bd964a65876a4e9492fa5a604218f562d58 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Sat, 4 Oct 2025 17:45:59 +0100 Subject: [PATCH 07/30] refactor: streamline database initialization by removing retry logic and simplifying index creation --- athena/app/services/database_service.py | 68 +++++-------------- .../semantic_memory_extraction_service.py | 13 ---- .../semantic_memory_storage_service.py | 11 --- 3 files changed, 16 insertions(+), 76 deletions(-) delete mode 100644 athena/app/services/semantic_memory_extraction_service.py delete mode 100644 athena/app/services/semantic_memory_storage_service.py diff --git a/athena/app/services/database_service.py b/athena/app/services/database_service.py index 07331e3..b8271f1 100644 --- a/athena/app/services/database_service.py +++ b/athena/app/services/database_service.py @@ -16,59 +16,33 @@ def __init__(self, DATABASE_URL: str, max_retries: int = 5, initial_backoff: flo self.engine, expire_on_commit=False, class_=AsyncSession ) self._logger = get_logger(__name__) - self._max_retries = max_retries - self._initial_backoff = initial_backoff # Create the database and tables - async def create_db_and_tables(self): # TODO: create dbs with custom names, "memories_xxx" + async def create_db_and_tables(self): async with self.engine.begin() as conn: - # Ensure pgvector extension exists (safe to ignore if unavailable) - try: - await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) - except Exception: - pass + # Ensure pgvector extension exists + await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + + # Create all tables await conn.run_sync(SQLModel.metadata.create_all) # Create ivfflat indexes for vector columns (if extension present) - try: - lists = settings.EMBEDDING_IVFFLAT_LISTS or 100 - await conn.exec_driver_sql( - f"CREATE INDEX IF NOT EXISTS idx_memory_units_task_embedding ON memory_units USING ivfflat (task_embedding vector_cosine_ops) WITH (lists = {lists})" - ) - await conn.exec_driver_sql( - f"CREATE INDEX IF NOT EXISTS idx_memory_units_state_embedding ON memory_units USING ivfflat (state_embedding vector_cosine_ops) WITH (lists = {lists})" - ) - await conn.exec_driver_sql( - f"CREATE INDEX IF NOT EXISTS idx_memory_units_task_state_embedding ON memory_units USING ivfflat (task_state_embedding vector_cosine_ops) WITH (lists = {lists})" - ) - except Exception: - # Index creation failed (likely no pgvector). Continue without indexes. - pass + lists = settings.EMBEDDING_IVFFLAT_LISTS or 100 + await conn.exec_driver_sql( + f"CREATE INDEX IF NOT EXISTS idx_memory_units_task_embedding ON memory_units USING ivfflat (task_embedding vector_cosine_ops) WITH (lists = {lists})" + ) + await conn.exec_driver_sql( + f"CREATE INDEX IF NOT EXISTS idx_memory_units_state_embedding ON memory_units USING ivfflat (state_embedding vector_cosine_ops) WITH (lists = {lists})" + ) + await conn.exec_driver_sql( + f"CREATE INDEX IF NOT EXISTS idx_memory_units_task_state_embedding ON memory_units USING ivfflat (task_state_embedding vector_cosine_ops) WITH (lists = {lists})" + ) async def start(self): """ Start the database service by creating the database and tables. This method is called when the service is initialized. """ - attempt = 0 - backoff = self._initial_backoff - while True: - try: - await self.create_db_and_tables() - self._logger.info("Database and tables created successfully.") - break - except Exception as exc: - attempt += 1 - if attempt > self._max_retries: - self._logger.error( - f"Database start failed after {self._max_retries} retries: {exc}" - ) - raise - self._logger.warning( - f"Database start failed (attempt {attempt}/{self._max_retries}): {exc}. " - f"Retrying in {backoff:.1f}s..." - ) - await asyncio.sleep(backoff) - backoff *= 2 + await self.create_db_and_tables() async def close(self): """ @@ -80,13 +54,3 @@ async def close(self): def get_sessionmaker(self) -> async_sessionmaker[AsyncSession]: """Return the async sessionmaker for dependency injection.""" return self.sessionmaker - - async def health_check(self) -> bool: - """Perform a lightweight connectivity check (SELECT 1).""" - try: - async with self.engine.connect() as conn: - await conn.exec_driver_sql("SELECT 1") - return True - except Exception as exc: - self._logger.warning(f"Database health_check failed: {exc}") - return False diff --git a/athena/app/services/semantic_memory_extraction_service.py b/athena/app/services/semantic_memory_extraction_service.py deleted file mode 100644 index 009390a..0000000 --- a/athena/app/services/semantic_memory_extraction_service.py +++ /dev/null @@ -1,13 +0,0 @@ -import asyncio -import html -import json -import re -import uuid -from datetime import datetime, timezone -from typing import Any, Dict, List, Optional - -from datasets import load_dataset -from tqdm import tqdm - -from athena.app.services.base_service import BaseService -from athena.app.services.llm_service import LLMService diff --git a/athena/app/services/semantic_memory_storage_service.py b/athena/app/services/semantic_memory_storage_service.py deleted file mode 100644 index 3815660..0000000 --- a/athena/app/services/semantic_memory_storage_service.py +++ /dev/null @@ -1,11 +0,0 @@ -from typing import List, Optional - -from sqlalchemy import select -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker -from sqlmodel import col - -from athena.app.services.base_service import BaseService -from athena.app.services.embedding_service import EmbeddingService -from athena.configuration.config import settings -from athena.entity.memory import SemanticMemoryUnitDB -from athena.models.memory import SemanticMemoryUnit From f21ac7ed5db5da3757ef0a6bdbab8ac1419a4454 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Sat, 4 Oct 2025 17:47:02 +0100 Subject: [PATCH 08/30] refactor: rename memory.py to episodic_memory.py for clarity --- athena/entity/{memory.py => episodic_memory.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename athena/entity/{memory.py => episodic_memory.py} (100%) diff --git a/athena/entity/memory.py b/athena/entity/episodic_memory.py similarity index 100% rename from athena/entity/memory.py rename to athena/entity/episodic_memory.py From 5f588ad3f9eeedc0c64fa5a9b9801a3f0505a019 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Sun, 5 Oct 2025 14:49:33 +0100 Subject: [PATCH 09/30] feat: add SemanticMemoryUnitDB model for persistent storage of semantic memory units --- athena/entity/__init__.py | 3 -- athena/entity/semantic_memory.py | 87 ++++++++++++++++++++++++++++++++ athena/models/__init__.py | 21 -------- 3 files changed, 87 insertions(+), 24 deletions(-) create mode 100644 athena/entity/semantic_memory.py diff --git a/athena/entity/__init__.py b/athena/entity/__init__.py index 51367cc..e69de29 100644 --- a/athena/entity/__init__.py +++ b/athena/entity/__init__.py @@ -1,3 +0,0 @@ -from .memory import EpisodicMemoryUnitDB - -__all__ = ["EpisodicMemoryUnitDB"] diff --git a/athena/entity/semantic_memory.py b/athena/entity/semantic_memory.py new file mode 100644 index 0000000..d77d21d --- /dev/null +++ b/athena/entity/semantic_memory.py @@ -0,0 +1,87 @@ +from datetime import datetime +from typing import Optional + +from pgvector.sqlalchemy import Vector +from sqlalchemy import Column, Text +from sqlmodel import Field, SQLModel + +from athena.configuration.config import settings +from athena.models.memory_context import MemoryContext +from athena.models.query import Query +from athena.models.semantic_memory import SemanticMemoryUnit + + +# Database models for persistent storage +class SemanticMemoryUnitDB(SQLModel, table=True): + """Database model for persistent storage of semantic memory units.""" + + __tablename__ = "semantic_memories" + + id: int = Field(primary_key=True, description="ID") + + # Source information + url: str = Field( + index=True, + description="Source URL of the repository" + ) + commit_id: str = Field( + index=True, + description="Git commit ID associated with the memory", + nullable=True, + ) + + # Query information (for retrieval) + query_essential_query: str = Field(sa_column=Column(Text, nullable=False)) + query_extra_requirements: str = Field(sa_column=Column(Text, nullable=False)) + query_purpose: str = Field(sa_column=Column(Text, nullable=False)) + + # Memory context + memory_context_overview: str = Field(sa_column=Column(Text, nullable=False)) + memory_context_content: str = Field(sa_column=Column(Text, nullable=False)) + + # Embeddings for semantic retrieval (pgvector) + _vec_dim: int = settings.EMBEDDING_DIM or 1024 # type: ignore[assignment] + + # Query embeddings for weighted retrieval + essential_query_embedding: Optional[list[float]] = Field( + sa_column=Column(Vector(_vec_dim), nullable=False) + ) + extra_requirements_embedding: Optional[list[float]] = Field( + sa_column=Column(Vector(_vec_dim), nullable=False) + ) + purpose_embedding: Optional[list[float]] = Field( + sa_column=Column(Vector(_vec_dim), nullable=False) + ) + + @classmethod + def from_semantic_memory_unit( + cls, + memory_unit: SemanticMemoryUnit, + url: Optional[str] = None, + commit_id: Optional[str] = None, + ) -> "SemanticMemoryUnitDB": + """Create a database model from a SemanticMemoryUnit.""" + + return cls( + url=url, + commit_id=commit_id, + query_essential_query=memory_unit.query.essential_query, + query_extra_requirements=memory_unit.query.extra_requirements, + query_purpose=memory_unit.query.purpose, + memory_context_overview=memory_unit.memory_context.overview, + memory_context_content=memory_unit.memory_context.content, + ) + + def to_semantic_memory_unit(self) -> SemanticMemoryUnit: + """Convert database model back to SemanticMemoryUnit.""" + return SemanticMemoryUnit( + query=Query( + essential_query=self.query_essential_query, + extra_requirements=self.query_extra_requirements, + purpose=self.query_purpose, + ), + memory_context=MemoryContext( + overview=self.memory_context_overview, + content=self.memory_context_content, + ), + ) \ No newline at end of file diff --git a/athena/models/__init__.py b/athena/models/__init__.py index 408202e..e69de29 100644 --- a/athena/models/__init__.py +++ b/athena/models/__init__.py @@ -1,21 +0,0 @@ -from .episodic_memory import ( - Action, - EpisodicMemorySource, - EpisodicMemoryTimestamp, - EpisodicMemoryUnit, - Result, - State, - Task, -) -from .message import Message - -__all__ = [ - "Message", - "EpisodicMemoryUnit", - "EpisodicMemorySource", - "EpisodicMemoryTimestamp", - "Task", - "State", - "Action", - "Result", -] From 9d752206806f5f13f6fb06c45de9745165435e3b Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Sun, 5 Oct 2025 14:56:25 +0100 Subject: [PATCH 10/30] refactor: clean up imports and improve code readability in memory services --- athena/app/services/database_service.py | 2 -- .../episodic_memory_extraction_service.py | 16 ++++++------ .../app/services/episodic_memory_service.py | 26 ++++++++++--------- .../episodic_memory_storage_service.py | 8 ++++-- athena/entity/episodic_memory.py | 10 ++++++- athena/entity/semantic_memory.py | 8 ++---- 6 files changed, 39 insertions(+), 31 deletions(-) diff --git a/athena/app/services/database_service.py b/athena/app/services/database_service.py index b8271f1..d60b8fb 100644 --- a/athena/app/services/database_service.py +++ b/athena/app/services/database_service.py @@ -1,5 +1,3 @@ -import asyncio - from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlmodel import SQLModel diff --git a/athena/app/services/episodic_memory_extraction_service.py b/athena/app/services/episodic_memory_extraction_service.py index 836f507..66b5936 100644 --- a/athena/app/services/episodic_memory_extraction_service.py +++ b/athena/app/services/episodic_memory_extraction_service.py @@ -2,7 +2,6 @@ import html import json import re -import uuid from datetime import datetime, timezone from typing import Any, Dict, List, Optional @@ -10,19 +9,19 @@ from tqdm import tqdm from athena.app.services.base_service import BaseService -from athena.app.services.llm_service import LLMService from athena.app.services.episodic_memory_storage_service import EpisodicMemoryStorageService -from athena.models import ( +from athena.app.services.llm_service import LLMService +from athena.models.episodic_memory import ( Action, EpisodicMemorySource, EpisodicMemoryTimestamp, EpisodicMemoryUnit, - Message, Result, State, Task, ) -from athena.prompts.memory_extraction import ( +from athena.models.message import Message +from athena.prompts.episodic_memory_extraction import ( ACTION_EXTRACTION_PROMPT, ACTION_JUDGE_PROMPT, RESULT_EXTRACTION_PROMPT, @@ -91,7 +90,7 @@ def __init__( self.llm_service = llm_service self.batch_size = batch_size self.max_retries = max_retries - self._extraction_cache: Dict[str, EpisodicEpisodicMemoryUnit] = {} + self._extraction_cache: Dict[str, EpisodicMemoryUnit] = {} self.memory_store = memory_store # self._logger = get_logger(__name__) @@ -302,7 +301,6 @@ def _create_memory_unit( state_todo = self._synthesize_state_todo_from_window(window_msgs, task, state_done, action) result = self._extract_result_from_window(window_msgs, action) return EpisodicMemoryUnit( - memory_id=str(uuid.uuid4()), timestamp=EpisodicMemoryTimestamp( created_at=datetime.now(timezone.utc), updated_at=None, @@ -398,7 +396,9 @@ def _synthesize_state_from_context( ) return State(done=state_done, todo=state_todo) # TODO: open_file, working_dir - def _synthesize_state_done_from_context(self, prior_units: List[EpisodicMemoryUnit], task: Task) -> str: + def _synthesize_state_done_from_context( + self, prior_units: List[EpisodicMemoryUnit], task: Task + ) -> str: """ Summarize previous context into a concise `state.done` string. Summary of what has ALREADY BEEN COMPLETED (no plans). diff --git a/athena/app/services/episodic_memory_service.py b/athena/app/services/episodic_memory_service.py index fd64356..bc98e43 100644 --- a/athena/app/services/episodic_memory_service.py +++ b/athena/app/services/episodic_memory_service.py @@ -3,7 +3,7 @@ from athena.app.services.base_service import BaseService from athena.app.services.episodic_memory_storage_service import EpisodicMemoryStorageService -from athena.models.memory import EpisodicMemoryUnit +from athena.models.episodic_memory import EpisodicMemoryUnit class EpisodicMemoryService(BaseService): @@ -17,7 +17,7 @@ class EpisodicMemoryService(BaseService): Design Principles: - Store complete execution trajectories as memory units - - Support semantic search for relevant past experiences + - Support semantic search for relevant experience - 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 @@ -27,7 +27,9 @@ class EpisodicMemoryService(BaseService): """ def __init__( - self, storage_backend: str = "in_memory", store: Optional[MemoryStorageService] = None + self, + storage_backend: str = "in_memory", + store: Optional[EpisodicMemoryStorageService] = None, ): """ Initialize the Memory Service with a specified storage backend. @@ -61,12 +63,12 @@ async def close(self): else: self._store.close() # type: ignore[misc] - async def store_memory(self, memory_unit: MemoryUnit) -> None: + async def store_memory(self, memory_unit: EpisodicMemoryUnit) -> None: """ Store a memory unit in the memory service. Args: - memory_unit: The MemoryUnit object containing task, state, + memory_unit: The EpisodicMemoryUnit object containing task, state, action, and result information to be stored. Deduplication is handled at the database layer on memory_id via upsert. @@ -78,7 +80,7 @@ async def store_memory(self, memory_unit: MemoryUnit) -> None: except Exception: return - async def store_memories(self, memory_units: Sequence[MemoryUnit]) -> None: + async def store_memories(self, memory_units: Sequence[EpisodicMemoryUnit]) -> None: """Bulk store multiple memory units efficiently.""" if self._store is None or not memory_units: return @@ -89,7 +91,7 @@ async def store_memories(self, memory_units: Sequence[MemoryUnit]) -> None: async def search_memory( self, query: str, limit: int = 10, field: str = "task_state" - ) -> List[MemoryUnit]: + ) -> List[EpisodicMemoryUnit]: """ Search for relevant memory units based on a query. @@ -98,7 +100,7 @@ async def search_memory( limit: Maximum number of results to return (default: 10) Returns: - List of MemoryUnit objects matching the search query, + List of EpisodicMemoryUnit objects matching the search query, ordered by relevance. This method should support semantic search across multiple @@ -112,7 +114,7 @@ async def search_memory( except Exception: return [] - async def get_memory_by_key(self, key: str) -> Optional[MemoryUnit]: + async def get_memory_by_key(self, key: str) -> Optional[EpisodicMemoryUnit]: """ Retrieve a specific memory unit by its memory id. @@ -120,7 +122,7 @@ async def get_memory_by_key(self, key: str) -> Optional[MemoryUnit]: key: The id of the memory unit Returns: - The MemoryUnit object if found, None otherwise. + The EpisodicMemoryUnit object if found, None otherwise. """ # Here we treat key as memory_id for simplicity if self._store is None: @@ -130,12 +132,12 @@ async def get_memory_by_key(self, key: str) -> Optional[MemoryUnit]: except Exception: return None - async def get_all_memories(self) -> List[MemoryUnit]: + async def get_all_memories(self) -> List[EpisodicMemoryUnit]: """ Retrieve all memory units stored in the service. Returns: - List of all MemoryUnit objects in the service. + List of all EpisodicMemoryUnit objects in the service. """ if self._store is None: return [] diff --git a/athena/app/services/episodic_memory_storage_service.py b/athena/app/services/episodic_memory_storage_service.py index a751bb4..d1bac05 100644 --- a/athena/app/services/episodic_memory_storage_service.py +++ b/athena/app/services/episodic_memory_storage_service.py @@ -66,7 +66,9 @@ async def _upsert_one(self, session: AsyncSession, unit: EpisodicMemoryUnit) -> # Check existing existing = await session.scalar( - select(EpisodicMemoryUnitDB).where(col(EpisodicMemoryUnitDB.memory_id) == unit.memory_id) + select(EpisodicMemoryUnitDB).where( + col(EpisodicMemoryUnitDB.memory_id) == unit.memory_id + ) ) if existing is None: @@ -177,7 +179,9 @@ def _serialize_state(u: EpisodicMemoryUnit) -> str: return "\n".join([str(p) for p in parts if p]) @staticmethod - def _serialize_task_state(u: EpisodicMemoryUnit) -> str: # TODO: text templates for static embeddings + def _serialize_task_state( + u: EpisodicMemoryUnit, + ) -> str: # TODO: text templates for static embeddings return ( EpisodicMemoryStorageService._serialize_task(u) + "\n\n" diff --git a/athena/entity/episodic_memory.py b/athena/entity/episodic_memory.py index f22aa2d..15957ac 100644 --- a/athena/entity/episodic_memory.py +++ b/athena/entity/episodic_memory.py @@ -7,7 +7,15 @@ from sqlmodel import Field, SQLModel, UniqueConstraint from athena.configuration.config import settings -from athena.models import Action, EpisodicMemorySource, EpisodicMemoryTimestamp, EpisodicMemoryUnit, Result, State, Task +from athena.models import ( + Action, + EpisodicMemorySource, + EpisodicMemoryTimestamp, + EpisodicMemoryUnit, + Result, + State, + Task, +) # Database models for persistent storage diff --git a/athena/entity/semantic_memory.py b/athena/entity/semantic_memory.py index d77d21d..cf8897c 100644 --- a/athena/entity/semantic_memory.py +++ b/athena/entity/semantic_memory.py @@ -1,4 +1,3 @@ -from datetime import datetime from typing import Optional from pgvector.sqlalchemy import Vector @@ -20,10 +19,7 @@ class SemanticMemoryUnitDB(SQLModel, table=True): id: int = Field(primary_key=True, description="ID") # Source information - url: str = Field( - index=True, - description="Source URL of the repository" - ) + url: str = Field(index=True, description="Source URL of the repository") commit_id: str = Field( index=True, description="Git commit ID associated with the memory", @@ -84,4 +80,4 @@ def to_semantic_memory_unit(self) -> SemanticMemoryUnit: overview=self.memory_context_overview, content=self.memory_context_content, ), - ) \ No newline at end of file + ) From 5c0442e24f49ce661d5b09d6e7ffa0b92554ee3b Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Sun, 5 Oct 2025 15:15:48 +0100 Subject: [PATCH 11/30] feat: implement vector indexing for episodic and semantic memory services --- athena/app/services/database_service.py | 22 ++++------- .../app/services/episodic_memory_service.py | 22 +++++++++++ .../app/services/semantic_memory_service.py | 37 ++++++++++++++++++- 3 files changed, 64 insertions(+), 17 deletions(-) diff --git a/athena/app/services/database_service.py b/athena/app/services/database_service.py index d60b8fb..30afe20 100644 --- a/athena/app/services/database_service.py +++ b/athena/app/services/database_service.py @@ -3,37 +3,29 @@ from sqlmodel import SQLModel from athena.app.services.base_service import BaseService -from athena.configuration.config import settings from athena.utils.logger_manager import get_logger class DatabaseService(BaseService): - def __init__(self, DATABASE_URL: str, max_retries: int = 5, initial_backoff: float = 1.0): + def __init__(self, DATABASE_URL: str): self.engine = create_async_engine(DATABASE_URL, echo=True) self.sessionmaker = async_sessionmaker( self.engine, expire_on_commit=False, class_=AsyncSession ) self._logger = get_logger(__name__) - # Create the database and tables - async def create_db_and_tables(self): + async def create_vector_extension( + self, + ): async with self.engine.begin() as conn: # Ensure pgvector extension exists await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + # Create the database and tables + async def create_db_and_tables(self): + async with self.engine.begin() as conn: # Create all tables await conn.run_sync(SQLModel.metadata.create_all) - # Create ivfflat indexes for vector columns (if extension present) - lists = settings.EMBEDDING_IVFFLAT_LISTS or 100 - await conn.exec_driver_sql( - f"CREATE INDEX IF NOT EXISTS idx_memory_units_task_embedding ON memory_units USING ivfflat (task_embedding vector_cosine_ops) WITH (lists = {lists})" - ) - await conn.exec_driver_sql( - f"CREATE INDEX IF NOT EXISTS idx_memory_units_state_embedding ON memory_units USING ivfflat (state_embedding vector_cosine_ops) WITH (lists = {lists})" - ) - await conn.exec_driver_sql( - f"CREATE INDEX IF NOT EXISTS idx_memory_units_task_state_embedding ON memory_units USING ivfflat (task_state_embedding vector_cosine_ops) WITH (lists = {lists})" - ) async def start(self): """ diff --git a/athena/app/services/episodic_memory_service.py b/athena/app/services/episodic_memory_service.py index bc98e43..8ef5452 100644 --- a/athena/app/services/episodic_memory_service.py +++ b/athena/app/services/episodic_memory_service.py @@ -2,7 +2,9 @@ from typing import List, Optional, Sequence from athena.app.services.base_service import BaseService +from athena.app.services.database_service import DatabaseService from athena.app.services.episodic_memory_storage_service import EpisodicMemoryStorageService +from athena.configuration.config import settings from athena.models.episodic_memory import EpisodicMemoryUnit @@ -28,6 +30,7 @@ class EpisodicMemoryService(BaseService): def __init__( self, + database_service: DatabaseService, storage_backend: str = "in_memory", store: Optional[EpisodicMemoryStorageService] = None, ): @@ -40,11 +43,30 @@ def __init__( - "database": Persistent database storage - "vector": Vector database for semantic search """ + self.database_service = database_service self.storage_backend = storage_backend self._store = store + async def create_vector_field_index(self): + async with self.database_service.engine.begin() as conn: + # Create ivfflat indexes for vector columns (if extension present) + lists = settings.EMBEDDING_IVFFLAT_LISTS or 100 + await conn.exec_driver_sql( + f"CREATE INDEX IF NOT EXISTS idx_memory_units_task_embedding ON episodic_memories USING ivfflat (task_embedding vector_cosine_ops) WITH (lists = {lists})" + ) + await conn.exec_driver_sql( + f"CREATE INDEX IF NOT EXISTS idx_memory_units_state_embedding ON episodic_memories USING ivfflat (state_embedding vector_cosine_ops) WITH (lists = {lists})" + ) + await conn.exec_driver_sql( + f"CREATE INDEX IF NOT EXISTS idx_memory_units_task_state_embedding ON episodic_memories USING ivfflat (task_state_embedding vector_cosine_ops) WITH (lists = {lists})" + ) + async def start(self): """Initialize the storage backend if needed and validate configuration.""" + # Create vector extension + await self.create_vector_field_index() + + # Validate storage backend configuration if self.storage_backend in {"database", "vector"} and self._store is None: raise RuntimeError( "EpisodicMemoryService requires a storage store for database/vector backends" diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index 0f025a8..ac9efad 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -1,13 +1,46 @@ from typing import List from athena.app.services.base_service import BaseService +from athena.app.services.database_service import DatabaseService +from athena.configuration.config import settings from athena.models.query import Query from athena.models.semantic_memory import SemanticMemoryUnit class SemanticMemoryService(BaseService): - def __init__(self): - pass + def __init__(self, database_service: DatabaseService): + self.database_service = database_service + + async def create_semantic_memory_vector_indexes(self): + """Create vector indexes for semantic memory table using IVFFlat algorithm.""" + async with self.database_service.engine.begin() as conn: + # Get lists parameter from settings (default 100) + lists = settings.EMBEDDING_IVFFLAT_LISTS or 100 + + # Create ivfflat indexes for each query component embedding + # Essential query embedding (highest priority, most important for retrieval) + await conn.exec_driver_sql( + f"CREATE INDEX IF NOT EXISTS idx_semantic_memories_essential_query_embedding " + f"ON semantic_memories USING ivfflat (essential_query_embedding vector_cosine_ops) " + f"WITH (lists = {lists})" + ) + + # Extra requirements embedding (medium priority) + await conn.exec_driver_sql( + f"CREATE INDEX IF NOT EXISTS idx_semantic_memories_extra_requirements_embedding " + f"ON semantic_memories USING ivfflat (extra_requirements_embedding vector_cosine_ops) " + f"WITH (lists = {lists})" + ) + + # Purpose embedding (medium priority) + await conn.exec_driver_sql( + f"CREATE INDEX IF NOT EXISTS idx_semantic_memories_purpose_embedding " + f"ON semantic_memories USING ivfflat (purpose_embedding vector_cosine_ops) " + f"WITH (lists = {lists})" + ) + + async def start(self): + await self.create_semantic_memory_vector_indexes() async def store_memory(self, url: str, commit_id: str, memory: SemanticMemoryUnit): # Logic to store memory From 7e13e794e2ee7b2630523263d0b158668315fd4f Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Sun, 5 Oct 2025 21:47:23 +0100 Subject: [PATCH 12/30] feat: add embedding index settings for semantic and episodic memory services --- athena/app/dependencies.py | 9 ++++++--- athena/app/services/episodic_memory_service.py | 2 +- .../app/services/episodic_memory_storage_service.py | 4 ++-- athena/app/services/llm_service.py | 13 +------------ athena/app/services/semantic_memory_service.py | 2 +- athena/configuration/config.py | 7 ++++--- athena/entity/episodic_memory.py | 2 +- 7 files changed, 16 insertions(+), 23 deletions(-) diff --git a/athena/app/dependencies.py b/athena/app/dependencies.py index 49511ab..d48dc15 100644 --- a/athena/app/dependencies.py +++ b/athena/app/dependencies.py @@ -8,6 +8,7 @@ from athena.app.services.episodic_memory_service import EpisodicMemoryService from athena.app.services.episodic_memory_storage_service import EpisodicMemoryStorageService from athena.app.services.llm_service import LLMService +from athena.app.services.semantic_memory_service import SemanticMemoryService from athena.configuration.config import settings @@ -33,8 +34,6 @@ def initialize_services() -> Dict[str, BaseService]: llm_service = LLMService( settings.MODEL_NAME, settings.MODEL_TEMPERATURE, - settings.MODEL_MAX_INPUT_TOKENS, - settings.MODEL_MAX_OUTPUT_TOKENS, settings.OPENAI_FORMAT_API_KEY, settings.OPENAI_FORMAT_BASE_URL, settings.ANTHROPIC_API_KEY, @@ -57,11 +56,15 @@ def initialize_services() -> Dict[str, BaseService]: ) episodic_memory_service = EpisodicMemoryService( - storage_backend=settings.MEMORY_STORAGE_BACKEND, store=episodic_memory_store + database_service=database_service, + storage_backend=settings.MEMORY_STORAGE_BACKEND, + store=episodic_memory_store, ) + semantic_memory_service = SemanticMemoryService(database_service=database_service) return { "llm_service": llm_service, "database_service": database_service, "episodic_memory_service": episodic_memory_service, + "semantic_memory_service": semantic_memory_service, } diff --git a/athena/app/services/episodic_memory_service.py b/athena/app/services/episodic_memory_service.py index 8ef5452..bd82736 100644 --- a/athena/app/services/episodic_memory_service.py +++ b/athena/app/services/episodic_memory_service.py @@ -50,7 +50,7 @@ def __init__( async def create_vector_field_index(self): async with self.database_service.engine.begin() as conn: # Create ivfflat indexes for vector columns (if extension present) - lists = settings.EMBEDDING_IVFFLAT_LISTS or 100 + lists = settings.EPISODE_MEMORY_EMBEDDING_IVFFLAT_LISTS await conn.exec_driver_sql( f"CREATE INDEX IF NOT EXISTS idx_memory_units_task_embedding ON episodic_memories USING ivfflat (task_embedding vector_cosine_ops) WITH (lists = {lists})" ) diff --git a/athena/app/services/episodic_memory_storage_service.py b/athena/app/services/episodic_memory_storage_service.py index d1bac05..d81e654 100644 --- a/athena/app/services/episodic_memory_storage_service.py +++ b/athena/app/services/episodic_memory_storage_service.py @@ -7,8 +7,8 @@ from athena.app.services.base_service import BaseService from athena.app.services.embedding_service import EmbeddingService from athena.configuration.config import settings -from athena.entity.memory import EpisodicMemoryUnitDB -from athena.models.memory import EpisodicMemoryUnit +from athena.entity.episodic_memory import EpisodicMemoryUnitDB +from athena.models.episodic_memory import EpisodicMemoryUnit def _ensure_dim(vec: List[float]) -> List[float]: diff --git a/athena/app/services/llm_service.py b/athena/app/services/llm_service.py index 419f016..bf899e0 100644 --- a/athena/app/services/llm_service.py +++ b/athena/app/services/llm_service.py @@ -14,8 +14,6 @@ def __init__( self, model_name: str, model_temperature: float, - model_max_input_tokens: int, - model_max_output_tokens: int, openai_format_api_key: Optional[str] = None, openai_format_base_url: Optional[str] = None, anthropic_api_key: Optional[str] = None, @@ -25,8 +23,6 @@ def __init__( self.model = get_model( model_name, temperature=model_temperature, - max_input_tokens=model_max_input_tokens, - max_output_tokens=model_max_output_tokens, openai_format_api_key=openai_format_api_key, openai_format_base_url=openai_format_base_url, anthropic_api_key=anthropic_api_key, @@ -38,8 +34,6 @@ def __init__( def get_model( model_name: str, temperature: float, - max_input_tokens: int, - max_output_tokens: int, openai_format_api_key: Optional[str] = None, openai_format_base_url: Optional[str] = None, anthropic_api_key: Optional[str] = None, @@ -51,7 +45,6 @@ def get_model( model_name=model_name, api_key=anthropic_api_key, temperature=temperature, - max_tokens_to_sample=max_output_tokens, max_retries=3, ) elif model_name.startswith("vertex:"): @@ -62,7 +55,6 @@ def get_model( project="prometheus-code-agent", location="us-central1", temperature=temperature, - max_output_tokens=max_output_tokens, max_retries=3, credentials=google_application_credentials, ) @@ -71,19 +63,16 @@ def get_model( model=model_name, api_key=gemini_api_key, temperature=temperature, - max_tokens=max_output_tokens, max_retries=3, ) else: """ - Use tiktoken_counter to ensure that the input messages do not exceed the maximum token limit. + Custom OpenAI chat model with specific configuration. """ return CustomChatOpenAI( - max_input_tokens=max_input_tokens, model=model_name, api_key=openai_format_api_key, base_url=openai_format_base_url, temperature=temperature, - max_tokens=max_output_tokens, max_retries=3, ) diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index ac9efad..8f27eb1 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -15,7 +15,7 @@ async def create_semantic_memory_vector_indexes(self): """Create vector indexes for semantic memory table using IVFFlat algorithm.""" async with self.database_service.engine.begin() as conn: # Get lists parameter from settings (default 100) - lists = settings.EMBEDDING_IVFFLAT_LISTS or 100 + lists = settings.SEMANTIC_MEMORY_EMBEDDING_IVFFLAT_LISTS # Create ivfflat indexes for each query component embedding # Essential query embedding (highest priority, most important for retrieval) diff --git a/athena/configuration/config.py b/athena/configuration/config.py index c0024e7..82f643c 100644 --- a/athena/configuration/config.py +++ b/athena/configuration/config.py @@ -30,9 +30,7 @@ class Settings(BaseSettings): GOOGLE_APPLICATION_CREDENTIALS: Optional[str] = None # Model parameters - MODEL_MAX_INPUT_TOKENS: int MODEL_TEMPERATURE: Optional[float] = None - MODEL_MAX_OUTPUT_TOKENS: Optional[int] = None # Database DATABASE_URL: str @@ -47,7 +45,10 @@ class Settings(BaseSettings): EMBEDDING_API_KEY: Optional[str] = None EMBEDDING_BASE_URL: Optional[str] = None EMBEDDING_DIM: Optional[int] = None - EMBEDDING_IVFFLAT_LISTS: Optional[int] = None + + # Embedding index settings + EPISODE_MEMORY_EMBEDDING_IVFFLAT_LISTS: Optional[int] = 100 + SEMANTIC_MEMORY_EMBEDDING_IVFFLAT_LISTS: Optional[int] = 10 settings = Settings() diff --git a/athena/entity/episodic_memory.py b/athena/entity/episodic_memory.py index 15957ac..659031e 100644 --- a/athena/entity/episodic_memory.py +++ b/athena/entity/episodic_memory.py @@ -7,7 +7,7 @@ from sqlmodel import Field, SQLModel, UniqueConstraint from athena.configuration.config import settings -from athena.models import ( +from athena.models.episodic_memory import ( Action, EpisodicMemorySource, EpisodicMemoryTimestamp, From 3552114b50961bdb0778e89ce5d623f19dced565 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Sun, 5 Oct 2025 21:49:29 +0100 Subject: [PATCH 13/30] feat: add vector extension creation during database service initialization --- athena/app/services/database_service.py | 1 + 1 file changed, 1 insertion(+) diff --git a/athena/app/services/database_service.py b/athena/app/services/database_service.py index 30afe20..0fdda57 100644 --- a/athena/app/services/database_service.py +++ b/athena/app/services/database_service.py @@ -32,6 +32,7 @@ async def start(self): Start the database service by creating the database and tables. This method is called when the service is initialized. """ + await self.create_vector_extension() await self.create_db_and_tables() async def close(self): From 79e4db8395e2494e6eacf526ea39d3f4cc2079a7 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Sun, 5 Oct 2025 22:22:14 +0100 Subject: [PATCH 14/30] feat: enhance database service with logging and refactor semantic memory service for improved session handling --- athena/app/services/database_service.py | 2 ++ athena/app/services/semantic_memory_service.py | 16 ++++++++++++---- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/athena/app/services/database_service.py b/athena/app/services/database_service.py index 0fdda57..4393e55 100644 --- a/athena/app/services/database_service.py +++ b/athena/app/services/database_service.py @@ -20,12 +20,14 @@ async def create_vector_extension( async with self.engine.begin() as conn: # Ensure pgvector extension exists await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + self._logger.info("pgvector extension ensured.") # Create the database and tables async def create_db_and_tables(self): async with self.engine.begin() as conn: # Create all tables await conn.run_sync(SQLModel.metadata.create_all) + self._logger.info("Database and tables created.") async def start(self): """ diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index 8f27eb1..2d73d7e 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -1,19 +1,22 @@ from typing import List +from sqlalchemy.ext.asyncio import AsyncSession + from athena.app.services.base_service import BaseService from athena.app.services.database_service import DatabaseService from athena.configuration.config import settings +from athena.entity.semantic_memory import SemanticMemoryUnitDB from athena.models.query import Query from athena.models.semantic_memory import SemanticMemoryUnit class SemanticMemoryService(BaseService): def __init__(self, database_service: DatabaseService): - self.database_service = database_service + self.engine = database_service.engine async def create_semantic_memory_vector_indexes(self): """Create vector indexes for semantic memory table using IVFFlat algorithm.""" - async with self.database_service.engine.begin() as conn: + async with self.engine.begin() as conn: # Get lists parameter from settings (default 100) lists = settings.SEMANTIC_MEMORY_EMBEDDING_IVFFLAT_LISTS @@ -43,8 +46,13 @@ async def start(self): await self.create_semantic_memory_vector_indexes() async def store_memory(self, url: str, commit_id: str, memory: SemanticMemoryUnit): - # Logic to store memory - pass + async with AsyncSession(self.engine) as session: + semantic_memory = SemanticMemoryUnitDB.from_semantic_memory_unit( + memory, url=url, commit_id=commit_id + ) + session.add(semantic_memory) + await session.commit() + await session.refresh(semantic_memory) async def retrieve_memory( self, url: str, commit_id: str, query: Query From 1cc0f192261535b91bd313703495c27026fc49a5 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Sun, 5 Oct 2025 22:49:49 +0100 Subject: [PATCH 15/30] feat: add semantic memory settings and refactor embedding service for async support --- athena/app/services/embedding_service.py | 11 +- .../episodic_memory_storage_service.py | 6 +- .../app/services/semantic_memory_service.py | 171 +++++++++++++++++- athena/configuration/config.py | 3 + athena/entity/semantic_memory.py | 46 +---- 5 files changed, 181 insertions(+), 56 deletions(-) diff --git a/athena/app/services/embedding_service.py b/athena/app/services/embedding_service.py index 872394d..01781b9 100644 --- a/athena/app/services/embedding_service.py +++ b/athena/app/services/embedding_service.py @@ -1,6 +1,6 @@ from typing import Iterable, List -import requests +import httpx from athena.app.services.base_service import BaseService @@ -18,16 +18,21 @@ def __init__(self, model: str, api_key: str, base_url: str, embed_dim: int): self.api_key = api_key self.base_url = base_url.rstrip("/") self.embed_dim = embed_dim + self.client = httpx.AsyncClient(timeout=60.0) - def embed(self, inputs: Iterable[str]) -> List[List[float]]: + async def embed(self, inputs: Iterable[str]) -> List[List[float]]: data = {"model": self.model, "input": list(inputs), "output_dimension": self.embed_dim} headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json", } url = f"{self.base_url}/v1/embeddings" - resp = requests.post(url, json=data, headers=headers, timeout=60) + resp = await self.client.post(url, json=data, headers=headers) resp.raise_for_status() payload = resp.json() vectors = [item["embedding"] for item in payload.get("data", [])] return vectors + + async def close(self): + """Close the httpx client.""" + await self.client.aclose() diff --git a/athena/app/services/episodic_memory_storage_service.py b/athena/app/services/episodic_memory_storage_service.py index d81e654..2a83f76 100644 --- a/athena/app/services/episodic_memory_storage_service.py +++ b/athena/app/services/episodic_memory_storage_service.py @@ -56,7 +56,7 @@ async def _upsert_one(self, session: AsyncSession, unit: EpisodicMemoryUnit) -> task_state_vec: Optional[List[float]] = None if self._embeddings is not None: combined_text = self._serialize_task_state(unit) - vecs = self._embeddings.embed([task_text, state_text, combined_text]) + vecs = await self._embeddings.embed([task_text, state_text, combined_text]) if len(vecs) >= 3: task_vec, state_vec, task_state_vec = ( _ensure_dim(vecs[0]), @@ -119,7 +119,9 @@ async def search_by_text( ) -> List[EpisodicMemoryUnit]: if self._embeddings is None: return [] - q_vec = _ensure_dim(self._embeddings.embed([text])[0]) + + vectors = await self._embeddings.embed([text]) + q_vec = _ensure_dim(vectors[0]) async with self._sessionmaker() as session: # Choose column by field col_expr = { diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index 2d73d7e..5fb48d8 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -1,9 +1,25 @@ -from typing import List +""" +Semantic Memory Service for storing and retrieving code knowledge. + +This service manages semantic memories which represent structured knowledge about code, +including queries and their corresponding memory contexts. It uses vector embeddings +for efficient semantic similarity search across three query components: +- Essential query: The core question or search intent +- Extra requirements: Additional constraints or specifications +- Purpose: The intended use case or goal + +The service leverages pgvector's IVFFlat indexing for fast approximate nearest neighbor +search on high-dimensional embeddings. +""" + +from typing import Sequence from sqlalchemy.ext.asyncio import AsyncSession +from sqlmodel import select from athena.app.services.base_service import BaseService from athena.app.services.database_service import DatabaseService +from athena.app.services.embedding_service import EmbeddingService from athena.configuration.config import settings from athena.entity.semantic_memory import SemanticMemoryUnitDB from athena.models.query import Query @@ -11,11 +27,42 @@ class SemanticMemoryService(BaseService): - def __init__(self, database_service: DatabaseService): + """ + Manages semantic memory storage and retrieval using vector embeddings. + + This service provides functionality to: + - Store semantic memory units with vector embeddings + - Retrieve relevant memories using weighted multi-vector similarity search + - Maintain vector indexes for efficient similarity queries + + Attributes: + engine: SQLAlchemy async engine for database operations + embedding_service: Service for generating vector embeddings from text + """ + + def __init__(self, database_service: DatabaseService, embedding_service: EmbeddingService): + """ + Initialize the semantic memory service. + + Args: + database_service: Database service providing engine and session management + embedding_service: Service for generating embeddings from text queries + """ self.engine = database_service.engine + self.embedding_service = embedding_service async def create_semantic_memory_vector_indexes(self): - """Create vector indexes for semantic memory table using IVFFlat algorithm.""" + """ + Create vector indexes for semantic memory table using IVFFlat algorithm. + + Creates three separate IVFFlat indexes for efficient cosine similarity search: + 1. Essential query embedding index (highest priority) + 2. Extra requirements embedding index (medium priority) + 3. Purpose embedding index (medium priority) + + The number of IVFFlat lists is configured via SEMANTIC_MEMORY_EMBEDDING_IVFFLAT_LISTS. + More lists improve query speed but increase index build time and memory usage. + """ async with self.engine.begin() as conn: # Get lists parameter from settings (default 100) lists = settings.SEMANTIC_MEMORY_EMBEDDING_IVFFLAT_LISTS @@ -43,12 +90,60 @@ async def create_semantic_memory_vector_indexes(self): ) async def start(self): + """Initialize the service by creating vector indexes.""" await self.create_semantic_memory_vector_indexes() + async def get_query_embeddings( + self, query: Query + ) -> tuple[list[float], list[float], list[float]]: + """ + Generate embeddings for all three components of a query. + + Args: + query: Query object containing essential_query, extra_requirements, and purpose + + Returns: + Tuple of three embedding vectors in order: + (essential_query_embedding, extra_requirements_embedding, purpose_embedding) + """ + texts = [query.essential_query, query.extra_requirements, query.purpose] + embeddings = await self.embedding_service.embed(texts) + return embeddings[0], embeddings[1], embeddings[2] + async def store_memory(self, url: str, commit_id: str, memory: SemanticMemoryUnit): + """ + Store a semantic memory unit with its vector embeddings. + + This method: + 1. Generates embeddings for the query components + 2. Creates a database record with all query and context information + 3. Persists the memory to the database + + Args: + url: Source URL of the repository + commit_id: Git commit ID associated with this memory + memory: SemanticMemoryUnit containing query and memory context + """ + # Get embeddings for the query + ( + essential_query_embedding, + extra_requirements_embedding, + purpose_embedding, + ) = await self.get_query_embeddings(memory.query) + + # Create a new SemanticMemoryUnitDB instance with embeddings async with AsyncSession(self.engine) as session: - semantic_memory = SemanticMemoryUnitDB.from_semantic_memory_unit( - memory, url=url, commit_id=commit_id + semantic_memory = SemanticMemoryUnitDB( + url=url, + commit_id=commit_id, + query_essential_query=memory.query_essential_query, + query_extra_requirements=memory.query_extra_requirements, + query_purpose=memory.query_purpose, + memory_context_overview=memory.memory_context_overview, + memory_context_content=memory.memory_context_content, + essential_query_embedding=essential_query_embedding, + extra_requirements_embedding=extra_requirements_embedding, + purpose_embedding=purpose_embedding, ) session.add(semantic_memory) await session.commit() @@ -56,6 +151,66 @@ async def store_memory(self, url: str, commit_id: str, memory: SemanticMemoryUni async def retrieve_memory( self, url: str, commit_id: str, query: Query - ) -> List[SemanticMemoryUnit]: - # Logic to retrieve memory based on a query - pass + ) -> Sequence[SemanticMemoryUnitDB]: + """ + Retrieve semantic memories most similar to the given query. + + Uses a weighted combination of cosine similarity across three query components: + - Essential query: 60% weight (most important for relevance) + - Extra requirements: 20% weight (filters and constraints) + - Purpose: 20% weight (intent and context) + + The combined similarity score ranks results, with higher scores indicating + better matches. + + Args: + url: Filter memories by this repository URL + commit_id: Filter memories by this commit ID + query: Query object to search for similar memories + + Returns: + List of SemanticMemoryUnitDB objects ordered by similarity (most similar first), + limited to SEMANTIC_MEMORY_MAX_RESULTS + """ + # Weights for similarity components: essential query gets highest weight + w1, w2, w3 = 0.6, 0.2, 0.2 + top_k = settings.SEMANTIC_MEMORY_MAX_RESULTS + + # Get embeddings for the query components + ( + essential_query_embedding, + extra_requirements_embedding, + purpose_embedding, + ) = await self.get_query_embeddings(query) + + # Compute weighted similarity score + # Cosine distance is converted to similarity (1 - distance) + similarity = ( + w1 + * ( + 1 + - SemanticMemoryUnitDB.essential_query_embedding.cosine_distance( + essential_query_embedding + ) + ) + + w2 + * ( + 1 + - SemanticMemoryUnitDB.extra_requirements_embedding.cosine_distance( + extra_requirements_embedding + ) + ) + + w3 * (1 - SemanticMemoryUnitDB.purpose_embedding.cosine_distance(purpose_embedding)) + ) + + # Build query: filter by url/commit and order by similarity + stmt = ( + select(SemanticMemoryUnitDB) + .order_by(similarity.desc()) + .limit(top_k) + .where(SemanticMemoryUnitDB.url == url, SemanticMemoryUnitDB.commit_id == commit_id) + ) + + async with AsyncSession(self.engine) as session: + result = await session.execute(stmt) + return result.scalars().all() diff --git a/athena/configuration/config.py b/athena/configuration/config.py index 82f643c..9daf8c4 100644 --- a/athena/configuration/config.py +++ b/athena/configuration/config.py @@ -50,5 +50,8 @@ class Settings(BaseSettings): EPISODE_MEMORY_EMBEDDING_IVFFLAT_LISTS: Optional[int] = 100 SEMANTIC_MEMORY_EMBEDDING_IVFFLAT_LISTS: Optional[int] = 10 + # Semantic Memory settings + SEMANTIC_MEMORY_MAX_RESULTS: int # Max semantic memory results to retrieve + settings = Settings() diff --git a/athena/entity/semantic_memory.py b/athena/entity/semantic_memory.py index cf8897c..473b832 100644 --- a/athena/entity/semantic_memory.py +++ b/athena/entity/semantic_memory.py @@ -1,13 +1,8 @@ -from typing import Optional - from pgvector.sqlalchemy import Vector from sqlalchemy import Column, Text from sqlmodel import Field, SQLModel from athena.configuration.config import settings -from athena.models.memory_context import MemoryContext -from athena.models.query import Query -from athena.models.semantic_memory import SemanticMemoryUnit # Database models for persistent storage @@ -39,45 +34,10 @@ class SemanticMemoryUnitDB(SQLModel, table=True): _vec_dim: int = settings.EMBEDDING_DIM or 1024 # type: ignore[assignment] # Query embeddings for weighted retrieval - essential_query_embedding: Optional[list[float]] = Field( - sa_column=Column(Vector(_vec_dim), nullable=False) - ) - extra_requirements_embedding: Optional[list[float]] = Field( + essential_query_embedding: list[float] = Field( sa_column=Column(Vector(_vec_dim), nullable=False) ) - purpose_embedding: Optional[list[float]] = Field( + extra_requirements_embedding: list[float] = Field( sa_column=Column(Vector(_vec_dim), nullable=False) ) - - @classmethod - def from_semantic_memory_unit( - cls, - memory_unit: SemanticMemoryUnit, - url: Optional[str] = None, - commit_id: Optional[str] = None, - ) -> "SemanticMemoryUnitDB": - """Create a database model from a SemanticMemoryUnit.""" - - return cls( - url=url, - commit_id=commit_id, - query_essential_query=memory_unit.query.essential_query, - query_extra_requirements=memory_unit.query.extra_requirements, - query_purpose=memory_unit.query.purpose, - memory_context_overview=memory_unit.memory_context.overview, - memory_context_content=memory_unit.memory_context.content, - ) - - def to_semantic_memory_unit(self) -> SemanticMemoryUnit: - """Convert database model back to SemanticMemoryUnit.""" - return SemanticMemoryUnit( - query=Query( - essential_query=self.query_essential_query, - extra_requirements=self.query_extra_requirements, - purpose=self.query_purpose, - ), - memory_context=MemoryContext( - overview=self.memory_context_overview, - content=self.memory_context_content, - ), - ) + purpose_embedding: list[float] = Field(sa_column=Column(Vector(_vec_dim), nullable=False)) From b825e700218182e78f7cbf367db9b37e83aaa9fc Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Sun, 5 Oct 2025 23:08:03 +0100 Subject: [PATCH 16/30] feat: implement episodic and semantic memory services with API endpoints --- athena/app/api/main.py | 39 ++------- athena/app/api/routes/__init__.py | 0 athena/app/api/routes/episodic_memory.py | 36 ++++++++ athena/app/api/routes/semantic_memory.py | 99 ++++++++++++++++++++++ athena/app/dependencies.py | 4 +- athena/app/main.py | 12 +-- athena/models/requests/__init__.py | 0 athena/models/requests/semantic_memory.py | 21 +++++ athena/models/responses/__init__.py | 0 athena/models/responses/response.py | 15 ++++ athena/models/responses/semantic_memory.py | 11 +++ 11 files changed, 195 insertions(+), 42 deletions(-) create mode 100644 athena/app/api/routes/__init__.py create mode 100644 athena/app/api/routes/episodic_memory.py create mode 100644 athena/app/api/routes/semantic_memory.py create mode 100644 athena/models/requests/__init__.py create mode 100644 athena/models/requests/semantic_memory.py create mode 100644 athena/models/responses/__init__.py create mode 100644 athena/models/responses/response.py create mode 100644 athena/models/responses/semantic_memory.py diff --git a/athena/app/api/main.py b/athena/app/api/main.py index f1e5356..73cf0b2 100644 --- a/athena/app/api/main.py +++ b/athena/app/api/main.py @@ -1,34 +1,11 @@ -from fastapi import APIRouter, HTTPException, Query, Request +from fastapi import APIRouter -from athena.configuration.config import settings +from athena.app.api.routes import episodic_memory, semantic_memory api_router = APIRouter() - - -@api_router.get("/episodic_memory/search", tags=["episodic_memory"]) -async def search_episodic_memory( - request: Request, - q: str, - field: str = "task_state", - limit: int = Query(default=settings.MEMORY_SEARCH_LIMIT, ge=1, le=100), -): - services = getattr(request.app.state, "service", {}) - episodic_memory_service = services.get("episodic_memory_service") - if episodic_memory_service is None: - raise HTTPException(status_code=503, detail="Episodic memory service not initialized") - if field not in {"task_state", "task", "state"}: - raise HTTPException(status_code=400, detail=f"Invalid field: {field}") - results = await episodic_memory_service.search_memory(q, limit=limit, field=field) - return [m.model_dump() for m in results] - - -@api_router.get("/episodic_memory/{memory_id}", tags=["episodic_memory"]) -async def get_episodic_memory(request: Request, memory_id: str): - services = getattr(request.app.state, "service", {}) - episodic_memory_service = services.get("episodic_memory_service") - if episodic_memory_service is None: - raise HTTPException(status_code=503, detail="Episodic memory service not initialized") - result = await episodic_memory_service.get_memory_by_key(memory_id) - if result is None: - raise HTTPException(status_code=404, detail="Episodic memory not found") - return result.model_dump() +api_router.include_router( + episodic_memory.router, prefix="/episodic-memory", tags=["episodic-memory"] +) +api_router.include_router( + semantic_memory.router, prefix="/semantic-memory", tags=["semantic-memory"] +) diff --git a/athena/app/api/routes/__init__.py b/athena/app/api/routes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/athena/app/api/routes/episodic_memory.py b/athena/app/api/routes/episodic_memory.py new file mode 100644 index 0000000..44e49de --- /dev/null +++ b/athena/app/api/routes/episodic_memory.py @@ -0,0 +1,36 @@ +from fastapi import APIRouter, HTTPException, Query, Request + +from athena.configuration.config import settings + +router = APIRouter() + + +@router.get("/search/", summary="Search Episodic Memory from database") +async def search_episodic_memory( + request: Request, + q: str, + field: str = "task_state", + limit: int = Query(default=settings.MEMORY_SEARCH_LIMIT, ge=1, le=100), +): + # Get the episodic memory service from the app state + episodic_memory_service = request.app.state.service["episodic_memory_service"] + + # Validate field + if field not in {"task_state", "task", "state"}: + raise HTTPException(status_code=400, detail=f"Invalid field: {field}") + + # Perform the search + results = await episodic_memory_service.search_memory(q, limit=limit, field=field) + return [m.model_dump() for m in results] + + +@router.get("/{memory_id}/", summary="Get Episodic Memory from database") +async def get_episodic_memory(request: Request, memory_id: str): + # Get the episodic memory service from the app state + episodic_memory_service = request.app.state.service["episodic_memory_service"] + + # Fetch the memory by ID + result = await episodic_memory_service.get_memory_by_key(memory_id) + if result is None: + raise HTTPException(status_code=404, detail="Episodic memory not found") + return result.model_dump() diff --git a/athena/app/api/routes/semantic_memory.py b/athena/app/api/routes/semantic_memory.py new file mode 100644 index 0000000..0086d63 --- /dev/null +++ b/athena/app/api/routes/semantic_memory.py @@ -0,0 +1,99 @@ +from typing import List + +from fastapi import APIRouter, HTTPException, Request + +from athena.models.requests.semantic_memory import ( + RetrieveSemanticMemoryRequest, + StoreSemanticMemoryRequest, +) +from athena.models.responses.response import Response +from athena.models.responses.semantic_memory import SemanticMemoryResponse +from athena.models.semantic_memory import SemanticMemoryUnit + +router = APIRouter() + + +@router.post("/store/", summary="Store semantic memory", response_model=Response) +async def store_semantic_memory( + request: Request, + body: StoreSemanticMemoryRequest, +): + """ + Store a semantic memory unit with vector embeddings. + + Args: + request: FastAPI request object + body: Request containing url, commit_id, query, and memory_context + + Returns: + Success message + """ + # Get the semantic memory service from the app state + semantic_memory_service = request.app.state.service["semantic_memory_service"] + + # Create SemanticMemoryUnit from request + memory_unit = SemanticMemoryUnit( + query=body.query, + memory_context=body.memory_context, + ) + + # Store the memory + await semantic_memory_service.store_memory( + url=body.url, + commit_id=body.commit_id, + memory=memory_unit, + ) + + return Response() + + +@router.post( + "/retrieve/", + response_model=Response[List[SemanticMemoryResponse]], + summary="Retrieve semantic memories", +) +async def retrieve_semantic_memory( + request: Request, + body: RetrieveSemanticMemoryRequest, +): + """ + Retrieve semantic memories similar to the given query. + + Uses weighted multi-vector similarity search across: + - Essential query (60% weight) + - Extra requirements (20% weight) + - Purpose (20% weight) + + Args: + request: FastAPI request object + body: Request containing url, commit_id, and query + + Returns: + List of semantic memory units ordered by similarity + """ + # Get the semantic memory service from the app state + semantic_memory_service = request.app.state.service["semantic_memory_service"] + + # Retrieve memories + results = await semantic_memory_service.retrieve_memory( + url=body.url, + commit_id=body.commit_id, + query=body.query, + ) + + if not results: + raise HTTPException(status_code=404, detail="No semantic memories found") + + # Convert to response models + return Response( + data=[ + SemanticMemoryResponse( + query_essential_query=memory.query_essential_query, + query_extra_requirements=memory.query_extra_requirements, + query_purpose=memory.query_purpose, + memory_context_overview=memory.memory_context_overview, + memory_context_content=memory.memory_context_content, + ) + for memory in results + ] + ) diff --git a/athena/app/dependencies.py b/athena/app/dependencies.py index d48dc15..da10017 100644 --- a/athena/app/dependencies.py +++ b/athena/app/dependencies.py @@ -60,7 +60,9 @@ def initialize_services() -> Dict[str, BaseService]: storage_backend=settings.MEMORY_STORAGE_BACKEND, store=episodic_memory_store, ) - semantic_memory_service = SemanticMemoryService(database_service=database_service) + semantic_memory_service = SemanticMemoryService( + database_service=database_service, embedding_service=embedding_service + ) return { "llm_service": llm_service, diff --git a/athena/app/main.py b/athena/app/main.py index c0a84d0..879184f 100644 --- a/athena/app/main.py +++ b/athena/app/main.py @@ -75,13 +75,5 @@ def custom_generate_unique_id(route: APIRoute) -> str: @app.get("/health", tags=["health"]) -async def health_check(): - services = getattr(app.state, "service", {}) - db = services.get("database_service") - db_ok = await db.health_check() if db is not None else False - status = "healthy" if db_ok else "degraded" - return { - "status": status, - "database": db_ok, - "timestamp": datetime.now(timezone.utc).isoformat(), - } +def health_check(): + return {"status": "healthy", "timestamp": datetime.now(timezone.utc).isoformat()} diff --git a/athena/models/requests/__init__.py b/athena/models/requests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/athena/models/requests/semantic_memory.py b/athena/models/requests/semantic_memory.py new file mode 100644 index 0000000..2a208cb --- /dev/null +++ b/athena/models/requests/semantic_memory.py @@ -0,0 +1,21 @@ +from pydantic import BaseModel + +from athena.models.memory_context import MemoryContext +from athena.models.query import Query + + +class StoreSemanticMemoryRequest(BaseModel): + """Request model for storing semantic memory.""" + + url: str + commit_id: str + query: Query + memory_context: MemoryContext + + +class RetrieveSemanticMemoryRequest(BaseModel): + """Request model for retrieving semantic memory.""" + + url: str + commit_id: str + query: Query diff --git a/athena/models/responses/__init__.py b/athena/models/responses/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/athena/models/responses/response.py b/athena/models/responses/response.py new file mode 100644 index 0000000..09c7e59 --- /dev/null +++ b/athena/models/responses/response.py @@ -0,0 +1,15 @@ +from typing import Generic, TypeVar + +from pydantic import BaseModel + +T = TypeVar("T") + + +class Response(BaseModel, Generic[T]): + """ + Generic response model for API responses. + """ + + code: int = 200 + message: str = "success" + data: T | None = None diff --git a/athena/models/responses/semantic_memory.py b/athena/models/responses/semantic_memory.py new file mode 100644 index 0000000..6f3d3f7 --- /dev/null +++ b/athena/models/responses/semantic_memory.py @@ -0,0 +1,11 @@ +from pydantic import BaseModel + + +class SemanticMemoryResponse(BaseModel): + """Response model for semantic memory retrieval.""" + + query_essential_query: str + query_extra_requirements: str + query_purpose: str + memory_context_overview: str + memory_context_content: str From 98ce71c889929202e297b4ac03aa78fca1b2eb43 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Mon, 6 Oct 2025 00:08:33 +0100 Subject: [PATCH 17/30] feat: implement episodic and semantic memory services with API endpoints --- athena/app/api/routes/semantic_memory.py | 24 +++++++++++++------ athena/app/dependencies.py | 15 +++++------- .../app/services/semantic_memory_service.py | 10 ++++---- athena/configuration/config.py | 8 +++---- athena/models/query.py | 4 ++-- athena/models/requests/semantic_memory.py | 20 +++++++--------- 6 files changed, 43 insertions(+), 38 deletions(-) diff --git a/athena/app/api/routes/semantic_memory.py b/athena/app/api/routes/semantic_memory.py index 0086d63..2434611 100644 --- a/athena/app/api/routes/semantic_memory.py +++ b/athena/app/api/routes/semantic_memory.py @@ -2,8 +2,8 @@ from fastapi import APIRouter, HTTPException, Request +from athena.models.query import Query from athena.models.requests.semantic_memory import ( - RetrieveSemanticMemoryRequest, StoreSemanticMemoryRequest, ) from athena.models.responses.response import Response @@ -47,14 +47,18 @@ async def store_semantic_memory( return Response() -@router.post( +@router.get( "/retrieve/", response_model=Response[List[SemanticMemoryResponse]], summary="Retrieve semantic memories", ) async def retrieve_semantic_memory( request: Request, - body: RetrieveSemanticMemoryRequest, + url: str, + commit_id: str, + essential_query: str, + extra_requirements: str, + purpose: str, ): """ Retrieve semantic memories similar to the given query. @@ -66,7 +70,11 @@ async def retrieve_semantic_memory( Args: request: FastAPI request object - body: Request containing url, commit_id, and query + url: Source URL of the repository + commit_id: Git commit ID associated with this memory + essential_query: The main query string + extra_requirements: Additional constraints or filters for the query + purpose: The intent or context behind the query Returns: List of semantic memory units ordered by similarity @@ -76,9 +84,11 @@ async def retrieve_semantic_memory( # Retrieve memories results = await semantic_memory_service.retrieve_memory( - url=body.url, - commit_id=body.commit_id, - query=body.query, + url=url, + commit_id=commit_id, + query=Query( + essential_query=essential_query, extra_requirements=extra_requirements, purpose=purpose + ), ) if not results: diff --git a/athena/app/dependencies.py b/athena/app/dependencies.py index da10017..c5a92a5 100644 --- a/athena/app/dependencies.py +++ b/athena/app/dependencies.py @@ -41,15 +41,12 @@ def initialize_services() -> Dict[str, BaseService]: settings.GOOGLE_APPLICATION_CREDENTIALS, ) - embedding_service = None - api_key = settings.EMBEDDING_API_KEY or settings.MISTRAL_API_KEY - if settings.EMBEDDING_MODEL and api_key and settings.EMBEDDING_BASE_URL: - embedding_service = EmbeddingService( - model=settings.EMBEDDING_MODEL, - api_key=api_key, - base_url=settings.EMBEDDING_BASE_URL, - embed_dim=settings.EMBEDDING_DIM or 1024, - ) + embedding_service = EmbeddingService( + model=settings.EMBEDDING_MODEL, + api_key=settings.EMBEDDING_API_KEY, + base_url=settings.EMBEDDING_BASE_URL, + embed_dim=settings.EMBEDDING_DIM, + ) episodic_memory_store = EpisodicMemoryStorageService( database_service.get_sessionmaker(), embedding_service diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index 5fb48d8..43bd947 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -136,11 +136,11 @@ async def store_memory(self, url: str, commit_id: str, memory: SemanticMemoryUni semantic_memory = SemanticMemoryUnitDB( url=url, commit_id=commit_id, - query_essential_query=memory.query_essential_query, - query_extra_requirements=memory.query_extra_requirements, - query_purpose=memory.query_purpose, - memory_context_overview=memory.memory_context_overview, - memory_context_content=memory.memory_context_content, + query_essential_query=memory.query.essential_query, + query_extra_requirements=memory.query.extra_requirements, + query_purpose=memory.query.purpose, + memory_context_overview=memory.memory_context.overview, + memory_context_content=memory.memory_context.content, essential_query_embedding=essential_query_embedding, extra_requirements_embedding=extra_requirements_embedding, purpose_embedding=purpose_embedding, diff --git a/athena/configuration/config.py b/athena/configuration/config.py index 9daf8c4..fe822a0 100644 --- a/athena/configuration/config.py +++ b/athena/configuration/config.py @@ -41,10 +41,10 @@ class Settings(BaseSettings): MEMORY_SEARCH_LIMIT: int = 10 # Default search result limit # Embeddings (OpenAI-format, works with Codestral embed via a compatible gateway) - EMBEDDING_MODEL: Optional[str] = None - EMBEDDING_API_KEY: Optional[str] = None - EMBEDDING_BASE_URL: Optional[str] = None - EMBEDDING_DIM: Optional[int] = None + EMBEDDING_MODEL: str + EMBEDDING_API_KEY: str + EMBEDDING_BASE_URL: str + EMBEDDING_DIM: Optional[int] = 1024 # Dimension of the embedding vectors # Embedding index settings EPISODE_MEMORY_EMBEDDING_IVFFLAT_LISTS: Optional[int] = 100 diff --git a/athena/models/query.py b/athena/models/query.py index b9c18b6..e58ea41 100644 --- a/athena/models/query.py +++ b/athena/models/query.py @@ -3,5 +3,5 @@ class Query(BaseModel): essential_query: str - extra_requirements: str = "" - purpose: str = "" + extra_requirements: str + purpose: str diff --git a/athena/models/requests/semantic_memory.py b/athena/models/requests/semantic_memory.py index 2a208cb..58d065e 100644 --- a/athena/models/requests/semantic_memory.py +++ b/athena/models/requests/semantic_memory.py @@ -1,4 +1,4 @@ -from pydantic import BaseModel +from pydantic import BaseModel, Field from athena.models.memory_context import MemoryContext from athena.models.query import Query @@ -7,15 +7,13 @@ class StoreSemanticMemoryRequest(BaseModel): """Request model for storing semantic memory.""" - url: str - commit_id: str + url: str = Field(description="The URL of the repository", max_length=100) + commit_id: str | None = Field( + default=None, + description="The commit id of the repository, " + "if not provided, the latest commit in the main branch will be used.", + min_length=40, + max_length=40, + ) query: Query memory_context: MemoryContext - - -class RetrieveSemanticMemoryRequest(BaseModel): - """Request model for retrieving semantic memory.""" - - url: str - commit_id: str - query: Query From 01cce63eea99403c463b2a457a7a5eecea027076 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Mon, 6 Oct 2025 00:30:26 +0100 Subject: [PATCH 18/30] fix semantic_memory_service --- athena/app/api/routes/semantic_memory.py | 16 ++++++++-------- athena/app/services/semantic_memory_service.py | 1 - 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/athena/app/api/routes/semantic_memory.py b/athena/app/api/routes/semantic_memory.py index 2434611..bffacb4 100644 --- a/athena/app/api/routes/semantic_memory.py +++ b/athena/app/api/routes/semantic_memory.py @@ -1,7 +1,8 @@ from typing import List -from fastapi import APIRouter, HTTPException, Request +from fastapi import APIRouter, Request +from athena.app.services.semantic_memory_service import SemanticMemoryService from athena.models.query import Query from athena.models.requests.semantic_memory import ( StoreSemanticMemoryRequest, @@ -29,7 +30,9 @@ async def store_semantic_memory( Success message """ # Get the semantic memory service from the app state - semantic_memory_service = request.app.state.service["semantic_memory_service"] + semantic_memory_service: SemanticMemoryService = request.app.state.service[ + "semantic_memory_service" + ] # Create SemanticMemoryUnit from request memory_unit = SemanticMemoryUnit( @@ -55,10 +58,10 @@ async def store_semantic_memory( async def retrieve_semantic_memory( request: Request, url: str, - commit_id: str, essential_query: str, - extra_requirements: str, - purpose: str, + extra_requirements: str = "", + purpose: str = "", + commit_id: str | None = None, ): """ Retrieve semantic memories similar to the given query. @@ -91,9 +94,6 @@ async def retrieve_semantic_memory( ), ) - if not results: - raise HTTPException(status_code=404, detail="No semantic memories found") - # Convert to response models return Response( data=[ diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index 43bd947..cc989e8 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -147,7 +147,6 @@ async def store_memory(self, url: str, commit_id: str, memory: SemanticMemoryUni ) session.add(semantic_memory) await session.commit() - await session.refresh(semantic_memory) async def retrieve_memory( self, url: str, commit_id: str, query: Query From ac22ee43fdb98021908f9e012dfde0601317bdc3 Mon Sep 17 00:00:00 2001 From: Zhaoyang-Chu Date: Mon, 6 Oct 2025 08:15:29 +0000 Subject: [PATCH 19/30] feat: debug and improve episodic memory services --- athena/app/services/llm_service.py | 14 +++++++++++++- athena/models/episodic_memory.py | 4 ++-- athena/scripts/offline_ingest_hf.py | 1 + 3 files changed, 16 insertions(+), 3 deletions(-) diff --git a/athena/app/services/llm_service.py b/athena/app/services/llm_service.py index bf899e0..7961240 100644 --- a/athena/app/services/llm_service.py +++ b/athena/app/services/llm_service.py @@ -4,6 +4,7 @@ from langchain_core.language_models.chat_models import BaseChatModel from langchain_google_genai import ChatGoogleGenerativeAI from langchain_google_vertexai import ChatVertexAI +from google.oauth2.service_account import Credentials from athena.app.services.base_service import BaseService from athena.chat_models.custom_chat_openai import CustomChatOpenAI @@ -50,13 +51,24 @@ def get_model( elif model_name.startswith("vertex:"): # example: model_name="vertex:gemini-2.5-pro" vertex_model = model_name.split("vertex:", 1)[1] + print(google_application_credentials) + if isinstance(google_application_credentials, dict): + creds = Credentials.from_service_account_info( + google_application_credentials, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ) + else: + creds = Credentials.from_service_account_file( + google_application_credentials, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ) return ChatVertexAI( model_name=vertex_model, project="prometheus-code-agent", location="us-central1", temperature=temperature, max_retries=3, - credentials=google_application_credentials, + credentials=creds, ) elif "gemini" in model_name: return ChatGoogleGenerativeAI( diff --git a/athena/models/episodic_memory.py b/athena/models/episodic_memory.py index 623fc98..cb8b0c9 100644 --- a/athena/models/episodic_memory.py +++ b/athena/models/episodic_memory.py @@ -98,7 +98,7 @@ class EpisodicMemoryUnit(BaseModel): Core episodic memory unit capturing one action of agent execution. This includes: - - The episodic memory id (episodic_memory_id) + - The episodic memory id (memory_id) - The episodic memory timestamp (timestamp) - The episodic memory source (source_name, run_id, metadata) - The task being worked on (issue and repository details) @@ -107,7 +107,7 @@ class EpisodicMemoryUnit(BaseModel): - The result of the action """ - episodic_memory_id: str = Field( + memory_id: str = Field( default_factory=lambda: str(uuid.uuid4()), description="Unique identifier for this episodic memory unit", ) diff --git a/athena/scripts/offline_ingest_hf.py b/athena/scripts/offline_ingest_hf.py index 9da2e1a..7c3cd14 100644 --- a/athena/scripts/offline_ingest_hf.py +++ b/athena/scripts/offline_ingest_hf.py @@ -40,6 +40,7 @@ async def main(): openai_format_base_url=settings.OPENAI_FORMAT_BASE_URL, anthropic_api_key=settings.ANTHROPIC_API_KEY, gemini_api_key=settings.GEMINI_API_KEY, + google_application_credentials=settings.GOOGLE_APPLICATION_CREDENTIALS, ) extractor = EpisodicMemoryExtractionService(llm_service=llm, memory_store=store) From 8fcf9e32fdebc3958fbd187262e2c4b2279de7e3 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Mon, 6 Oct 2025 12:05:04 +0100 Subject: [PATCH 20/30] add minimum similarity for semantic_memory_service --- athena/app/services/semantic_memory_service.py | 14 ++++++++++---- athena/configuration/config.py | 1 + 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index cc989e8..ece8dfa 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -160,7 +160,8 @@ async def retrieve_memory( - Purpose: 20% weight (intent and context) The combined similarity score ranks results, with higher scores indicating - better matches. + better matches. Only results with similarity >= SEMANTIC_MEMORY_MIN_SIMILARITY + are returned. Args: url: Filter memories by this repository URL @@ -169,11 +170,12 @@ async def retrieve_memory( Returns: List of SemanticMemoryUnitDB objects ordered by similarity (most similar first), - limited to SEMANTIC_MEMORY_MAX_RESULTS + limited to SEMANTIC_MEMORY_MAX_RESULTS and filtered by minimum similarity threshold """ # Weights for similarity components: essential query gets highest weight w1, w2, w3 = 0.6, 0.2, 0.2 top_k = settings.SEMANTIC_MEMORY_MAX_RESULTS + min_similarity = settings.SEMANTIC_MEMORY_MIN_SIMILARITY # Get embeddings for the query components ( @@ -202,12 +204,16 @@ async def retrieve_memory( + w3 * (1 - SemanticMemoryUnitDB.purpose_embedding.cosine_distance(purpose_embedding)) ) - # Build query: filter by url/commit and order by similarity + # Build query: filter by url/commit, similarity threshold, and order by similarity stmt = ( select(SemanticMemoryUnitDB) .order_by(similarity.desc()) .limit(top_k) - .where(SemanticMemoryUnitDB.url == url, SemanticMemoryUnitDB.commit_id == commit_id) + .where( + SemanticMemoryUnitDB.url == url, + SemanticMemoryUnitDB.commit_id == commit_id, + similarity >= min_similarity, + ) ) async with AsyncSession(self.engine) as session: diff --git a/athena/configuration/config.py b/athena/configuration/config.py index fe822a0..9a7a0d3 100644 --- a/athena/configuration/config.py +++ b/athena/configuration/config.py @@ -52,6 +52,7 @@ class Settings(BaseSettings): # Semantic Memory settings SEMANTIC_MEMORY_MAX_RESULTS: int # Max semantic memory results to retrieve + SEMANTIC_MEMORY_MIN_SIMILARITY: float # Minimum similarity threshold (0-1) settings = Settings() From e5b02c883a32b690ec24661b27b9afe16fe3549b Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Mon, 6 Oct 2025 12:23:21 +0100 Subject: [PATCH 21/30] add minimum similarity for semantic_memory_service --- athena/app/api/routes/semantic_memory.py | 14 +++++--------- athena/app/services/semantic_memory_service.py | 18 +++++++----------- athena/entity/semantic_memory.py | 7 +------ athena/models/requests/semantic_memory.py | 9 +-------- 4 files changed, 14 insertions(+), 34 deletions(-) diff --git a/athena/app/api/routes/semantic_memory.py b/athena/app/api/routes/semantic_memory.py index bffacb4..50faf76 100644 --- a/athena/app/api/routes/semantic_memory.py +++ b/athena/app/api/routes/semantic_memory.py @@ -24,7 +24,7 @@ async def store_semantic_memory( Args: request: FastAPI request object - body: Request containing url, commit_id, query, and memory_context + body: Request containing repository_id, query, and memory_context Returns: Success message @@ -42,8 +42,7 @@ async def store_semantic_memory( # Store the memory await semantic_memory_service.store_memory( - url=body.url, - commit_id=body.commit_id, + repository_id=body.repository_id, memory=memory_unit, ) @@ -57,11 +56,10 @@ async def store_semantic_memory( ) async def retrieve_semantic_memory( request: Request, - url: str, + repository_id: int, essential_query: str, extra_requirements: str = "", purpose: str = "", - commit_id: str | None = None, ): """ Retrieve semantic memories similar to the given query. @@ -73,8 +71,7 @@ async def retrieve_semantic_memory( Args: request: FastAPI request object - url: Source URL of the repository - commit_id: Git commit ID associated with this memory + repository_id: Repository identifier essential_query: The main query string extra_requirements: Additional constraints or filters for the query purpose: The intent or context behind the query @@ -87,8 +84,7 @@ async def retrieve_semantic_memory( # Retrieve memories results = await semantic_memory_service.retrieve_memory( - url=url, - commit_id=commit_id, + repository_id=repository_id, query=Query( essential_query=essential_query, extra_requirements=extra_requirements, purpose=purpose ), diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index ece8dfa..8973329 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -110,7 +110,7 @@ async def get_query_embeddings( embeddings = await self.embedding_service.embed(texts) return embeddings[0], embeddings[1], embeddings[2] - async def store_memory(self, url: str, commit_id: str, memory: SemanticMemoryUnit): + async def store_memory(self, repository_id: int, memory: SemanticMemoryUnit): """ Store a semantic memory unit with its vector embeddings. @@ -120,8 +120,7 @@ async def store_memory(self, url: str, commit_id: str, memory: SemanticMemoryUni 3. Persists the memory to the database Args: - url: Source URL of the repository - commit_id: Git commit ID associated with this memory + repository_id: Repository identifier memory: SemanticMemoryUnit containing query and memory context """ # Get embeddings for the query @@ -134,8 +133,7 @@ async def store_memory(self, url: str, commit_id: str, memory: SemanticMemoryUni # Create a new SemanticMemoryUnitDB instance with embeddings async with AsyncSession(self.engine) as session: semantic_memory = SemanticMemoryUnitDB( - url=url, - commit_id=commit_id, + repository_id=repository_id, query_essential_query=memory.query.essential_query, query_extra_requirements=memory.query.extra_requirements, query_purpose=memory.query.purpose, @@ -149,7 +147,7 @@ async def store_memory(self, url: str, commit_id: str, memory: SemanticMemoryUni await session.commit() async def retrieve_memory( - self, url: str, commit_id: str, query: Query + self, repository_id: int, query: Query ) -> Sequence[SemanticMemoryUnitDB]: """ Retrieve semantic memories most similar to the given query. @@ -164,8 +162,7 @@ async def retrieve_memory( are returned. Args: - url: Filter memories by this repository URL - commit_id: Filter memories by this commit ID + repository_id: Filter memories by this repository identifier query: Query object to search for similar memories Returns: @@ -204,14 +201,13 @@ async def retrieve_memory( + w3 * (1 - SemanticMemoryUnitDB.purpose_embedding.cosine_distance(purpose_embedding)) ) - # Build query: filter by url/commit, similarity threshold, and order by similarity + # Build query: filter by repository_id, similarity threshold, and order by similarity stmt = ( select(SemanticMemoryUnitDB) .order_by(similarity.desc()) .limit(top_k) .where( - SemanticMemoryUnitDB.url == url, - SemanticMemoryUnitDB.commit_id == commit_id, + SemanticMemoryUnitDB.repository_id == repository_id, similarity >= min_similarity, ) ) diff --git a/athena/entity/semantic_memory.py b/athena/entity/semantic_memory.py index 473b832..af52426 100644 --- a/athena/entity/semantic_memory.py +++ b/athena/entity/semantic_memory.py @@ -14,12 +14,7 @@ class SemanticMemoryUnitDB(SQLModel, table=True): id: int = Field(primary_key=True, description="ID") # Source information - url: str = Field(index=True, description="Source URL of the repository") - commit_id: str = Field( - index=True, - description="Git commit ID associated with the memory", - nullable=True, - ) + repository_id: int = Field(index=True, description="Repository identifier") # Query information (for retrieval) query_essential_query: str = Field(sa_column=Column(Text, nullable=False)) diff --git a/athena/models/requests/semantic_memory.py b/athena/models/requests/semantic_memory.py index 58d065e..269f998 100644 --- a/athena/models/requests/semantic_memory.py +++ b/athena/models/requests/semantic_memory.py @@ -7,13 +7,6 @@ class StoreSemanticMemoryRequest(BaseModel): """Request model for storing semantic memory.""" - url: str = Field(description="The URL of the repository", max_length=100) - commit_id: str | None = Field( - default=None, - description="The commit id of the repository, " - "if not provided, the latest commit in the main branch will be used.", - min_length=40, - max_length=40, - ) + repository_id: int = Field(description="Repository identifier") query: Query memory_context: MemoryContext From eb9d38090bdc48abc49fdb09de75aff0abd1626a Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Mon, 6 Oct 2025 12:38:23 +0100 Subject: [PATCH 22/30] add delete semantic_memory API service --- athena/app/api/routes/semantic_memory.py | 34 ++++++++++++++++--- .../app/services/semantic_memory_service.py | 28 ++++++++++++--- 2 files changed, 53 insertions(+), 9 deletions(-) diff --git a/athena/app/api/routes/semantic_memory.py b/athena/app/api/routes/semantic_memory.py index 50faf76..7ca3bf2 100644 --- a/athena/app/api/routes/semantic_memory.py +++ b/athena/app/api/routes/semantic_memory.py @@ -50,7 +50,7 @@ async def store_semantic_memory( @router.get( - "/retrieve/", + "/retrieve/{repository_id}/", response_model=Response[List[SemanticMemoryResponse]], summary="Retrieve semantic memories", ) @@ -65,9 +65,9 @@ async def retrieve_semantic_memory( Retrieve semantic memories similar to the given query. Uses weighted multi-vector similarity search across: - - Essential query (60% weight) - - Extra requirements (20% weight) - - Purpose (20% weight) + - Essential query (70% weight) + - Extra requirements (15% weight) + - Purpose (15% weight) Args: request: FastAPI request object @@ -103,3 +103,29 @@ async def retrieve_semantic_memory( for memory in results ] ) + + +@router.delete("/{repository_id}/", summary="Delete semantic memories by repository", response_model=Response) +async def delete_semantic_memories_by_repository( + request: Request, + repository_id: int, +): + """ + Delete all semantic memories for a given repository. + + Args: + request: FastAPI request object + repository_id: Repository identifier to filter memories for deletion + + Returns: + Success message + """ + # Get the semantic memory service from the app state + semantic_memory_service: SemanticMemoryService = request.app.state.service[ + "semantic_memory_service" + ] + + # Delete memories + await semantic_memory_service.delete_memories_by_repository(repository_id=repository_id) + + return Response() diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index 8973329..2cc7153 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -15,7 +15,7 @@ from typing import Sequence from sqlalchemy.ext.asyncio import AsyncSession -from sqlmodel import select +from sqlmodel import delete, select from athena.app.services.base_service import BaseService from athena.app.services.database_service import DatabaseService @@ -153,9 +153,9 @@ async def retrieve_memory( Retrieve semantic memories most similar to the given query. Uses a weighted combination of cosine similarity across three query components: - - Essential query: 60% weight (most important for relevance) - - Extra requirements: 20% weight (filters and constraints) - - Purpose: 20% weight (intent and context) + - Essential query: 70% weight (most important for relevance) + - Extra requirements: 15% weight (filters and constraints) + - Purpose: 15% weight (intent and context) The combined similarity score ranks results, with higher scores indicating better matches. Only results with similarity >= SEMANTIC_MEMORY_MIN_SIMILARITY @@ -170,7 +170,7 @@ async def retrieve_memory( limited to SEMANTIC_MEMORY_MAX_RESULTS and filtered by minimum similarity threshold """ # Weights for similarity components: essential query gets highest weight - w1, w2, w3 = 0.6, 0.2, 0.2 + w1, w2, w3 = 0.7, 0.15, 0.15 top_k = settings.SEMANTIC_MEMORY_MAX_RESULTS min_similarity = settings.SEMANTIC_MEMORY_MIN_SIMILARITY @@ -215,3 +215,21 @@ async def retrieve_memory( async with AsyncSession(self.engine) as session: result = await session.execute(stmt) return result.scalars().all() + + async def delete_memories_by_repository(self, repository_id: int): + """ + Delete all semantic memories for a given repository. + + Args: + repository_id: Repository identifier to filter memories for deletion + + Returns: + Number of memories deleted + """ + stmt = delete(SemanticMemoryUnitDB).where( + SemanticMemoryUnitDB.repository_id == repository_id + ) + + async with AsyncSession(self.engine) as session: + await session.execute(stmt) + await session.commit() From 33e08ef5dbdea809f490f8a367d1ab279681d65b Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Mon, 6 Oct 2025 12:38:43 +0100 Subject: [PATCH 23/30] style formating --- athena/app/api/routes/semantic_memory.py | 4 +++- athena/app/services/llm_service.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/athena/app/api/routes/semantic_memory.py b/athena/app/api/routes/semantic_memory.py index 7ca3bf2..c7a4e53 100644 --- a/athena/app/api/routes/semantic_memory.py +++ b/athena/app/api/routes/semantic_memory.py @@ -105,7 +105,9 @@ async def retrieve_semantic_memory( ) -@router.delete("/{repository_id}/", summary="Delete semantic memories by repository", response_model=Response) +@router.delete( + "/{repository_id}/", summary="Delete semantic memories by repository", response_model=Response +) async def delete_semantic_memories_by_repository( request: Request, repository_id: int, diff --git a/athena/app/services/llm_service.py b/athena/app/services/llm_service.py index 7961240..3926ef7 100644 --- a/athena/app/services/llm_service.py +++ b/athena/app/services/llm_service.py @@ -1,10 +1,10 @@ from typing import Optional +from google.oauth2.service_account import Credentials from langchain_anthropic import ChatAnthropic from langchain_core.language_models.chat_models import BaseChatModel from langchain_google_genai import ChatGoogleGenerativeAI from langchain_google_vertexai import ChatVertexAI -from google.oauth2.service_account import Credentials from athena.app.services.base_service import BaseService from athena.chat_models.custom_chat_openai import CustomChatOpenAI From 7abe6b40a62d02ce001be49af9bffbde6988742f Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Tue, 7 Oct 2025 16:13:03 +0100 Subject: [PATCH 24/30] fix storing format --- athena/app/api/routes/semantic_memory.py | 10 +-- .../app/services/semantic_memory_service.py | 67 ++++++++++++++++--- athena/entity/context.py | 34 ++++++++++ athena/entity/semantic_memory.py | 1 - athena/models/context.py | 12 ++++ athena/models/memory_context.py | 6 +- athena/models/responses/semantic_memory.py | 6 +- 7 files changed, 120 insertions(+), 16 deletions(-) create mode 100644 athena/entity/context.py create mode 100644 athena/models/context.py diff --git a/athena/app/api/routes/semantic_memory.py b/athena/app/api/routes/semantic_memory.py index c7a4e53..e54e734 100644 --- a/athena/app/api/routes/semantic_memory.py +++ b/athena/app/api/routes/semantic_memory.py @@ -94,11 +94,11 @@ async def retrieve_semantic_memory( return Response( data=[ SemanticMemoryResponse( - query_essential_query=memory.query_essential_query, - query_extra_requirements=memory.query_extra_requirements, - query_purpose=memory.query_purpose, - memory_context_overview=memory.memory_context_overview, - memory_context_content=memory.memory_context_content, + query_essential_query=memory.query.essential_query, + query_extra_requirements=memory.query.extra_requirements, + query_purpose=memory.query.purpose, + memory_context_overview=memory.memory_context.overview, + memory_context_contexts=memory.memory_context.contexts, ) for memory in results ] diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index 2cc7153..5e3793a 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -21,7 +21,10 @@ from athena.app.services.database_service import DatabaseService from athena.app.services.embedding_service import EmbeddingService from athena.configuration.config import settings +from athena.entity.context import ContextDB from athena.entity.semantic_memory import SemanticMemoryUnitDB +from athena.models.context import Context +from athena.models.memory_context import MemoryContext from athena.models.query import Query from athena.models.semantic_memory import SemanticMemoryUnit @@ -112,12 +115,13 @@ async def get_query_embeddings( async def store_memory(self, repository_id: int, memory: SemanticMemoryUnit): """ - Store a semantic memory unit with its vector embeddings. + Store a semantic memory unit with its vector embeddings and contexts. This method: 1. Generates embeddings for the query components - 2. Creates a database record with all query and context information - 3. Persists the memory to the database + 2. Creates a database record with query and overview information + 3. Stores each context separately in the contexts table + 4. Persists everything to the database in a transaction Args: repository_id: Repository identifier @@ -138,17 +142,29 @@ async def store_memory(self, repository_id: int, memory: SemanticMemoryUnit): query_extra_requirements=memory.query.extra_requirements, query_purpose=memory.query.purpose, memory_context_overview=memory.memory_context.overview, - memory_context_content=memory.memory_context.content, essential_query_embedding=essential_query_embedding, extra_requirements_embedding=extra_requirements_embedding, purpose_embedding=purpose_embedding, ) session.add(semantic_memory) + await session.flush() # Flush to get the semantic_memory.id + + # Store each context separately + for context in memory.memory_context.contexts: + context_db = ContextDB( + semantic_memory_id=semantic_memory.id, + relative_path=context.relative_path, + content=context.content, + start_line_number=context.start_line_number, + end_line_number=context.end_line_number, + ) + session.add(context_db) + await session.commit() async def retrieve_memory( self, repository_id: int, query: Query - ) -> Sequence[SemanticMemoryUnitDB]: + ) -> Sequence[SemanticMemoryUnit]: """ Retrieve semantic memories most similar to the given query. @@ -166,8 +182,9 @@ async def retrieve_memory( query: Query object to search for similar memories Returns: - List of SemanticMemoryUnitDB objects ordered by similarity (most similar first), - limited to SEMANTIC_MEMORY_MAX_RESULTS and filtered by minimum similarity threshold + List of SemanticMemoryUnit objects with contexts loaded, ordered by similarity + (most similar first), limited to SEMANTIC_MEMORY_MAX_RESULTS and filtered by + minimum similarity threshold """ # Weights for similarity components: essential query gets highest weight w1, w2, w3 = 0.7, 0.15, 0.15 @@ -214,7 +231,41 @@ async def retrieve_memory( async with AsyncSession(self.engine) as session: result = await session.execute(stmt) - return result.scalars().all() + semantic_memories = result.scalars().all() + + # Load contexts for each semantic memory + memory_units = [] + for sem_mem in semantic_memories: + # Query contexts for this semantic memory + contexts_stmt = select(ContextDB).where(ContextDB.semantic_memory_id == sem_mem.id) + contexts_result = await session.execute(contexts_stmt) + contexts_db = contexts_result.scalars().all() + + # Convert to Context model objects + contexts = [ + Context( + relative_path=ctx.relative_path, + content=ctx.content, + start_line_number=ctx.start_line_number, + end_line_number=ctx.end_line_number, + ) + for ctx in contexts_db + ] + + # Create SemanticMemoryUnit + memory_unit = SemanticMemoryUnit( + query=Query( + essential_query=sem_mem.query_essential_query, + extra_requirements=sem_mem.query_extra_requirements, + purpose=sem_mem.query_purpose, + ), + memory_context=MemoryContext( + overview=sem_mem.memory_context_overview, contexts=contexts + ), + ) + memory_units.append(memory_unit) + + return memory_units async def delete_memories_by_repository(self, repository_id: int): """ diff --git a/athena/entity/context.py b/athena/entity/context.py new file mode 100644 index 0000000..f0b9cdc --- /dev/null +++ b/athena/entity/context.py @@ -0,0 +1,34 @@ +from typing import Optional + +from sqlalchemy import Column, ForeignKey, Integer, Text +from sqlmodel import Field, SQLModel + + +class ContextDB(SQLModel, table=True): + """Database model for storing code context information.""" + + __tablename__ = "contexts" + + id: int = Field(primary_key=True, description="ID") + + # Foreign key to semantic_memories + semantic_memory_id: int = Field( + sa_column=Column( + Integer, + ForeignKey("semantic_memories.id", ondelete="CASCADE"), + nullable=False, + index=True, + ), + description="Foreign key to semantic memory", + ) + + relative_path: str = Field( + sa_column=Column(Text, nullable=False), description="Relative file path" + ) + content: str = Field(sa_column=Column(Text, nullable=False), description="Code content") + start_line_number: Optional[int] = Field( + default=None, sa_column=Column(Integer), description="Starting line number" + ) + end_line_number: Optional[int] = Field( + default=None, sa_column=Column(Integer), description="Ending line number" + ) diff --git a/athena/entity/semantic_memory.py b/athena/entity/semantic_memory.py index af52426..3404c2a 100644 --- a/athena/entity/semantic_memory.py +++ b/athena/entity/semantic_memory.py @@ -23,7 +23,6 @@ class SemanticMemoryUnitDB(SQLModel, table=True): # Memory context memory_context_overview: str = Field(sa_column=Column(Text, nullable=False)) - memory_context_content: str = Field(sa_column=Column(Text, nullable=False)) # Embeddings for semantic retrieval (pgvector) _vec_dim: int = settings.EMBEDDING_DIM or 1024 # type: ignore[assignment] diff --git a/athena/models/context.py b/athena/models/context.py new file mode 100644 index 0000000..20b0107 --- /dev/null +++ b/athena/models/context.py @@ -0,0 +1,12 @@ +from typing import Optional + +from pydantic import BaseModel + + +class Context(BaseModel): + """Represents a code context with file location and content.""" + + relative_path: str + content: str + start_line_number: Optional[int] = None + end_line_number: Optional[int] = None diff --git a/athena/models/memory_context.py b/athena/models/memory_context.py index ce4afee..746e4ca 100644 --- a/athena/models/memory_context.py +++ b/athena/models/memory_context.py @@ -1,6 +1,10 @@ +from typing import List + from pydantic import BaseModel +from athena.models.context import Context + class MemoryContext(BaseModel): overview: str - content: str + contexts: List[Context] diff --git a/athena/models/responses/semantic_memory.py b/athena/models/responses/semantic_memory.py index 6f3d3f7..711315f 100644 --- a/athena/models/responses/semantic_memory.py +++ b/athena/models/responses/semantic_memory.py @@ -1,5 +1,9 @@ +from typing import List + from pydantic import BaseModel +from athena.models.context import Context + class SemanticMemoryResponse(BaseModel): """Response model for semantic memory retrieval.""" @@ -8,4 +12,4 @@ class SemanticMemoryResponse(BaseModel): query_extra_requirements: str query_purpose: str memory_context_overview: str - memory_context_content: str + memory_context_contexts: List[Context] From bb51e05f8256581677ff4be957cbd184e0006d08 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Tue, 7 Oct 2025 21:31:39 +0100 Subject: [PATCH 25/30] delete context overview field --- athena/app/api/routes/semantic_memory.py | 1 - athena/app/services/semantic_memory_service.py | 5 +---- athena/entity/semantic_memory.py | 3 --- athena/models/memory_context.py | 1 - athena/models/responses/semantic_memory.py | 1 - 5 files changed, 1 insertion(+), 10 deletions(-) diff --git a/athena/app/api/routes/semantic_memory.py b/athena/app/api/routes/semantic_memory.py index e54e734..4e872de 100644 --- a/athena/app/api/routes/semantic_memory.py +++ b/athena/app/api/routes/semantic_memory.py @@ -97,7 +97,6 @@ async def retrieve_semantic_memory( query_essential_query=memory.query.essential_query, query_extra_requirements=memory.query.extra_requirements, query_purpose=memory.query.purpose, - memory_context_overview=memory.memory_context.overview, memory_context_contexts=memory.memory_context.contexts, ) for memory in results diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index 5e3793a..25ec88b 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -141,7 +141,6 @@ async def store_memory(self, repository_id: int, memory: SemanticMemoryUnit): query_essential_query=memory.query.essential_query, query_extra_requirements=memory.query.extra_requirements, query_purpose=memory.query.purpose, - memory_context_overview=memory.memory_context.overview, essential_query_embedding=essential_query_embedding, extra_requirements_embedding=extra_requirements_embedding, purpose_embedding=purpose_embedding, @@ -259,9 +258,7 @@ async def retrieve_memory( extra_requirements=sem_mem.query_extra_requirements, purpose=sem_mem.query_purpose, ), - memory_context=MemoryContext( - overview=sem_mem.memory_context_overview, contexts=contexts - ), + memory_context=MemoryContext(contexts=contexts), ) memory_units.append(memory_unit) diff --git a/athena/entity/semantic_memory.py b/athena/entity/semantic_memory.py index 3404c2a..062addf 100644 --- a/athena/entity/semantic_memory.py +++ b/athena/entity/semantic_memory.py @@ -21,9 +21,6 @@ class SemanticMemoryUnitDB(SQLModel, table=True): query_extra_requirements: str = Field(sa_column=Column(Text, nullable=False)) query_purpose: str = Field(sa_column=Column(Text, nullable=False)) - # Memory context - memory_context_overview: str = Field(sa_column=Column(Text, nullable=False)) - # Embeddings for semantic retrieval (pgvector) _vec_dim: int = settings.EMBEDDING_DIM or 1024 # type: ignore[assignment] diff --git a/athena/models/memory_context.py b/athena/models/memory_context.py index 746e4ca..6669fae 100644 --- a/athena/models/memory_context.py +++ b/athena/models/memory_context.py @@ -6,5 +6,4 @@ class MemoryContext(BaseModel): - overview: str contexts: List[Context] diff --git a/athena/models/responses/semantic_memory.py b/athena/models/responses/semantic_memory.py index 711315f..07c8121 100644 --- a/athena/models/responses/semantic_memory.py +++ b/athena/models/responses/semantic_memory.py @@ -11,5 +11,4 @@ class SemanticMemoryResponse(BaseModel): query_essential_query: str query_extra_requirements: str query_purpose: str - memory_context_overview: str memory_context_contexts: List[Context] From a8a3c719b79d0ef89cd68f89a7e6e7abe068fd53 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Tue, 7 Oct 2025 21:48:15 +0100 Subject: [PATCH 26/30] fix --- athena/app/api/routes/semantic_memory.py | 4 ++-- athena/app/services/semantic_memory_service.py | 5 ++--- athena/models/memory_context.py | 9 --------- athena/models/requests/semantic_memory.py | 6 ++++-- athena/models/semantic_memory.py | 6 ++++-- 5 files changed, 12 insertions(+), 18 deletions(-) delete mode 100644 athena/models/memory_context.py diff --git a/athena/app/api/routes/semantic_memory.py b/athena/app/api/routes/semantic_memory.py index 4e872de..258160b 100644 --- a/athena/app/api/routes/semantic_memory.py +++ b/athena/app/api/routes/semantic_memory.py @@ -37,7 +37,7 @@ async def store_semantic_memory( # Create SemanticMemoryUnit from request memory_unit = SemanticMemoryUnit( query=body.query, - memory_context=body.memory_context, + contexts=body.contexts ) # Store the memory @@ -97,7 +97,7 @@ async def retrieve_semantic_memory( query_essential_query=memory.query.essential_query, query_extra_requirements=memory.query.extra_requirements, query_purpose=memory.query.purpose, - memory_context_contexts=memory.memory_context.contexts, + memory_context_contexts=memory.contexts, ) for memory in results ] diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index 25ec88b..c011a75 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -24,7 +24,6 @@ from athena.entity.context import ContextDB from athena.entity.semantic_memory import SemanticMemoryUnitDB from athena.models.context import Context -from athena.models.memory_context import MemoryContext from athena.models.query import Query from athena.models.semantic_memory import SemanticMemoryUnit @@ -149,7 +148,7 @@ async def store_memory(self, repository_id: int, memory: SemanticMemoryUnit): await session.flush() # Flush to get the semantic_memory.id # Store each context separately - for context in memory.memory_context.contexts: + for context in memory.contexts: context_db = ContextDB( semantic_memory_id=semantic_memory.id, relative_path=context.relative_path, @@ -258,7 +257,7 @@ async def retrieve_memory( extra_requirements=sem_mem.query_extra_requirements, purpose=sem_mem.query_purpose, ), - memory_context=MemoryContext(contexts=contexts), + contexts=contexts, ) memory_units.append(memory_unit) diff --git a/athena/models/memory_context.py b/athena/models/memory_context.py deleted file mode 100644 index 6669fae..0000000 --- a/athena/models/memory_context.py +++ /dev/null @@ -1,9 +0,0 @@ -from typing import List - -from pydantic import BaseModel - -from athena.models.context import Context - - -class MemoryContext(BaseModel): - contexts: List[Context] diff --git a/athena/models/requests/semantic_memory.py b/athena/models/requests/semantic_memory.py index 269f998..84b153c 100644 --- a/athena/models/requests/semantic_memory.py +++ b/athena/models/requests/semantic_memory.py @@ -1,6 +1,8 @@ +from typing import List + from pydantic import BaseModel, Field -from athena.models.memory_context import MemoryContext +from athena.models.context import Context from athena.models.query import Query @@ -9,4 +11,4 @@ class StoreSemanticMemoryRequest(BaseModel): repository_id: int = Field(description="Repository identifier") query: Query - memory_context: MemoryContext + contexts: List[Context] diff --git a/athena/models/semantic_memory.py b/athena/models/semantic_memory.py index 62db680..ec240a7 100644 --- a/athena/models/semantic_memory.py +++ b/athena/models/semantic_memory.py @@ -1,9 +1,11 @@ +from typing import List + from pydantic import BaseModel -from athena.models.memory_context import MemoryContext +from athena.models.context import Context from athena.models.query import Query class SemanticMemoryUnit(BaseModel): query: Query - memory_context: MemoryContext + contexts: List[Context] From 6682cdb299e096a6e58641d12aed320daf68fd76 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Wed, 8 Oct 2025 17:36:46 +0100 Subject: [PATCH 27/30] update semantic memory settings and adjust weight distribution for query components --- athena/app/api/routes/semantic_memory.py | 6 +++--- athena/app/services/semantic_memory_service.py | 8 ++++---- example.env | 4 ++++ 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/athena/app/api/routes/semantic_memory.py b/athena/app/api/routes/semantic_memory.py index 258160b..2d484e0 100644 --- a/athena/app/api/routes/semantic_memory.py +++ b/athena/app/api/routes/semantic_memory.py @@ -65,9 +65,9 @@ async def retrieve_semantic_memory( Retrieve semantic memories similar to the given query. Uses weighted multi-vector similarity search across: - - Essential query (70% weight) - - Extra requirements (15% weight) - - Purpose (15% weight) + - Essential query (50% weight) + - Extra requirements (25% weight) + - Purpose (25% weight) Args: request: FastAPI request object diff --git a/athena/app/services/semantic_memory_service.py b/athena/app/services/semantic_memory_service.py index c011a75..5466fa5 100644 --- a/athena/app/services/semantic_memory_service.py +++ b/athena/app/services/semantic_memory_service.py @@ -167,9 +167,9 @@ async def retrieve_memory( Retrieve semantic memories most similar to the given query. Uses a weighted combination of cosine similarity across three query components: - - Essential query: 70% weight (most important for relevance) - - Extra requirements: 15% weight (filters and constraints) - - Purpose: 15% weight (intent and context) + - Essential query: 50% weight (most important for relevance) + - Extra requirements: 25% weight (filters and constraints) + - Purpose: 25% weight (intent and context) The combined similarity score ranks results, with higher scores indicating better matches. Only results with similarity >= SEMANTIC_MEMORY_MIN_SIMILARITY @@ -185,7 +185,7 @@ async def retrieve_memory( minimum similarity threshold """ # Weights for similarity components: essential query gets highest weight - w1, w2, w3 = 0.7, 0.15, 0.15 + w1, w2, w3 = 0.5, 0.25, 0.5 top_k = settings.SEMANTIC_MEMORY_MAX_RESULTS min_similarity = settings.SEMANTIC_MEMORY_MIN_SIMILARITY diff --git a/example.env b/example.env index ab66400..9b65715 100644 --- a/example.env +++ b/example.env @@ -30,3 +30,7 @@ ATHENA_EMBEDDING_IVFFLAT_LISTS=100 # Database settings ATHENA_DATABASE_URL=postgresql+asyncpg://postgres:password@postgres:5432/postgres + +# Semantic memory settings +ATHENA_SEMANTIC_MEMORY_MAX_RESULTS=5 +ATHENA_SEMANTIC_MEMORY_MIN_SIMILARITY=0.8 From 89adec3ce85d7a58e3ba1118409e1ca15177a11c Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Wed, 8 Oct 2025 17:38:09 +0100 Subject: [PATCH 28/30] add semantic memory settings to docker-compose --- docker-compose.yml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index b7d03c5..ba2860d 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -47,13 +47,15 @@ services: - ATHENA_OPENAI_FORMAT_BASE_URL=${ATHENA_OPENAI_FORMAT_BASE_URL} # Model settings - - ATHENA_MODEL_MAX_INPUT_TOKENS=${ATHENA_MODEL_MAX_INPUT_TOKENS} - ATHENA_MODEL_TEMPERATURE=${ATHENA_MODEL_TEMPERATURE} - - ATHENA_MODEL_MAX_OUTPUT_TOKENS=${ATHENA_MODEL_MAX_OUTPUT_TOKENS} # Database settings - ATHENA_DATABASE_URL=${ATHENA_DATABASE_URL} + # Semantic memory settings + - ATHENA_SEMANTIC_MEMORY_MAX_RESULTS=${ATHENA_SEMANTIC_MEMORY_MAX_RESULTS} + - ATHENA_SEMANTIC_MEMORY_MIN_SIMILARITY=${ATHENA_SEMANTIC_MEMORY_MIN_SIMILARITY} + volumes: - .:/app - /var/run/docker.sock:/var/run/docker.sock From 7d3a1e0eda5684b997c3e6bb5bac67ad92577cb2 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Wed, 8 Oct 2025 17:45:55 +0100 Subject: [PATCH 29/30] update postgres container name and adjust semantic memory similarity threshold --- docker-compose.yml | 2 +- example.env | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index ba2860d..6a440b3 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -5,7 +5,7 @@ networks: services: postgres: image: pgvector/pgvector:pg16 - container_name: postgres_container + container_name: athena_postgres_container networks: - athena_network environment: diff --git a/example.env b/example.env index 9b65715..d9637fa 100644 --- a/example.env +++ b/example.env @@ -33,4 +33,4 @@ ATHENA_DATABASE_URL=postgresql+asyncpg://postgres:password@postgres:5432/postgre # Semantic memory settings ATHENA_SEMANTIC_MEMORY_MAX_RESULTS=5 -ATHENA_SEMANTIC_MEMORY_MIN_SIMILARITY=0.8 +ATHENA_SEMANTIC_MEMORY_MIN_SIMILARITY=0.85 From 6c2ad4a931ccfe680ab20bb22015de0c545d7523 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Wed, 8 Oct 2025 17:48:10 +0100 Subject: [PATCH 30/30] refactor: streamline SemanticMemoryUnit instantiation --- athena/app/api/routes/semantic_memory.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/athena/app/api/routes/semantic_memory.py b/athena/app/api/routes/semantic_memory.py index 2d484e0..a6a1bf3 100644 --- a/athena/app/api/routes/semantic_memory.py +++ b/athena/app/api/routes/semantic_memory.py @@ -35,10 +35,7 @@ async def store_semantic_memory( ] # Create SemanticMemoryUnit from request - memory_unit = SemanticMemoryUnit( - query=body.query, - contexts=body.contexts - ) + memory_unit = SemanticMemoryUnit(query=body.query, contexts=body.contexts) # Store the memory await semantic_memory_service.store_memory(