r/src/tools/tool_research.c

// retoor <retoor@molodetz.nl>

#include "tool.h"
#include "db.h"
#include "http_client.h"
#include "agent.h"
#include "r_config.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/sha.h>
#include <sqlite3.h>

extern db_handle db_open(const char *path);
extern void db_close(db_handle db);
extern r_status_t db_execute(db_handle db, const char *sql, struct json_object **result);

static void compute_url_hash(const char *url, char *output) {
    unsigned char hash[SHA256_DIGEST_LENGTH];
    SHA256((unsigned char*)url, strlen(url), hash);
    for(int i = 0; i < SHA256_DIGEST_LENGTH; i++) {
        snprintf(output + (i * 2), 3, "%02x", hash[i]);
    }
    output[64] = '\0';
}

static char *research_dispatcher_execute(tool_t *self, struct json_object *args) {
    (void)self;
    struct json_object *urls_obj, *batch_id_obj;
    if (!json_object_object_get_ex(args, "urls", &urls_obj) ||
        !json_object_object_get_ex(args, "batch_id", &batch_id_obj)) return strdup("Error: missing urls or batch_id");

    const char *batch_id = json_object_get_string(batch_id_obj);
    db_handle db = db_open(NULL);
    int count = json_object_array_length(urls_obj);
    int dispatched = 0;

    for (int i = 0; i < count; i++) {
        const char *url = json_object_get_string(json_object_array_get_idx(urls_obj, i));
        char hash[65];
        compute_url_hash(url, hash);

        char *sql = sqlite3_mprintf("INSERT OR IGNORE INTO research_tasks (url_hash, url, status, batch_id) VALUES (%Q, %Q, 'pending', %Q)",
                                    hash, url, batch_id);
        db_execute(db, sql, NULL);
        sqlite3_free(sql);
        dispatched++;
    }

    db_close(db);
    char buf[128];
    snprintf(buf, sizeof(buf), "Dispatched %d URLs to research_tasks table for batch %s.", dispatched, batch_id);
    return strdup(buf);
}

static char *fetch_and_scrape_execute(tool_t *self, struct json_object *args) {
    (void)self;
    struct json_object *url_obj;
    if (!json_object_object_get_ex(args, "url", &url_obj)) return strdup("Error: missing url");
    const char *url = json_object_get_string(url_obj);
    char hash[65];
    compute_url_hash(url, hash);

    db_handle db = db_open(NULL);
    
    // Atomic Lock / Claim
    char *lock_sql = sqlite3_mprintf("UPDATE research_tasks SET status = 'processing', updated_at = CURRENT_TIMESTAMP WHERE url_hash = %Q AND (status = 'pending' OR (status = 'processing' AND updated_at < datetime('now', '-10 minutes')))", hash);
    db_execute(db, lock_sql, NULL);
    sqlite3_free(lock_sql);

    http_client_handle http = http_client_create(NULL);
    char *content = NULL;
    r_status_t st = http_get(http, url, &content);
    http_client_destroy(http);

    if (st != R_SUCCESS || !content) {
        char *fail_sql = sqlite3_mprintf("UPDATE research_tasks SET status = 'failed' WHERE url_hash = %Q", hash);
        db_execute(db, fail_sql, NULL);
        sqlite3_free(fail_sql);
        db_close(db);
        return strdup("Error: Failed to fetch URL content.");
    }

    // Truncate for summary (simplified for MVP)
    char summary[2048];
    strncpy(summary, content, 2047);
    summary[2047] = '\0';

    char *update_sql = sqlite3_mprintf("UPDATE research_tasks SET status = 'completed', summary = %Q, updated_at = CURRENT_TIMESTAMP WHERE url_hash = %Q", summary, hash);
    db_execute(db, update_sql, NULL);
    sqlite3_free(update_sql);
    
    db_close(db);
    free(content);
    return strdup("URL fetched, scraped, and summary stored in research_tasks.");
}

static char *suggest_subtask_execute(tool_t *self, struct json_object *args) {
    (void)self;
    struct json_object *url_obj, *reason_obj;
    if (!json_object_object_get_ex(args, "url", &url_obj) ||
        !json_object_object_get_ex(args, "reason", &reason_obj)) return strdup("Error: missing url or reason");

    fprintf(stderr, "\033[1;34m[Escalation] Worker suggests new subtask: %s (Reason: %s)\033[0m\n", 
            json_object_get_string(url_obj), json_object_get_string(reason_obj));
    
    return strdup("Subtask suggestion logged and escalated to Manager.");
}

static struct json_object *dispatch_desc(void) {
    struct json_object *root = json_object_new_object();
    json_object_object_add(root, "type", json_object_new_string("function"));
    struct json_object *f = json_object_new_object();
    json_object_object_add(f, "name", json_object_new_string("research_dispatcher"));
    json_object_object_add(f, "description", json_object_new_string("Manager tool: Fan-out a list of URLs into the research_tasks queue."));
    struct json_object *p = json_object_new_object();
    json_object_object_add(p, "type", json_object_new_string("object"));
    struct json_object *props = json_object_new_object();
    struct json_object *u = json_object_new_object();
    json_object_object_add(u, "type", json_object_new_string("array"));
    json_object_object_add(u, "items", json_object_new_object());
    json_object_object_add(json_object_object_get(u, "items"), "type", json_object_new_string("string"));
    json_object_object_add(props, "urls", u);
    struct json_object *b = json_object_new_object();
    json_object_object_add(b, "type", json_object_new_string("string"));
    json_object_object_add(props, "batch_id", b);
    json_object_object_add(p, "properties", props);
    struct json_object *req = json_object_new_array();
    json_object_array_add(req, json_object_new_string("urls"));
    json_object_array_add(req, json_object_new_string("batch_id"));
    json_object_object_add(p, "required", req);
    json_object_object_add(p, "additionalProperties", json_object_new_boolean(0));
    json_object_object_add(f, "parameters", p);
    r_config_handle cfg = r_config_get_instance();
    if (r_config_use_strict(cfg)) json_object_object_add(f, "strict", json_object_new_boolean(1));
    json_object_object_add(root, "function", f);
    return root;
}

static struct json_object *fetch_desc(void) {
    struct json_object *root = json_object_new_object();
    json_object_object_add(root, "type", json_object_new_string("function"));
    struct json_object *f = json_object_new_object();
    json_object_object_add(f, "name", json_object_new_string("fetch_and_scrape"));
    json_object_object_add(f, "description", json_object_new_string("Worker tool: Fetch a single URL, scrape it, and save the summary to the database."));
    struct json_object *p = json_object_new_object();
    json_object_object_add(p, "type", json_object_new_string("object"));
    struct json_object *props = json_object_new_object();
    struct json_object *u = json_object_new_object();
    json_object_object_add(u, "type", json_object_new_string("string"));
    json_object_object_add(props, "url", u);
    json_object_object_add(p, "properties", props);
    struct json_object *req = json_object_new_array();
    json_object_array_add(req, json_object_new_string("url"));
    json_object_object_add(p, "required", req);
    json_object_object_add(p, "additionalProperties", json_object_new_boolean(0));
    json_object_object_add(f, "parameters", p);
    r_config_handle cfg = r_config_get_instance();
    if (r_config_use_strict(cfg)) json_object_object_add(f, "strict", json_object_new_boolean(1));
    json_object_object_add(root, "function", f);
    return root;
}

static struct json_object *suggest_desc(void) {
    struct json_object *root = json_object_new_object();
    json_object_object_add(root, "type", json_object_new_string("function"));
    struct json_object *f = json_object_new_object();
    json_object_object_add(f, "name", json_object_new_string("suggest_subtask"));
    json_object_object_add(f, "description", json_object_new_string("Worker tool: Suggest a new URL discovered within a page for future research. Escalates to Manager."));
    struct json_object *p = json_object_new_object();
    json_object_object_add(p, "type", json_object_new_string("object"));
    struct json_object *props = json_object_new_object();
    struct json_object *u = json_object_new_object();
    json_object_object_add(u, "type", json_object_new_string("string"));
    json_object_object_add(props, "url", u);
    struct json_object *r = json_object_new_object();
    json_object_object_add(r, "type", json_object_new_string("string"));
    json_object_object_add(props, "reason", r);
    json_object_object_add(p, "properties", props);
    struct json_object *req = json_object_new_array();
    json_object_array_add(req, json_object_new_string("url"));
    json_object_array_add(req, json_object_new_string("reason"));
    json_object_object_add(p, "required", req);
    json_object_object_add(p, "additionalProperties", json_object_new_boolean(0));
    json_object_object_add(f, "parameters", p);
    r_config_handle cfg = r_config_get_instance();
    if (r_config_use_strict(cfg)) json_object_object_add(f, "strict", json_object_new_boolean(1));
    json_object_object_add(root, "function", f);
    return root;
}

static const tool_vtable_t dispatch_vtable = { .get_description = dispatch_desc, .execute = research_dispatcher_execute };
static const tool_vtable_t fetch_vtable = { .get_description = fetch_desc, .execute = fetch_and_scrape_execute };
static const tool_vtable_t suggest_vtable = { .get_description = suggest_desc, .execute = suggest_subtask_execute };

static tool_t dispatch_tool = { .vtable = &dispatch_vtable, .name = "research_dispatcher" };
static tool_t fetch_tool = { .vtable = &fetch_vtable, .name = "fetch_and_scrape" };
static tool_t suggest_tool = { .vtable = &suggest_vtable, .name = "suggest_subtask" };

tool_t *tool_research_dispatcher_create(void) { return &dispatch_tool; }
tool_t *tool_fetch_and_scrape_create(void) { return &fetch_tool; }
tool_t *tool_suggest_subtask_create(void) { return &suggest_tool; }