mirror of
https://github.com/crawlab-team/crawlab.git
synced 2026-01-22 17:31:03 +01:00
2d0508a0c2
- Implemented new API methods in CrawlabClient for managing LLM providers and AutoProbes. - Added prompts for setting up AI-powered web scraping with AutoProbe and configuring LLM providers. - Created a new module for AI tools, integrating various functionalities including listing, creating, updating, and deleting LLM providers and AutoProbes. - Enhanced error handling and response formatting for better user experience.
610 lines
17 KiB
Go
610 lines
17 KiB
Go
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
|
|
import { CrawlabClient } from "../client.js";
|
|
import { z } from "zod";
|
|
|
|
const AI_TOOLS = {
|
|
// LLM Provider Management
|
|
list_llm_providers: "crawlab_list_llm_providers",
|
|
get_llm_provider: "crawlab_get_llm_provider",
|
|
create_llm_provider: "crawlab_create_llm_provider",
|
|
update_llm_provider: "crawlab_update_llm_provider",
|
|
delete_llm_provider: "crawlab_delete_llm_provider",
|
|
|
|
// Chat Conversations
|
|
list_conversations: "crawlab_list_conversations",
|
|
get_conversation: "crawlab_get_conversation",
|
|
create_conversation: "crawlab_create_conversation",
|
|
update_conversation: "crawlab_update_conversation",
|
|
delete_conversation: "crawlab_delete_conversation",
|
|
get_conversation_messages: "crawlab_get_conversation_messages",
|
|
get_chat_message: "crawlab_get_chat_message",
|
|
|
|
// AutoProbe (AI Web Scraping)
|
|
list_autoprobes: "crawlab_list_autoprobes",
|
|
get_autoprobe: "crawlab_get_autoprobe",
|
|
create_autoprobe: "crawlab_create_autoprobe",
|
|
update_autoprobe: "crawlab_update_autoprobe",
|
|
delete_autoprobe: "crawlab_delete_autoprobe",
|
|
run_autoprobe_task: "crawlab_run_autoprobe_task",
|
|
get_autoprobe_tasks: "crawlab_get_autoprobe_tasks",
|
|
get_autoprobe_preview: "crawlab_get_autoprobe_preview",
|
|
get_autoprobe_pattern: "crawlab_get_autoprobe_pattern",
|
|
|
|
// AutoProbe V2
|
|
list_autoprobes_v2: "crawlab_list_autoprobes_v2",
|
|
get_autoprobe_v2: "crawlab_get_autoprobe_v2",
|
|
create_autoprobe_v2: "crawlab_create_autoprobe_v2",
|
|
update_autoprobe_v2: "crawlab_update_autoprobe_v2",
|
|
delete_autoprobe_v2: "crawlab_delete_autoprobe_v2",
|
|
run_autoprobe_v2_task: "crawlab_run_autoprobe_v2_task",
|
|
get_autoprobe_v2_tasks: "crawlab_get_autoprobe_v2_tasks",
|
|
get_autoprobe_v2_preview: "crawlab_get_autoprobe_v2_preview",
|
|
get_autoprobe_v2_pattern: "crawlab_get_autoprobe_v2_pattern",
|
|
get_autoprobe_v2_pattern_results: "crawlab_get_autoprobe_v2_pattern_results",
|
|
};
|
|
|
|
export function configureAITools(server: McpServer, client: CrawlabClient) {
|
|
// LLM Provider Management
|
|
server.tool(
|
|
AI_TOOLS.list_llm_providers,
|
|
"List all LLM providers configured in Crawlab",
|
|
{
|
|
page: z.number().optional().describe("Page number for pagination (default: 1)"),
|
|
page_size: z.number().optional().describe("Number of providers per page (default: 10)"),
|
|
filter: z.string().optional().describe("Filter providers by name or type"),
|
|
},
|
|
async (args) => {
|
|
try {
|
|
const result = await client.getLLMProviders(args);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error listing LLM providers: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
server.tool(
|
|
AI_TOOLS.get_llm_provider,
|
|
"Get details of a specific LLM provider",
|
|
{
|
|
id: z.string().describe("ID of the LLM provider"),
|
|
},
|
|
async ({ id }) => {
|
|
try {
|
|
const result = await client.getLLMProvider(id);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error getting LLM provider: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
server.tool(
|
|
AI_TOOLS.create_llm_provider,
|
|
"Create a new LLM provider configuration",
|
|
{
|
|
name: z.string().describe("Display name for the provider"),
|
|
type: z.string().describe("Provider type (e.g., 'openai', 'azure-openai', 'anthropic', 'gemini')"),
|
|
api_key: z.string().describe("API key for the provider"),
|
|
api_base_url: z.string().optional().describe("Custom API base URL"),
|
|
api_version: z.string().optional().describe("API version"),
|
|
default_model: z.string().optional().describe("Default model for this provider"),
|
|
models: z.array(z.string()).optional().describe("List of supported models"),
|
|
deployment_name: z.string().optional().describe("Deployment name (for Azure)"),
|
|
},
|
|
async (args) => {
|
|
try {
|
|
const result = await client.createLLMProvider(args);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error creating LLM provider: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
server.tool(
|
|
AI_TOOLS.update_llm_provider,
|
|
"Update an existing LLM provider configuration",
|
|
{
|
|
id: z.string().describe("ID of the LLM provider to update"),
|
|
name: z.string().optional().describe("Display name for the provider"),
|
|
type: z.string().optional().describe("Provider type"),
|
|
api_key: z.string().optional().describe("API key for the provider"),
|
|
api_base_url: z.string().optional().describe("Custom API base URL"),
|
|
api_version: z.string().optional().describe("API version"),
|
|
default_model: z.string().optional().describe("Default model for this provider"),
|
|
models: z.array(z.string()).optional().describe("List of supported models"),
|
|
deployment_name: z.string().optional().describe("Deployment name (for Azure)"),
|
|
},
|
|
async ({ id, ...updateData }) => {
|
|
try {
|
|
const result = await client.updateLLMProvider(id, updateData);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error updating LLM provider: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
server.tool(
|
|
AI_TOOLS.delete_llm_provider,
|
|
"Delete an LLM provider configuration",
|
|
{
|
|
id: z.string().describe("ID of the LLM provider to delete"),
|
|
},
|
|
async ({ id }) => {
|
|
try {
|
|
await client.deleteLLMProvider(id);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `LLM provider ${id} deleted successfully`,
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error deleting LLM provider: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
// Chat Conversations
|
|
server.tool(
|
|
AI_TOOLS.list_conversations,
|
|
"List all AI chat conversations",
|
|
{
|
|
page: z.number().optional().describe("Page number for pagination"),
|
|
page_size: z.number().optional().describe("Number of conversations per page"),
|
|
filter: z.string().optional().describe("Filter conversations by title or content"),
|
|
},
|
|
async (args) => {
|
|
try {
|
|
const result = await client.getConversations(args);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error listing conversations: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
server.tool(
|
|
AI_TOOLS.get_conversation,
|
|
"Get details of a specific conversation",
|
|
{
|
|
id: z.string().describe("ID of the conversation"),
|
|
},
|
|
async ({ id }) => {
|
|
try {
|
|
const result = await client.getConversation(id);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error getting conversation: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
server.tool(
|
|
AI_TOOLS.get_conversation_messages,
|
|
"Get all messages from a conversation",
|
|
{
|
|
id: z.string().describe("ID of the conversation"),
|
|
},
|
|
async ({ id }) => {
|
|
try {
|
|
const result = await client.getConversationMessages(id);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error getting conversation messages: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
// AutoProbe V2 (Latest)
|
|
server.tool(
|
|
AI_TOOLS.list_autoprobes_v2,
|
|
"List all AutoProbe V2 configurations (AI-powered web scraping)",
|
|
{
|
|
page: z.number().optional().describe("Page number for pagination"),
|
|
page_size: z.number().optional().describe("Number of AutoProbes per page"),
|
|
filter: z.string().optional().describe("Filter AutoProbes by name or URL"),
|
|
},
|
|
async (args) => {
|
|
try {
|
|
const result = await client.getAutoProbesV2(args);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error listing AutoProbes V2: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
server.tool(
|
|
AI_TOOLS.get_autoprobe_v2,
|
|
"Get details of a specific AutoProbe V2 configuration",
|
|
{
|
|
id: z.string().describe("ID of the AutoProbe"),
|
|
},
|
|
async ({ id }) => {
|
|
try {
|
|
const result = await client.getAutoProbeV2(id);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error getting AutoProbe V2: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
server.tool(
|
|
AI_TOOLS.create_autoprobe_v2,
|
|
"Create a new AutoProbe V2 configuration for AI-powered web scraping",
|
|
{
|
|
name: z.string().describe("Name of the AutoProbe"),
|
|
url: z.string().describe("Target URL to scrape"),
|
|
description: z.string().optional().describe("Description of what to extract"),
|
|
query: z.string().optional().describe("Natural language query describing the data to extract"),
|
|
settings: z.object({}).optional().describe("Additional configuration settings"),
|
|
},
|
|
async (args) => {
|
|
try {
|
|
const result = await client.createAutoProbeV2(args);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error creating AutoProbe V2: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
server.tool(
|
|
AI_TOOLS.run_autoprobe_v2_task,
|
|
"Run an AutoProbe V2 scraping task",
|
|
{
|
|
id: z.string().describe("ID of the AutoProbe"),
|
|
query: z.string().optional().describe("Custom query for this run"),
|
|
view_port: z.object({}).optional().describe("Custom viewport settings"),
|
|
},
|
|
async ({ id, ...params }) => {
|
|
try {
|
|
const result = await client.runAutoProbeV2Task(id, params);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error running AutoProbe V2 task: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
server.tool(
|
|
AI_TOOLS.get_autoprobe_v2_preview,
|
|
"Get a preview of what an AutoProbe V2 would extract",
|
|
{
|
|
id: z.string().describe("ID of the AutoProbe"),
|
|
},
|
|
async ({ id }) => {
|
|
try {
|
|
const result = await client.getAutoProbeV2Preview(id);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error getting AutoProbe V2 preview: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
server.tool(
|
|
AI_TOOLS.get_autoprobe_v2_pattern,
|
|
"Get the extraction pattern for an AutoProbe V2",
|
|
{
|
|
id: z.string().describe("ID of the AutoProbe"),
|
|
},
|
|
async ({ id }) => {
|
|
try {
|
|
const result = await client.getAutoProbeV2Pattern(id);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error getting AutoProbe V2 pattern: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
// AutoProbe V1 (Legacy Support)
|
|
server.tool(
|
|
AI_TOOLS.list_autoprobes,
|
|
"List all AutoProbe V1 configurations (legacy)",
|
|
{
|
|
page: z.number().optional().describe("Page number for pagination"),
|
|
page_size: z.number().optional().describe("Number of AutoProbes per page"),
|
|
filter: z.string().optional().describe("Filter AutoProbes by name or URL"),
|
|
},
|
|
async (args) => {
|
|
try {
|
|
const result = await client.getAutoProbes(args);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error listing AutoProbes V1: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
server.tool(
|
|
AI_TOOLS.get_autoprobe,
|
|
"Get details of a specific AutoProbe V1 configuration",
|
|
{
|
|
id: z.string().describe("ID of the AutoProbe"),
|
|
},
|
|
async ({ id }) => {
|
|
try {
|
|
const result = await client.getAutoProbe(id);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error getting AutoProbe V1: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
|
|
server.tool(
|
|
AI_TOOLS.create_autoprobe,
|
|
"Create a new AutoProbe V1 configuration (legacy)",
|
|
{
|
|
name: z.string().describe("Name of the AutoProbe"),
|
|
url: z.string().describe("Target URL to scrape"),
|
|
description: z.string().optional().describe("Description of what to extract"),
|
|
query: z.string().optional().describe("Natural language query describing the data to extract"),
|
|
},
|
|
async (args) => {
|
|
try {
|
|
const result = await client.createAutoProbe(args);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: JSON.stringify(result, null, 2),
|
|
},
|
|
],
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Error creating AutoProbe V1: ${error instanceof Error ? error.message : String(error)}`,
|
|
},
|
|
],
|
|
isError: true,
|
|
};
|
|
}
|
|
}
|
|
);
|
|
}
|