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

subagents/knowledge_base: wire KB specialist into orchestrator (renderer/projector split, FS middleware stack, cloud-mode gating)

Browse source
This commit is contained in:
CREDO23 2026-05-11 20:43:44 +02:00
parent 09fc99c435
commit df2afed18d
10 changed files with 230 additions and 41 deletions
Download patch file Download diff file Expand all files Collapse all files

27
surfsense_backend/app/agents/new_chat/middleware/knowledge_search.py
View file

@ -584,6 +584,7 @@ class KnowledgePriorityMiddleware(AgentMiddleware): # type: ignore[type-arg]
available_document_types: list[str] | None = None,
top_k: int = 10,
mentioned_document_ids: list[int] | None = None,
inject_system_message: bool = True, # For backwards compatibility
) -> None:
self.llm = llm
self.search_space_id = search_space_id
@ -592,6 +593,7 @@ class KnowledgePriorityMiddleware(AgentMiddleware): # type: ignore[type-arg]
self.available_document_types = available_document_types
self.top_k = top_k
self.mentioned_document_ids = mentioned_document_ids or []
self.inject_system_message = inject_system_message
# Build the kb-planner private Runnable ONCE here so we don't pay
# the ``create_agent`` compile cost (50-200ms) on every turn.
# Disabled by default behind ``enable_kb_planner_runnable``; when
@ -772,14 +774,16 @@ class KnowledgePriorityMiddleware(AgentMiddleware): # type: ignore[type-arg]
"mentioned": True,
}
]
new_messages = list(state.get("messages") or [])
insert_at = max(len(new_messages) - 1, 0)
new_messages.insert(insert_at, _render_priority_message(priority))
return {
update: dict[str, Any] = {
"kb_priority": priority,
"kb_matched_chunk_ids": {},
"messages": new_messages,
}
if self.inject_system_message:
new_messages = list(state.get("messages") or [])
insert_at = max(len(new_messages) - 1, 0)
new_messages.insert(insert_at, _render_priority_message(priority))
update["messages"] = new_messages
return update
async def _authenticated_priority(
self,
@ -876,10 +880,6 @@ class KnowledgePriorityMiddleware(AgentMiddleware): # type: ignore[type-arg]
priority, matched_chunk_ids = await self._materialize_priority(merged)
new_messages = list(messages)
insert_at = max(len(new_messages) - 1, 0)
new_messages.insert(insert_at, _render_priority_message(priority))
_perf_log.info(
"[kb_priority] completed in %.3fs query=%r priority=%d mentioned=%d",
asyncio.get_event_loop().time() - t0,
@ -888,11 +888,16 @@ class KnowledgePriorityMiddleware(AgentMiddleware): # type: ignore[type-arg]
len(mentioned_results),
)
return {
update: dict[str, Any] = {
"kb_priority": priority,
"kb_matched_chunk_ids": matched_chunk_ids,
"messages": new_messages,
}
if self.inject_system_message:
new_messages = list(messages)
insert_at = max(len(new_messages) - 1, 0)
new_messages.insert(insert_at, _render_priority_message(priority))
update["messages"] = new_messages
return update
async def _materialize_priority(
self, merged: list[dict[str, Any]]

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

@ -105,12 +105,14 @@ class KnowledgeTreeMiddleware(AgentMiddleware): # type: ignore[type-arg]
llm: BaseChatModel | None = None,
max_entries: int = MAX_TREE_ENTRIES,
max_tokens: int = MAX_TREE_TOKENS,
inject_system_message: bool = True, # For backwards compatibility
) -> None:
self.search_space_id = search_space_id
self.filesystem_mode = filesystem_mode
self.llm = llm
self.max_entries = max_entries
self.max_tokens = max_tokens
self.inject_system_message = inject_system_message
self._cache: dict[tuple[int, int, bool], str] = {}
async def abefore_agent( # type: ignore[override]
@ -132,10 +134,13 @@ class KnowledgeTreeMiddleware(AgentMiddleware): # type: ignore[type-arg]
else:
tree_msg = await self._render_kb_tree(state)
messages = list(state.get("messages") or [])
insert_at = max(len(messages) - 1, 0)
messages.insert(insert_at, SystemMessage(content=tree_msg))
update["messages"] = messages
update["workspace_tree_text"] = tree_msg
if self.inject_system_message:
messages = list(state.get("messages") or [])
insert_at = max(len(messages) - 1, 0)
messages.insert(insert_at, SystemMessage(content=tree_msg))
update["messages"] = messages
return update
def before_agent( # type: ignore[override]