agent_base#

Classes

AgentBase(task, job_id, task_id, config[, ...])

Base class for agents to work on a task.

BuildAndTestHandler()

Handler for evaluating agent-produced kernel folders.

EvaluateFunctionResult(tool_response, ...)

class kernelfoundry.algorithm.agent_base.EvaluateFunctionResult(tool_response, eval_result, program)[source]#
tool_response: dict#

Alias for field number 0

eval_result: EvalResult | None#

Alias for field number 1

program: Program | None#

Alias for field number 2

class kernelfoundry.algorithm.agent_base.BuildAndTestHandler[source]#

Handler for evaluating agent-produced kernel folders.

Subclass and override call() to customise evaluation logic, or compose a separate object for any post-session teardown logic.

Example:

from kernelfoundry.algorithm.agent_base import AgentBase, BuildAndTestHandler

class MyHandler(BuildAndTestHandler):
    def call(self, task, folder_path, job_id, task_id, prompt,
             iteration, branch, llm_model, session_log,
             previous_program=None):
        result = super().call(
            task, folder_path, job_id, task_id, prompt,
            iteration, branch, llm_model, session_log,
            previous_program,
        )
        # Optionally modify result here.
        return result

agent = MyAgent(task, job_id, build_test_handler=MyHandler())
call(task: Task, folder_path: str | Path, job_id: int, task_id: str, prompt: str, iteration: int, branch: int, llm_model: str, session_log: str, previous_program: Program | None = None, agent_session_id: str | None = None) → EvaluateFunctionResult[source]#

Evaluate an agent-produced folder and return an EvaluateFunctionResult.

Extracts the EVOLVE block from folder_path, combines it with task via with_blocks(), then runs the Evaluator on the result.

Parameters:

  • task – The base task whose EVOLVE block is replaced by the content in folder_path.

  • folder_path – Path to the folder produced by the agent (as passed to the build_and_test MCP tool).

  • job_id – Job ID associated with the current agent run.

  • task_id – Task identifier for the current run.

  • prompt – The prompt that was given to the agent.

  • iteration – The current iteration number passed to the agent’s run function.

  • branch – The branch identifier.

  • llm_model – The language model used by the agent, for logging purposes.

  • session_log – The session log for the current agent run.

  • previous_program – The previously generated program, if any.

  • agent_session_id – The session identifier of the agent that produced this kernel, if any.

Returns:

An EvaluateFunctionResult whose tool_response dict contains the keys expected by the build_and_test tool:

  • success (bool) — whether the kernel passed the correctness check

  • job_id (int) — the supplied job_id

  • eval_log (str) — the condensed evaluation log

  • runtime_stats (dict) — detailed runtime statistics from the evaluation

  • speedup (float) — runtime improvement compared to the reference implementation

class kernelfoundry.algorithm.agent_base.AgentBase(task: Task, job_id: int, task_id: str, config: dict, container_image: Image | None = None, initial_session_state: dict | None = None, build_test_handler: BuildAndTestHandler | None = None, branch: int = 0, parent_session_uuid: str | None = None, parent_program: Program | None = None, skills: list[Skill] | None = None)[source]#

Base class for agents to work on a task.

The agent autonomously generates solutions to the given task and evaluates them using the build_and_test tool provided by the MCP server.

Example usage pattern of the agent:

from kernelfoundry.algorithm.agent_base import AgentBase, BuildAndTestHandler

# Subclass BuildAndTestHandler to customise evaluation or post-session behaviour
class MyHandler(BuildAndTestHandler):
    def call(self, task, folder_path, job_id, task_id, prompt,
             iteration, branch, llm_model, session_log,
             previous_program=None):
        result = super().call(
            task, folder_path, job_id, task_id, prompt,
            iteration, branch, llm_model, session_log,
            previous_program=previous_program,
        )
        # Optionally modify result here.
        return result

    def session_end(self, session_log):
        print(session_log)

# Initialize the agent with a task and job ID
agent = MyAgent(task, job_id, build_test_handler=MyHandler())

ans = agent.run("Optimize the kernel for better performance.")

agent2 = agent.fork()  # Create a new instance continuing from the same session state
# These can run in separate threads to explore in parallel.
agent.run("Continue your optimization efforts.")
agent2.run("Focus on reducing memory bandwidth usage.")
__init__(task: Task, job_id: int, task_id: str, config: dict, container_image: Image | None = None, initial_session_state: dict | None = None, build_test_handler: BuildAndTestHandler | None = None, branch: int = 0, parent_session_uuid: str | None = None, parent_program: Program | None = None, skills: list[Skill] | None = None)[source]#

Initialize the agent with a starting point task and job ID.

Parameters:

  • task (Task) – The task for the agent to work on.

  • job_id (int) – The job ID associated with this agent.

  • config (dict) – The main configuration dictionary for the job.

  • container_image (Image | None) – Optional image for running the agent in a containerized environment.

  • initial_session_state (dict | None) – Optional session state to restore from a previous run, as returned by session_state(). When provided, the agent continues from that state rather than starting a fresh session.

  • build_test_handler – A BuildAndTestHandler instance whose call() method is invoked after every build_and_test tool call. Defaults to a plain BuildAndTestHandler instance.

  • branch – An integer identifier for the branch used for logging.

  • parent_session_uuid – The session UUID of the parent agent, if any.

  • parent_program – The program used as the parent for the next evaluation, if any.

  • skills – Optional list of Skill instances to make available to the agent. Any filtering is expected to be done by the caller.

set_parent_program(program: Program | None) → None[source]#

Set the parent program used as the parent for the next evaluation.

abstract run(prompt: str, iteration: int) → list[tuple[Program, EvalResult]][source]#

Run the agent with the given prompt and return a list of (Program, EvalResult) tuples.

Parameters:

  • prompt (str) – The input prompt for the agent to process.

  • iteration (int) – The current iteration number, starting from 0.

Returns:

A list of tuples containing the generated

program for each build_and_test call and the respective evaluation result.

Return type:

list[tuple[Program, EvalResult]]

abstract fork(branch: int, parent_program: Program | None = None) → AgentBase[source]#

Create a new instance of the agent with the same session state.

Parameters:

  • branch (int) – An integer identifier for the branch used for logging.

  • parent_program (Program | None) – The program to set as the parent for the forked agent’s next evaluation, if any.

This function must not be called while the agent is running.

abstract session_state() → dict[source]#

Return the current session state of the agent as a dictionary.

The returned dict is JSON serializable and should contain all necessary information to restore the agent’s state in a new instance. Note that the state is implementation-specific and may include binary data using Base64 encoding or similar approaches.