From 61794aa8326eb5a9f2b335ef96c3de2d65fdcee7 Mon Sep 17 00:00:00 2001 From: Zhaoyang-Chu Date: Tue, 14 Oct 2025 10:23:13 +0000 Subject: [PATCH] fix --- .gitignore | 1 + .../episodic_memory_storage_service.py | 30 ++++++++++--------- athena/app/services/llm_service.py | 1 - docker-compose.yml | 2 ++ example.env | 2 +- 5 files changed, 20 insertions(+), 16 deletions(-) diff --git a/.gitignore b/.gitignore index 68cc726..19c6e68 100644 --- a/.gitignore +++ b/.gitignore @@ -312,3 +312,4 @@ cython_debug/ # postgres data data_postgres/ +secrets/ diff --git a/athena/app/services/episodic_memory_storage_service.py b/athena/app/services/episodic_memory_storage_service.py index 2a83f76..1d71312 100644 --- a/athena/app/services/episodic_memory_storage_service.py +++ b/athena/app/services/episodic_memory_storage_service.py @@ -67,26 +67,26 @@ 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 + col(EpisodicMemoryUnitDB.episodic_memory_id) == unit.memory_id ) ) if existing is None: - row = EpisodicMemoryUnitDB.from_memory_unit(unit) + row = EpisodicMemoryUnitDB.from_episodic_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 = EpisodicMemoryUnitDB.from_memory_unit(unit) + fresh = EpisodicMemoryUnitDB.from_episodic_memory_unit(unit) for attr in ( - "memory_created_at", - "memory_updated_at", - "memory_invalid_at", - "memory_source_name", - "memory_run_id", - "memory_metadata", + "episodic_memory_created_at", + "episodic_memory_updated_at", + "episodic_memory_invalid_at", + "episodic_memory_source_name", + "episodic_memory_run_id", + "episodic_memory_metadata", "task_issue_title", "task_issue_body", "task_issue_comments", @@ -138,14 +138,16 @@ async def search_by_text( .limit(limit) ) rows: List[EpisodicMemoryUnitDB] = list(res.scalars()) - return [r.to_memory_unit() for r in rows] + return [r.to_episodic_memory_unit() for r in rows] async def get_by_memory_id(self, memory_id: str) -> Optional[EpisodicMemoryUnit]: async with self._sessionmaker() as session: row = await session.scalar( - select(EpisodicMemoryUnitDB).where(col(EpisodicMemoryUnitDB.memory_id) == memory_id) + select(EpisodicMemoryUnitDB).where( + col(EpisodicMemoryUnitDB.episodic_memory_id) == memory_id + ) ) - return None if row is None else row.to_memory_unit() + return None if row is None else row.to_episodic_memory_unit() async def list_all(self, limit: Optional[int] = None) -> List[EpisodicMemoryUnit]: async with self._sessionmaker() as session: @@ -154,12 +156,12 @@ async def list_all(self, limit: Optional[int] = None) -> List[EpisodicMemoryUnit stmt = stmt.limit(limit) res = await session.execute(stmt) rows: List[EpisodicMemoryUnitDB] = list(res.scalars()) - return [r.to_memory_unit() for r in rows] + return [r.to_episodic_memory_unit() for r in rows] async def clear_all(self) -> None: async with self._sessionmaker() as session: # Delete all rows from table - await session.execute("DELETE FROM memory_units") + await session.execute("DELETE FROM episodic_memories") await session.commit() @staticmethod diff --git a/athena/app/services/llm_service.py b/athena/app/services/llm_service.py index 3926ef7..f704131 100644 --- a/athena/app/services/llm_service.py +++ b/athena/app/services/llm_service.py @@ -51,7 +51,6 @@ 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, diff --git a/docker-compose.yml b/docker-compose.yml index 7b0ce36..cb911bb 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -43,6 +43,7 @@ services: - ATHENA_GEMINI_API_KEY=${ATHENA_GEMINI_API_KEY} - ATHENA_OPENAI_FORMAT_API_KEY=${ATHENA_OPENAI_FORMAT_API_KEY} - ATHENA_OPENAI_FORMAT_BASE_URL=${ATHENA_OPENAI_FORMAT_BASE_URL} + - ATHENA_GOOGLE_APPLICATION_CREDENTIALS=${ATHENA_GOOGLE_APPLICATION_CREDENTIALS} # Model settings - ATHENA_MODEL_TEMPERATURE=${ATHENA_MODEL_TEMPERATURE} @@ -57,6 +58,7 @@ services: volumes: - .:/app - /var/run/docker.sock:/var/run/docker.sock + - ./secrets/prometheus-code-agent.json:/secrets/prometheus-code-agent.json:ro depends_on: postgres: condition: service_healthy \ No newline at end of file diff --git a/example.env b/example.env index 621689c..dc3a8a4 100644 --- a/example.env +++ b/example.env @@ -6,7 +6,7 @@ ATHENA_ENVIRONMENT=local ATHENA_BACKEND_CORS_ORIGINS=["*"] # LLM model settings -ATHENA_MODEL_NAME=gpt-4o +ATHENA_MODEL_NAME=vertex:gemini-2.5-flash # API keys for various LLM providers ATHENA_ANTHROPIC_API_KEY=anthropic_api_key