diff --git a/surfsense_web/components/tool-ui/index.ts b/surfsense_web/components/tool-ui/index.ts index 6371554ae..4d885a38c 100644 --- a/surfsense_web/components/tool-ui/index.ts +++ b/surfsense_web/components/tool-ui/index.ts @@ -51,17 +51,11 @@ export { SandboxExecuteToolUI, } from "./sandbox-execute"; export { - type MemoryItem, - type RecallMemoryArgs, - RecallMemoryArgsSchema, - type RecallMemoryResult, - RecallMemoryResultSchema, - RecallMemoryToolUI, - type SaveMemoryArgs, - SaveMemoryArgsSchema, - type SaveMemoryResult, - SaveMemoryResultSchema, - SaveMemoryToolUI, + type UpdateMemoryArgs, + UpdateMemoryArgsSchema, + type UpdateMemoryResult, + UpdateMemoryResultSchema, + UpdateMemoryToolUI, } from "./user-memory"; export { GenerateVideoPresentationToolUI } from "./video-presentation"; export { type WriteTodosData, WriteTodosSchema, WriteTodosToolUI } from "./write-todos"; diff --git a/surfsense_web/components/tool-ui/user-memory.tsx b/surfsense_web/components/tool-ui/user-memory.tsx index e232bdcc7..800b9e601 100644 --- a/surfsense_web/components/tool-ui/user-memory.tsx +++ b/surfsense_web/components/tool-ui/user-memory.tsx @@ -1,100 +1,38 @@ "use client"; import type { ToolCallMessagePartProps } from "@assistant-ui/react"; -import { BrainIcon, CheckIcon, Loader2Icon, SearchIcon, XIcon } from "lucide-react"; +import { AlertTriangleIcon, BrainIcon, CheckIcon, Loader2Icon, XIcon } from "lucide-react"; import { z } from "zod"; // ============================================================================ -// Zod Schemas for save_memory tool +// Zod Schemas for update_memory tool // ============================================================================ -const SaveMemoryArgsSchema = z.object({ - content: z.string(), - category: z.string().default("fact"), +const UpdateMemoryArgsSchema = z.object({ + updated_memory: z.string(), }); -const SaveMemoryResultSchema = z.object({ +const UpdateMemoryResultSchema = z.object({ status: z.enum(["saved", "error"]), - memory_id: z.number().nullish(), - memory_text: z.string().nullish(), - category: z.string().nullish(), message: z.string().nullish(), - error: z.string().nullish(), + warning: z.string().nullish(), }); -type SaveMemoryArgs = z.infer; -type SaveMemoryResult = z.infer; +type UpdateMemoryArgs = z.infer; +type UpdateMemoryResult = z.infer; // ============================================================================ -// Zod Schemas for recall_memory tool +// Update Memory Tool UI // ============================================================================ -const RecallMemoryArgsSchema = z.object({ - query: z.string().nullish(), - category: z.string().nullish(), - top_k: z.number().default(5), -}); - -const MemoryItemSchema = z.object({ - id: z.number(), - memory_text: z.string(), - category: z.string(), - updated_at: z.string().nullish(), -}); - -const RecallMemoryResultSchema = z.object({ - status: z.enum(["success", "error"]), - count: z.number().nullish(), - memories: z.array(MemoryItemSchema).nullish(), - formatted_context: z.string().nullish(), - error: z.string().nullish(), -}); - -type RecallMemoryArgs = z.infer; -type RecallMemoryResult = z.infer; -type MemoryItem = z.infer; - -// ============================================================================ -// Category badge colors -// ============================================================================ - -const categoryColors: Record = { - preference: "bg-blue-500/10 text-blue-600 dark:text-blue-400", - fact: "bg-green-500/10 text-green-600 dark:text-green-400", - instruction: "bg-purple-500/10 text-purple-600 dark:text-purple-400", - context: "bg-orange-500/10 text-orange-600 dark:text-orange-400", -}; - -function CategoryBadge({ category }: { category: string }) { - const colorClass = categoryColors[category] || "bg-gray-500/10 text-gray-600 dark:text-gray-400"; - return ( - - {category} - - ); -} - -// ============================================================================ -// Save Memory Tool UI -// ============================================================================ - -export const SaveMemoryToolUI = ({ - args, +export const UpdateMemoryToolUI = ({ result, status, -}: ToolCallMessagePartProps) => { +}: ToolCallMessagePartProps) => { const isRunning = status.type === "running" || status.type === "requires-action"; const isComplete = status.type === "complete"; const isError = result?.status === "error"; - // Parse args safely - const parsedArgs = SaveMemoryArgsSchema.safeParse(args); - const content = parsedArgs.success ? parsedArgs.data.content : ""; - const category = parsedArgs.success ? parsedArgs.data.category : "fact"; - - // Loading state if (isRunning) { return (
@@ -102,13 +40,12 @@ export const SaveMemoryToolUI = ({
- Saving to memory... + Updating memory...
); } - // Error state if (isError) { return (
@@ -116,14 +53,15 @@ export const SaveMemoryToolUI = ({
- Failed to save memory - {result?.error &&

{result.error}

} + Failed to update memory + {result?.message && ( +

{result.message}

+ )}
); } - // Success state if (isComplete && result?.status === "saved") { return (
@@ -133,138 +71,19 @@ export const SaveMemoryToolUI = ({
- Memory saved - + Memory updated
-

{content}

-
-
- ); - } - - // Default/incomplete state - show what's being saved - if (content) { - return ( -
-
- -
-
-
- Saving memory - -
-

{content}

-
-
- ); - } - - return null; -}; - -// ============================================================================ -// Recall Memory Tool UI -// ============================================================================ - -export const RecallMemoryToolUI = ({ - args, - result, - status, -}: ToolCallMessagePartProps) => { - const isRunning = status.type === "running" || status.type === "requires-action"; - const isComplete = status.type === "complete"; - const isError = result?.status === "error"; - - // Parse args safely - const parsedArgs = RecallMemoryArgsSchema.safeParse(args); - const query = parsedArgs.success ? parsedArgs.data.query : null; - - // Loading state - if (isRunning) { - return ( -
-
- -
-
- - {query ? `Searching memories for "${query}"...` : "Recalling memories..."} - -
-
- ); - } - - // Error state - if (isError) { - return ( -
-
- -
-
- Failed to recall memories - {result?.error &&

{result.error}

} -
-
- ); - } - - // Success state with memories - if (isComplete && result?.status === "success") { - const memories = result.memories || []; - const count = result.count || 0; - - if (count === 0) { - return ( -
-
- -
- No memories found -
- ); - } - - return ( -
-
- - - Recalled {count} {count === 1 ? "memory" : "memories"} - -
-
- {memories.slice(0, 5).map((memory: MemoryItem) => ( -
- - {memory.memory_text} + {result.warning && ( +
+ +

{result.warning}

- ))} - {memories.length > 5 && ( -

...and {memories.length - 5} more

)}
); } - // Default/incomplete state - if (query) { - return ( -
-
- -
- Searching memories for "{query}" -
- ); - } - return null; }; @@ -273,13 +92,8 @@ export const RecallMemoryToolUI = ({ // ============================================================================ export { - SaveMemoryArgsSchema, - SaveMemoryResultSchema, - RecallMemoryArgsSchema, - RecallMemoryResultSchema, - type SaveMemoryArgs, - type SaveMemoryResult, - type RecallMemoryArgs, - type RecallMemoryResult, - type MemoryItem, + UpdateMemoryArgsSchema, + UpdateMemoryResultSchema, + type UpdateMemoryArgs, + type UpdateMemoryResult, }; diff --git a/surfsense_web/contracts/enums/toolIcons.tsx b/surfsense_web/contracts/enums/toolIcons.tsx index 8e5d1e452..fd12aaa9c 100644 --- a/surfsense_web/contracts/enums/toolIcons.tsx +++ b/surfsense_web/contracts/enums/toolIcons.tsx @@ -19,8 +19,7 @@ const TOOL_ICONS: Record = { scrape_webpage: ScanLine, web_search: Globe, search_surfsense_docs: BookOpen, - save_memory: Brain, - recall_memory: Brain, + update_memory: Brain, }; export function getToolIcon(name: string): LucideIcon { diff --git a/surfsense_web/contracts/types/search-space.types.ts b/surfsense_web/contracts/types/search-space.types.ts index 8a0a2fb4c..7b4fefb62 100644 --- a/surfsense_web/contracts/types/search-space.types.ts +++ b/surfsense_web/contracts/types/search-space.types.ts @@ -9,6 +9,7 @@ export const searchSpace = z.object({ user_id: z.string(), citations_enabled: z.boolean(), qna_custom_instructions: z.string().nullable(), + shared_memory_md: z.string().nullable().optional(), member_count: z.number(), is_owner: z.boolean(), }); @@ -54,6 +55,7 @@ export const updateSearchSpaceRequest = z.object({ description: true, citations_enabled: true, qna_custom_instructions: true, + shared_memory_md: true, }) .partial(), });