Skip to content

Refactor: Introduce useActionHandler pattern and standardize logging #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
plvo wants to merge 1 commit into
base: 1.x
Choose a base branch
from
Open

Refactor: Introduce useActionHandler pattern and standardize logging #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
116 changes: 39 additions & 77 deletions src/actions/callToolAction.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,20 +1,22 @@
import {
type Action,
type HandlerCallback,
type IAgentRuntime,
type Memory,
type State,
logger,
import type {
Action,
HandlerCallback,
IAgentRuntime,
Memory,
State,
} from "@elizaos/core";
import type { McpService } from "../service";
import { MCP_SERVICE_NAME } from "../types";
import { handleMcpError } from "../utils/error";
import { handleToolResponse, processToolResult } from "../utils/processing";
import { createToolSelectionArgument, createToolSelectionName } from "../utils/selection";
import { handleNoToolAvailable } from "../utils/handler";
import { mcpLogger } from "@/utils/mcp-logger";
import { handleMcpError } from "@/utils/error";
import { handleToolResponse, processToolResult } from "@/utils/processing";
import { createToolSelectionArgument, createToolSelectionName } from "@/utils/selection";
import { useActionHandler } from "@/utils/use-action";
import { validateAction } from "@/utils/validation";
import { handleNoToolAvailable } from "@/utils/handlers";

const ACTION_NAME = "CALL_TOOL";

export const callToolAction: Action = {
name: "CALL_TOOL",
name: ACTION_NAME,
similes: [
"CALL_MCP_TOOL",
"USE_TOOL",
Expand All @@ -28,101 +30,61 @@ export const callToolAction: Action = {
],
description: "Calls a tool from an MCP server to perform a specific task",

validate: async (runtime: IAgentRuntime, _message: Memory, _state?: State): Promise<boolean> => {
const mcpService = runtime.getService<McpService>(MCP_SERVICE_NAME);
if (!mcpService) return false;

const servers = mcpService.getServers();
return (
servers.length > 0 &&
servers.some(
(server) => server.status === "connected" && server.tools && server.tools.length > 0
)
);
validate: async (runtime: IAgentRuntime, message: Memory, state?: State): Promise<boolean> => {
return await validateAction(ACTION_NAME, runtime, message, state);
},

handler: async (
runtime: IAgentRuntime,
message: Memory,
_state?: State,
_options?: { [key: string]: unknown },
state?: State,
options?: { [key: string]: unknown },
callback?: HandlerCallback
): Promise<boolean> => {
const composedState = await runtime.composeState(message, ["RECENT_MESSAGES", "MCP"]);
const mcpService = runtime.getService<McpService>(MCP_SERVICE_NAME);
if (!mcpService) {
throw new Error("MCP service not available");
}
const mcpProvider = mcpService.getProviderData();
const context = await useActionHandler({ actionName: ACTION_NAME, runtime, message, state, options, callback });

try {
// Select the tool with this servername and toolname
const toolSelectionName = await createToolSelectionName({
runtime,
state: composedState,
message,
callback,
mcpProvider,
});
const toolSelectionName = await createToolSelectionName({...context});
if (!toolSelectionName || toolSelectionName.noToolAvailable) {
logger.warn("[NO_TOOL_AVAILABLE] No appropriate tool available for the request");
mcpLogger.warn("[NO_TOOL_AVAILABLE] No appropriate tool available for the request");
return handleNoToolAvailable(callback, toolSelectionName);
}
const { serverName, toolName, reasoning } = toolSelectionName;
logger.info(
`[CALLING] Calling tool "${serverName}/${toolName}" on server with reasoning: "${reasoning}"`
);
mcpLogger.info(`[CALLING] Calling tool "${serverName}/${toolName}" on server with reasoning: "${reasoning}"`);

// Create the tool selection "argument" based on the selected tool name
const toolSelectionArgument = await createToolSelectionArgument({
runtime,
state: composedState,
message,
callback,
mcpProvider,
toolSelectionName,
});
const toolSelectionArgument = await createToolSelectionArgument({ ...context, toolSelectionName });
if (!toolSelectionArgument) {
logger.warn(
"[NO_TOOL_SELECTION_ARGUMENT] No appropriate tool selection argument available"
);
mcpLogger.warn("[NO_TOOL_SELECTION_ARGUMENT] No appropriate tool selection argument available");
return handleNoToolAvailable(callback, toolSelectionName);
}
logger.info(
`[SELECTED] Tool Selection result:\n${JSON.stringify(toolSelectionArgument, null, 2)}`
);
mcpLogger.info(`[SELECTED] Tool Selection result:\n${JSON.stringify(toolSelectionArgument, null, 2)}`);

const result = await mcpService.callTool(
serverName,
toolName,
toolSelectionArgument.toolArguments
);
const result = await context.mcpService.callTool(serverName, toolName, toolSelectionArgument.toolArguments);
mcpLogger.info(`[CALLED] Tool "${serverName}/${toolName}" result:\n"${JSON.stringify(result, null, 2)}"`);

const { toolOutput, hasAttachments, attachments } = processToolResult(
const { toolOutput, hasAttachments, attachments } = processToolResult({
...context,
result,
serverName,
toolName,
runtime,
message.entityId
);
messageEntityId: context.message.entityId,
});

await handleToolResponse(
runtime,
message,
mcpLogger.info('[HANDLE] Handling tool response...');
await handleToolResponse({
...context,
serverName,
toolName,
toolSelectionArgument.toolArguments,
toolArguments: toolSelectionArgument.toolArguments,
toolOutput,
hasAttachments,
attachments,
composedState,
mcpProvider,
callback
);
});

return true;
} catch (error) {
return handleMcpError(composedState, mcpProvider, error, runtime, message, "tool", callback);
return await handleMcpError({ ...context, type: 'tool', error });
}
},

Expand Down
171 changes: 34 additions & 137 deletions src/actions/readResourceAction.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,67 +1,26 @@
import {
type Action,
type HandlerCallback,
type IAgentRuntime,
type Memory,
ModelType,
type State,
composePromptFromState,
logger,
import type {
Action,
HandlerCallback,
IAgentRuntime,
Memory,
State,
} from "@elizaos/core";
import type { McpService } from "../service";
import { resourceSelectionTemplate } from "../templates/resourceSelectionTemplate";
import { MCP_SERVICE_NAME } from "../types";
import { handleMcpError } from "../utils/error";
import { mcpLogger } from "@/utils/mcp-logger";
import { handleMcpError } from "@/utils/error";
import {
handleResourceAnalysis,
processResourceResult,
sendInitialResponse,
} from "../utils/processing";
import {
createResourceSelectionFeedbackPrompt,
validateResourceSelection,
} from "../utils/validation";
import type { ResourceSelection } from "../utils/validation";
import { withModelRetry } from "../utils/wrapper";
} from "@/utils/processing";
import { useActionHandler } from "@/utils/use-action";
import { createResourceSelection } from "@/utils/selection";
import { handleNoResourceAvailable, handleResourceAnalysis } from "@/utils/handlers";
import { validateAction } from "@/utils/validation";

function createResourceSelectionPrompt(composedState: State, userMessage: string): string {
const mcpData = composedState.values.mcp || {};
const serverNames = Object.keys(mcpData);

let resourcesDescription = "";
for (const serverName of serverNames) {
const server = mcpData[serverName];
if (server.status !== "connected") continue;

const resourceUris = Object.keys(server.resources || {});
for (const uri of resourceUris) {
const resource = server.resources[uri];
resourcesDescription += `Resource: ${uri} (Server: ${serverName})\n`;
resourcesDescription += `Name: ${resource.name || "No name available"}\n`;
resourcesDescription += `Description: ${
resource.description || "No description available"
}\n`;
resourcesDescription += `MIME Type: ${resource.mimeType || "Not specified"}\n\n`;
}
}

const enhancedState: State = {
...composedState,
values: {
...composedState.values,
resourcesDescription,
userMessage,
},
};

return composePromptFromState({
state: enhancedState,
template: resourceSelectionTemplate,
});
}
const ACTION_NAME = 'READ_RESOURCE';

export const readResourceAction: Action = {
name: "READ_RESOURCE",
name: ACTION_NAME,
similes: [
"READ_MCP_RESOURCE",
"GET_RESOURCE",
Expand All @@ -73,107 +32,45 @@ export const readResourceAction: Action = {
],
description: "Reads a resource from an MCP server",

validate: async (runtime: IAgentRuntime, _message: Memory, _state?: State): Promise<boolean> => {
const mcpService = runtime.getService<McpService>(MCP_SERVICE_NAME);
if (!mcpService) return false;

const servers = mcpService.getServers();
return (
servers.length > 0 &&
servers.some(
(server) => server.status === "connected" && server.resources && server.resources.length > 0
)
);
validate: async (runtime: IAgentRuntime, message: Memory, state?: State): Promise<boolean> => {
return await validateAction(ACTION_NAME, runtime, message, state);
},

handler: async (
runtime: IAgentRuntime,
message: Memory,
_state?: State,
_options?: { [key: string]: unknown },
state?: State,
options?: { [key: string]: unknown },
callback?: HandlerCallback
): Promise<boolean> => {
const composedState = await runtime.composeState(message, ["RECENT_MESSAGES", "MCP"]);

const mcpService = runtime.getService<McpService>(MCP_SERVICE_NAME);
if (!mcpService) {
throw new Error("MCP service not available");
}

const mcpProvider = mcpService.getProviderData();
const context = await useActionHandler({ actionName: ACTION_NAME, runtime, message, state, options, callback });

try {
mcpLogger.info('[INITIAL_RESPONSE] Sending initial response...');
await sendInitialResponse(callback);

const resourceSelectionPrompt = createResourceSelectionPrompt(
composedState,
message.content.text || ""
);

const resourceSelection = await runtime.useModel(ModelType.TEXT_SMALL, {
prompt: resourceSelectionPrompt,
});
const resourceSelection = await createResourceSelection({ ...context });
mcpLogger.info(`[SELECTED] Resource Selection response:\n${JSON.stringify(resourceSelection, null, 2)}`);

const parsedSelection = await withModelRetry<ResourceSelection>({
runtime,
state: composedState,
message,
callback,
input: resourceSelection,
validationFn: (data) => validateResourceSelection(data),
createFeedbackPromptFn: (originalResponse, errorMessage, state, userMessage) =>
createResourceSelectionFeedbackPrompt(
originalResponse as string,
errorMessage,
state,
userMessage
),
failureMsg: `I'm having trouble finding the resource you're looking for. Could you provide more details about what you need?`,
retryCount: 0,
});

if (!parsedSelection || parsedSelection.noResourceAvailable) {
if (callback && parsedSelection?.noResourceAvailable) {
await callback({
text: "I don't have a specific resource that contains the information you're looking for. Let me try to assist you directly instead.",
thought:
"No appropriate MCP resource available for this request. Falling back to direct assistance.",
actions: ["REPLY"],
});
}
return true;
if (!resourceSelection || resourceSelection.noResourceAvailable) {
mcpLogger.info('[NO_RESOURCE_AVAILABLE] No appropriate resource available for the request');
return handleNoResourceAvailable(callback);
}

const { serverName, uri, reasoning } = parsedSelection;

logger.debug(`Selected resource "${uri}" on server "${serverName}" because: ${reasoning}`);

const result = await mcpService.readResource(serverName, uri);
logger.debug(`Read resource ${uri} from server ${serverName}`);
const { serverName, uri, reasoning } = resourceSelection;
mcpLogger.info(`[FETCHING] Fetching resource "${serverName}/${uri}" with reasoning: "${reasoning}"`);

const result = await context.mcpService.readResource(serverName, uri);
mcpLogger.info(`[FETCHED] Resource "${serverName}/${uri}" result: \n"${JSON.stringify(result, null, 2)}"`);

const { resourceContent, resourceMeta } = processResourceResult(result, uri);

await handleResourceAnalysis(
runtime,
message,
uri,
serverName,
resourceContent,
resourceMeta,
callback
);
mcpLogger.info('[HANDLE] Handling resource response...');
await handleResourceAnalysis({ ...context, serverName, uri, resourceContent, resourceMeta });

return true;
} catch (error) {
return handleMcpError(
composedState,
mcpProvider,
error,
runtime,
message,
"resource",
callback
);
return await handleMcpError({ ...context, error, type: 'resource' });
}
},

Expand Down
5 changes: 3 additions & 2 deletions src/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { type IAgentRuntime, type Plugin, logger } from "@elizaos/core";
import type { IAgentRuntime, Plugin } from "@elizaos/core";
import { mcpLogger } from "@/utils/mcp-logger";
import { callToolAction } from "./actions/callToolAction";
import { readResourceAction } from "./actions/readResourceAction";
import { provider } from "./provider";
Expand All @@ -9,7 +10,7 @@ const mcpPlugin: Plugin = {
description: "Plugin for connecting to MCP (Model Context Protocol) servers",

init: async (_config: Record<string, string>, _runtime: IAgentRuntime) => {
logger.info("Initializing MCP plugin...");
mcpLogger.info("Initializing MCP plugin...");
},

services: [McpService],
Expand Down
Loading