Spaces:
Running
Running
File size: 4,949 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 |
"""
Finalyzer Agent: Creates final polished solution with clear output.
This agent runs when the verifier confirms the plan is sufficient.
It generates a final version of the code with improved formatting and output.
"""
from langchain_core.messages import AIMessage
from ..utils.code_execution import execute_code_safely
from ..utils.formatters import extract_code, format_data_descriptions, gemini_text
from ..utils.state import DSStarState
def finalyzer_node(state: DSStarState) -> dict:
"""
Finalyzer Agent Node: Creates final polished solution.
Takes the working code and creates a final version with:
- Clear answer to the original question
- Proper output formatting
- Self-contained executable code
Args:
state: Current DSStarState
Returns:
Dictionary with updated state fields:
- current_code: Final polished code
- execution_result: Final execution output
- messages: Agent communication messages
- next: "__end__" (workflow complete)
"""
print("=" * 60)
print("FINALYZER AGENT STARTING...")
print("=" * 60)
data_context = format_data_descriptions(state["data_descriptions"])
prompt = f"""You are an expert data analyst creating final solutions.
Original Question: {state["query"]}
Available Data:
{data_context}
Working Code:
{state["current_code"]}
Execution Result:
{state["execution_result"]}
Task: Create a final version of the code that:
1. Clearly prints the answer to the question
2. Includes proper formatting of the output
3. Is self-contained and executable
4. Has clear print statements
Provide ONLY the final 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):
response_text = gemini_text(response)
elif hasattr(response, "content"):
response_text = response.content
else:
response_text = str(response)
final_code = extract_code(response_text)
print("\nFinal Code Generated:")
print("-" * 60)
print(final_code[:300] + "..." if len(final_code) > 300 else final_code)
print("-" * 60)
# Execute final code
print("\nExecuting final code...")
success, stdout, stderr = execute_code_safely(final_code)
if success:
final_result = stdout
print("\n✓ Final execution successful")
else:
# If final execution fails, use previous result
print("\n⚠ Final execution failed, using previous result")
final_result = state["execution_result"]
print("\nFinal Result:")
print("-" * 60)
print(final_result[:300] + "..." if len(final_result) > 300 else final_result)
print("-" * 60)
print("=" * 60)
print("SOLUTION COMPLETE ✓")
print("=" * 60)
return {
"current_code": final_code,
"execution_result": final_result,
"messages": [AIMessage(content="Solution finalized")],
"next": "__end__",
}
except Exception as e:
# On error, return current state
print(f"\n✗ Finalyzer error: {str(e)}")
print("Using current solution as final")
return {
"messages": [
AIMessage(content=f"Finalyzer error: {str(e)}, using current solution")
],
"next": "__end__",
}
# Standalone test function
def test_finalyzer(
llm, query: str, data_descriptions: dict, current_code: str, execution_result: str
):
"""
Test the finalyzer agent independently.
Args:
llm: LLM instance
query: User query
data_descriptions: Dict of filename -> description
current_code: Working code to finalize
execution_result: Current execution result
Returns:
Dictionary with finalyzer results
"""
# Create minimal test state
test_state = {
"llm": llm,
"query": query,
"data_descriptions": data_descriptions,
"plan": [],
"current_code": current_code,
"execution_result": execution_result,
"is_sufficient": True,
"router_decision": "",
"iteration": 0,
"max_iterations": 20,
"messages": [],
"next": "finalyzer",
}
result = finalyzer_node(test_state)
print("\n" + "=" * 60)
print("FINALYZER TEST RESULTS")
print("=" * 60)
print("Final Code:")
print(result.get("current_code", "No code"))
print("\nFinal Result:")
print(result.get("execution_result", "No result"))
return result
|