apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-05-27 19:25:15 +02:00
Code Issues Projects Releases Packages Wiki Activity

perf(new-chat): add memory_injection middleware timing log

Browse source
This commit is contained in:
CREDO23 2026-05-19 21:30:19 +02:00
parent 1df40fbe31
commit b3b66e4c48
1 changed files with 22 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

22
surfsense_backend/app/agents/new_chat/middleware/memory_injection.py
View file

@ -8,6 +8,7 @@ Injects memory markdown into the system prompt on every turn:
from __future__ import annotations from __future__ import annotations
import logging import logging
import time
from typing import Any from typing import Any
from uuid import UUID from uuid import UUID
@ -19,8 +20,10 @@ from sqlalchemy.ext.asyncio import AsyncSession
from app.agents.new_chat.tools.update_memory import MEMORY_HARD_LIMIT, MEMORY_SOFT_LIMIT from app.agents.new_chat.tools.update_memory import MEMORY_HARD_LIMIT, MEMORY_SOFT_LIMIT
from app.db import ChatVisibility, SearchSpace, User, shielded_async_session from app.db import ChatVisibility, SearchSpace, User, shielded_async_session
from app.utils.perf import get_perf_logger
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
_perf_log = get_perf_logger()
class MemoryInjectionMiddleware(AgentMiddleware): # type: ignore[type-arg] class MemoryInjectionMiddleware(AgentMiddleware): # type: ignore[type-arg]
@ -53,9 +56,13 @@ class MemoryInjectionMiddleware(AgentMiddleware): # type: ignore[type-arg]
if not isinstance(last_message, HumanMessage): if not isinstance(last_message, HumanMessage):
return None return None
start = time.perf_counter()
db_elapsed = 0.0
memory_blocks: list[str] = [] memory_blocks: list[str] = []
scope = "team" if self.visibility == ChatVisibility.SEARCH_SPACE else "user"
async with shielded_async_session() as session: async with shielded_async_session() as session:
db_start = time.perf_counter()
if self.visibility == ChatVisibility.SEARCH_SPACE: if self.visibility == ChatVisibility.SEARCH_SPACE:
team_memory = await self._load_team_memory(session) team_memory = await self._load_team_memory(session)
if team_memory: if team_memory:
@ -96,7 +103,15 @@ class MemoryInjectionMiddleware(AgentMiddleware): # type: ignore[type-arg]
f"</memory_warning>" f"</memory_warning>"
) )
db_elapsed = time.perf_counter() - db_start
if not memory_blocks: if not memory_blocks:
_perf_log.info(
"[memory_injection] scope=%s injected=0 db=%.3fs total=%.3fs",
scope,
db_elapsed,
time.perf_counter() - start,
)
return None return None
memory_text = "\n\n".join(memory_blocks) memory_text = "\n\n".join(memory_blocks)
@ -106,6 +121,13 @@ class MemoryInjectionMiddleware(AgentMiddleware): # type: ignore[type-arg]
insert_idx = 1 if len(new_messages) > 1 else 0 insert_idx = 1 if len(new_messages) > 1 else 0
new_messages.insert(insert_idx, memory_msg) new_messages.insert(insert_idx, memory_msg)
_perf_log.info(
"[memory_injection] scope=%s injected=1 chars=%d db=%.3fs total=%.3fs",
scope,
len(memory_text),
db_elapsed,
time.perf_counter() - start,
)
return {"messages": new_messages} return {"messages": new_messages}
async def _load_user_memory( async def _load_user_memory(