gateway/gwserver/test.py

"""
Test script for AgentService workflow with real implementations.
This script tests a workflow execution with a user query using actual AI service.
"""

import asyncio
import logging
import json
import os
import sys
import uuid
from datetime import datetime
from typing import Dict, Any, List, Optional

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)

logger = logging.getLogger(__name__)

# Define a custom WorkflowManager class for testing
class TestWorkflowManager:
    """
    Custom implementation of WorkflowManager for testing purposes.
    """
    
    def __init__(self, mandate_id: int = None, user_id: int = None, 
                ai_service = None, lucydom_interface = None):
        """Initialize the workflow manager"""
        self.mandate_id = mandate_id
        self.user_id = user_id
        self.ai_service = ai_service
        self.lucydom_interface = lucydom_interface
        self.workflows = {}
        self.results_dir = "./_results"
        os.makedirs(self.results_dir, exist_ok=True)
    
    async def create_workflow(self, message, files=None):
        """Create a new workflow with the given message and files."""
        workflow_id = f"wf_{datetime.now().strftime('%Y%m%d%H%M%S')}"
        
        # Initialize workflow
        workflow = self._initialize_workflow(workflow_id)
        logger.info(f"Created workflow with ID: {workflow_id}")
        
        # Create workflow execution
        from modules.agentservice_workflow_execution import WorkflowExecution
        execution = WorkflowExecution(
            self, workflow_id, self.mandate_id, self.user_id, 
            self.ai_service, self.lucydom_interface
        )
        
        # Execute workflow
        result = await execution.execute(message, workflow, files, True)
        return result
    
    def _initialize_workflow(self, workflow_id: str) -> Dict[str, Any]:
        """Initialize a new workflow."""
        current_time = datetime.now().isoformat()
        
        workflow = {
            "id": workflow_id,
            "name": f"Test Workflow {workflow_id}",
            "mandate_id": self.mandate_id,
            "user_id": self.user_id,
            "status": "running",
            "started_at": current_time,
            "last_activity": current_time,
            "messages": [],
            "logs": []
        }
        self.workflows[workflow_id] = workflow
        return workflow
    
    def _save_workflow(self, workflow: Dict[str, Any]) -> None:
        """Save workflow."""
        self.workflows[workflow["id"]] = workflow
        # Also save to LucyDOM
        if self.lucydom_interface:
            self.lucydom_interface.save_workflow_state(workflow)
    
    def _add_log(self, workflow_id, message, log_type, agent_id=None, agent_name=None):
        """Add log to workflow."""
        if workflow_id in self.workflows:
            if "logs" not in self.workflows[workflow_id]:
                self.workflows[workflow_id]["logs"] = []
            
            log_entry = {
                "id": f"log_{uuid.uuid4()}",
                "message": message,
                "type": log_type,
                "timestamp": datetime.now().isoformat(),
                "agent_id": agent_id,
                "agent_name": agent_name
            }
            
            self.workflows[workflow_id]["logs"].append(log_entry)
            logger.info(f"Log [{log_type}]: {message}")

async def run_test_workflow():
    """Run a test workflow with a user prompt."""
    try:
        # Import necessary modules
        from modules.agentservice_workflow_manager import get_workflow_manager
        from connectors.connector_aichat_openai import ChatService
        from modules.lucydom_interface import get_lucydom_interface
        
        # Initialize services
        # Set up mandate_id and user_id for testing
        mandate_id = 1
        user_id = 1
        
        # Initialize the AI service
        ai_service = ChatService()
        
        # Initialize the LucyDOM interface for database access
        lucydom_interface = get_lucydom_interface(mandate_id, user_id)
        
        # Initialize our custom workflow manager for testing
        workflow_manager = TestWorkflowManager(
            mandate_id=mandate_id,
            user_id=user_id,
            ai_service=ai_service,
            lucydom_interface=lucydom_interface
        )
        
        # Create a test message
        test_message = {
            "role": "user",
            "content": "Please analyze the CSV file and give me a summary of the data."
        }
        
        # Add a sample CSV file
        # For testing, let's create a real CSV file
        csv_content = b"name,age,location\nJohn,30,New York\nAlice,25,London\nBob,35,Paris\nEmma,28,Berlin"
        file_meta = lucydom_interface.save_uploaded_file(csv_content, "data.csv")
        logger.info(f"Uploaded test CSV file: {file_meta}")
        
        # List of files to include in the workflow
        files = [file_meta]
        
        # Execute the workflow with real implementation
        logger.info("Executing workflow...")
        result = await workflow_manager.create_workflow(test_message, files)
        
        # Print the result
        logger.info(f"Workflow execution completed with status: {result.get('status')}")
        logger.info(f"Workflow ID: {result.get('workflow_id')}")
        
        # Print messages in the workflow
        messages = result.get("messages", [])
        logger.info(f"Number of messages in workflow: {len(messages)}")
        
        for i, msg in enumerate(messages, 1):
            role = msg.get("role", "unknown")
            agent_type = msg.get("agent_type", "")
            agent_name = msg.get("agent_name", "")
            content = msg.get("content", "")
            
            agent_info = f" ({agent_name})" if agent_name else ""
            logger.info(f"Message {i} - {role}{agent_info}:")
            
            # Print first 100 chars of content
            if content:
                preview = content[:100] + "..." if len(content) > 100 else content
                logger.info(f"Content preview: {preview}")
            else:
                logger.info("No content")
            
            logger.info("-" * 40)
        
        # Save the complete result to a JSON file for inspection
        with open("workflow_test_result.json", "w") as f:
            json.dump(result, f, indent=2)
        logger.info("Workflow result saved to 'workflow_test_result.json'")
        
        # Clean up - close AI service
        await ai_service.close()
        
        return result
        
    except ImportError as e:
        logger.error(f"Import error: {e}")
        logger.error("This test script requires the actual modules to be in the Python path.")
        logger.error("You may need to adjust your PYTHONPATH to include the project directory.")
        return None
    except Exception as e:
        logger.error(f"Error running test workflow: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return None

if __name__ == "__main__":
    # Run the async function
    asyncio.run(run_test_workflow())