From 70544e0f4999784e51971f8d529291a47ba957dd Mon Sep 17 00:00:00 2001 From: makeitdouble Date: Thu, 28 May 2026 18:35:56 +0300 Subject: [PATCH] feature: message memory - add summerized message memory to the runtime context --- clients/brain_client.py | 30 +++- memory/local_json_memory.py | 98 ------------- memory/memory_base.py | 37 ----- memory/message_memory.py | 260 +++++++++++++++++++++++++++++++++++ runtime/runtime_context.py | 10 ++ tests/test_message_memory.py | 224 ++++++++++++++++++++++++++++++ websocket.py | 10 ++ 7 files changed, 530 insertions(+), 139 deletions(-) delete mode 100644 memory/local_json_memory.py delete mode 100644 memory/memory_base.py create mode 100644 memory/message_memory.py create mode 100644 tests/test_message_memory.py diff --git a/clients/brain_client.py b/clients/brain_client.py index 65dfb8a7..931ca273 100644 --- a/clients/brain_client.py +++ b/clients/brain_client.py @@ -482,10 +482,16 @@ def build_brain_runtime_context( .to_runtime_xml() ) + runtime_memory = "" search_result = "" search_result_id = "" if context is not None: + runtime_memory = getattr( + context, + "runtime_memory", + "", + ) search_result = getattr( context, "runtime_search_result", @@ -497,8 +503,21 @@ def build_brain_runtime_context( "", ) + runtime_context_parts = [ + runtime_xml + ] + + if runtime_memory.strip(): + runtime_context_parts.append( + "\n" + f"{indent_xml(escape(runtime_memory))}\n" + "" + ) + if not search_result: - return runtime_xml + return "\n".join( + runtime_context_parts + ) search_result = strip_empty_results_xml( search_result @@ -523,9 +542,12 @@ def build_brain_runtime_context( ) return ( - runtime_xml - + "\n" - + tool_results_xml + "\n".join( + runtime_context_parts + + [ + tool_results_xml + ] + ) ) diff --git a/memory/local_json_memory.py b/memory/local_json_memory.py deleted file mode 100644 index e60faac9..00000000 --- a/memory/local_json_memory.py +++ /dev/null @@ -1,98 +0,0 @@ -import os -import json -import uuid -import asyncio -from typing import List, Dict, Any, Optional -from memory_base import MemoryBase - -class LocalJsonMemory(MemoryBase): - def __init__(self, file_path: str = "jin_memory.json"): - self.file_path = file_path - self.lock = asyncio.Lock() # Защита от одновременной записи из разных асинхронных задач - self._ensure_file_exists() - - def _ensure_file_exists(self): - if not os.path.exists(self.file_path): - with open(self.file_path, "w", encoding="utf-8") as f: - json.dump({}, f, ensure_ascii=False, indent=4) - - async def _read_db(self) -> Dict[str, Dict[str, Any]]: - async with self.lock: - with open(self.file_path, "r", encoding="utf-8") as f: - return json.load(f) - - async def _write_db(self, data: Dict[str, Dict[str, Any]]): - async with self.lock: - with open(self.file_path, "w", encoding="utf-8") as f: - json.dump(data, f, ensure_ascii=False, indent=4) - - async def add_fact( - self, - text: str, - importance: float, - decay_rate: float, - evidence_id: Optional[str] = None - ) -> str: - db = await self._read_db() - fact_id = str(uuid.uuid4()) - - db[fact_id] = { - "id": fact_id, - "text": text, - "weight": max(0.0, min(1.0, importance)), # Ограничение 0.0 - 1.0 - "decay_rate": max(0.0, decay_rate), - "evidence_id": evidence_id, - "hits": 1 # Сколько раз этот факт подтверждался/вспоминался - } - - await self._write_db(db) - return fact_id - - async def get_relevant_facts(self, query: str, limit: int = 10) -> List[Dict[str, Any]]: - db = await self._read_db() - query_words = set(query.lower().split()) - scored_facts = [] - - for fact in db.values(): - # Простейший поиск по пересечению слов (заглушка до графов/эмбеддингов) - fact_words = set(fact["text"].lower().split()) - matches = len(query_words.intersection(fact_words)) - - # Релевантность = совпадение слов * текущий вес факта в памяти - relevance = matches * fact["weight"] - - if relevance > 0 or fact["decay_rate"] == 0.0: # Пропускаем константы (имя и т.д.) всегда - scored_facts.append((relevance, fact)) - - # Сортируем по релевантности и возвращаем чистый список фактов - scored_facts.sort(key=lambda x: x[0], reverse=True) - return [item[1] for item in scored_facts[:limit]] - - async def apply_night_decay(self) -> Dict[str, int]: - """ - Линейно снижает вес каждого факта на его decay_rate. - Если вес падает ниже порогового (0.15) и это не константа — факт удаляется. - """ - db = await self._read_db() - updated_count = 0 - deleted_count = 0 - new_db = {} - - for fact_id, fact in db.items(): - # Если decay_rate == 0, факт бессмертен (например, имя пользователя) - if fact["decay_rate"] == 0.0: - new_db[fact_id] = fact - continue - - # Формула затухания: уменьшаем вес - fact["weight"] = round(fact["weight"] - fact["decay_rate"], 3) - - # Порог утилизации факта (например, 0.15) - if fact["weight"] <= 0.15: - deleted_count += 1 - else: - new_db[fact_id] = fact - updated_count += 1 - - await self._write_db(new_db) - return {"updated": updated_count, "deleted": deleted_count} diff --git a/memory/memory_base.py b/memory/memory_base.py deleted file mode 100644 index 54d22293..00000000 --- a/memory/memory_base.py +++ /dev/null @@ -1,37 +0,0 @@ -from abc import ABC, abstractmethod -from typing import List, Dict, Any, Optional - -class MemoryBase(ABC): - @abstractmethod - async def add_fact( - self, - text: str, - importance: float, - decay_rate: float, - evidence_id: Optional[str] = None - ) -> str: - """ - Добавить новый факт в память. - :param text: Содержание факта (например, "Пользователь пьет кофе в 8 утра") - :param importance: Стартовая важность факта (от 0.0 до 1.0) - :param decay_rate: Скорость затухания (0.0 — никогда не забывать, 0.5 — быстро) - :param evidence_id: ID связанного скриншота/доказательства (опционально) - :return: Уникальный ID созданного факта - """ - pass - - @abstractmethod - async def get_relevant_facts(self, query: str, limit: int = 10) -> List[Dict[str, Any]]: - """ - Извлечь факты, релевантные текущему контексту (для сборки XML-пакета). - """ - pass - - @abstractmethod - async def apply_night_decay(self) -> Dict[str, int]: - """ - Запустить ночной полураспад. Уменьшает веса фактов на их decay_rate. - Удаляет затухшие факты. - :return: Статистика (сколько фактов обновлено, сколько удалено) - """ - pass diff --git a/memory/message_memory.py b/memory/message_memory.py new file mode 100644 index 00000000..46ae60ff --- /dev/null +++ b/memory/message_memory.py @@ -0,0 +1,260 @@ +import asyncio +import contextlib + +from clients.service_client import ( + ask_service_model, +) +from settings.config_loader import ( + config, +) +from utils.response_extractor import ( + ResponseExtractor, +) + + +async def safe_call( + call, + *args, + **kwargs, +): + + if call is None: + return + + with contextlib.suppress(Exception): + await call( + *args, + **kwargs, + ) + + +def build_runtime_memory_system_prompt() -> str: + + return ( + "You are JIN's runtime memory summarizer.\n" + "Return only the new compressed memory state as plain text.\n" + "Do not output JSON.\n" + "Do not use Markdown headings.\n" + "Do not explain your reasoning or the summarization process.\n" + "Preserve still-relevant existing memory. Update it instead of replacing it blindly.\n" + "Drop old details only when they are clearly obsolete, duplicated, or no longer useful.\n" + "Decide the summary depth from the signal in the latest turn.\n" + "Use shallow summarization for simple factual, isolated, or low-signal turns: " + "keep only the dry fact, topic, or unresolved reference that could help " + "the next answer.\n" + "Use deep summarization for turns that reveal user intent, project direction, " + "preferences, recurring patterns, decisions, emotional tone, or a meaningful " + "shift in the conversation trajectory.\n" + "If the user asks a follow-up that depends on prior context, preserve the " + "referent clearly enough for the next brain prompt to resolve it.\n" + "If the user switches topic, keep the new topic without forcing it into the " + "previous one; only note a pattern if the sequence itself is meaningful.\n" + "Do not infer stable user traits from a single turn.\n" + "Do not over-interpret jokes, tests, or casual topic changes.\n" + "Prefer compact continuity over transcript-like detail.\n" + "Remove noise, implementation chatter, and one-off details unless they change " + "what JIN should understand next.\n" + "Keep the memory under 1200 characters unless the latest turn is highly important.\n" + "The final memory should feel like live state, not chat history.\n" + ) + + +def build_runtime_memory_user_prompt( + *, + current_memory: str, + user_message: str, + assistant_message: str, +) -> str: + + return ( + "Current runtime memory:\n" + f"{current_memory.strip() or 'User and JIN just started interacting.'}\n\n" + "Latest user message:\n" + f"{user_message.strip()}\n\n" + "Latest JIN answer:\n" + f"{assistant_message.strip()}\n\n" + "Rewrite the runtime memory now." + ) + + +def extract_runtime_memory_text( + response: dict, +) -> str: + + text = ( + ResponseExtractor.extract_content_text( + response + ) + or ResponseExtractor.extract_reasoning_text( + response + ) + ) + + return text.strip() + + +async def summarize_runtime_memory( + *, + context, + user_message: str, + assistant_message: str, +) -> str: + + if not assistant_message.strip(): + return getattr( + context, + "runtime_memory", + "", + ) + + service_client = ( + getattr( + context, + "clients", + {}, + ) + .get( + "service" + ) + ) + + if service_client is None: + return getattr( + context, + "runtime_memory", + "", + ) + + current_memory = getattr( + context, + "runtime_memory", + "", + ) + + try: + response = await ask_service_model( + client=service_client, + system_prompt=( + build_runtime_memory_system_prompt() + ), + user_prompt=( + build_runtime_memory_user_prompt( + current_memory=current_memory, + user_message=user_message, + assistant_message=assistant_message, + ) + ), + temperature=( + config.SERVICE_TEMPERATURE + ), + max_tokens=( + min( + config.SERVICE_MAX_TOKENS, + 512, + ) + ), + ) + + updated_memory = extract_runtime_memory_text( + response + ) + + if updated_memory: + context.runtime_memory = updated_memory + context.runtime_memory_updates = ( + getattr( + context, + "runtime_memory_updates", + 0, + ) + + 1 + ) + + logger = getattr( + context, + "logger", + None, + ) + log_service = getattr( + logger, + "log_service", + None, + ) + + await safe_call( + log_service, + "[MEMORY] runtime memory updated", + ) + + return getattr( + context, + "runtime_memory", + "", + ) + + except asyncio.CancelledError: + raise + + except Exception as error: + logger = getattr( + context, + "logger", + None, + ) + log_error = getattr( + logger, + "log_error", + None, + ) + + await safe_call( + log_error, + "[MEMORY] runtime memory update failed", + details=str( + error + ), + ) + + return getattr( + context, + "runtime_memory", + "", + ) + + +def schedule_runtime_memory_update( + *, + context, + user_message: str, + assistant_message: str, +) -> asyncio.Task | None: + + if not assistant_message.strip(): + return None + + task = asyncio.create_task( + summarize_runtime_memory( + context=context, + user_message=user_message, + assistant_message=assistant_message, + ) + ) + + background_tasks = getattr( + context, + "background_tasks", + None, + ) + + if background_tasks is None: + background_tasks = set() + context.background_tasks = background_tasks + + background_tasks.add( + task + ) + task.add_done_callback( + background_tasks.discard + ) + + return task diff --git a/runtime/runtime_context.py b/runtime/runtime_context.py index bf9c2c42..9241fc00 100644 --- a/runtime/runtime_context.py +++ b/runtime/runtime_context.py @@ -43,3 +43,13 @@ class RuntimeContext: runtime_usage_events: list[dict] = field( default_factory=list ) + + runtime_memory: str = ( + "User and JIN just started interacting." + ) + + runtime_memory_updates: int = 0 + + background_tasks: set = field( + default_factory=set + ) diff --git a/tests/test_message_memory.py b/tests/test_message_memory.py new file mode 100644 index 00000000..27699e25 --- /dev/null +++ b/tests/test_message_memory.py @@ -0,0 +1,224 @@ +import unittest +from types import SimpleNamespace + +from clients.brain_client import ( + build_brain_system_prompt, +) +from memory.message_memory import ( + build_runtime_memory_system_prompt, + schedule_runtime_memory_update, + summarize_runtime_memory, +) + + +class FakeServiceClient: + + def __init__( + self, + response_text: str, + ): + + self.response_text = response_text + self.calls = [] + + async def ask( + self, + *, + system_prompt: str, + user_prompt: str, + temperature: float, + max_tokens: int, + ): + + self.calls.append({ + "system_prompt": system_prompt, + "user_prompt": user_prompt, + "temperature": temperature, + "max_tokens": max_tokens, + }) + + return { + "choices": [ + { + "message": { + "content": self.response_text, + }, + }, + ], + } + + +class FakeLogger: + + def __init__(self): + self.service_logs = [] + self.errors = [] + + async def log_service( + self, + message: str, + ): + + self.service_logs.append( + message + ) + + async def log_error( + self, + message: str, + details: str | None = None, + ): + + self.errors.append( + ( + message, + details, + ) + ) + + +class MessageMemoryTests( + unittest.IsolatedAsyncioTestCase +): + + def test_runtime_memory_prompt_focuses_on_summary_depth(self): + + prompt = build_runtime_memory_system_prompt() + + self.assertIn( + "Decide the summary depth", + prompt, + ) + self.assertIn( + "Use shallow summarization", + prompt, + ) + self.assertIn( + "Use deep summarization", + prompt, + ) + self.assertNotIn( + "after one completed user/JIN turn", + prompt, + ) + + def test_brain_prompt_includes_runtime_memory(self): + + context = SimpleNamespace( + runtime_memory=( + "The user recently asked about Lamborghini pricing." + ), + deep_thought_count=0, + runtime_search_result="", + runtime_search_result_id="", + ) + + prompt = build_brain_system_prompt( + context=context, + runtime_actions={ + "CAN_SEARCH": False, + "CAN_DEEP_THOUGHT": False, + }, + ) + + self.assertIn( + "", + prompt, + ) + self.assertIn( + "Lamborghini pricing", + prompt, + ) + + async def test_summarizer_updates_runtime_memory(self): + + service_client = FakeServiceClient( + "The user is testing live runtime memory." + ) + logger = FakeLogger() + context = SimpleNamespace( + clients={ + "service": service_client, + }, + logger=logger, + runtime_memory="User and JIN just started interacting.", + runtime_memory_updates=0, + ) + + updated_memory = await summarize_runtime_memory( + context=context, + user_message="Do you remember this?", + assistant_message="Yes, I can keep the live context updated.", + ) + + self.assertEqual( + updated_memory, + "The user is testing live runtime memory.", + ) + self.assertEqual( + context.runtime_memory, + "The user is testing live runtime memory.", + ) + self.assertEqual( + context.runtime_memory_updates, + 1, + ) + self.assertIn( + "Do you remember this?", + service_client.calls[0]["user_prompt"], + ) + self.assertEqual( + service_client.calls[0]["system_prompt"], + build_runtime_memory_system_prompt(), + ) + self.assertIn( + "[MEMORY] runtime memory updated", + logger.service_logs, + ) + + async def test_scheduled_update_is_background_task(self): + + service_client = FakeServiceClient( + "Updated background memory." + ) + context = SimpleNamespace( + clients={ + "service": service_client, + }, + logger=FakeLogger(), + runtime_memory="Initial memory.", + runtime_memory_updates=0, + ) + + task = schedule_runtime_memory_update( + context=context, + user_message="First message", + assistant_message="First answer", + ) + + self.assertIsNotNone( + task + ) + self.assertTrue( + hasattr( + context, + "background_tasks", + ) + ) + + await task + + self.assertEqual( + context.runtime_memory, + "Updated background memory.", + ) + self.assertEqual( + len( + context.background_tasks + ), + 0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/websocket.py b/websocket.py index cb43932e..ecf614d7 100644 --- a/websocket.py +++ b/websocket.py @@ -28,6 +28,10 @@ AgentState, ) +from memory.message_memory import ( + schedule_runtime_memory_update, +) + from utils.brain import ( get_brain_runtime_config, ) @@ -245,6 +249,12 @@ async def process_message( "type": "agent_runtime_end", }) + schedule_runtime_memory_update( + context=context, + user_message=user_text, + assistant_message=state.brain_response, + ) + except asyncio.CancelledError: await logger.log_runtime(