Spaces:
Running
Running
File size: 5,010 Bytes
8ff817c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
"""
Coder Agent: Implements the plan as executable Python code.
This agent generates Python code that implements all steps in the current plan.
"""
from langchain_core.messages import AIMessage
from ..utils.code_execution import execute_with_debug
from ..utils.formatters import extract_code, format_data_descriptions, format_plan
from ..utils.state import DSStarState
def coder_node(state: DSStarState) -> dict:
"""
Coder Agent Node: Generates and executes Python code for the plan.
On first call: Generates code implementing all plan steps
On subsequent calls: Updates code to include new plan steps
Args:
state: Current DSStarState
Returns:
Dictionary with updated state fields:
- current_code: Generated Python code
- execution_result: Output from code execution
- messages: Agent communication messages
- next: Next node to visit ("verifier")
"""
print("=" * 60)
print("CODER AGENT STARTING...")
print("=" * 60)
data_context = format_data_descriptions(state["data_descriptions"])
plan_text = format_plan(state["plan"])
is_initial = state["current_code"] == ""
if is_initial:
print("Generating INITIAL code implementation...")
prompt = f"""You are an expert Python developer for data science.
Available Data Files:
{data_context}
Plan to Implement:
{plan_text}
Task: Write a Python script that implements ALL steps in the plan.
Requirements:
- Use pandas for data manipulation
- All files are in 'data/' directory
- Print intermediate results for each step
- No try-except blocks
- Clean, readable code
Provide ONLY the Python code in a markdown code block."""
else:
print(f"Updating code to implement {len(state['plan'])} steps...")
prompt = f"""You are an expert Python developer for data science.
Available Data Files:
{data_context}
Complete Plan:
{plan_text}
Previous Code:
{state["current_code"]}
Task: Update the code to implement the COMPLETE current plan.
Build upon the previous code, extending it to include all plan steps.
Requirements:
- Use pandas for data manipulation
- All files are in 'data/' directory
- Print intermediate and final results
- No try-except blocks
- Clean, readable code
Provide ONLY the updated Python code in a markdown code block."""
try:
# Get LLM response
response = state["llm"].invoke(prompt)
# Handle different response formats
if hasattr(response, "content") and isinstance(response.content, list):
from ..utils.formatters import gemini_text
response_text = gemini_text(response)
elif hasattr(response, "content"):
response_text = response.content
else:
response_text = str(response)
code = extract_code(response_text)
print("\nGenerated Code:")
print("-" * 60)
print(code[:200] + "..." if len(code) > 200 else code)
print("-" * 60)
print("\nExecuting code...")
# Execute with debugging
result = execute_with_debug(
code, state["llm"], is_analysis=False, data_context=data_context
)
print("\nExecution Result:")
print("-" * 60)
print(result[:200] + "..." if len(result) > 200 else result)
print("-" * 60)
print("=" * 60)
return {
"current_code": code,
"execution_result": result,
"messages": [AIMessage(content="Code executed")],
"next": "verifier",
}
except Exception as e:
print(f"\n✗ Coder error: {str(e)}")
return {
"messages": [AIMessage(content=f"Coder error: {str(e)}")],
"next": "__end__",
}
# Standalone test function
def test_coder(llm, query: str, data_descriptions: dict, plan: list):
"""
Test the coder agent independently.
Args:
llm: LLM instance
query: User query
data_descriptions: Dict of filename -> description
plan: List of plan steps to implement
Returns:
Dictionary with coder results
"""
# Create minimal test state
test_state = {
"llm": llm,
"query": query,
"data_descriptions": data_descriptions,
"plan": plan,
"current_code": "",
"execution_result": "",
"is_sufficient": False,
"router_decision": "",
"iteration": 0,
"max_iterations": 20,
"messages": [],
"next": "coder",
}
result = coder_node(test_state)
print("\n" + "=" * 60)
print("CODER TEST RESULTS")
print("=" * 60)
print("Code:")
print(result.get("current_code", "No code generated"))
print("\nExecution Result:")
print(result.get("execution_result", "No result"))
return result
|