r/src/agent.c

// retoor <retoor@molodetz.nl>

#include "agent.h"
#include "http_client.h"
#include "r_config.h"
#include "tool.h"
#include "context_manager.h"
#include <json-c/json.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

struct agent_t {
    char *goal;
    int iteration_count;
    int max_iterations;
    int tool_retry_count;
    int max_tool_retries;
    agent_state_t state;
    time_t start_time;
    char *last_error;
    bool verbose;
    messages_handle messages;
    bool owns_messages;
    http_client_handle http;
    tool_registry_t *tools;
};

static const char *incomplete_phrases[] = {
    "I'll ",      "I will ",    "Let me ",     "I'm going to ",
    "Next, I",    "Now I'll",   "Now I will",  "I'll now",
    "I need to",  "I should",   "I can ",      "Going to ",
    "Will now",   "Proceeding", "Starting to", "About to ",
    "First, I",   "Then I",     "After that",  "Following that",
    "Now let me", "Let's check", "Let me check",
    NULL
};

static const char *incomplete_endings[] = {
    "...", ":", "files:", "content:", "implementation:",
    NULL
};

extern tool_registry_t *tools_get_registry(void);

agent_handle agent_create(const char *goal, messages_handle messages) {
    struct agent_t *agent = calloc(1, sizeof(struct agent_t));
    if (!agent) return NULL;

    if (goal) {
        agent->goal = strdup(goal);
        if (!agent->goal) {
            free(agent);
            return NULL;
        }
    }

    r_config_handle cfg = r_config_get_instance();

    agent->iteration_count = 0;
    agent->max_iterations = AGENT_MAX_ITERATIONS;
    agent->tool_retry_count = 0;
    agent->max_tool_retries = AGENT_MAX_TOOL_RETRIES;
    agent->state = AGENT_STATE_IDLE;
    agent->start_time = time(NULL);
    agent->verbose = r_config_is_verbose(cfg);

    if (messages) {
        agent->messages = messages;
        agent->owns_messages = false;
    } else {
        agent->messages = messages_create(r_config_get_session_id(cfg));
        agent->owns_messages = true;
    }

    if (!agent->messages) {
        free(agent->goal);
        free(agent);
        return NULL;
    }

    agent->http = http_client_create(r_config_get_api_key(cfg));
    if (!agent->http) {
        if (agent->owns_messages) {
            messages_destroy(agent->messages);
        }
        free(agent->goal);
        free(agent);
        return NULL;
    }

    agent->tools = tools_get_registry();

    return agent;
}

void agent_destroy(agent_handle agent) {
    if (!agent) return;
    if (agent->http) http_client_destroy(agent->http);
    if (agent->messages && agent->owns_messages) messages_destroy(agent->messages);
    free(agent->goal);
    free(agent->last_error);
    free(agent);
}

void agent_set_max_iterations(agent_handle agent, int max) {
    if (agent) agent->max_iterations = max;
}

void agent_set_verbose(agent_handle agent, bool verbose) {
    if (agent) agent->verbose = verbose;
}

void agent_set_tool_registry(agent_handle agent, tool_registry_t *registry) {
    if (agent && registry) agent->tools = registry;
}

agent_state_t agent_get_state(agent_handle agent) {
    return agent ? agent->state : AGENT_STATE_ERROR;
}

const char *agent_get_error(agent_handle agent) {
    return agent ? agent->last_error : NULL;
}

int agent_get_iteration_count(agent_handle agent) {
    return agent ? agent->iteration_count : 0;
}

static void agent_set_error(agent_handle agent, const char *error) {
    if (!agent) return;
    free(agent->last_error);
    agent->last_error = error ? strdup(error) : NULL;
}

static char *agent_build_request(agent_handle agent, const char *role, const char *message) {
    r_config_handle cfg = r_config_get_instance();

    struct json_object *root = json_object_new_object();
    if (!root) return NULL;

    json_object_object_add(root, "model",
        json_object_new_string(r_config_get_model(cfg)));

    if (role && message) {
        messages_add(agent->messages, role, message);
        if (r_config_use_tools(cfg) && agent->tools) {
            json_object_object_add(root, "tools",
                tool_registry_get_descriptions(agent->tools));
        }
    }

    json_object_object_add(root, "messages",
        json_object_get(messages_to_json(agent->messages)));
    json_object_object_add(root, "temperature",
        json_object_new_double(r_config_get_temperature(cfg)));

    char *result = strdup(json_object_to_json_string_ext(root, JSON_C_TO_STRING_PRETTY));
    json_object_put(root);
    return result;
}

static struct json_object *agent_process_response(agent_handle agent, const char *json_data) {
    r_config_handle cfg = r_config_get_instance();

    char *response = NULL;
    r_status_t status = http_post(agent->http, r_config_get_api_url(cfg), json_data, &response);

    if (status != R_SUCCESS || !response) {
        return NULL;
    }

    struct json_object *parsed = json_tokener_parse(response);
    free(response);

    if (!parsed) return NULL;

    struct json_object *error_obj;
    if (json_object_object_get_ex(parsed, "error", &error_obj)) {
        const char *err_str = json_object_to_json_string(error_obj);
        
        // Smart error detection for context overflow
        if (strstr(err_str, "too long") || 
            strstr(err_str, "context_length_exceeded") ||
            strstr(err_str, "Input is too long")) {
            agent_set_error(agent, "CONTEXT_OVERFLOW");
        } else {
            fprintf(stderr, "API Error: %s\n", err_str);
        }
        
        json_object_put(parsed);
        return NULL;
    }

    struct json_object *choices;
    if (!json_object_object_get_ex(parsed, "choices", &choices)) {
        json_object_put(parsed);
        return NULL;
    }

    struct json_object *first_choice = json_object_array_get_idx(choices, 0);
    if (!first_choice) {
        json_object_put(parsed);
        return NULL;
    }

    return first_choice;
}

static bool agent_has_tool_calls(struct json_object *choice) {
    struct json_object *message_obj;
    if (!json_object_object_get_ex(choice, "message", &message_obj)) return false;

    struct json_object *tool_calls;
    if (!json_object_object_get_ex(message_obj, "tool_calls", &tool_calls)) return false;

    return json_object_array_length(tool_calls) > 0;
}

static struct json_object *agent_get_tool_calls(struct json_object *choice) {
    struct json_object *message_obj;
    if (!json_object_object_get_ex(choice, "message", &message_obj)) return NULL;

    struct json_object *tool_calls;
    if (!json_object_object_get_ex(message_obj, "tool_calls", &tool_calls)) return NULL;

    return tool_calls;
}

static struct json_object *agent_get_message(struct json_object *choice) {
    struct json_object *message_obj;
    if (json_object_object_get_ex(choice, "message", &message_obj)) {
        return message_obj;
    }
    return NULL;
}

static char *agent_get_content(struct json_object *choice) {
    struct json_object *message_obj;
    if (!json_object_object_get_ex(choice, "message", &message_obj)) return NULL;

    struct json_object *content_obj;
    if (!json_object_object_get_ex(message_obj, "content", &content_obj)) return NULL;

    const char *content = json_object_get_string(content_obj);
    return content ? strdup(content) : NULL;
}

static bool agent_response_indicates_incomplete(const char *content) {
    if (!content) return false;

    for (int i = 0; incomplete_phrases[i]; i++) {
        if (strstr(content, incomplete_phrases[i])) return true;
    }

    size_t len = strlen(content);
    if (len > 3) {
        for (int i = 0; incomplete_endings[i]; i++) {
            size_t end_len = strlen(incomplete_endings[i]);
            if (len >= end_len && strcmp(content + len - end_len, incomplete_endings[i]) == 0) {
                return true;
            }
        }
    }

    return false;
}

char *agent_run(agent_handle agent, const char *user_message) {
    if (!agent) return NULL;

    agent->state = AGENT_STATE_RUNNING;
    agent->iteration_count = 0;
    agent->tool_retry_count = 0;

    if (!user_message || !*user_message) {
        agent->state = AGENT_STATE_ERROR;
        agent_set_error(agent, "Empty user message");
        return NULL;
    }

    messages_load(agent->messages);

    char *json_data = agent_build_request(agent, "user", user_message);
    if (!json_data) {
        agent->state = AGENT_STATE_ERROR;
        agent_set_error(agent, "Failed to create chat JSON");
        return NULL;
    }

    char *accumulated_response = NULL;
    size_t accumulated_len = 0;

    while (agent->state == AGENT_STATE_RUNNING || agent->state == AGENT_STATE_EXECUTING_TOOLS) {
        agent->iteration_count++;

        if (agent->iteration_count > agent->max_iterations) {
            agent->state = AGENT_STATE_MAX_ITERATIONS;
            agent_set_error(agent, "Maximum iterations reached");
            if (agent->verbose) {
                fprintf(stderr, "[Agent] Max iterations (%d) reached\n", agent->max_iterations);
            }
            free(json_data);
            break;
        }

        if (agent->verbose) {
            fprintf(stderr, "[Agent] Iteration %d/%d\n",
                agent->iteration_count, agent->max_iterations);
        }

        struct json_object *choice = agent_process_response(agent, json_data);
        
        if (!choice && agent->last_error && strcmp(agent->last_error, "CONTEXT_OVERFLOW") == 0) {
            if (context_manager_shrink(agent->messages) == R_SUCCESS) {
                // Retry with shrunk history
                free(json_data);
                json_data = agent_build_request(agent, NULL, NULL);
                agent->state = AGENT_STATE_RUNNING;
                // Don't increment iteration_count for retries due to context
                agent->iteration_count--;
                continue;
            } else {
                agent_set_error(agent, "Context limit reached and cannot be shrunk further.");
                free(json_data);
                break;
            }
        }

        free(json_data);
        json_data = NULL;

        if (!choice) {
            agent->tool_retry_count++;
            if (agent->tool_retry_count >= agent->max_tool_retries) {
                agent->state = AGENT_STATE_ERROR;
                agent_set_error(agent, "API request failed after retries");
                break;
            }
            if (agent->verbose) {
                fprintf(stderr, "[Agent] API error, retry %d/%d\n",
                    agent->tool_retry_count, agent->max_tool_retries);
            }
            json_data = agent_build_request(agent, NULL, NULL);
            agent->state = AGENT_STATE_RUNNING;
            continue;
        }

        agent->tool_retry_count = 0;

        struct json_object *message_obj = agent_get_message(choice);
        if (message_obj) {
            messages_add_object(agent->messages, json_object_get(message_obj));
        }

        char *content = agent_get_content(choice);
        if (content && *content) {
            // Print content immediately to the user
            extern void parse_markdown_to_ansi(const char *content);
            parse_markdown_to_ansi(content);
            printf("\n");

            size_t content_len = strlen(content);
            char *new_acc = realloc(accumulated_response, accumulated_len + content_len + 2);
            if (new_acc) {
                accumulated_response = new_acc;
                if (accumulated_len > 0) {
                    strcat(accumulated_response, "\n");
                    accumulated_len += 1;
                }
                strcpy(accumulated_response + accumulated_len, content);
                accumulated_len += content_len;
            }
        }

        bool has_tools = agent_has_tool_calls(choice);

        if (agent->verbose) {
            fprintf(stderr, "[Agent] has_tool_calls=%s\n", has_tools ? "true" : "false");
        }

        if (has_tools) {
            agent->state = AGENT_STATE_EXECUTING_TOOLS;

            struct json_object *tool_calls = agent_get_tool_calls(choice);

            if (agent->verbose) {
                int num_tools = json_object_array_length(tool_calls);
                fprintf(stderr, "[Agent] Executing %d tool(s)\n", num_tools);
            }

            struct json_object *results = tool_registry_execute(agent->tools, tool_calls, agent->verbose);

            int count = json_object_array_length(results);
            for (int i = 0; i < count; i++) {
                struct json_object *result = json_object_array_get_idx(results, i);
                messages_add_tool_call(agent->messages, json_object_get(result));
            }

            agent->state = AGENT_STATE_RUNNING;
            json_data = agent_build_request(agent, NULL, NULL);
            if (!json_data) {
                agent->state = AGENT_STATE_ERROR;
                agent_set_error(agent, "Failed to create follow-up JSON");
                free(content);
                break;
            }

        } else if (content && agent_response_indicates_incomplete(content)) {
            if (agent->verbose) {
                fprintf(stderr, "[Agent] Response indicates incomplete work, auto-continuing\n");
            }

            json_data = agent_build_request(agent, "user",
                "Continue. Execute the necessary actions to complete the task.");
            agent->state = AGENT_STATE_RUNNING;
            if (!json_data) {
                agent->state = AGENT_STATE_ERROR;
                agent_set_error(agent, "Failed to create continue JSON");
                free(content);
                break;
            }
        } else {
            agent->state = AGENT_STATE_COMPLETED;
            if (agent->verbose) {
                fprintf(stderr, "[Agent] Completed in %d iteration(s)\n",
                    agent->iteration_count);
            }
        }
        free(content);
    }

    free(json_data);
    return accumulated_response;
}

char *agent_chat(const char *user_message, messages_handle messages) {
    agent_handle agent = agent_create(user_message, messages);
    if (!agent) return NULL;

    char *response = agent_run(agent, user_message);

    if (agent->verbose && agent->state != AGENT_STATE_COMPLETED && agent->last_error) {
        fprintf(stderr, "[Agent] Error: %s\n", agent->last_error);
    }

    agent_destroy(agent);
    return response;
}

char *agent_chat_with_limit(const char *user_message, int max_iterations, messages_handle messages) {
    agent_handle agent = agent_create(user_message, messages);
    if (!agent) return NULL;

    agent_set_max_iterations(agent, max_iterations);
    char *response = agent_run(agent, user_message);

    if (agent->verbose && agent->state != AGENT_STATE_COMPLETED && agent->last_error) {
        fprintf(stderr, "[Agent] Error: %s\n", agent->last_error);
    }

    agent_destroy(agent);
    return response;
}