mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-04-26 01:06:23 +02:00
refactor: improve write_todos tool and UI components
- Refactored the write_todos tool to enhance argument and result schemas using Zod for better validation and type safety. - Updated the WriteTodosToolUI to streamline the rendering logic and improve loading states, ensuring a smoother user experience. - Enhanced the Plan and TodoItem components to better handle streaming states and display progress, providing clearer feedback during task management. - Cleaned up code formatting and structure for improved readability and maintainability.
This commit is contained in:
Anish Sarkar
parent
2c86287264
commit
ebc04f590e
18 changed files with 833 additions and 751 deletions
|
|
@ -10,20 +10,20 @@
|
|||
import { atom } from "jotai";
|
||||
|
||||
export interface PlanTodo {
|
||||
id: string;
|
||||
label: string;
|
||||
status: "pending" | "in_progress" | "completed" | "cancelled";
|
||||
description?: string;
|
||||
id: string;
|
||||
label: string;
|
||||
status: "pending" | "in_progress" | "completed" | "cancelled";
|
||||
description?: string;
|
||||
}
|
||||
|
||||
export interface PlanState {
|
||||
id: string;
|
||||
title: string;
|
||||
description?: string;
|
||||
todos: PlanTodo[];
|
||||
lastUpdated: number;
|
||||
/** The toolCallId of the first component that rendered this plan */
|
||||
ownerToolCallId: string;
|
||||
id: string;
|
||||
title: string;
|
||||
description?: string;
|
||||
todos: PlanTodo[];
|
||||
lastUpdated: number;
|
||||
/** The toolCallId of the first component that rendered this plan */
|
||||
ownerToolCallId: string;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -38,14 +38,14 @@ let firstPlanOwner: { toolCallId: string; title: string } | null = null;
|
|||
* All subsequent calls update the state but don't render their own card.
|
||||
*/
|
||||
export function registerPlanOwner(title: string, toolCallId: string): boolean {
|
||||
if (!firstPlanOwner) {
|
||||
// First plan in this conversation - claim ownership
|
||||
firstPlanOwner = { toolCallId, title };
|
||||
return true;
|
||||
}
|
||||
|
||||
// Check if we're the owner
|
||||
return firstPlanOwner.toolCallId === toolCallId;
|
||||
if (!firstPlanOwner) {
|
||||
// First plan in this conversation - claim ownership
|
||||
firstPlanOwner = { toolCallId, title };
|
||||
return true;
|
||||
}
|
||||
|
||||
// Check if we're the owner
|
||||
return firstPlanOwner.toolCallId === toolCallId;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -53,35 +53,35 @@ export function registerPlanOwner(title: string, toolCallId: string): boolean {
|
|||
* Returns the first plan's title if one exists, otherwise the provided title
|
||||
*/
|
||||
export function getCanonicalPlanTitle(title: string): string {
|
||||
return firstPlanOwner?.title || title;
|
||||
return firstPlanOwner?.title || title;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a plan already exists in this conversation
|
||||
*/
|
||||
export function hasPlan(): boolean {
|
||||
return firstPlanOwner !== null;
|
||||
return firstPlanOwner !== null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the first plan's info
|
||||
*/
|
||||
export function getFirstPlanInfo(): { toolCallId: string; title: string } | null {
|
||||
return firstPlanOwner;
|
||||
return firstPlanOwner;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a toolCallId is the owner of the plan SYNCHRONOUSLY
|
||||
*/
|
||||
export function isPlanOwner(toolCallId: string): boolean {
|
||||
return !firstPlanOwner || firstPlanOwner.toolCallId === toolCallId;
|
||||
return !firstPlanOwner || firstPlanOwner.toolCallId === toolCallId;
|
||||
}
|
||||
|
||||
/**
|
||||
* Clear ownership registry (call when starting a new chat)
|
||||
*/
|
||||
export function clearPlanOwnerRegistry(): void {
|
||||
firstPlanOwner = null;
|
||||
firstPlanOwner = null;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -94,56 +94,53 @@ export const planStatesAtom = atom<Map<string, PlanState>>(new Map());
|
|||
* Input type for updating plan state
|
||||
*/
|
||||
export interface UpdatePlanInput {
|
||||
id: string;
|
||||
title: string;
|
||||
description?: string;
|
||||
todos: PlanTodo[];
|
||||
toolCallId: string;
|
||||
id: string;
|
||||
title: string;
|
||||
description?: string;
|
||||
todos: PlanTodo[];
|
||||
toolCallId: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Helper atom to update a plan state
|
||||
*/
|
||||
export const updatePlanStateAtom = atom(
|
||||
null,
|
||||
(get, set, plan: UpdatePlanInput) => {
|
||||
const states = new Map(get(planStatesAtom));
|
||||
|
||||
// Register ownership synchronously if not already done
|
||||
registerPlanOwner(plan.title, plan.toolCallId);
|
||||
|
||||
// Get the actual owner from the first plan
|
||||
const ownerToolCallId = firstPlanOwner?.toolCallId || plan.toolCallId;
|
||||
|
||||
// Always use the canonical (first) title for the plan key
|
||||
const canonicalTitle = getCanonicalPlanTitle(plan.title);
|
||||
|
||||
states.set(canonicalTitle, {
|
||||
id: plan.id,
|
||||
title: canonicalTitle,
|
||||
description: plan.description,
|
||||
todos: plan.todos,
|
||||
lastUpdated: Date.now(),
|
||||
ownerToolCallId,
|
||||
});
|
||||
set(planStatesAtom, states);
|
||||
}
|
||||
);
|
||||
export const updatePlanStateAtom = atom(null, (get, set, plan: UpdatePlanInput) => {
|
||||
const states = new Map(get(planStatesAtom));
|
||||
|
||||
// Register ownership synchronously if not already done
|
||||
registerPlanOwner(plan.title, plan.toolCallId);
|
||||
|
||||
// Get the actual owner from the first plan
|
||||
const ownerToolCallId = firstPlanOwner?.toolCallId || plan.toolCallId;
|
||||
|
||||
// Always use the canonical (first) title for the plan key
|
||||
const canonicalTitle = getCanonicalPlanTitle(plan.title);
|
||||
|
||||
states.set(canonicalTitle, {
|
||||
id: plan.id,
|
||||
title: canonicalTitle,
|
||||
description: plan.description,
|
||||
todos: plan.todos,
|
||||
lastUpdated: Date.now(),
|
||||
ownerToolCallId,
|
||||
});
|
||||
set(planStatesAtom, states);
|
||||
});
|
||||
|
||||
/**
|
||||
* Helper atom to get the latest plan state by title
|
||||
*/
|
||||
export const getPlanStateAtom = atom((get) => {
|
||||
const states = get(planStatesAtom);
|
||||
return (title: string) => states.get(title);
|
||||
const states = get(planStatesAtom);
|
||||
return (title: string) => states.get(title);
|
||||
});
|
||||
|
||||
/**
|
||||
* Helper atom to clear all plan states (useful when starting a new chat)
|
||||
*/
|
||||
export const clearPlanStatesAtom = atom(null, (get, set) => {
|
||||
clearPlanOwnerRegistry();
|
||||
set(planStatesAtom, new Map());
|
||||
clearPlanOwnerRegistry();
|
||||
set(planStatesAtom, new Map());
|
||||
});
|
||||
|
||||
/**
|
||||
|
|
@ -151,84 +148,80 @@ export const clearPlanStatesAtom = atom(null, (get, set) => {
|
|||
* Call this when loading messages from the database to restore plan state
|
||||
*/
|
||||
export interface HydratePlanInput {
|
||||
toolCallId: string;
|
||||
result: {
|
||||
id?: string;
|
||||
title?: string;
|
||||
description?: string;
|
||||
todos?: Array<{
|
||||
id: string;
|
||||
label: string;
|
||||
status: "pending" | "in_progress" | "completed" | "cancelled";
|
||||
description?: string;
|
||||
}>;
|
||||
};
|
||||
toolCallId: string;
|
||||
result: {
|
||||
id?: string;
|
||||
title?: string;
|
||||
description?: string;
|
||||
todos?: Array<{
|
||||
id: string;
|
||||
label: string;
|
||||
status: "pending" | "in_progress" | "completed" | "cancelled";
|
||||
description?: string;
|
||||
}>;
|
||||
};
|
||||
}
|
||||
|
||||
export const hydratePlanStateAtom = atom(
|
||||
null,
|
||||
(get, set, plan: HydratePlanInput) => {
|
||||
if (!plan.result?.todos || plan.result.todos.length === 0) return;
|
||||
|
||||
const states = new Map(get(planStatesAtom));
|
||||
const title = plan.result.title || "Planning Approach";
|
||||
|
||||
// Register this as the owner if no plan exists yet
|
||||
registerPlanOwner(title, plan.toolCallId);
|
||||
|
||||
// Get the canonical title
|
||||
const canonicalTitle = getCanonicalPlanTitle(title);
|
||||
const ownerToolCallId = firstPlanOwner?.toolCallId || plan.toolCallId;
|
||||
|
||||
// Only set if this is newer or doesn't exist
|
||||
const existing = states.get(canonicalTitle);
|
||||
if (!existing) {
|
||||
states.set(canonicalTitle, {
|
||||
id: plan.result.id || `plan-${Date.now()}`,
|
||||
title: canonicalTitle,
|
||||
description: plan.result.description,
|
||||
todos: plan.result.todos,
|
||||
lastUpdated: Date.now(),
|
||||
ownerToolCallId,
|
||||
});
|
||||
set(planStatesAtom, states);
|
||||
}
|
||||
}
|
||||
);
|
||||
export const hydratePlanStateAtom = atom(null, (get, set, plan: HydratePlanInput) => {
|
||||
if (!plan.result?.todos || plan.result.todos.length === 0) return;
|
||||
|
||||
const states = new Map(get(planStatesAtom));
|
||||
const title = plan.result.title || "Planning Approach";
|
||||
|
||||
// Register this as the owner if no plan exists yet
|
||||
registerPlanOwner(title, plan.toolCallId);
|
||||
|
||||
// Get the canonical title
|
||||
const canonicalTitle = getCanonicalPlanTitle(title);
|
||||
const ownerToolCallId = firstPlanOwner?.toolCallId || plan.toolCallId;
|
||||
|
||||
// Only set if this is newer or doesn't exist
|
||||
const existing = states.get(canonicalTitle);
|
||||
if (!existing) {
|
||||
states.set(canonicalTitle, {
|
||||
id: plan.result.id || `plan-${Date.now()}`,
|
||||
title: canonicalTitle,
|
||||
description: plan.result.description,
|
||||
todos: plan.result.todos,
|
||||
lastUpdated: Date.now(),
|
||||
ownerToolCallId,
|
||||
});
|
||||
set(planStatesAtom, states);
|
||||
}
|
||||
});
|
||||
|
||||
/**
|
||||
* Extract write_todos tool call data from message content
|
||||
* Returns an array of { toolCallId, result } for each write_todos call found
|
||||
*/
|
||||
export function extractWriteTodosFromContent(content: unknown): HydratePlanInput[] {
|
||||
if (!Array.isArray(content)) return [];
|
||||
|
||||
const results: HydratePlanInput[] = [];
|
||||
|
||||
for (const part of content) {
|
||||
if (
|
||||
typeof part === "object" &&
|
||||
part !== null &&
|
||||
"type" in part &&
|
||||
(part as { type: string }).type === "tool-call" &&
|
||||
"toolName" in part &&
|
||||
(part as { toolName: string }).toolName === "write_todos" &&
|
||||
"toolCallId" in part &&
|
||||
"result" in part
|
||||
) {
|
||||
const toolCall = part as {
|
||||
toolCallId: string;
|
||||
result: HydratePlanInput["result"];
|
||||
};
|
||||
if (toolCall.result) {
|
||||
results.push({
|
||||
toolCallId: toolCall.toolCallId,
|
||||
result: toolCall.result,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
if (!Array.isArray(content)) return [];
|
||||
|
||||
const results: HydratePlanInput[] = [];
|
||||
|
||||
for (const part of content) {
|
||||
if (
|
||||
typeof part === "object" &&
|
||||
part !== null &&
|
||||
"type" in part &&
|
||||
(part as { type: string }).type === "tool-call" &&
|
||||
"toolName" in part &&
|
||||
(part as { toolName: string }).toolName === "write_todos" &&
|
||||
"toolCallId" in part &&
|
||||
"result" in part
|
||||
) {
|
||||
const toolCall = part as {
|
||||
toolCallId: string;
|
||||
result: HydratePlanInput["result"];
|
||||
};
|
||||
if (toolCall.result) {
|
||||
results.push({
|
||||
toolCallId: toolCall.toolCallId,
|
||||
result: toolCall.result,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue