Code Representation Schema

The Code Representation Schema provides a structured and standardized way to define and manage executable code within agentic workflows. This schema facilitates clear separation of concerns, enabling robust parsing, manipulation, and execution of generated or provided code.

The `Code` Schema

The Code schema defines the fundamental structure for representing a code solution. It breaks down a complete code block into logical components, which aids in both human readability and programmatic processing.

Structure:

The Code schema is a Pydantic BaseModel with the following fields:

prefix (str): A descriptive string outlining the problem, the approach taken, or any contextual information relevant to the code. This field is useful for documentation, logging, or providing context to subsequent processing steps.
imports (str): A dedicated block for all necessary import statements. Separating imports ensures they are consistently placed and can be easily extracted or validated independently of the main logic.
code (str): The core executable code block, excluding any import statements. This field contains the primary logic of the solution.

Usage:

Developers use the Code schema to encapsulate generated or pre-defined code snippets. This structured approach is particularly beneficial in scenarios where code needs to be dynamically constructed, analyzed, or executed.

Example:

from code_runner_agent import Code

# Representing a simple Python function
my_code_solution = Code(
    prefix="This code defines a function to add two numbers.",
    imports="import math\nfrom typing import Union",
    code="""
def add_numbers(a: Union[int, float], b: Union[int, float]) -> Union[int, float]:
    return a + b
"""
)

print(f"Code Prefix: {my_code_solution.prefix}")
print(f"Code Imports:\n{my_code_solution.imports}")
print(f"Main Code Block:\n{my_code_solution.code}")

Agent State Management with `AgentState`

The AgentState schema captures the complete operational state of an agent at any given point in its execution. This includes communication history, the current code generation, and iteration details. The AgentState schema integrates the Code schema to manage the agent's generated solution.

Structure:

The AgentState schema is a Pydantic BaseModel with the following fields:

messages (list[dict[str, str]]): A chronological list of messages exchanged during the agent's operation. Each message is typically a dictionary containing role (e.g., "user", "assistant") and content.
generation (Code): An instance of the Code schema, representing the agent's current or most recent code solution. This field directly leverages the structured code representation.
iterations (int): A counter tracking the number of iterations or steps the agent has performed.
error (str): A string indicating any error encountered during the agent's process. A default value of "no" signifies no error.
output (Optional[str]): The result or output produced by executing the generation code, if applicable.

Integration and Usage:

The AgentState schema is central to managing the lifecycle of an agent. It provides a snapshot of the agent's progress, allowing for introspection, debugging, and continuation of workflows. The generation field, holding a Code object, is a key integration point, enabling agents to track and refine their code solutions.

Example:

from code_runner_agent import AgentState, Code

# Initialize an empty agent state
initial_state = AgentState()
print(f"Initial Agent State: {initial_state.model_dump_json(indent=2)}")

# Simulate an agent generating code
generated_code = Code(
    prefix="Initial attempt to greet the user.",
    imports="",
    code="""
def greet(name: str) -> str:
    return f"Hello, {name}!"
"""
)

# Update the agent state with the generated code and a message
updated_state = AgentState(
    messages=[{"role": "assistant", "content": "Here is my first code draft."}],
    generation=generated_code,
    iterations=1
)

print(f"\nUpdated Agent State:\n{updated_state.model_dump_json(indent=2)}")

# Accessing the generated code from the agent state
print(f"\nAgent's current code generation:\n{updated_state.generation.code}")

Practical Considerations and Best Practices

Serialization and Deserialization: Both Code and AgentState are Pydantic BaseModel instances. This allows for seamless serialization to and deserialization from JSON, making them ideal for persistence, inter-process communication, and API payloads. Use model_dump_json() for JSON output and parse_obj() or model_validate() for parsing.
Iterative Development: The AgentState schema is designed to support iterative agent workflows. Agents can update their generation field with refined Code objects, increment iterations, and log messages to track progress and changes over time.
Error Handling: The error field in AgentState provides a standardized way to communicate issues. Best practice involves updating this field with descriptive error messages when code execution fails or validation issues arise.
Modularity of Code: The separation of prefix, imports, and code in the Code schema encourages modularity. This structure simplifies tasks like linting, dependency analysis, or dynamic execution where imports might need to be handled separately from the main logic.
Contextual Information: Leverage the prefix field in the Code schema to provide crucial context. This is especially important in multi-turn agent interactions where the rationale behind a code solution might not be immediately obvious from the code itself.