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 knowledge_tree middleware timing log

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

17
surfsense_backend/app/agents/new_chat/middleware/knowledge_tree.py
View file

@ -24,6 +24,7 @@ from __future__ import annotations
import asyncio import asyncio
import logging import logging
import time
from typing import Any from typing import Any
from langchain.agents.middleware import AgentMiddleware, AgentState from langchain.agents.middleware import AgentMiddleware, AgentState
@ -41,6 +42,9 @@ from app.agents.new_chat.path_resolver import (
doc_to_virtual_path, doc_to_virtual_path,
) )
from app.db import Document, shielded_async_session from app.db import Document, shielded_async_session
from app.utils.perf import get_perf_logger
_perf_log = get_perf_logger()
try: try:
from litellm import token_counter from litellm import token_counter
@ -124,6 +128,7 @@ class KnowledgeTreeMiddleware(AgentMiddleware): # type: ignore[type-arg]
if self.filesystem_mode != FilesystemMode.CLOUD: if self.filesystem_mode != FilesystemMode.CLOUD:
return None return None
start = time.perf_counter()
update: dict[str, Any] = {} update: dict[str, Any] = {}
if not state.get("cwd"): if not state.get("cwd"):
update["cwd"] = DOCUMENTS_ROOT update["cwd"] = DOCUMENTS_ROOT
@ -131,7 +136,11 @@ class KnowledgeTreeMiddleware(AgentMiddleware): # type: ignore[type-arg]
anon_doc = state.get("kb_anon_doc") anon_doc = state.get("kb_anon_doc")
if anon_doc: if anon_doc:
tree_msg = self._render_anon_tree(anon_doc) tree_msg = self._render_anon_tree(anon_doc)
cache_outcome = "anon"
else: else:
version = int(state.get("tree_version") or 0)
cache_key = (self.search_space_id, version, False)
cache_outcome = "hit" if cache_key in self._cache else "miss"
tree_msg = await self._render_kb_tree(state) tree_msg = await self._render_kb_tree(state)
update["workspace_tree_text"] = tree_msg update["workspace_tree_text"] = tree_msg
@ -141,6 +150,14 @@ class KnowledgeTreeMiddleware(AgentMiddleware): # type: ignore[type-arg]
insert_at = max(len(messages) - 1, 0) insert_at = max(len(messages) - 1, 0)
messages.insert(insert_at, SystemMessage(content=tree_msg)) messages.insert(insert_at, SystemMessage(content=tree_msg))
update["messages"] = messages update["messages"] = messages
_perf_log.info(
"[knowledge_tree] cache=%s chars=%d elapsed=%.3fs space=%d",
cache_outcome,
len(tree_msg),
time.perf_counter() - start,
self.search_space_id,
)
return update return update
def before_agent( # type: ignore[override] def before_agent( # type: ignore[override]