"""
Test Feature Code Generation Pipeline
This test demonstrates the Research Agent's ability to:
1. Learn from a user-provided example (XML material file)
2. Extract schema and patterns
3. Design a feature specification
4. Generate working Python code from the learned template
5. Save the generated code to a file
Author: Atomizer Development Team
Version: 0.1.0 (Phase 2 Week 2)
Last Updated: 2025-01-16
"""
import sys
from pathlib import Path
# Set UTF-8 encoding for Windows console
if sys.platform == 'win32':
import codecs
sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer, errors='replace')
sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer, errors='replace')
# Add project root to path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))
from optimization_engine.research_agent import (
ResearchAgent,
ResearchFindings,
CONFIDENCE_LEVELS
)
def test_code_generation():
"""Test complete code generation workflow from example to working code."""
print("\n" + "="*80)
print("FEATURE CODE GENERATION TEST")
print("="*80)
agent = ResearchAgent()
# Step 1: User provides material XML example
print("\n" + "-"*80)
print("[Step 1] User Provides Example Material XML")
print("-"*80)
example_xml = """
7850
200
0.29
1.17e-05
295
"""
print("\n Example XML (steel material):")
for line in example_xml.split('\n')[:4]:
print(f" {line}")
print(" ...")
# Step 2: Agent learns from example
print("\n" + "-"*80)
print("[Step 2] Agent Learns Schema from Example")
print("-"*80)
findings = ResearchFindings(
sources={'user_example': 'steel_material.xml'},
raw_data={'user_example': example_xml},
confidence_scores={'user_example': CONFIDENCE_LEVELS['user_validated']}
)
knowledge = agent.synthesize_knowledge(findings)
print(f"\n Learned schema:")
if knowledge.schema and 'xml_structure' in knowledge.schema:
xml_schema = knowledge.schema['xml_structure']
print(f" Root element: {xml_schema['root_element']}")
print(f" Attributes: {xml_schema.get('attributes', {})}")
print(f" Required fields ({len(xml_schema['required_fields'])}):")
for field in xml_schema['required_fields']:
print(f" - {field}")
print(f"\n Confidence: {knowledge.confidence:.2f}")
# Step 3: Design feature specification
print("\n" + "-"*80)
print("[Step 3] Design Feature Specification")
print("-"*80)
feature_name = "nx_material_generator"
feature_spec = agent.design_feature(knowledge, feature_name)
print(f"\n Feature designed:")
print(f" Feature ID: {feature_spec['feature_id']}")
print(f" Category: {feature_spec['category']}")
print(f" Subcategory: {feature_spec['subcategory']}")
print(f" Lifecycle stage: {feature_spec['lifecycle_stage']}")
print(f" Implementation file: {feature_spec['implementation']['file_path']}")
print(f" Number of inputs: {len(feature_spec['interface']['inputs'])}")
print(f"\n Input parameters:")
for input_param in feature_spec['interface']['inputs']:
print(f" - {input_param['name']}: {input_param['type']}")
# Step 4: Generate Python code
print("\n" + "-"*80)
print("[Step 4] Generate Python Code from Learned Template")
print("-"*80)
generated_code = agent.generate_feature_code(feature_spec, knowledge)
print(f"\n Generated {len(generated_code)} characters of Python code")
print(f"\n Code preview (first 20 lines):")
print(" " + "-"*76)
for i, line in enumerate(generated_code.split('\n')[:20]):
print(f" {line}")
print(" " + "-"*76)
print(f" ... ({len(generated_code.split(chr(10)))} total lines)")
# Step 5: Validate generated code
print("\n" + "-"*80)
print("[Step 5] Validate Generated Code")
print("-"*80)
# Check that code has necessary components
validations = [
('Function definition', f'def {feature_name}(' in generated_code),
('Docstring', '"""' in generated_code),
('Type hints', ('-> Dict[str, Any]' in generated_code or ': float' in generated_code)),
('XML Element handling', 'ET.Element' in generated_code),
('Return statement', 'return {' in generated_code),
('Example usage', 'if __name__ == "__main__":' in generated_code)
]
all_valid = True
print("\n Code validation:")
for check_name, passed in validations:
status = "✓" if passed else "✗"
print(f" {status} {check_name}")
if not passed:
all_valid = False
assert all_valid, "Generated code is missing required components"
# Step 6: Save generated code to file
print("\n" + "-"*80)
print("[Step 6] Save Generated Code")
print("-"*80)
# Create custom_functions directory if it doesn't exist
custom_functions_dir = project_root / "optimization_engine" / "custom_functions"
custom_functions_dir.mkdir(parents=True, exist_ok=True)
output_file = custom_functions_dir / f"{feature_name}.py"
output_file.write_text(generated_code, encoding='utf-8')
print(f"\n Code saved to: {output_file}")
print(f" File size: {output_file.stat().st_size} bytes")
print(f" Lines of code: {len(generated_code.split(chr(10)))}")
# Step 7: Test that code is syntactically valid Python
print("\n" + "-"*80)
print("[Step 7] Verify Code is Valid Python")
print("-"*80)
try:
compile(generated_code, '', 'exec')
print("\n ✓ Code compiles successfully!")
print(" Generated code is syntactically valid Python")
except SyntaxError as e:
print(f"\n ✗ Syntax error: {e}")
assert False, "Generated code has syntax errors"
# Summary
print("\n" + "="*80)
print("CODE GENERATION TEST SUMMARY")
print("="*80)
print("\n Workflow Completed:")
print(" ✓ User provided example XML")
print(" ✓ Agent learned schema (5 fields)")
print(" ✓ Feature specification designed")
print(f" ✓ Python code generated ({len(generated_code)} chars)")
print(f" ✓ Code saved to {output_file.name}")
print(" ✓ Code is syntactically valid Python")
print("\n What This Demonstrates:")
print(" - Agent can learn from a single example")
print(" - Schema extraction works correctly")
print(" - Code generation follows learned patterns")
print(" - Generated code has proper structure (docstrings, type hints, examples)")
print(" - Output is ready to use (valid Python)")
print("\n Next Steps (in real usage):")
print(" 1. User tests the generated function")
print(" 2. User provides feedback if adjustments needed")
print(" 3. Agent refines code based on feedback")
print(" 4. Feature gets added to feature registry")
print(" 5. Future requests use this template automatically")
print("\n" + "="*80)
print("Code Generation: SUCCESS! ✓")
print("="*80 + "\n")
return True
if __name__ == '__main__':
try:
success = test_code_generation()
sys.exit(0 if success else 1)
except Exception as e:
print(f"\n[ERROR] {e}")
import traceback
traceback.print_exc()
sys.exit(1)