aiecs 1.5.1__py3-none-any.whl

aiecs/tools/docs/ai_document_orchestrator.py

@@ -0,0 +1,599 @@
+import asyncio
+import logging
+from typing import Dict, Any, List, Optional, Callable
+from enum import Enum
+from datetime import datetime
+
+from pydantic import BaseModel, Field, ConfigDict
+
+from aiecs.tools.base_tool import BaseTool
+from aiecs.tools import register_tool
+
+
+class ProcessingMode(str, Enum):
+    """AI document processing modes"""
+
+    SUMMARIZE = "summarize"
+    EXTRACT_INFO = "extract_info"
+    ANALYZE = "analyze"
+    TRANSLATE = "translate"
+    CLASSIFY = "classify"
+    ANSWER_QUESTIONS = "answer_questions"
+    CUSTOM = "custom"
+
+
+class AIProvider(str, Enum):
+    """Supported AI providers"""
+
+    OPENAI = "openai"
+    VERTEX_AI = "vertex_ai"
+    XAI = "xai"
+    LOCAL = "local"
+
+
+class AIDocumentOrchestratorError(Exception):
+    """Base exception for AI Document Orchestrator errors"""
+
+
+class AIProviderError(AIDocumentOrchestratorError):
+    """Raised when AI provider operations fail"""
+
+
+class ProcessingError(AIDocumentOrchestratorError):
+    """Raised when document processing fails"""
+
+
+@register_tool("ai_document_orchestrator")
+class AIDocumentOrchestrator(BaseTool):
+    """
+    AI-powered document processing orchestrator that:
+    1. Coordinates document parsing with AI analysis
+    2. Manages AI provider interactions
+    3. Handles complex document processing workflows
+    4. Provides intelligent content analysis and extraction
+
+    Integrates with:
+    - DocumentParserTool for document parsing
+    - Various AI providers for content analysis
+    - Existing AIECS infrastructure
+    """
+
+    # Configuration schema
+    class Config(BaseModel):
+        """Configuration for the AI document orchestrator tool"""
+
+        model_config = ConfigDict(env_prefix="AI_DOC_ORCHESTRATOR_")
+
+        default_ai_provider: str = Field(default="openai", description="Default AI provider to use")
+        max_chunk_size: int = Field(
+            default=4000, description="Maximum chunk size for AI processing"
+        )
+        max_concurrent_requests: int = Field(
+            default=5, description="Maximum concurrent AI requests"
+        )
+        default_temperature: float = Field(
+            default=0.1, description="Default temperature for AI model"
+        )
+        max_tokens: int = Field(default=2000, description="Maximum tokens for AI response")
+        timeout: int = Field(default=60, description="Timeout in seconds for AI operations")
+
+    def __init__(self, config: Optional[Dict] = None):
+        """Initialize AI Document Orchestrator with settings"""
+        super().__init__(config)
+
+        # Parse configuration
+        self.config = self.Config(**(config or {}))
+
+        self.logger = logging.getLogger(__name__)
+
+        # Initialize document parser
+        self._init_document_parser()
+
+        # Initialize AI providers
+        self._init_ai_providers()
+
+        # Processing templates
+        self._init_processing_templates()
+
+    def _init_document_parser(self):
+        """Initialize document parser tool"""
+        try:
+            from aiecs.tools.docs.document_parser_tool import (
+                DocumentParserTool,
+            )
+
+            self.document_parser = DocumentParserTool()
+        except ImportError:
+            self.logger.error("DocumentParserTool not available")
+            self.document_parser = None
+
+    def _init_ai_providers(self):
+        """Initialize AI providers"""
+        self.ai_providers = {}
+
+        try:
+            # Initialize AIECS client for AI operations
+            from aiecs import AIECS
+
+            self.aiecs_client = AIECS()
+            self.ai_providers["aiecs"] = self.aiecs_client
+        except ImportError:
+            self.logger.warning("AIECS client not available")
+            self.aiecs_client = None
+
+    def _init_processing_templates(self):
+        """Initialize processing templates for different AI tasks"""
+        self.processing_templates = {
+            ProcessingMode.SUMMARIZE: {
+                "system_prompt": "You are an expert document summarizer. Create concise, informative summaries.",
+                "user_prompt_template": "Please summarize the following document content:\n\n{content}\n\nProvide a clear, structured summary highlighting the key points.",
+            },
+            ProcessingMode.EXTRACT_INFO: {
+                "system_prompt": "You are an expert information extractor. Extract specific information from documents.",
+                "user_prompt_template": "Extract the following information from the document:\n{extraction_criteria}\n\nDocument content:\n{content}\n\nProvide the extracted information in a structured format.",
+            },
+            ProcessingMode.ANALYZE: {
+                "system_prompt": "You are an expert document analyzer. Provide thorough analysis of document content.",
+                "user_prompt_template": "Analyze the following document content and provide insights:\n\n{content}\n\nInclude analysis of:\n- Main themes and topics\n- Key findings\n- Important details\n- Overall structure and organization",
+            },
+            ProcessingMode.TRANSLATE: {
+                "system_prompt": "You are an expert translator. Provide accurate translations while preserving meaning and context.",
+                "user_prompt_template": "Translate the following document content to {target_language}:\n\n{content}\n\nMaintain the original structure and formatting where possible.",
+            },
+            ProcessingMode.CLASSIFY: {
+                "system_prompt": "You are an expert document classifier. Classify documents accurately based on their content.",
+                "user_prompt_template": "Classify the following document content into the appropriate categories:\n\nCategories: {categories}\n\nDocument content:\n{content}\n\nProvide the classification with confidence scores and reasoning.",
+            },
+            ProcessingMode.ANSWER_QUESTIONS: {
+                "system_prompt": "You are an expert document analyst. Answer questions based on document content accurately.",
+                "user_prompt_template": "Based on the following document content, answer these questions:\n\nQuestions:\n{questions}\n\nDocument content:\n{content}\n\nProvide clear, accurate answers with references to the relevant parts of the document.",
+            },
+            ProcessingMode.CUSTOM: {
+                "system_prompt": "You are an expert document analyst. Follow the custom instructions provided.",
+                "user_prompt_template": "{custom_prompt}\n\nDocument content:\n{content}\n\nPlease provide your analysis based on the custom instructions above.",
+            },
+        }
+
+    # Schema definitions
+    class ProcessDocumentSchema(BaseModel):
+        """Schema for process_document operation"""
+
+        source: str = Field(description="URL or file path to the document")
+        processing_mode: ProcessingMode = Field(description="AI processing mode to apply")
+        ai_provider: Optional[AIProvider] = Field(default=None, description="AI provider to use")
+        processing_params: Optional[Dict[str, Any]] = Field(
+            default=None, description="Additional processing parameters"
+        )
+        parse_params: Optional[Dict[str, Any]] = Field(
+            default=None, description="Document parsing parameters"
+        )
+        ai_params: Optional[Dict[str, Any]] = Field(
+            default=None, description="AI provider parameters"
+        )
+
+    class BatchProcessSchema(BaseModel):
+        """Schema for batch_process_documents operation"""
+
+        sources: List[str] = Field(description="List of URLs or file paths")
+        processing_mode: ProcessingMode = Field(description="AI processing mode to apply")
+        ai_provider: Optional[AIProvider] = Field(default=None, description="AI provider to use")
+        processing_params: Optional[Dict[str, Any]] = Field(
+            default=None, description="Additional processing parameters"
+        )
+        max_concurrent: Optional[int] = Field(
+            default=None, description="Maximum concurrent processing"
+        )
+
+    class AnalyzeDocumentSchema(BaseModel):
+        """Schema for analyze_document operation (AI-first approach)"""
+
+        source: str = Field(description="URL or file path to the document")
+        analysis_type: str = Field(description="Type of analysis to perform")
+        custom_prompt: Optional[str] = Field(
+            default=None, description="Custom AI prompt for analysis"
+        )
+        ai_provider: Optional[AIProvider] = Field(default=None, description="AI provider to use")
+
+    def process_document(
+        self,
+        source: str,
+        processing_mode: ProcessingMode,
+        ai_provider: Optional[AIProvider] = None,
+        processing_params: Optional[Dict[str, Any]] = None,
+        parse_params: Optional[Dict[str, Any]] = None,
+        ai_params: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Process a document using AI with intelligent orchestration
+
+        Args:
+            source: URL or file path to document
+            processing_mode: AI processing mode to apply
+            ai_provider: AI provider to use (optional)
+            processing_params: Additional processing parameters
+            parse_params: Document parsing parameters
+            ai_params: AI provider parameters
+
+        Returns:
+            Dict containing processed results and metadata
+        """
+        try:
+            start_time = datetime.now()
+
+            # Step 1: Parse the document
+            self.logger.info(f"Starting document processing: {source}")
+            parsed_result = self._parse_document(source, parse_params or {})
+
+            # Step 2: Prepare content for AI processing
+            content = self._prepare_content_for_ai(parsed_result, processing_mode)
+
+            # Step 3: Process with AI
+            ai_result = self._process_with_ai(
+                content,
+                processing_mode,
+                ai_provider or self.config.default_ai_provider,
+                processing_params or {},
+                ai_params or {},
+            )
+
+            # Step 4: Combine results
+            result = {
+                "source": source,
+                "processing_mode": processing_mode,
+                "ai_provider": ai_provider or self.config.default_ai_provider,
+                "document_info": {
+                    "type": parsed_result.get("document_type"),
+                    "detection_confidence": parsed_result.get("detection_confidence"),
+                    "content_stats": parsed_result.get("content_stats"),
+                },
+                "ai_result": ai_result,
+                "processing_metadata": {
+                    "start_time": start_time.isoformat(),
+                    "end_time": datetime.now().isoformat(),
+                    "processing_duration": (datetime.now() - start_time).total_seconds(),
+                },
+            }
+
+            # Step 5: Post-process if needed
+            result = self._post_process_result(result, processing_mode, processing_params or {})
+
+            return result
+
+        except Exception as e:
+            raise ProcessingError(f"Document processing failed: {str(e)}")
+
+    async def process_document_async(
+        self,
+        source: str,
+        processing_mode: ProcessingMode,
+        ai_provider: Optional[AIProvider] = None,
+        processing_params: Optional[Dict[str, Any]] = None,
+        parse_params: Optional[Dict[str, Any]] = None,
+        ai_params: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """Async version of process_document"""
+        return await asyncio.to_thread(
+            self.process_document,
+            source=source,
+            processing_mode=processing_mode,
+            ai_provider=ai_provider,
+            processing_params=processing_params,
+            parse_params=parse_params,
+            ai_params=ai_params,
+        )
+
+    def batch_process_documents(
+        self,
+        sources: List[str],
+        processing_mode: ProcessingMode,
+        ai_provider: Optional[AIProvider] = None,
+        processing_params: Optional[Dict[str, Any]] = None,
+        max_concurrent: Optional[int] = None,
+    ) -> Dict[str, Any]:
+        """
+        Process multiple documents in batch with intelligent orchestration
+
+        Args:
+            sources: List of URLs or file paths
+            processing_mode: AI processing mode to apply
+            ai_provider: AI provider to use
+            processing_params: Additional processing parameters
+            max_concurrent: Maximum concurrent processing
+
+        Returns:
+            Dict containing batch processing results
+        """
+        try:
+            start_time = datetime.now()
+            max_concurrent = max_concurrent or self.config.max_concurrent_requests
+
+            # Process documents in batches
+            results = asyncio.run(
+                self._batch_process_async(
+                    sources,
+                    processing_mode,
+                    ai_provider,
+                    processing_params,
+                    max_concurrent,
+                )
+            )
+
+            # Aggregate results
+            batch_result = {
+                "sources": sources,
+                "processing_mode": processing_mode,
+                "ai_provider": ai_provider or self.config.default_ai_provider,
+                "total_documents": len(sources),
+                "successful_documents": len([r for r in results if r.get("status") == "success"]),
+                "failed_documents": len([r for r in results if r.get("status") == "error"]),
+                "results": results,
+                "batch_metadata": {
+                    "start_time": start_time.isoformat(),
+                    "end_time": datetime.now().isoformat(),
+                    "total_duration": (datetime.now() - start_time).total_seconds(),
+                },
+            }
+
+            return batch_result
+
+        except Exception as e:
+            raise ProcessingError(f"Batch processing failed: {str(e)}")
+
+    def analyze_document(
+        self,
+        source: str,
+        analysis_type: str,
+        custom_prompt: Optional[str] = None,
+        ai_provider: Optional[AIProvider] = None,
+    ) -> Dict[str, Any]:
+        """
+        Perform AI-first document analysis
+
+        Args:
+            source: URL or file path to document
+            analysis_type: Type of analysis to perform
+            custom_prompt: Custom AI prompt for analysis
+            ai_provider: AI provider to use
+
+        Returns:
+            Dict containing analysis results
+        """
+        try:
+            # Parse document first
+            parsed_result = self._parse_document(source, {})
+            content = parsed_result.get("content", "")
+
+            # Prepare AI prompt
+            if custom_prompt:
+                prompt = custom_prompt.format(content=content, analysis_type=analysis_type)
+            else:
+                prompt = f"Perform {analysis_type} analysis on the following document:\n\n{content}"
+
+            # Process with AI
+            ai_result = self._call_ai_provider(
+                prompt, ai_provider or self.config.default_ai_provider, {}
+            )
+
+            return {
+                "source": source,
+                "analysis_type": analysis_type,
+                "document_info": {
+                    "type": parsed_result.get("document_type"),
+                    "content_stats": parsed_result.get("content_stats"),
+                },
+                "analysis_result": ai_result,
+                "timestamp": datetime.now().isoformat(),
+            }
+
+        except Exception as e:
+            raise ProcessingError(f"Document analysis failed: {str(e)}")
+
+    def _parse_document(self, source: str, parse_params: Dict[str, Any]) -> Dict[str, Any]:
+        """Parse document using DocumentParserTool"""
+        if not self.document_parser:
+            raise ProcessingError("DocumentParserTool not available")
+
+        try:
+            return self.document_parser.parse_document(source, **parse_params)
+        except Exception as e:
+            raise ProcessingError(f"Document parsing failed: {str(e)}")
+
+    def _prepare_content_for_ai(
+        self, parsed_result: Dict[str, Any], processing_mode: ProcessingMode
+    ) -> str:
+        """Prepare parsed content for AI processing"""
+        content = parsed_result.get("content", "")
+
+        if isinstance(content, dict):
+            # Extract text from structured content
+            text_content = content.get("text", str(content))
+        else:
+            text_content = str(content)
+
+        # Chunk content if too large
+        max_size = self.config.max_chunk_size
+        if len(text_content) > max_size:
+            # For now, truncate - could implement smart chunking
+            text_content = text_content[:max_size] + "\n\n[Content truncated...]"
+
+        return text_content
+
+    def _process_with_ai(
+        self,
+        content: str,
+        processing_mode: ProcessingMode,
+        ai_provider: AIProvider,
+        processing_params: Dict[str, Any],
+        ai_params: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        """Process content with AI based on processing mode"""
+        try:
+            # Get processing template
+            template = self.processing_templates.get(processing_mode)
+            if not template:
+                raise ProcessingError(f"No template found for processing mode: {processing_mode}")
+
+            # Format prompt
+            prompt = self._format_prompt(template, content, processing_params)
+
+            # Call AI provider
+            ai_result = self._call_ai_provider(prompt, ai_provider, ai_params)
+
+            return {
+                "processing_mode": processing_mode,
+                "prompt_used": prompt,
+                "ai_response": ai_result,
+                "ai_provider": ai_provider,
+            }
+
+        except Exception as e:
+            raise AIProviderError(f"AI processing failed: {str(e)}")
+
+    def _format_prompt(self, template: Dict[str, str], content: str, params: Dict[str, Any]) -> str:
+        """Format AI prompt using template and parameters"""
+        user_prompt = template["user_prompt_template"]
+
+        # Replace content placeholder
+        formatted_prompt = user_prompt.replace("{content}", content)
+
+        # Replace other parameters
+        for key, value in params.items():
+            placeholder = f"{{{key}}}"
+            if placeholder in formatted_prompt:
+                formatted_prompt = formatted_prompt.replace(placeholder, str(value))
+
+        return formatted_prompt
+
+    def _call_ai_provider(
+        self, prompt: str, ai_provider: AIProvider, ai_params: Dict[str, Any]
+    ) -> str:
+        """Call AI provider with prompt"""
+        try:
+            if self.aiecs_client:
+                # Use AIECS client for AI operations
+                from aiecs.domain.task.task_context import TaskContext
+
+                task_context = TaskContext(
+                    task_id=f"doc_processing_{datetime.now().timestamp()}",
+                    task_type="document_processing",
+                    input_data={"prompt": prompt},
+                    metadata=ai_params,
+                )
+
+                # This would need to be adapted based on actual AIECS API
+                result = self.aiecs_client.process_task(task_context)
+                return result.get("response", "")
+            else:
+                # Fallback to direct AI provider calls
+                return self._direct_ai_call(prompt, ai_provider, ai_params)
+
+        except Exception as e:
+            raise AIProviderError(f"AI provider call failed: {str(e)}")
+
+    def _direct_ai_call(
+        self, prompt: str, ai_provider: AIProvider, ai_params: Dict[str, Any]
+    ) -> str:
+        """Direct AI provider call (fallback)"""
+        # This is a placeholder for direct AI provider integration
+        # In a real implementation, you would integrate with specific AI APIs
+        self.logger.warning("Using mock AI response - implement actual AI provider integration")
+        return f"Mock AI response for prompt: {prompt[:100]}..."
+
+    async def _batch_process_async(
+        self,
+        sources: List[str],
+        processing_mode: ProcessingMode,
+        ai_provider: Optional[AIProvider],
+        processing_params: Optional[Dict[str, Any]],
+        max_concurrent: int,
+    ) -> List[Dict[str, Any]]:
+        """Process documents in parallel with concurrency control"""
+        semaphore = asyncio.Semaphore(max_concurrent)
+
+        async def process_single(source: str) -> Dict[str, Any]:
+            async with semaphore:
+                try:
+                    result = await self.process_document_async(
+                        source=source,
+                        processing_mode=processing_mode,
+                        ai_provider=ai_provider,
+                        processing_params=processing_params,
+                    )
+                    return {
+                        "source": source,
+                        "status": "success",
+                        "result": result,
+                    }
+                except Exception as e:
+                    return {
+                        "source": source,
+                        "status": "error",
+                        "error": str(e),
+                    }
+
+        tasks = [process_single(source) for source in sources]
+        return await asyncio.gather(*tasks)
+
+    def _post_process_result(
+        self,
+        result: Dict[str, Any],
+        processing_mode: ProcessingMode,
+        params: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        """Post-process results based on processing mode"""
+        # Add any post-processing logic here
+        # For example, formatting, validation, additional analysis
+
+        if processing_mode == ProcessingMode.EXTRACT_INFO:
+            # Validate extracted information
+            result["validation"] = self._validate_extracted_info(result, params)
+        elif processing_mode == ProcessingMode.CLASSIFY:
+            # Add confidence scoring
+            result["confidence_analysis"] = self._analyze_classification_confidence(result)
+
+        return result
+
+    def _validate_extracted_info(
+        self, result: Dict[str, Any], params: Dict[str, Any]
+    ) -> Dict[str, str]:
+        """Validate extracted information"""
+        # Placeholder for validation logic
+        return {"status": "validated", "notes": "Validation completed"}
+
+    def _analyze_classification_confidence(self, result: Dict[str, Any]) -> Dict[str, Any]:
+        """Analyze classification confidence"""
+        # Placeholder for confidence analysis
+        return {
+            "overall_confidence": 0.85,
+            "factors": ["content_quality", "model_certainty"],
+        }
+
+    # Utility methods for custom processing
+    def create_custom_processor(self, system_prompt: str, user_prompt_template: str) -> Callable:
+        """Create a custom processing function"""
+
+        def custom_processor(source: str, **kwargs) -> Dict[str, Any]:
+            # Add custom template
+            self.processing_templates[ProcessingMode.CUSTOM] = {
+                "system_prompt": system_prompt,
+                "user_prompt_template": user_prompt_template,
+            }
+
+            return self.process_document(
+                source=source,
+                processing_mode=ProcessingMode.CUSTOM,
+                processing_params=kwargs,
+            )
+
+        return custom_processor
+
+    def get_processing_stats(self) -> Dict[str, Any]:
+        """Get processing statistics"""
+        # Placeholder for statistics tracking
+        return {
+            "total_documents_processed": 0,
+            "average_processing_time": 0,
+            "success_rate": 1.0,
+            "most_common_document_types": [],
+            "ai_provider_usage": {},
+        }