Atomizer/examples/interactive_research_session.py

"""
Interactive Research Agent Session

This example demonstrates real-time learning and interaction with the Research Agent.
Users can make requests, provide examples, and see the agent learn and generate code.

Author: Atomizer Development Team
Version: 0.1.0 (Phase 3)
Last Updated: 2025-01-16
"""

import sys
from pathlib import Path
from typing import Optional, Dict, Any

# Set UTF-8 encoding for Windows console
if sys.platform == 'win32':
    import codecs
    # Only wrap if not already wrapped
    if not isinstance(sys.stdout, codecs.StreamWriter):
        if hasattr(sys.stdout, 'buffer'):
            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
)


class InteractiveResearchSession:
    """Interactive session manager for Research Agent conversations."""

    def __init__(self, auto_mode: bool = False):
        self.agent = ResearchAgent()
        self.conversation_history = []
        self.current_gap: Optional[KnowledgeGap] = None
        self.current_findings: Optional[ResearchFindings] = None
        self.auto_mode = auto_mode  # For automated testing

    def print_header(self, text: str, char: str = "="):
        """Print formatted header."""
        print(f"\n{char * 80}")
        print(text)
        print(f"{char * 80}\n")

    def print_section(self, text: str):
        """Print section divider."""
        print(f"\n{'-' * 80}")
        print(text)
        print(f"{'-' * 80}\n")

    def display_knowledge_gap(self, gap: KnowledgeGap):
        """Display detected knowledge gap in user-friendly format."""
        print("  Knowledge Gap Analysis:")
        print(f"\n  Missing Features ({len(gap.missing_features)}):")
        for feature in gap.missing_features:
            print(f"    - {feature}")

        print(f"\n  Missing Knowledge ({len(gap.missing_knowledge)}):")
        for knowledge in gap.missing_knowledge:
            print(f"    - {knowledge}")

        print(f"\n  Confidence Level: {gap.confidence:.0%}")

        if gap.confidence < 0.5:
            print("  Status: New domain - Learning required")
        elif gap.confidence < 0.8:
            print("  Status: Partial knowledge - Some research needed")
        else:
            print("  Status: Known domain - Can reuse existing knowledge")

    def display_research_plan(self, plan):
        """Display research plan in user-friendly format."""
        # Handle both ResearchPlan objects and lists
        steps = plan.steps if hasattr(plan, 'steps') else plan

        print("  Research Plan Created:")
        print(f"\n  Will gather knowledge in {len(steps)} steps:\n")

        for i, step in enumerate(steps, 1):
            action = step['action'].replace('_', ' ').title()
            confidence = step['expected_confidence']

            print(f"  Step {i}: {action}")
            print(f"     Expected confidence: {confidence:.0%}")

            if 'details' in step:
                if 'prompt' in step['details']:
                    print(f"     What I'll ask: \"{step['details']['prompt'][:60]}...\"")
                elif 'query' in step['details']:
                    print(f"     Search query: \"{step['details']['query']}\"")
            print()

    def ask_for_example(self, prompt: str, file_types: list) -> Optional[str]:
        """Ask user for an example file or content."""
        print(f"  {prompt}\n")
        print(f"  Suggested file types: {', '.join(file_types)}\n")
        print("  Options:")
        print("    1. Enter file path to existing example")
        print("    2. Paste example content directly")
        print("    3. Skip (type 'skip')\n")

        user_input = input("  Your choice: ").strip()

        if user_input.lower() == 'skip':
            return None

        # Check if it's a file path
        file_path = Path(user_input)
        if file_path.exists() and file_path.is_file():
            try:
                content = file_path.read_text(encoding='utf-8')
                print(f"\n  Loaded {len(content)} characters from {file_path.name}")
                return content
            except Exception as e:
                print(f"\n  Error reading file: {e}")
                return None

        # Otherwise, treat as direct content
        if len(user_input) > 10:  # Minimum reasonable example size
            print(f"\n  Received {len(user_input)} characters of example content")
            return user_input

        print("\n  Input too short to be a valid example")
        return None

    def execute_research_plan(self, gap: KnowledgeGap) -> ResearchFindings:
        """Execute research plan interactively."""
        plan = self.agent.create_research_plan(gap)
        self.display_research_plan(plan)

        # Handle both ResearchPlan objects and lists
        steps = plan.steps if hasattr(plan, 'steps') else plan

        sources = {}
        raw_data = {}
        confidence_scores = {}

        for i, step in enumerate(steps, 1):
            action = step['action']

            print(f"\n  Executing Step {i}/{len(steps)}: {action.replace('_', ' ').title()}")
            print("  " + "-" * 76)

            if action == 'ask_user_for_example':
                prompt = step['details']['prompt']
                file_types = step['details'].get('suggested_file_types', ['.xml', '.py'])

                example_content = self.ask_for_example(prompt, file_types)

                if example_content:
                    sources['user_example'] = 'user_provided_example'
                    raw_data['user_example'] = example_content
                    confidence_scores['user_example'] = CONFIDENCE_LEVELS['user_validated']
                    print(f"  Step {i} completed with high confidence ({CONFIDENCE_LEVELS['user_validated']:.0%})")
                else:
                    print(f"  Step {i} skipped by user")

            elif action == 'search_knowledge_base':
                query = step['details']['query']
                print(f"  Searching knowledge base for: \"{query}\"")

                result = self.agent.search_knowledge_base(query)

                if result and result['confidence'] > 0.7:
                    sources['knowledge_base'] = result['session_id']
                    raw_data['knowledge_base'] = result
                    confidence_scores['knowledge_base'] = result['confidence']
                    print(f"  Found existing knowledge! Session: {result['session_id']}")
                    print(f"    Confidence: {result['confidence']:.0%}, Relevance: {result['relevance_score']:.0%}")
                else:
                    print(f"  No reliable existing knowledge found")

            elif action == 'query_nx_mcp':
                query = step['details']['query']
                print(f"  Would query NX MCP server: \"{query}\"")
                print(f"  ℹ️  (MCP integration pending - Phase 3)")
                confidence_scores['nx_mcp'] = 0.0  # Not yet implemented

            elif action == 'web_search':
                query = step['details']['query']
                print(f"  Would search web: \"{query}\"")
                print(f"  ℹ️  (Web search integration pending - Phase 3)")
                confidence_scores['web_search'] = 0.0  # Not yet implemented

            elif action == 'search_nxopen_tse':
                query = step['details']['query']
                print(f"  Would search NXOpen TSE: \"{query}\"")
                print(f"  ℹ️  (TSE search pending - Phase 3)")
                confidence_scores['tse_search'] = 0.0  # Not yet implemented

        return ResearchFindings(
            sources=sources,
            raw_data=raw_data,
            confidence_scores=confidence_scores
        )

    def display_learning_results(self, knowledge):
        """Display what the agent learned."""
        print("  Knowledge Synthesized:")
        print(f"\n  Overall Confidence: {knowledge.confidence:.0%}\n")

        if knowledge.schema:
            if 'xml_structure' in knowledge.schema:
                xml_schema = knowledge.schema['xml_structure']
                print(f"  Learned XML Structure:")
                print(f"    Root element: <{xml_schema['root_element']}>")

                if xml_schema.get('attributes'):
                    print(f"    Attributes: {xml_schema['attributes']}")

                print(f"    Required fields ({len(xml_schema['required_fields'])}):")
                for field in xml_schema['required_fields']:
                    print(f"      • {field}")

                if xml_schema.get('optional_fields'):
                    print(f"    Optional fields ({len(xml_schema['optional_fields'])}):")
                    for field in xml_schema['optional_fields']:
                        print(f"      • {field}")

        if knowledge.patterns:
            print(f"\n  Patterns Identified: {len(knowledge.patterns)}")
            if isinstance(knowledge.patterns, dict):
                for pattern_type, pattern_list in knowledge.patterns.items():
                    print(f"    {pattern_type}: {len(pattern_list)} found")
            else:
                print(f"    Total patterns: {len(knowledge.patterns)}")

    def generate_and_save_feature(self, feature_name: str, knowledge) -> Optional[Path]:
        """Generate feature code and save to file."""
        print(f"\n   Designing feature: {feature_name}")

        feature_spec = self.agent.design_feature(knowledge, feature_name)

        print(f"    Category: {feature_spec['category']}")
        print(f"    Lifecycle stage: {feature_spec['lifecycle_stage']}")
        print(f"    Input parameters: {len(feature_spec['interface']['inputs'])}")

        print(f"\n  Generating Python code...")

        generated_code = self.agent.generate_feature_code(feature_spec, knowledge)

        print(f"    Generated {len(generated_code)} characters ({len(generated_code.split(chr(10)))} lines)")

        # Validate syntax
        try:
            compile(generated_code, '<generated>', 'exec')
            print(f"    Code is syntactically valid Python")
        except SyntaxError as e:
            print(f"    Syntax error: {e}")
            return None

        # Save to file
        output_file = feature_spec['implementation']['file_path']
        output_path = project_root / output_file
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(generated_code, encoding='utf-8')

        print(f"\n  Saved to: {output_file}")

        return output_path

    def handle_request(self, user_request: str):
        """Handle a user request through the full research workflow."""
        self.print_header(f"Processing Request: {user_request[:60]}...")

        # Step 1: Identify knowledge gap
        self.print_section("[Step 1] Analyzing Knowledge Gap")
        gap = self.agent.identify_knowledge_gap(user_request)
        self.display_knowledge_gap(gap)

        self.current_gap = gap

        # Check if we can skip research
        if not gap.research_needed:
            print("\n  I already have the knowledge to handle this!")
            print("   Proceeding directly to generation...\n")
            # In a full implementation, would generate directly here
            return

        # Step 2: Execute research plan
        self.print_section("[Step 2] Executing Research Plan")
        findings = self.execute_research_plan(gap)
        self.current_findings = findings

        # Step 3: Synthesize knowledge
        self.print_section("[Step 3] Synthesizing Knowledge")
        knowledge = self.agent.synthesize_knowledge(findings)
        self.display_learning_results(knowledge)

        # Step 4: Generate feature
        if knowledge.confidence > 0.5:
            self.print_section("[Step 4] Generating Feature Code")

            # Extract feature name from request
            feature_name = user_request.lower().replace(' ', '_')[:30]
            if not feature_name.isidentifier():
                feature_name = "generated_feature"

            output_file = self.generate_and_save_feature(feature_name, knowledge)

            if output_file:
                # Step 5: Document session
                self.print_section("[Step 5] Documenting Research Session")

                topic = feature_name
                session_path = self.agent.document_session(
                    topic=topic,
                    knowledge_gap=gap,
                    findings=findings,
                    knowledge=knowledge,
                    generated_files=[str(output_file)]
                )

                print(f"   Session documented: {session_path.name}")
                print(f"    Files created:")
                for file in session_path.iterdir():
                    if file.is_file():
                        print(f"      • {file.name}")

                self.print_header("Request Completed Successfully!", "=")
                print(f"  Generated file: {output_file.relative_to(project_root)}")
                print(f"  Knowledge confidence: {knowledge.confidence:.0%}")
                print(f"  Session saved: {session_path.name}\n")
        else:
            print(f"\n  ️  Confidence too low ({knowledge.confidence:.0%}) to generate reliable code")
            print(f"   Try providing more examples or information\n")

    def run(self):
        """Run interactive session."""
        self.print_header("Interactive Research Agent Session", "=")

        print("  Welcome! I'm your Research Agent. I can learn from examples and")
        print("  generate code for optimization features.\n")
        print("  Commands:")
        print("    • Type your request in natural language")
        print("    • Type 'demo' for a demonstration")
        print("    • Type 'quit' to exit\n")

        while True:
            try:
                user_input = input("\nYour request: ").strip()

                if not user_input:
                    continue

                if user_input.lower() in ['quit', 'exit', 'q']:
                    print("\n   Goodbye! Session ended.\n")
                    break

                if user_input.lower() == 'demo':
                    self.run_demo()
                    continue

                # Process the request
                self.handle_request(user_input)

            except KeyboardInterrupt:
                print("\n\n   Goodbye! Session ended.\n")
                break
            except Exception as e:
                print(f"\n  Error: {e}")
                import traceback
                traceback.print_exc()

    def run_demo(self):
        """Run a demonstration of the Research Agent capabilities."""
        self.print_header("Research Agent Demonstration", "=")

        print("  This demo will show:")
        print("    1. Learning from a user example (material XML)")
        print("    2. Generating Python code from learned pattern")
        print("    3. Reusing knowledge for a second request\n")

        if not self.auto_mode:
            input("  Press Enter to start demo...")

        # Demo request 1: Learn from steel example
        demo_request_1 = "Create an NX material XML generator for steel"

        print(f"\n   Demo Request 1: \"{demo_request_1}\"\n")

        # Provide example automatically for demo
        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>
</PhysicalMaterial>"""

        print("  [Auto-providing example for demo]\n")

        gap1 = self.agent.identify_knowledge_gap(demo_request_1)
        self.display_knowledge_gap(gap1)

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

        knowledge1 = self.agent.synthesize_knowledge(findings1)
        self.display_learning_results(knowledge1)

        output_file1 = self.generate_and_save_feature("nx_material_generator_demo", knowledge1)

        if output_file1:
            print(f"\n  First request completed!")
            print(f"  Generated: {output_file1.name}\n")

        if not self.auto_mode:
            input("  Press Enter for second request (knowledge reuse demo)...")

        # Demo request 2: Reuse learned knowledge
        demo_request_2 = "Create aluminum 6061-T6 material XML"

        print(f"\n   Demo Request 2: \"{demo_request_2}\"\n")

        gap2 = self.agent.identify_knowledge_gap(demo_request_2)
        self.display_knowledge_gap(gap2)

        if gap2.confidence > 0.7:
            print("\n  Knowledge Reuse Success!")
            print("  I already learned the material XML structure from your first request.")
            print("  No need to ask for another example!\n")

        print("\n  Demo completed! Notice how:")
        print("    • First request: Low confidence, asked for example")
        print("    • Second request: High confidence, reused learned template")
        print("    • This is the power of learning and knowledge accumulation!\n")


def main():
    """Main entry point for interactive research session."""
    session = InteractiveResearchSession()
    session.run()


if __name__ == '__main__':
    main()