r/src/tools/tool_agent.c

// retoor <retoor@molodetz.nl>
#include "tool.h"
#include "agent.h"
#include "messages.h"
#include "r_config.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
typedef struct {
    tool_t tool;
} tool_agent_t;
static struct json_object *tool_spawn_agent_get_description(void) {
    struct json_object *obj = json_object_new_object();
    json_object_object_add(obj, "name", json_object_new_string("spawn_agent"));
    json_object_object_add(obj, "description", json_object_new_string("Spawn a specialized sub-agent to handle a specific task."));
    struct json_object *params = json_object_new_object();
    json_object_object_add(params, "type", json_object_new_string("object"));
    struct json_object *props = json_object_new_object();
    struct json_object *persona = json_object_new_object();
    json_object_object_add(persona, "type", json_object_new_string("string"));
    json_object_object_add(persona, "description", json_object_new_string("The persona of the agent (researcher, developer, security, fetcher)."));
    struct json_object *persona_enum = json_object_new_array();
    json_object_array_add(persona_enum, json_object_new_string("researcher"));
    json_object_array_add(persona_enum, json_object_new_string("developer"));
    json_object_array_add(persona_enum, json_object_new_string("security"));
    json_object_array_add(persona_enum, json_object_new_string("fetcher"));
    json_object_object_add(persona, "enum", persona_enum);
    json_object_object_add(props, "persona", persona);
    struct json_object *goal = json_object_new_object();
    json_object_object_add(goal, "type", json_object_new_string("string"));
    json_object_object_add(goal, "description", json_object_new_string("The specific task or goal for the sub-agent."));
    json_object_object_add(props, "goal", goal);
    struct json_object *max_subagents = json_object_new_object();
    json_object_object_add(max_subagents, "type", json_object_new_string("integer"));
    json_object_object_add(max_subagents, "description", json_object_new_string("Remaining budget for spawning recursive sub-agents. Decrement this by 1 when spawning a sub-agent. Default is 2."));
    json_object_object_add(max_subagents, "default", json_object_new_int(2));
    json_object_object_add(props, "max_subagents", max_subagents);
    json_object_object_add(params, "properties", props);
    struct json_object *required = json_object_new_array();
    json_object_array_add(required, json_object_new_string("persona"));
    json_object_array_add(required, json_object_new_string("goal"));
    json_object_array_add(required, json_object_new_string("max_subagents"));
    json_object_object_add(params, "required", required);
    json_object_object_add(params, "additionalProperties", json_object_new_boolean(0));
    json_object_object_add(obj, "parameters", params);
    r_config_handle cfg = r_config_get_instance();
    if (r_config_use_strict(cfg)) {
        json_object_object_add(obj, "strict", json_object_new_boolean(1));
    }
    struct json_object *full_obj = json_object_new_object();
    json_object_object_add(full_obj, "type", json_object_new_string("function"));
    json_object_object_add(full_obj, "function", obj);
    return full_obj;
}
static char *tool_spawn_agent_execute(tool_t *self, struct json_object *args) {
    (void)self;
    struct json_object *persona_obj, *goal_obj, *max_subagents_obj;
    if (!json_object_object_get_ex(args, "persona", &persona_obj) ||
        !json_object_object_get_ex(args, "goal", &goal_obj)) {
        return strdup("Error: Missing persona or goal");
    }
    int max_subagents = 2;
    if (json_object_object_get_ex(args, "max_subagents", &max_subagents_obj)) {
        max_subagents = json_object_get_int(max_subagents_obj);
    }
    if (max_subagents <= 0) {
        return strdup("Error: Spawning limit reached. You are not allowed to spawn more sub-agents. Perform the task yourself using existing tools.");
    }
    const char *persona_str = json_object_get_string(persona_obj);
    const char *goal_str = json_object_get_string(goal_obj);
    tool_registry_type_t type = TOOL_TYPE_ALL;
    const char *system_prompt_base = NULL;
    if (strcmp(persona_str, "researcher") == 0) {
        type = TOOL_TYPE_RESEARCHER;
        system_prompt_base = "You are a specialized Research Agent. Your goal is to find, extract, and summarize information. "
                        "Do not attempt to write or execute code unless it's for data analysis. "
                        "Focus on using web search, http fetch, and reading files. "
                        "## Sub-Agent Rules\n"
                        "- Your output is for a master agent, not the final user.\n"
                        "- DO NOT ask questions or for permission.\n"
                        "- Provide RAW DATA and summaries.\n"
                        "- Do not say 'task complete'.\n"
                        "## Hierarchical Research Workflow\n"
                        "When web_search returns URLs with content:\n"
                        "1. Spawn 'fetcher' agents in parallel to fetch URL contents\n"
                        "2. Each fetcher should use http_fetch tool for individual URLs\n"
                        "3. Aggregate results and synthesize with citations\n"
                        "Citation format: [Source N] Title (URL)\n"
                        "Use spawn_agent extensively for URL fetching from search results.";
    } else if (strcmp(persona_str, "developer") == 0) {
        type = TOOL_TYPE_DEVELOPER;
        system_prompt_base = "You are a specialized Developer Agent. Your goal is to write, test, and debug code. "
                        "## Sub-Agent Rules\n"
                        "- Your output is for a master agent, not the final user.\n"
                        "- DO NOT ask questions or for permission.\n"
                        "- Just perform the requested development task and report results.\n"
                        "Use the terminal, file editing tools, and python execution to fulfill your task. "
                        "Always verify your changes by running tests or the code itself.";
    } else if (strcmp(persona_str, "security") == 0) {
        type = TOOL_TYPE_SECURITY;
        system_prompt_base = "You are a specialized Security Auditor Agent. Your goal is to find vulnerabilities and perform security analysis. "
                        "## Sub-Agent Rules\n"
                        "- Your output is for a master agent, not the final user.\n"
                        "- DO NOT ask questions or for permission.\n"
                        "Be pedantic and thorough. Use fuzzing, port scanning, and code analysis tools. "
                        "Report any findings clearly with potential impact.";
    } else if (strcmp(persona_str, "fetcher") == 0) {
        type = TOOL_TYPE_ALL;
        system_prompt_base = "You are a specialized URL Fetcher Agent. Your goal is to fetch and extract content from multiple URLs. "
                        "## Sub-Agent Rules\n"
                        "- Just return the content of the URLs.\n"
                        "- Do not provide commentary unless necessary.\n"
                        "Use http_fetch tool to retrieve content from each URL. "
                        "Handle errors gracefully and continue with other URLs. "
                        "Clean HTML/extract text suitable for LLM analysis. "
                        "Truncate content to ~10K chars per URL to stay within token limits.";
    } else {
        return strdup("Error: Invalid persona");
    }
    time_t now = time(NULL);
    struct tm *tm_info = localtime(&now);
    char datetime[64];
    strftime(datetime, sizeof(datetime), "%Y-%m-%d %H:%M:%S %Z", tm_info);
    size_t prompt_size = strlen(system_prompt_base) + 1024;
    char *system_prompt = malloc(prompt_size);
    if (!system_prompt) return strdup("Error: Out of memory");
    snprintf(system_prompt, prompt_size, 
             "Current date/time: %s\n\n%s\n\n"
             "CRITICAL: It is currently %s.\n"
             "ORCHESTRATION BUDGET: You are allowed to spawn up to %d more levels of sub-agents. "
             "When using spawn_agent, you MUST pass 'max_subagents' as %d.\n", 
             datetime, system_prompt_base, datetime, max_subagents, max_subagents - 1);
    char session_id[256];
    snprintf(session_id, sizeof(session_id), "subagent-%s-%u", persona_str, (unsigned int)time(NULL));
    messages_handle msgs = messages_create(session_id);
    messages_add(msgs, "system", system_prompt);
    free(system_prompt);
    agent_handle agent = agent_create(goal_str, msgs);
    if (!agent) {
        messages_destroy(msgs);
        return strdup("Error: Failed to create sub-agent");
    }
    agent_set_is_subagent(agent, true);
    tool_registry_t *specialized_tools = tool_registry_get_specialized(type);
    if (specialized_tools) {
        agent_set_tool_registry(agent, specialized_tools);
    }
    
    agent_set_max_iterations(agent, 50); // Sub-agents have lower limit
    char *agent_response = agent_run(agent, goal_str);
    char *result = messages_to_json_string(msgs);
    agent_destroy(agent);
    free(agent_response);
    
    if (specialized_tools) {
        tool_registry_destroy(specialized_tools);
    }
    if (!result) {
        return strdup("Sub-agent failed to provide a result.");
    }
    return result;
}
static void tool_spawn_agent_print_action(const char *name, struct json_object *args) {
    struct json_object *persona_obj, *goal_obj;
    const char *persona = "unknown";
    const char *goal = "unknown";
    if (json_object_object_get_ex(args, "persona", &persona_obj)) persona = json_object_get_string(persona_obj);
    if (json_object_object_get_ex(args, "goal", &goal_obj)) goal = json_object_get_string(goal_obj);
    printf("\033[1;34m[Agent] Spawning %s agent for: %s\033[0m\n", persona, goal);
}
static const tool_vtable_t tool_spawn_agent_vtable = {
    .get_description = tool_spawn_agent_get_description,
    .execute = tool_spawn_agent_execute,
    .print_action = tool_spawn_agent_print_action
};
tool_t *tool_spawn_agent_create(void) {
    tool_agent_t *tool = calloc(1, sizeof(tool_agent_t));
    if (!tool) return NULL;
    tool->tool.vtable = &tool_spawn_agent_vtable;
    tool->tool.name = "spawn_agent";
    return &tool->tool;
}