Atomizer/tests/demo_research_agent.py

"""
Quick Interactive Demo of Research Agent

This demo shows the Research Agent learning from a material XML example
and documenting the research session.

Run this to see Phase 2 in action!
"""

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,
    KnowledgeGap,
    CONFIDENCE_LEVELS
)


def main():
    print("\n" + "="*70)
    print("  RESEARCH AGENT DEMO - Phase 2 Self-Learning System")
    print("="*70)

    # Initialize agent
    agent = ResearchAgent()
    print("\n[1] Research Agent initialized")
    print(f"    Feature registry loaded: {agent.feature_registry_path}")
    print(f"    Knowledge base: {agent.knowledge_base_path}")

    # Test 1: Detect knowledge gap
    print("\n" + "-"*70)
    print("[2] Testing Knowledge Gap Detection")
    print("-"*70)

    request = "Create NX material XML for titanium Ti-6Al-4V"
    print(f"\nUser request: \"{request}\"")

    gap = agent.identify_knowledge_gap(request)
    print(f"\n  Analysis:")
    print(f"    Missing features: {gap.missing_features}")
    print(f"    Missing knowledge: {gap.missing_knowledge}")
    print(f"    Confidence: {gap.confidence:.2f}")
    print(f"    Research needed: {gap.research_needed}")

    # Test 2: Learn from example
    print("\n" + "-"*70)
    print("[3] Learning from User Example")
    print("-"*70)

    # Simulated user provides this example
    example_xml = """<?xml version="1.0" encoding="UTF-8"?>
<PhysicalMaterial name="Steel_AISI_1020" version="1.0">
    <Density units="kg/m3">7850</Density>
    <YoungModulus units="GPa">200</YoungModulus>
    <PoissonRatio>0.29</PoissonRatio>
    <ThermalExpansion units="1/K">1.17e-05</ThermalExpansion>
    <YieldStrength units="MPa">295</YieldStrength>
    <UltimateTensileStrength units="MPa">420</UltimateTensileStrength>
</PhysicalMaterial>"""

    print("\nUser provides example: steel_material.xml")
    print("  (Simulating user uploading a file)")

    # Create research findings
    findings = ResearchFindings(
        sources={'user_example': 'steel_material.xml'},
        raw_data={'user_example': example_xml},
        confidence_scores={'user_example': CONFIDENCE_LEVELS['user_validated']}
    )

    print(f"\n  Source: user_example")
    print(f"  Confidence: {CONFIDENCE_LEVELS['user_validated']:.2f} (user-validated)")

    # Test 3: Synthesize knowledge
    print("\n" + "-"*70)
    print("[4] Synthesizing Knowledge")
    print("-"*70)

    knowledge = agent.synthesize_knowledge(findings)

    print(f"\n  {knowledge.synthesis_notes}")

    if knowledge.schema and 'xml_structure' in knowledge.schema:
        xml_schema = knowledge.schema['xml_structure']
        print(f"\n  Learned Schema:")
        print(f"    Root element: {xml_schema['root_element']}")
        print(f"    Required fields: {len(xml_schema['required_fields'])}")
        for field in xml_schema['required_fields'][:3]:
            print(f"      - {field}")
        if len(xml_schema['required_fields']) > 3:
            print(f"      ... and {len(xml_schema['required_fields']) - 3} more")

    # Test 4: Document session
    print("\n" + "-"*70)
    print("[5] Documenting Research Session")
    print("-"*70)

    session_path = agent.document_session(
        topic='nx_materials_demo',
        knowledge_gap=gap,
        findings=findings,
        knowledge=knowledge,
        generated_files=[
            'optimization_engine/custom_functions/nx_material_generator.py',
            'knowledge_base/templates/material_xml_template.py'
        ]
    )

    print(f"\n  Session saved to:")
    print(f"    {session_path}")

    print(f"\n  Files created:")
    for file in ['user_question.txt', 'sources_consulted.txt', 'findings.md', 'decision_rationale.md']:
        file_path = session_path / file
        if file_path.exists():
            print(f"    [OK] {file}")
        else:
            print(f"    [MISSING] {file}")

    # Show content of findings
    print("\n  Preview of findings.md:")
    findings_path = session_path / 'findings.md'
    if findings_path.exists():
        content = findings_path.read_text(encoding='utf-8')
        for i, line in enumerate(content.split('\n')[:12]):
            print(f"    {line}")
        print("    ...")

    # Test 5: Now agent can generate materials
    print("\n" + "-"*70)
    print("[6] Agent is Now Ready to Generate Materials!")
    print("-"*70)

    print("\n  Next time you request a material XML, the agent will:")
    print("    1. Search knowledge base and find this research session")
    print("    2. Retrieve the learned schema")
    print("    3. Generate new material XML following the pattern")
    print("    4. Confidence: HIGH (based on user-validated example)")

    print("\n  Example usage:")
    print('    User: "Create aluminum alloy 6061-T6 material XML"')
    print('    Agent: "I know how to do this! Using learned schema..."')
    print('           [Generates XML with Al 6061-T6 properties]')

    # Summary
    print("\n" + "="*70)
    print("  DEMO COMPLETE - Research Agent Successfully Learned!")
    print("="*70)

    print("\n  What was accomplished:")
    print("    [OK] Detected knowledge gap (material XML generation)")
    print("    [OK] Learned XML schema from user example")
    print("    [OK] Extracted reusable patterns")
    print("    [OK] Documented research session for future reference")
    print("    [OK] Ready to generate similar features autonomously")

    print("\n  Knowledge persisted in:")
    print(f"    {session_path}")

    print("\n  This demonstrates Phase 2: Self-Extending Research System")
    print("  The agent can now learn ANY new capability from examples!\n")


if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        print(f"\n[ERROR] {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)