From 79bc3c2d981bae15681d6f2228dacd4cb2de32c0 Mon Sep 17 00:00:00 2001 From: Cameron Date: Wed, 5 Nov 2025 18:11:51 -0800 Subject: [PATCH] Add /link and /unlink commands for managing agent tools (#59) Co-authored-by: Letta --- examples/send-image-simple.ts | 47 ++++++ examples/send-image.ts | 81 ++++++++++ src/agent/modify.ts | 147 ++++++++++++++++++ src/cli/App.tsx | 89 +++++++++++ src/cli/commands/registry.ts | 14 ++ src/cli/components/CommandPreview.tsx | 7 +- src/cli/components/WelcomeScreen.tsx | 8 + src/index.ts | 47 +++++- src/tests/agent/link-unlink.test.ts | 205 ++++++++++++++++++++++++++ test-image-send.ts | 60 ++++++++ 10 files changed, 701 insertions(+), 4 deletions(-) create mode 100644 examples/send-image-simple.ts create mode 100644 examples/send-image.ts create mode 100644 src/tests/agent/link-unlink.test.ts create mode 100644 test-image-send.ts diff --git a/examples/send-image-simple.ts b/examples/send-image-simple.ts new file mode 100644 index 0000000..071f8db --- /dev/null +++ b/examples/send-image-simple.ts @@ -0,0 +1,47 @@ +#!/usr/bin/env bun +/** + * Simplest possible example: Send an image to a Letta agent with streaming + */ + +import { readFileSync } from "node:fs"; +import { sendMessageStream } from "../src/agent/message"; + +async function main() { + const agentId = "agent-YOUR-AGENT-ID"; // Replace with your agent ID + const imagePath = "./screenshot.png"; // Replace with your image path + + // 1. Read and encode image + const imageData = readFileSync(imagePath).toString("base64"); + + // 2. Send message with streaming + const stream = await sendMessageStream(agentId, [ + { + role: "user", + content: [ + { + type: "text", + text: "What's in this image?", + }, + { + type: "image", + source: { + type: "base64", + mediaType: "image/png", + data: imageData, + }, + }, + ], + }, + ]); + + // 3. Print streaming response + console.log("Agent response:"); + for await (const chunk of stream) { + if (chunk.messageType === "assistant_message") { + process.stdout.write(chunk.content || ""); + } + } + console.log("\n"); +} + +main().catch(console.error); diff --git a/examples/send-image.ts b/examples/send-image.ts new file mode 100644 index 0000000..e9310e0 --- /dev/null +++ b/examples/send-image.ts @@ -0,0 +1,81 @@ +#!/usr/bin/env bun +/** + * Minimal example: Send an image to a Letta agent + * + * Usage: + * bun examples/send-image.ts + * + * Example: + * bun examples/send-image.ts agent-123abc screenshot.png + */ + +import { readFileSync } from "node:fs"; +import { getClient } from "../src/agent/client"; +import { sendMessageStream } from "../src/agent/message"; + +async function main() { + const agentId = process.argv[2]; + const imagePath = process.argv[3]; + + if (!agentId || !imagePath) { + console.error("Usage: bun send-image.ts "); + process.exit(1); + } + + // Step 1: Read image file and convert to base64 + const imageBuffer = readFileSync(imagePath); + const base64Data = imageBuffer.toString("base64"); + + // Determine media type from file extension + const ext = imagePath.toLowerCase(); + const mediaType = ext.endsWith(".png") + ? "image/png" + : ext.endsWith(".jpg") || ext.endsWith(".jpeg") + ? "image/jpeg" + : ext.endsWith(".gif") + ? "image/gif" + : ext.endsWith(".webp") + ? "image/webp" + : "image/jpeg"; // default + + console.log(`Sending ${imagePath} to agent ${agentId}...`); + + // Step 2: Create message content with text and image + const messageContent = [ + { type: "text" as const, text: "What do you see in this image?" }, + { + type: "image" as const, + source: { + type: "base64" as const, + mediaType, + data: base64Data, + }, + }, + ]; + + // Step 3: Send message to agent + const stream = await sendMessageStream(agentId, [ + { + role: "user", + content: messageContent, + }, + ]); + + // Step 4: Process the streaming response + console.log("\nAgent response:"); + for await (const chunk of stream) { + if (chunk.messageType === "assistant_message") { + process.stdout.write(chunk.content || ""); + } else if (chunk.messageType === "reasoning_message") { + // Optionally show internal monologue + // console.log(`[thinking] ${chunk.reasoning}`); + } + } + + console.log("\n"); +} + +main().catch((err) => { + console.error("Error:", err); + process.exit(1); +}); diff --git a/src/agent/modify.ts b/src/agent/modify.ts index dac5fbf..f9055ad 100644 --- a/src/agent/modify.ts +++ b/src/agent/modify.ts @@ -2,6 +2,7 @@ // Utilities for modifying agent configuration import type { LlmConfig } from "@letta-ai/letta-client/resources/models/models"; +import { getToolNames } from "../tools/manager"; import { getClient } from "./client"; /** @@ -44,3 +45,149 @@ export async function updateAgentLLMConfig( return finalConfig; } + +export interface LinkResult { + success: boolean; + message: string; + addedCount?: number; +} + +export interface UnlinkResult { + success: boolean; + message: string; + removedCount?: number; +} + +/** + * Attach all Letta Code tools to an agent. + * + * @param agentId - The agent ID + * @returns Result with success status and message + */ +export async function linkToolsToAgent(agentId: string): Promise { + try { + const client = await getClient(); + + // Get ALL agent tools from agent state + const agent = await client.agents.retrieve(agentId); + const currentTools = agent.tools || []; + const currentToolIds = currentTools + .map((t) => t.id) + .filter((id): id is string => typeof id === "string"); + const currentToolNames = new Set( + currentTools + .map((t) => t.name) + .filter((name): name is string => typeof name === "string"), + ); + + // Get Letta Code tool names + const lettaCodeToolNames = getToolNames(); + + // Find tools to add (tools that aren't already attached) + const toolsToAdd = lettaCodeToolNames.filter( + (name) => !currentToolNames.has(name), + ); + + if (toolsToAdd.length === 0) { + return { + success: true, + message: "All Letta Code tools already attached", + addedCount: 0, + }; + } + + // Look up tool IDs from global tool list + const toolsToAddIds: string[] = []; + for (const toolName of toolsToAdd) { + const tools = await client.tools.list({ name: toolName }); + const tool = tools[0]; + if (tool?.id) { + toolsToAddIds.push(tool.id); + } + } + + // Combine current tools with new tools + const newToolIds = [...currentToolIds, ...toolsToAddIds]; + + // Get current tool_rules and add requires_approval rules for new tools + const currentToolRules = agent.tool_rules || []; + const newToolRules = [ + ...currentToolRules, + ...toolsToAdd.map((toolName) => ({ + tool_name: toolName, + type: "requires_approval" as const, + prompt_template: null, + })), + ]; + + await client.agents.modify(agentId, { + tool_ids: newToolIds, + tool_rules: newToolRules, + }); + + return { + success: true, + message: `Attached ${toolsToAddIds.length} Letta Code tool(s) to agent`, + addedCount: toolsToAddIds.length, + }; + } catch (error) { + return { + success: false, + message: `Failed: ${error instanceof Error ? error.message : String(error)}`, + }; + } +} + +/** + * Remove all Letta Code tools from an agent. + * + * @param agentId - The agent ID + * @returns Result with success status and message + */ +export async function unlinkToolsFromAgent( + agentId: string, +): Promise { + try { + const client = await getClient(); + + // Get ALL agent tools from agent state (not tools.list which may be incomplete) + const agent = await client.agents.retrieve(agentId); + const allTools = agent.tools || []; + const lettaCodeToolNames = new Set(getToolNames()); + + // Filter out Letta Code tools, keep everything else + const remainingTools = allTools.filter( + (t) => t.name && !lettaCodeToolNames.has(t.name), + ); + const removedCount = allTools.length - remainingTools.length; + + // Extract IDs from remaining tools (filter out any undefined IDs) + const remainingToolIds = remainingTools + .map((t) => t.id) + .filter((id): id is string => typeof id === "string"); + + // Remove approval rules for Letta Code tools being unlinked + const currentToolRules = agent.tool_rules || []; + const remainingToolRules = currentToolRules.filter( + (rule) => + rule.type !== "requires_approval" || + !lettaCodeToolNames.has(rule.tool_name), + ); + + await client.agents.modify(agentId, { + tool_ids: remainingToolIds, + tool_rules: remainingToolRules, + }); + + return { + success: true, + message: `Removed ${removedCount} Letta Code tool(s) from agent`, + removedCount, + }; + } catch (error) { + return { + success: false, + message: `Failed: ${error instanceof Error ? error.message : String(error)}`, + }; + } +} diff --git a/src/cli/App.tsx b/src/cli/App.tsx index c745ec0..1981975 100644 --- a/src/cli/App.tsx +++ b/src/cli/App.tsx @@ -15,6 +15,7 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { getResumeData } from "../agent/check-approval"; import { getClient } from "../agent/client"; import { sendMessageStream } from "../agent/message"; +import { linkToolsToAgent, unlinkToolsFromAgent } from "../agent/modify"; import { SessionStats } from "../agent/stats"; import type { ApprovalContext } from "../permissions/analyzer"; import { permissionMode } from "../permissions/mode"; @@ -111,6 +112,8 @@ export default function App({ loadingState?: | "assembling" | "upserting" + | "linking" + | "unlinking" | "initializing" | "checking" | "ready"; @@ -837,6 +840,92 @@ export default function App({ return { submitted: true }; } + // Special handling for /link command - attach Letta Code tools + if (msg.trim() === "/link") { + const cmdId = uid("cmd"); + buffersRef.current.byId.set(cmdId, { + kind: "command", + id: cmdId, + input: msg, + output: "Attaching Letta Code tools to agent...", + phase: "running", + }); + buffersRef.current.order.push(cmdId); + refreshDerived(); + + setCommandRunning(true); + + try { + const result = await linkToolsToAgent(agentId); + + buffersRef.current.byId.set(cmdId, { + kind: "command", + id: cmdId, + input: msg, + output: result.message, + phase: "finished", + success: result.success, + }); + refreshDerived(); + } catch (error) { + buffersRef.current.byId.set(cmdId, { + kind: "command", + id: cmdId, + input: msg, + output: `Failed: ${error instanceof Error ? error.message : String(error)}`, + phase: "finished", + success: false, + }); + refreshDerived(); + } finally { + setCommandRunning(false); + } + return { submitted: true }; + } + + // Special handling for /unlink command - remove Letta Code tools + if (msg.trim() === "/unlink") { + const cmdId = uid("cmd"); + buffersRef.current.byId.set(cmdId, { + kind: "command", + id: cmdId, + input: msg, + output: "Removing Letta Code tools from agent...", + phase: "running", + }); + buffersRef.current.order.push(cmdId); + refreshDerived(); + + setCommandRunning(true); + + try { + const result = await unlinkToolsFromAgent(agentId); + + buffersRef.current.byId.set(cmdId, { + kind: "command", + id: cmdId, + input: msg, + output: result.message, + phase: "finished", + success: result.success, + }); + refreshDerived(); + } catch (error) { + buffersRef.current.byId.set(cmdId, { + kind: "command", + id: cmdId, + input: msg, + output: `Failed: ${error instanceof Error ? error.message : String(error)}`, + phase: "finished", + success: false, + }); + refreshDerived(); + } finally { + setCommandRunning(false); + } + return { submitted: true }; + } + // Immediately add command to transcript with "running" phase const cmdId = uid("cmd"); buffersRef.current.byId.set(cmdId, { diff --git a/src/cli/commands/registry.ts b/src/cli/commands/registry.ts index cacec0e..5700e1f 100644 --- a/src/cli/commands/registry.ts +++ b/src/cli/commands/registry.ts @@ -50,6 +50,20 @@ export const commands: Record = { return "Clearing credentials..."; }, }, + "/link": { + desc: "Attach Letta Code tools to current agent", + handler: () => { + // Handled specially in App.tsx to access agent ID and client + return "Attaching tools..."; + }, + }, + "/unlink": { + desc: "Remove Letta Code tools from current agent", + handler: () => { + // Handled specially in App.tsx to access agent ID and client + return "Removing tools..."; + }, + }, }; /** diff --git a/src/cli/components/CommandPreview.tsx b/src/cli/components/CommandPreview.tsx index ab2fc9e..3acdc6a 100644 --- a/src/cli/components/CommandPreview.tsx +++ b/src/cli/components/CommandPreview.tsx @@ -34,9 +34,10 @@ export function CommandPreview({ paddingX={1} > {commandList.map((item) => ( - - {item.cmd} - {item.desc} + + + {item.cmd.padEnd(15)} {item.desc} + ))} {showAgentUrl && ( diff --git a/src/cli/components/WelcomeScreen.tsx b/src/cli/components/WelcomeScreen.tsx index 91b63d2..fc7e50b 100644 --- a/src/cli/components/WelcomeScreen.tsx +++ b/src/cli/components/WelcomeScreen.tsx @@ -9,6 +9,8 @@ import { colors } from "./colors"; type LoadingState = | "assembling" | "upserting" + | "linking" + | "unlinking" | "initializing" | "checking" | "ready"; @@ -82,6 +84,12 @@ export function WelcomeScreen({ if (loadingState === "upserting") { return "Upserting tools..."; } + if (loadingState === "linking") { + return "Attaching Letta Code tools..."; + } + if (loadingState === "unlinking") { + return "Removing Letta Code tools..."; + } if (loadingState === "checking") { return "Checking for pending approvals..."; } diff --git a/src/index.ts b/src/index.ts index 1395b60..355f05e 100755 --- a/src/index.ts +++ b/src/index.ts @@ -91,6 +91,8 @@ async function main() { "permission-mode": { type: "string" }, yolo: { type: "boolean" }, "output-format": { type: "string" }, + link: { type: "boolean" }, + unlink: { type: "boolean" }, }, strict: true, allowPositionals: true, @@ -233,6 +235,21 @@ async function main() { } } + // Handle --link and --unlink flags (modify tools before starting session) + const shouldLink = values.link as boolean | undefined; + const shouldUnlink = values.unlink as boolean | undefined; + + // Validate --link/--unlink flags require --agent + if (shouldLink || shouldUnlink) { + if (!specifiedAgentId) { + console.error( + `Error: --${shouldLink ? "link" : "unlink"} requires --agent `, + ); + process.exit(1); + } + // Implementation is in InteractiveSession init() + } + if (isHeadless) { // For headless mode, load tools synchronously await loadTools(); @@ -263,7 +280,13 @@ async function main() { model?: string; }) { const [loadingState, setLoadingState] = useState< - "assembling" | "upserting" | "initializing" | "checking" | "ready" + | "assembling" + | "upserting" + | "linking" + | "unlinking" + | "initializing" + | "checking" + | "ready" >("assembling"); const [agentId, setAgentId] = useState(null); const [agentState, setAgentState] = useState(null); @@ -279,6 +302,28 @@ async function main() { const client = await getClient(); await upsertToolsToServer(client); + // Handle --link/--unlink after upserting tools + if (shouldLink || shouldUnlink) { + if (!agentIdArg) { + console.error("Error: --link/--unlink requires --agent "); + process.exit(1); + } + + setLoadingState(shouldLink ? "linking" : "unlinking"); + const { linkToolsToAgent, unlinkToolsFromAgent } = await import( + "./agent/modify" + ); + + const result = shouldLink + ? await linkToolsToAgent(agentIdArg) + : await unlinkToolsFromAgent(agentIdArg); + + if (!result.success) { + console.error(`āœ— ${result.message}`); + process.exit(1); + } + } + setLoadingState("initializing"); const { createAgent } = await import("./agent/create"); const { getModelUpdateArgs } = await import("./agent/model"); diff --git a/src/tests/agent/link-unlink.test.ts b/src/tests/agent/link-unlink.test.ts new file mode 100644 index 0000000..11483bd --- /dev/null +++ b/src/tests/agent/link-unlink.test.ts @@ -0,0 +1,205 @@ +import { afterAll, beforeAll, describe, expect, test } from "bun:test"; +import { Letta } from "@letta-ai/letta-client"; +import { linkToolsToAgent, unlinkToolsFromAgent } from "../../agent/modify"; +import { settingsManager } from "../../settings-manager"; +import { getToolNames, loadTools } from "../../tools/manager"; + +// Skip these integration tests if LETTA_API_KEY is not set +const shouldSkip = !process.env.LETTA_API_KEY; +const describeOrSkip = shouldSkip ? describe.skip : describe; + +describeOrSkip("Link/Unlink Tools", () => { + let client: Letta; + let testAgentId: string; + + beforeAll(async () => { + // Initialize settings and load tools + await settingsManager.initialize(); + await loadTools(); + + // Create a test agent + const apiKey = process.env.LETTA_API_KEY; + if (!apiKey) { + throw new Error("LETTA_API_KEY required for tests"); + } + + client = new Letta({ apiKey }); + + const agent = await client.agents.create({ + model: "openai/gpt-4o-mini", + embedding: "openai/text-embedding-3-small", + memory_blocks: [ + { label: "human", value: "Test user" }, + { label: "persona", value: "Test agent" }, + ], + tools: [], + }); + + testAgentId = agent.id; + }); + + afterAll(async () => { + // Cleanup: delete test agent + if (testAgentId) { + try { + await client.agents.delete(testAgentId); + } catch (_error) { + // Ignore cleanup errors + } + } + }); + + test("linkToolsToAgent attaches all Letta Code tools", async () => { + // Reset: ensure tools are not already attached + await unlinkToolsFromAgent(testAgentId); + + const result = await linkToolsToAgent(testAgentId); + + expect(result.success).toBe(true); + expect(result.addedCount).toBeGreaterThan(0); + + // Verify tools were attached + const agent = await client.agents.retrieve(testAgentId); + const toolNames = agent.tools?.map((t) => t.name) || []; + const lettaCodeTools = getToolNames(); + + for (const toolName of lettaCodeTools) { + expect(toolNames).toContain(toolName); + } + }, 30000); + + test("linkToolsToAgent adds approval rules for all tools", async () => { + // First unlink to reset + await unlinkToolsFromAgent(testAgentId); + + // Link tools + await linkToolsToAgent(testAgentId); + + // Verify approval rules were added + const agent = await client.agents.retrieve(testAgentId); + const approvalRules = agent.tool_rules?.filter( + (rule) => rule.type === "requires_approval", + ); + + const lettaCodeTools = getToolNames(); + expect(approvalRules?.length).toBe(lettaCodeTools.length); + + // Check all Letta Code tools have approval rules + const rulesToolNames = approvalRules?.map((r) => r.tool_name) || []; + for (const toolName of lettaCodeTools) { + expect(rulesToolNames).toContain(toolName); + } + }, 30000); + + test("linkToolsToAgent returns success when tools already attached", async () => { + // Reset and link once + await unlinkToolsFromAgent(testAgentId); + await linkToolsToAgent(testAgentId); + + // Link again + const result = await linkToolsToAgent(testAgentId); + + expect(result.success).toBe(true); + expect(result.addedCount).toBe(0); + expect(result.message).toContain("already attached"); + }, 30000); + + test("unlinkToolsFromAgent removes all Letta Code tools", async () => { + // First link tools + await linkToolsToAgent(testAgentId); + + // Then unlink + const result = await unlinkToolsFromAgent(testAgentId); + + expect(result.success).toBe(true); + expect(result.removedCount).toBeGreaterThan(0); + + // Verify tools were removed + const agent = await client.agents.retrieve(testAgentId); + const toolNames = agent.tools?.map((t) => t.name) || []; + const lettaCodeTools = getToolNames(); + + for (const toolName of lettaCodeTools) { + expect(toolNames).not.toContain(toolName); + } + }, 30000); + + test("unlinkToolsFromAgent removes approval rules", async () => { + // First link tools + await linkToolsToAgent(testAgentId); + + // Then unlink + await unlinkToolsFromAgent(testAgentId); + + // Verify approval rules were removed + const agent = await client.agents.retrieve(testAgentId); + const approvalRules = agent.tool_rules?.filter( + (rule) => rule.type === "requires_approval", + ); + + const lettaCodeTools = new Set(getToolNames()); + const remainingApprovalRules = approvalRules?.filter((r) => + lettaCodeTools.has(r.tool_name), + ); + + expect(remainingApprovalRules?.length || 0).toBe(0); + }, 30000); + + test("unlinkToolsFromAgent preserves non-Letta-Code tools", async () => { + // Link Letta Code tools + await linkToolsToAgent(testAgentId); + + // Attach memory tool + const memoryTools = await client.tools.list({ name: "memory" }); + const memoryTool = memoryTools[0]; + if (memoryTool?.id) { + await client.agents.tools.attach(memoryTool.id, { + agent_id: testAgentId, + }); + } + + // Unlink Letta Code tools + await unlinkToolsFromAgent(testAgentId); + + // Verify memory tool is still there + const agent = await client.agents.retrieve(testAgentId); + const toolNames = agent.tools?.map((t) => t.name) || []; + + expect(toolNames).toContain("memory"); + + // Verify Letta Code tools are gone + const lettaCodeTools = getToolNames(); + for (const toolName of lettaCodeTools) { + expect(toolNames).not.toContain(toolName); + } + }, 30000); + + test("unlinkToolsFromAgent preserves non-approval tool_rules", async () => { + // Link tools + await linkToolsToAgent(testAgentId); + + // Add a continue_loop rule manually + const agent = await client.agents.retrieve(testAgentId); + const newToolRules = [ + ...(agent.tool_rules || []), + { + tool_name: "memory", + type: "continue_loop" as const, + prompt_template: "Test rule", + }, + ]; + + await client.agents.modify(testAgentId, { tool_rules: newToolRules }); + + // Unlink Letta Code tools + await unlinkToolsFromAgent(testAgentId); + + // Verify continue_loop rule is still there + const updatedAgent = await client.agents.retrieve(testAgentId); + const continueLoopRules = updatedAgent.tool_rules?.filter( + (r) => r.type === "continue_loop" && r.tool_name === "memory", + ); + + expect(continueLoopRules?.length).toBe(1); + }, 30000); +}); diff --git a/test-image-send.ts b/test-image-send.ts new file mode 100644 index 0000000..dd00c42 --- /dev/null +++ b/test-image-send.ts @@ -0,0 +1,60 @@ +import { createAgent } from "./src/agent/create"; +import { sendMessageStream } from "./src/agent/message"; +import { readFileSync, writeFileSync } from "node:fs"; + +async function main() { + // Create a simple test image (1x1 red PNG) + const testImageBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="; + const testImagePath = "/tmp/test.png"; + writeFileSync(testImagePath, Buffer.from(testImageBase64, "base64")); + console.log("Created test image at", testImagePath); + + // Create agent + console.log("\nCreating test agent..."); + const agent = await createAgent("image-test-agent"); + console.log("Agent created:", agent.id); + + // Read image + const imageData = readFileSync(testImagePath).toString("base64"); + + // Send message with image + console.log("\nSending image to agent..."); + const stream = await sendMessageStream(agent.id, [ + { + role: "user", + content: [ + { + type: "text", + text: "What do you see in this image?", + }, + { + type: "image", + source: { + type: "base64", + media_type: "image/png", + data: imageData, + }, + }, + ], + }, + ]); + + // Print response + console.log("\nAgent response:"); + let fullResponse = ""; + for await (const chunk of stream) { + if (chunk.message_type === "assistant_message" && chunk.content) { + fullResponse += chunk.content; + process.stdout.write(chunk.content); + } + } + if (!fullResponse) { + console.log("(no assistant message received)"); + } + console.log("\n\nāœ… Done!"); +} + +main().catch((err) => { + console.error("āŒ Error:", err); + process.exit(1); +});