apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-05-02 20:32:39 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

refactor: use ChainOfThoughtItem for thinking step display and update dependency lock file.

Browse source
This commit is contained in:
Anish Sarkar 2025-12-27 16:40:53 +05:30
parent 70383931bb
commit ad47d94b40
2 changed files with 159 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

163
surfsense_backend/app/tasks/chat/stream_new_chat.py
View file

@ -246,7 +246,8 @@ async def stream_new_chat(
config = {
"configurable": {
"thread_id": str(chat_id),
}
},
"recursion_limit": 80, # Increase from default 25 to allow more tool iterations
}
# Start the message stream
@ -268,6 +269,8 @@ async def stream_new_chat(
completed_step_ids: set[str] = set()
# Track if we just finished a tool (text flows silently after tools)
just_finished_tool: bool = False
# Track write_todos calls to show "Creating plan" vs "Updating plan"
write_todos_call_count: int = 0
def next_thinking_step_id() -> str:
nonlocal thinking_step_counter
@ -476,6 +479,60 @@ async def stream_new_chat(
status="in_progress",
items=last_active_step_items,
)
elif tool_name == "write_todos":
# Track write_todos calls for better messaging
write_todos_call_count += 1
todos = (
tool_input.get("todos", [])
if isinstance(tool_input, dict)
else []
)
todo_count = len(todos) if isinstance(todos, list) else 0
if write_todos_call_count == 1:
# First call - creating the plan
last_active_step_title = "Creating plan"
last_active_step_items = [f"Defining {todo_count} tasks..."]
else:
# Subsequent calls - updating the plan
# Try to provide context about what's being updated
in_progress_count = (
sum(
1
for t in todos
if isinstance(t, dict)
and t.get("status") == "in_progress"
)
if isinstance(todos, list)
else 0
)
completed_count = (
sum(
1
for t in todos
if isinstance(t, dict)
and t.get("status") == "completed"
)
if isinstance(todos, list)
else 0
)
last_active_step_title = "Updating progress"
last_active_step_items = (
[
f"Progress: {completed_count}/{todo_count} completed",
f"In progress: {in_progress_count} tasks",
]
if completed_count > 0
else [f"Working on {todo_count} tasks"]
)
yield streaming_service.format_thinking_step(
step_id=tool_step_id,
title=last_active_step_title,
status="in_progress",
items=last_active_step_items,
)
elif tool_name == "generate_podcast":
podcast_title = (
tool_input.get("podcast_title", "SurfSense Podcast")
@ -500,6 +557,15 @@ async def stream_new_chat(
status="in_progress",
items=last_active_step_items,
)
elif tool_name == "ls":
last_active_step_title = "Exploring files"
last_active_step_items = []
yield streaming_service.format_thinking_step(
step_id=tool_step_id,
title="Exploring files",
status="in_progress",
items=None,
)
else:
last_active_step_title = f"Using {tool_name.replace('_', ' ')}"
last_active_step_items = []
@ -745,19 +811,100 @@ async def stream_new_chat(
items=completed_items,
)
elif tool_name == "write_todos":
# Build completion items for planning
# Build completion items for planning/updating
if isinstance(tool_output, dict):
todos = tool_output.get("todos", [])
todo_count = len(todos) if isinstance(todos, list) else 0
completed_items = [
*last_active_step_items,
f"Tasks: {todo_count} steps defined",
]
completed_count = (
sum(
1
for t in todos
if isinstance(t, dict)
and t.get("status") == "completed"
)
if isinstance(todos, list)
else 0
)
in_progress_count = (
sum(
1
for t in todos
if isinstance(t, dict)
and t.get("status") == "in_progress"
)
if isinstance(todos, list)
else 0
)
# Use context-aware completion message
if last_active_step_title == "Creating plan":
completed_items = [f"Created {todo_count} tasks"]
else:
# Updating progress - show stats
completed_items = [
f"Progress: {completed_count}/{todo_count} completed",
]
if in_progress_count > 0:
# Find the currently in-progress task name
in_progress_task = next(
(
t.get("content", "")[:40]
for t in todos
if isinstance(t, dict)
and t.get("status") == "in_progress"
),
None,
)
if in_progress_task:
completed_items.append(
f"Current: {in_progress_task}..."
)
else:
completed_items = [*last_active_step_items, "Plan created"]
completed_items = ["Plan updated"]
yield streaming_service.format_thinking_step(
step_id=original_step_id,
title="Creating plan",
title=last_active_step_title,
status="completed",
items=completed_items,
)
elif tool_name == "ls":
# Build completion items showing file names found
if isinstance(tool_output, dict):
result = tool_output.get("result", "")
elif isinstance(tool_output, str):
result = tool_output
else:
result = str(tool_output) if tool_output else ""
# Parse file paths and extract just the file names
file_names = []
if result:
# The ls tool returns paths, extract just the file/folder names
for line in result.strip().split("\n"):
line = line.strip()
if line:
# Get just the filename from the path
name = line.rstrip("/").split("/")[-1]
if name and len(name) <= 40:
file_names.append(name)
elif name:
file_names.append(name[:37] + "...")
# Build display items - wrap file names in brackets for icon rendering
if file_names:
if len(file_names) <= 5:
# Wrap each file name in brackets for styled tile rendering
completed_items = [f"[{name}]" for name in file_names]
else:
# Show first few with brackets and count
completed_items = [f"[{name}]" for name in file_names[:4]]
completed_items.append(f"(+{len(file_names) - 4} more)")
else:
completed_items = ["No files found"]
yield streaming_service.format_thinking_step(
step_id=original_step_id,
title="Exploring files",
status="completed",
items=completed_items,
)