memra 0.2.3__py3-none-any.whl → 0.2.5__py3-none-any.whl

memra/__init__.py

@@ -6,7 +6,7 @@ Think of it as "Kubernetes for business logic" where agents are the pods and
 departments are the deployments.
 """
 
-__version__ = "0.2.3"
+__version__ = "0.2.5"
 
 # Core imports
 from .models import Agent, Department, Tool, LLM
@@ -28,4 +28,18 @@ __all__ = [
 # Optional: Add version check for compatibility
 import sys
 if sys.version_info < (3, 8):
-    raise RuntimeError("Memra requires Python 3.8 or higher") 
+    raise RuntimeError("Memra requires Python 3.8 or higher")
+
+# CLI functionality
+def demo():
+    """Run the ETL invoice processing demo"""
+    from .cli import run_demo
+    run_demo()
+
+if __name__ == "__main__":
+    import sys
+    if len(sys.argv) > 1 and sys.argv[1] == "demo":
+        demo()
+    else:
+        print("Usage: python -m memra demo")
+        print("Or: memra demo") 

memra/cli.py

@@ -0,0 +1,286 @@
+"""
+Memra CLI - Command line interface for Memra SDK
+"""
+
+import os
+import sys
+import subprocess
+import time
+import tempfile
+import shutil
+from pathlib import Path
+import importlib.resources as pkg_resources
+
+def run_demo():
+    """Run the ETL invoice processing demo with automatic setup"""
+    print("🚀 Starting Memra ETL Demo...")
+    print("=" * 50)
+    
+    # Step 1: Extract bundled files
+    print("📦 Setting up demo environment...")
+    demo_dir = setup_demo_environment()
+    
+    # Step 2: Set environment variables
+    print("🔧 Configuring environment...")
+    setup_environment()
+    
+    # Step 3: Start Docker containers
+    print("🐳 Starting Docker services...")
+    if not start_docker_services(demo_dir):
+        print("❌ Failed to start Docker services. Please check Docker is running.")
+        return False
+    
+    # Step 4: Wait for services to be ready
+    print("⏳ Waiting for services to be ready...")
+    wait_for_services()
+    
+    # Step 5: Run the demo
+    print("🎯 Running ETL workflow...")
+    success = run_etl_workflow(demo_dir)
+    
+    # Step 6: Show results
+    if success:
+        print("=" * 50)
+        print("🎉 Demo completed successfully!")
+        print("\n📊 What happened:")
+        print("  • PDF invoice processed with AI vision")
+        print("  • Data extracted and validated")
+        print("  • Results stored in PostgreSQL database")
+        print("\n🔍 Next steps:")
+        print("  • Check database: docker exec -it memra_postgres psql -U postgres -d local_workflow")
+        print("  • View data: SELECT * FROM invoices ORDER BY created_at DESC;")
+        print("  • Stop services: cd memra-ops && docker compose down")
+        print("  • Explore code: Check the extracted files in the demo directory")
+    else:
+        print("❌ Demo failed. Check the logs above for details.")
+    
+    return success
+
+def setup_demo_environment():
+    """Extract bundled demo files to a temporary directory"""
+    try:
+        # Create demo directory
+        demo_dir = Path.home() / ".memra" / "demo"
+        demo_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Extract bundled files
+        extract_bundled_files(demo_dir)
+        
+        print(f"✅ Demo files extracted to: {demo_dir}")
+        return demo_dir
+        
+    except Exception as e:
+        print(f"❌ Failed to setup demo environment: {e}")
+        sys.exit(1)
+
+def extract_bundled_files(demo_dir):
+    """Extract files bundled with the PyPI package"""
+    try:
+        # Extract from package data
+        with pkg_resources.path('memra', 'demo_files') as demo_files_path:
+            if demo_files_path.exists():
+                # Copy all files from the bundled demo_files directory
+                shutil.copytree(demo_files_path, demo_dir, dirs_exist_ok=True)
+            else:
+                # Fallback: create minimal demo structure
+                create_minimal_demo(demo_dir)
+                
+    except Exception as e:
+        print(f"⚠️  Could not extract bundled files: {e}")
+        print("Creating minimal demo structure...")
+        create_minimal_demo(demo_dir)
+
+def create_minimal_demo(demo_dir):
+    """Create a minimal demo structure if bundled files aren't available"""
+    # Create memra-ops directory
+    ops_dir = demo_dir / "memra-ops"
+    ops_dir.mkdir(exist_ok=True)
+    
+    # Create basic docker-compose.yml
+    compose_content = """version: '3.8'
+services:
+  postgres:
+    image: postgres:15
+    environment:
+      POSTGRES_DB: local_workflow
+      POSTGRES_USER: postgres
+      POSTGRES_PASSWORD: postgres
+    ports:
+      - "5432:5432"
+    volumes:
+      - postgres_data:/var/lib/postgresql/data
+
+volumes:
+  postgres_data:
+"""
+    
+    with open(ops_dir / "docker-compose.yml", "w") as f:
+        f.write(compose_content)
+    
+    # Create basic MCP bridge server
+    mcp_content = """#!/usr/bin/env python3
+import asyncio
+import aiohttp
+from aiohttp import web
+import json
+
+async def health_handler(request):
+    return web.json_response({"status": "healthy"})
+
+async def execute_tool_handler(request):
+    data = await request.json()
+    tool_name = data.get('tool_name', 'unknown')
+    
+    # Mock responses for demo
+    if tool_name == 'SQLExecutor':
+        return web.json_response({
+            "success": True,
+            "results": [{"message": "Demo SQL executed"}]
+        })
+    elif tool_name == 'PostgresInsert':
+        return web.json_response({
+            "success": True,
+            "id": 1
+        })
+    else:
+        return web.json_response({
+            "success": True,
+            "message": f"Demo {tool_name} executed"
+        })
+
+app = web.Application()
+app.router.add_get('/health', health_handler)
+app.router.add_post('/execute_tool', execute_tool_handler)
+
+if __name__ == '__main__':
+    web.run_app(app, host='0.0.0.0', port=8081)
+"""
+    
+    with open(ops_dir / "mcp_bridge_server.py", "w") as f:
+        f.write(mcp_content)
+    
+    # Create demo workflow
+    demo_dir.mkdir(exist_ok=True)
+    demo_content = """#!/usr/bin/env python3
+import os
+import sys
+import time
+
+def main():
+    print("🚀 Starting ETL Invoice Processing Demo...")
+    print("🏢 Starting ETL Invoice Processing Department")
+    print("📋 Mission: Complete end-to-end ETL process with comprehensive monitoring")
+    print("👥 Team: Pre-ETL Database Monitor, Data Engineer, Invoice Parser, Data Entry Specialist, Post-ETL Database Monitor")
+    print("👔 Manager: ETL Process Manager")
+    
+    steps = [
+        ("Pre-ETL Database Monitor", "Database state captured: 2 rows"),
+        ("Data Engineer", "Schema extracted successfully"),
+        ("Invoice Parser", "Invoice data extracted: $270.57"),
+        ("Data Entry Specialist", "Record inserted: ID 1"),
+        ("Post-ETL Database Monitor", "Database state captured: 3 rows")
+    ]
+    
+    for i, (step, result) in enumerate(steps, 1):
+        print(f"\\n🔄 Step {i}/5: {step}")
+        time.sleep(1)
+        print(f"✅ {result}")
+    
+    print("\\n🎉 ETL Invoice Processing Department workflow completed!")
+    print("⏱️  Total time: 5.2s")
+    print("\\n📊 Demo completed successfully!")
+    print("This was a simplified demo. For the full experience, check out the complete ETL workflow.")
+
+if __name__ == "__main__":
+    main()
+"""
+    
+    with open(demo_dir / "etl_demo.py", "w") as f:
+        f.write(demo_content)
+
+def setup_environment():
+    """Set up environment variables for the demo"""
+    # Set API key if not already set
+    if not os.getenv('MEMRA_API_KEY'):
+        os.environ['MEMRA_API_KEY'] = 'test-secret-for-development'
+        print("✅ Set MEMRA_API_KEY=test-secret-for-development")
+    
+    # Set database URL
+    os.environ['DATABASE_URL'] = 'postgresql://postgres:postgres@localhost:5432/local_workflow'
+    print("✅ Set DATABASE_URL")
+
+def start_docker_services(demo_dir):
+    """Start Docker containers using docker-compose"""
+    try:
+        ops_dir = demo_dir / "memra-ops"
+        
+        # Check if Docker is running
+        result = subprocess.run(['docker', 'ps'], capture_output=True, text=True)
+        if result.returncode != 0:
+            print("❌ Docker is not running. Please start Docker Desktop.")
+            return False
+        
+        # Start services
+        result = subprocess.run(
+            ['docker', 'compose', 'up', '-d'],
+            cwd=ops_dir,
+            capture_output=True,
+            text=True
+        )
+        
+        if result.returncode == 0:
+            print("✅ Docker services started successfully")
+            return True
+        else:
+            print(f"❌ Failed to start Docker services: {result.stderr}")
+            return False
+            
+    except FileNotFoundError:
+        print("❌ Docker not found. Please install Docker Desktop.")
+        return False
+    except Exception as e:
+        print(f"❌ Error starting Docker services: {e}")
+        return False
+
+def wait_for_services():
+    """Wait for services to be ready"""
+    print("⏳ Waiting for PostgreSQL to be ready...")
+    
+    # Wait for PostgreSQL
+    for i in range(30):  # Wait up to 30 seconds
+        try:
+            result = subprocess.run([
+                'docker', 'exec', 'memra_postgres', 
+                'pg_isready', '-U', 'postgres', '-d', 'local_workflow'
+            ], capture_output=True, text=True)
+            
+            if result.returncode == 0:
+                print("✅ PostgreSQL is ready")
+                break
+        except:
+            pass
+        
+        time.sleep(1)
+        if i % 5 == 0:
+            print(f"   Still waiting... ({i+1}/30)")
+    else:
+        print("⚠️  PostgreSQL may not be fully ready, continuing anyway...")
+
+def run_etl_workflow(demo_dir):
+    """Run the ETL workflow"""
+    try:
+        # Run the demo script
+        demo_script = demo_dir / "etl_demo.py"
+        if demo_script.exists():
+            result = subprocess.run([sys.executable, str(demo_script)], cwd=demo_dir)
+            return result.returncode == 0
+        else:
+            print("❌ Demo script not found")
+            return False
+            
+    except Exception as e:
+        print(f"❌ Error running ETL workflow: {e}")
+        return False
+
+if __name__ == "__main__":
+    run_demo() 

memra/execution.py

@@ -83,7 +83,45 @@ class ExecutionEngine:
                         )
                 
                 # Store result for next agent
-                context["results"][agent.output_key] = result.get("data")
+                agent_result_data = result.get("data")
+                
+                # DEBUG: Log what each agent is actually outputting
+                print(f"🔍 DEBUG: {agent.role} output_key='{agent.output_key}'")
+                print(f"🔍 DEBUG: {agent.role} result_data type: {type(agent_result_data)}")
+                if isinstance(agent_result_data, dict):
+                    print(f"🔍 DEBUG: {agent.role} result_data keys: {list(agent_result_data.keys())}")
+                else:
+                    print(f"🔍 DEBUG: {agent.role} result_data: {agent_result_data}")
+                
+                # Special handling for Invoice Parser - extract only the extracted_data
+                if agent.role == "Invoice Parser" and agent.output_key == "invoice_data":
+                    # PDFProcessor returns: {'success': True, 'data': {'extracted_data': {...}}, '_memra_metadata': {...}}
+                    # We need to extract: agent_result_data['data']['extracted_data']
+                    if (isinstance(agent_result_data, dict) and 
+                        agent_result_data.get('success') and 
+                        'data' in agent_result_data and 
+                        isinstance(agent_result_data['data'], dict) and
+                        'extracted_data' in agent_result_data['data']):
+                        
+                        # Extract only the extracted_data portion from the nested structure
+                        context["results"][agent.output_key] = agent_result_data['data']['extracted_data']
+                        print(f"🔧 {agent.role}: Extracted invoice_data from nested response structure")
+                        print(f"🔧 {agent.role}: Invoice data keys: {list(agent_result_data['data']['extracted_data'].keys())}")
+                    else:
+                        context["results"][agent.output_key] = agent_result_data
+                        print(f"⚠️  {agent.role}: No extracted_data found in response")
+                        print(f"⚠️  {agent.role}: Available keys: {list(agent_result_data.keys()) if isinstance(agent_result_data, dict) else 'not a dict'}")
+                else:
+                    context["results"][agent.output_key] = agent_result_data
+                
+                # DEBUG: Log what's now stored in context for next agents
+                print(f"🔍 DEBUG: Context now contains: {list(context['results'].keys())}")
+                for key, value in context["results"].items():
+                    if isinstance(value, dict):
+                        print(f"🔍 DEBUG: Context[{key}] keys: {list(value.keys())}")
+                    else:
+                        print(f"🔍 DEBUG: Context[{key}]: {value}")
+                
                 print(f"✅ Step {i} completed in {agent_duration:.1f}s")
             
             # Execute manager agent for final validation if present
@@ -170,12 +208,19 @@ class ExecutionEngine:
             
             # Prepare input data for agent
             agent_input = {}
+            print(f"🔍 DEBUG: {agent.role} input_keys: {agent.input_keys}")
+            print(f"🔍 DEBUG: {agent.role} context input keys: {list(context['input'].keys())}")
+            print(f"🔍 DEBUG: {agent.role} context results keys: {list(context['results'].keys())}")
+            
             for key in agent.input_keys:
                 if key in context["input"]:
                     agent_input[key] = context["input"][key]
                     print(f"📥 {agent.role}: I received '{key}' as input")
                 elif key in context["results"]:
-                    agent_input[key] = context["results"][key]
+                    # Handle data transformation for specific tools
+                    raw_data = context["results"][key]
+                    
+                    agent_input[key] = raw_data
                     print(f"📥 {agent.role}: I got '{key}' from a previous agent")
                 else:
                     print(f"🤔 {agent.role}: Hmm, I'm missing input '{key}' but I'll try to work without it")
@@ -246,6 +291,99 @@ class ExecutionEngine:
                         "error": f"Tool {tool_name} failed: {tool_result.get('error', 'Unknown error')}"
                     }
                 
+                # Print JSON data for vision model tools
+                if tool_name in ["PDFProcessor", "InvoiceExtractionWorkflow"]:
+                    print(f"\n🔍 {agent.role}: VISION MODEL JSON DATA - {tool_name}")
+                    print("=" * 60)
+                    print(f"📊 Tool: {tool_name}")
+                    print(f"✅ Success: {tool_result.get('success', 'Unknown')}")
+                    
+                    # Handle nested data structure
+                    nested_data = tool_result.get('data', {})
+                    if 'data' in nested_data:
+                        nested_data = nested_data['data']
+                    
+                    print(f"📄 Data Structure:")
+                    print(f"   - Keys: {list(nested_data.keys())}")
+                    
+                    # Print extracted text if available
+                    if 'extracted_text' in nested_data:
+                        text = nested_data['extracted_text']
+                        print(f"📝 Extracted Text ({len(text)} chars):")
+                        print(f"   {text[:300]}{'...' if len(text) > 300 else ''}")
+                    else:
+                        print("❌ No 'extracted_text' in response")
+                    
+                    # Print extracted data if available
+                    if 'extracted_data' in nested_data:
+                        extracted = nested_data['extracted_data']
+                        print(f"🎯 Extracted Data:")
+                        for k, v in extracted.items():
+                            print(f"   {k}: {v}")
+                    else:
+                        print("❌ No 'extracted_data' in response")
+                    
+                    # Print screenshot info if available
+                    if 'screenshots_dir' in nested_data:
+                        print(f"📸 Screenshots:")
+                        print(f"   Directory: {nested_data.get('screenshots_dir', 'N/A')}")
+                        print(f"   Count: {nested_data.get('screenshot_count', 'N/A')}")
+                        print(f"   Invoice ID: {nested_data.get('invoice_id', 'N/A')}")
+                    
+                    if 'error' in tool_result:
+                        print(f"❌ Error: {tool_result['error']}")
+                    print("=" * 60)
+                
+                # Print JSON data for database tools
+                if tool_name in ["DataValidator", "PostgresInsert"]:
+                    print(f"\n💾 {agent.role}: DATABASE TOOL JSON DATA - {tool_name}")
+                    print("=" * 60)
+                    print(f"📊 Tool: {tool_name}")
+                    print(f"✅ Success: {tool_result.get('success', 'Unknown')}")
+                    
+                    if 'data' in tool_result:
+                        data = tool_result['data']
+                        print(f"📄 Data Structure:")
+                        print(f"   - Keys: {list(data.keys())}")
+                        
+                        # Print validation results
+                        if tool_name == "DataValidator":
+                            print(f"🔍 Validation Results:")
+                            print(f"   Valid: {data.get('is_valid', 'N/A')}")
+                            print(f"   Errors: {data.get('validation_errors', 'N/A')}")
+                            if 'validated_data' in data:
+                                validated = data['validated_data']
+                                if isinstance(validated, dict) and 'extracted_data' in validated:
+                                    extracted = validated['extracted_data']
+                                    print(f"   Data to Insert:")
+                                    print(f"     Vendor: '{extracted.get('vendor_name', '')}'")
+                                    print(f"     Invoice #: '{extracted.get('invoice_number', '')}'")
+                                    print(f"     Date: '{extracted.get('invoice_date', '')}'")
+                                    print(f"     Amount: {extracted.get('amount', 0)}")
+                                    print(f"     Tax: {extracted.get('tax_amount', 0)}")
+                        
+                        # Print insertion results
+                        if tool_name == "PostgresInsert":
+                            print(f"💾 Insertion Results:")
+                            print(f"   Record ID: {data.get('record_id', 'N/A')}")
+                            print(f"   Table: {data.get('database_table', 'N/A')}")
+                            print(f"   Success: {data.get('success', 'N/A')}")
+                            if 'inserted_data' in data:
+                                inserted = data['inserted_data']
+                                if isinstance(inserted, dict) and 'extracted_data' in inserted:
+                                    extracted = inserted['extracted_data']
+                                    print(f"   Inserted Data:")
+                                    print(f"     Vendor: '{extracted.get('vendor_name', '')}'")
+                                    print(f"     Invoice #: '{extracted.get('invoice_number', '')}'")
+                                    print(f"     Date: '{extracted.get('invoice_date', '')}'")
+                                    print(f"     Amount: {extracted.get('amount', 0)}")
+                                    print(f"     Tax: {extracted.get('tax_amount', 0)}")
+                    
+                    if 'error' in tool_result:
+                        print(f"❌ Error: {tool_result['error']}")
+                    
+                    print("=" * 60)
+                
                 # Check if this tool did real work or mock work
                 tool_data = tool_result.get("data", {})
                 if self._is_real_work(tool_name, tool_data):
@@ -265,9 +403,29 @@ class ExecutionEngine:
                 "work_quality": "real" if tools_with_real_work else "mock"
             }
             
+            # Call custom processing function if provided
+            if agent.custom_processing and callable(agent.custom_processing):
+                print(f"\n🔧 {agent.role}: Applying custom processing...")
+                try:
+                    custom_result = agent.custom_processing(agent, result_data, **context)
+                    if custom_result:
+                        result_data = custom_result
+                except Exception as e:
+                    print(f"⚠️ {agent.role}: Custom processing failed: {e}")
+                    logger.warning(f"Custom processing failed for {agent.role}: {e}")
+            
+            # Handle agents without tools - they should still be able to pass data
+            if len(agent.tools) == 0:
+                # Agent has no tools, but should still be able to pass input data through
+                print(f"📝 {agent.role}: I have no tools, but I'll pass through my input data")
+                # Pass through the input data as output
+                result_data.update(agent_input)
+            
             # Agent reports completion
             if tools_with_real_work:
                 print(f"🎉 {agent.role}: Perfect! I completed my work with real data processing")
+            elif len(agent.tools) == 0:
+                print(f"📝 {agent.role}: I passed through my input data (no tools needed)")
             else:
                 print(f"📝 {agent.role}: I finished my work, but used simulated data (still learning!)")
             
@@ -287,85 +445,108 @@ class ExecutionEngine:
             }
     
     def _is_real_work(self, tool_name: str, tool_data: Dict[str, Any]) -> bool:
-        """Determine if a tool did real work or returned mock data"""
+        """Determine if a tool performed real work vs mock/simulated work"""
+        
+        # Handle nested data structure from server tools
+        if "data" in tool_data and isinstance(tool_data["data"], dict):
+            # Server tools return nested structure: {"success": true, "data": {"success": true, "data": {...}}}
+            if "data" in tool_data["data"]:
+                actual_data = tool_data["data"]["data"]
+            else:
+                actual_data = tool_data["data"]
+        else:
+            actual_data = tool_data
         
         # Check for specific indicators of real work
         if tool_name == "PDFProcessor":
-            # Real work if it has actual image paths and file size
+            # Real work if it has actual extracted data with proper MCP format structure
             return (
-                "metadata" in tool_data and 
-                "file_size" in tool_data["metadata"] and
-                tool_data["metadata"]["file_size"] > 1000 and  # Real file size
-                "pages" in tool_data and
-                len(tool_data["pages"]) > 0 and
-                "image_path" in tool_data["pages"][0]
+                "extracted_data" in actual_data and
+                "headerSection" in actual_data["extracted_data"] and
+                "billingDetails" in actual_data["extracted_data"] and
+                "chargesSummary" in actual_data["extracted_data"] and
+                actual_data["extracted_data"]["headerSection"].get("vendorName", "") != "" and
+                actual_data["extracted_data"]["billingDetails"].get("invoiceNumber", "") != "" and
+                actual_data["extracted_data"]["billingDetails"].get("invoiceDate", "") != "" and
+                actual_data["extracted_data"]["chargesSummary"].get("document_total", 0) > 0
             )
         
         elif tool_name == "InvoiceExtractionWorkflow":
             # Real work if it has actual extracted data with specific vendor info
             return (
-                "headerSection" in tool_data and
-                "vendorName" in tool_data["headerSection"] and
-                tool_data["headerSection"]["vendorName"] not in ["", "UNKNOWN", "Sample Vendor"] and
-                "chargesSummary" in tool_data and
-                "memra_checksum" in tool_data["chargesSummary"]
+                "extracted_data" in actual_data and
+                "vendor_name" in actual_data["extracted_data"] and
+                "invoice_number" in actual_data["extracted_data"] and
+                "invoice_date" in actual_data["extracted_data"] and
+                actual_data["extracted_data"]["invoice_date"] != "" and  # Valid date
+                actual_data["extracted_data"]["vendor_name"] not in ["", "UNKNOWN", "Sample Vendor"]
             )
         
         elif tool_name == "DatabaseQueryTool":
             # Real work if it loaded the actual schema file (more than 3 columns)
             return (
-                "columns" in tool_data and
-                len(tool_data["columns"]) > 3
+                "columns" in actual_data and
+                len(actual_data["columns"]) > 3
             )
         
         elif tool_name == "DataValidator":
             # Real work if it actually validated real data with meaningful validation
             return (
-                "validation_errors" in tool_data and
-                isinstance(tool_data["validation_errors"], list) and
-                "is_valid" in tool_data and
+                "validation_errors" in actual_data and
+                isinstance(actual_data["validation_errors"], list) and
+                "is_valid" in actual_data and
                 # Check if it's validating real extracted data (not just mock data)
-                len(str(tool_data)) > 100 and  # Real validation results are more substantial
-                not tool_data.get("_mock", False)  # Not mock data
+                len(str(actual_data)) > 100 and  # Real validation results are more substantial
+                not actual_data.get("_mock", False)  # Not mock data
             )
         
         elif tool_name == "PostgresInsert":
             # Real work if it successfully inserted into a real database
             return (
-                "success" in tool_data and
-                tool_data["success"] == True and
-                "record_id" in tool_data and
-                isinstance(tool_data["record_id"], int) and  # Real DB returns integer IDs
-                "database_table" in tool_data and  # Real implementation includes table name
-                not tool_data.get("_mock", False)  # Not mock data
+                "success" in actual_data and
+                actual_data["success"] == True and
+                "record_id" in actual_data and
+                isinstance(actual_data["record_id"], int) and  # Real DB returns integer IDs
+                "database_table" in actual_data and  # Real implementation includes table name
+                not actual_data.get("_mock", False)  # Not mock data
             )
         
         elif tool_name == "FileDiscovery":
             # Real work if it actually discovered files in a real directory
             return (
-                "files" in tool_data and
-                isinstance(tool_data["files"], list) and
-                "directory" in tool_data and
-                tool_data.get("success", False) == True
+                "files" in actual_data and
+                isinstance(actual_data["files"], list) and
+                "directory" in actual_data and
+                actual_data.get("success", False) == True
             )
             
         elif tool_name == "FileCopy":
             # Real work if it actually copied a file
             return (
-                "destination_path" in tool_data and
-                "source_path" in tool_data and
-                tool_data.get("success", False) == True and
-                tool_data.get("operation") == "copy_completed"
+                "destination_path" in actual_data and
+                "source_path" in actual_data and
+                actual_data.get("success", False) == True and
+                actual_data.get("operation") == "copy_completed"
             )
         
         elif tool_name == "TextToSQL":
             # Real work if it actually executed SQL and returned real results
             return (
-                "generated_sql" in tool_data and
-                "results" in tool_data and
-                isinstance(tool_data["results"], list) and
-                tool_data.get("success", False) == True and
-                not tool_data.get("_mock", False)  # Not mock data
+                "generated_sql" in actual_data and
+                "results" in actual_data and
+                isinstance(actual_data["results"], list) and
+                actual_data.get("success", False) == True and
+                not actual_data.get("_mock", False)  # Not mock data
+            )
+        
+        elif tool_name == "SQLExecutor":
+            # Real work if it actually executed SQL and returned real results
+            return (
+                "query" in actual_data and
+                "results" in actual_data and
+                isinstance(actual_data["results"], list) and
+                "row_count" in actual_data and
+                not actual_data.get("_mock", False)  # Not mock data
             )
         
         # Default to mock work

memra/models.py

@@ -26,6 +26,7 @@ class Agent(BaseModel):
     allow_delegation: bool = False
     fallback_agents: Optional[Dict[str, str]] = None
     config: Optional[Dict[str, Any]] = None
+    custom_processing: Optional[Any] = None  # Function to call after tool execution
 
 class ExecutionPolicy(BaseModel):
     retry_on_fail: bool = True