PraisonAI 3.0.0__py3-none-any.whl

praisonai/capabilities/messages.py

@@ -0,0 +1,270 @@
+"""
+Messages Capabilities Module
+
+Provides Anthropic-style messages API functionality via LiteLLM.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional, Any, Dict, List
+
+
+@dataclass
+class MessageResult:
+    """Result from message operations."""
+    id: str
+    content: Optional[List[Dict[str, Any]]] = None
+    role: str = "assistant"
+    model: Optional[str] = None
+    stop_reason: Optional[str] = None
+    usage: Optional[Dict[str, int]] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class TokenCountResult:
+    """Result from token counting."""
+    input_tokens: int
+    model: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+def messages_create(
+    messages: List[Dict[str, Any]],
+    model: str = "claude-3-5-sonnet-20241022",
+    max_tokens: int = 1024,
+    system: Optional[str] = None,
+    temperature: float = 1.0,
+    tools: Optional[List[Dict[str, Any]]] = None,
+    timeout: float = 600.0,
+    api_key: Optional[str] = None,
+    api_base: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **kwargs
+) -> MessageResult:
+    """
+    Create a message using Anthropic-style API.
+    
+    Args:
+        messages: List of messages
+        model: Model to use (e.g., "claude-3-5-sonnet-20241022")
+        max_tokens: Maximum tokens to generate
+        system: System prompt
+        temperature: Sampling temperature
+        tools: List of tools
+        timeout: Request timeout in seconds
+        api_key: Optional API key override
+        api_base: Optional API base URL override
+        metadata: Optional metadata for tracing
+        
+    Returns:
+        MessageResult with response
+        
+    Example:
+        >>> result = messages_create([{"role": "user", "content": "Hello"}])
+        >>> print(result.content)
+    """
+    import litellm
+    
+    # Convert to LiteLLM format
+    litellm_messages = []
+    if system:
+        litellm_messages.append({"role": "system", "content": system})
+    litellm_messages.extend(messages)
+    
+    call_kwargs = {
+        'model': model,
+        'messages': litellm_messages,
+        'max_tokens': max_tokens,
+        'temperature': temperature,
+        'timeout': timeout,
+    }
+    
+    if tools:
+        call_kwargs['tools'] = tools
+    if api_key:
+        call_kwargs['api_key'] = api_key
+    if api_base:
+        call_kwargs['api_base'] = api_base
+    
+    call_kwargs.update(kwargs)
+    
+    if metadata:
+        call_kwargs['metadata'] = metadata
+    
+    response = litellm.completion(**call_kwargs)
+    
+    choice = response.choices[0] if response.choices else None
+    message = choice.message if choice else None
+    
+    # Convert to Anthropic-style content blocks
+    content = None
+    if message and message.content:
+        content = [{"type": "text", "text": message.content}]
+    
+    usage = None
+    if hasattr(response, 'usage') and response.usage:
+        usage = {
+            'input_tokens': getattr(response.usage, 'prompt_tokens', 0),
+            'output_tokens': getattr(response.usage, 'completion_tokens', 0),
+        }
+    
+    return MessageResult(
+        id=getattr(response, 'id', ''),
+        content=content,
+        role=getattr(message, 'role', 'assistant') if message else 'assistant',
+        model=getattr(response, 'model', model),
+        stop_reason=getattr(choice, 'finish_reason', None) if choice else None,
+        usage=usage,
+        metadata=metadata or {},
+    )
+
+
+async def amessages_create(
+    messages: List[Dict[str, Any]],
+    model: str = "claude-3-5-sonnet-20241022",
+    max_tokens: int = 1024,
+    system: Optional[str] = None,
+    temperature: float = 1.0,
+    tools: Optional[List[Dict[str, Any]]] = None,
+    timeout: float = 600.0,
+    api_key: Optional[str] = None,
+    api_base: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **kwargs
+) -> MessageResult:
+    """
+    Async: Create a message using Anthropic-style API.
+    
+    See messages_create() for full documentation.
+    """
+    import litellm
+    
+    litellm_messages = []
+    if system:
+        litellm_messages.append({"role": "system", "content": system})
+    litellm_messages.extend(messages)
+    
+    call_kwargs = {
+        'model': model,
+        'messages': litellm_messages,
+        'max_tokens': max_tokens,
+        'temperature': temperature,
+        'timeout': timeout,
+    }
+    
+    if tools:
+        call_kwargs['tools'] = tools
+    if api_key:
+        call_kwargs['api_key'] = api_key
+    if api_base:
+        call_kwargs['api_base'] = api_base
+    
+    call_kwargs.update(kwargs)
+    
+    if metadata:
+        call_kwargs['metadata'] = metadata
+    
+    response = await litellm.acompletion(**call_kwargs)
+    
+    choice = response.choices[0] if response.choices else None
+    message = choice.message if choice else None
+    
+    content = None
+    if message and message.content:
+        content = [{"type": "text", "text": message.content}]
+    
+    usage = None
+    if hasattr(response, 'usage') and response.usage:
+        usage = {
+            'input_tokens': getattr(response.usage, 'prompt_tokens', 0),
+            'output_tokens': getattr(response.usage, 'completion_tokens', 0),
+        }
+    
+    return MessageResult(
+        id=getattr(response, 'id', ''),
+        content=content,
+        role=getattr(message, 'role', 'assistant') if message else 'assistant',
+        model=getattr(response, 'model', model),
+        stop_reason=getattr(choice, 'finish_reason', None) if choice else None,
+        usage=usage,
+        metadata=metadata or {},
+    )
+
+
+def count_tokens(
+    messages: List[Dict[str, Any]],
+    model: str = "gpt-4o-mini",
+    system: Optional[str] = None,
+    api_key: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **kwargs
+) -> TokenCountResult:
+    """
+    Count tokens in messages.
+    
+    Args:
+        messages: List of messages
+        model: Model to use for tokenization
+        system: System prompt
+        api_key: Optional API key override
+        metadata: Optional metadata for tracing
+        
+    Returns:
+        TokenCountResult with token count
+        
+    Example:
+        >>> result = count_tokens([{"role": "user", "content": "Hello"}])
+        >>> print(result.input_tokens)
+    """
+    import litellm
+    
+    all_messages = []
+    if system:
+        all_messages.append({"role": "system", "content": system})
+    all_messages.extend(messages)
+    
+    try:
+        token_count = litellm.token_counter(model=model, messages=all_messages)
+    except Exception:
+        # Fallback: estimate tokens
+        text = ""
+        for msg in all_messages:
+            content = msg.get("content", "")
+            if isinstance(content, str):
+                text += content
+            elif isinstance(content, list):
+                for item in content:
+                    if isinstance(item, dict) and item.get("type") == "text":
+                        text += item.get("text", "")
+        # Rough estimate: 4 chars per token
+        token_count = len(text) // 4
+    
+    return TokenCountResult(
+        input_tokens=token_count,
+        model=model,
+        metadata=metadata or {},
+    )
+
+
+async def acount_tokens(
+    messages: List[Dict[str, Any]],
+    model: str = "gpt-4o-mini",
+    system: Optional[str] = None,
+    api_key: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **kwargs
+) -> TokenCountResult:
+    """
+    Async: Count tokens in messages.
+    
+    See count_tokens() for full documentation.
+    """
+    # Token counting is synchronous, just wrap it
+    return count_tokens(
+        messages=messages,
+        model=model,
+        system=system,
+        api_key=api_key,
+        metadata=metadata,
+        **kwargs
+    )

praisonai/capabilities/moderations.py

@@ -0,0 +1,154 @@
+"""
+Moderations Capabilities Module
+
+Provides content moderation functionality via LiteLLM.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional, Union, List, Any, Dict
+
+
+@dataclass
+class ModerationResult:
+    """Result from content moderation."""
+    flagged: bool
+    categories: Dict[str, bool] = field(default_factory=dict)
+    category_scores: Dict[str, float] = field(default_factory=dict)
+    model: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+def moderate(
+    input: Union[str, List[str]],
+    model: str = "omni-moderation-latest",
+    timeout: float = 600.0,
+    api_key: Optional[str] = None,
+    api_base: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **kwargs
+) -> List[ModerationResult]:
+    """
+    Check content for policy violations.
+    
+    Args:
+        input: Text or list of texts to moderate
+        model: Model name (e.g., "text-moderation-latest", "text-moderation-stable")
+        timeout: Request timeout in seconds
+        api_key: Optional API key override
+        api_base: Optional API base URL override
+        metadata: Optional metadata for tracing
+        
+    Returns:
+        List of ModerationResult objects
+        
+    Example:
+        >>> results = moderate("Some text to check")
+        >>> if results[0].flagged:
+        ...     print("Content flagged!")
+    """
+    import litellm
+    
+    call_kwargs = {
+        'input': input,
+        'model': model,
+        'timeout': timeout,
+    }
+    
+    if api_key:
+        call_kwargs['api_key'] = api_key
+    if api_base:
+        call_kwargs['api_base'] = api_base
+    
+    call_kwargs.update(kwargs)
+    
+    if metadata:
+        call_kwargs['metadata'] = metadata
+    
+    response = litellm.moderation(**call_kwargs)
+    
+    results = []
+    if hasattr(response, 'results'):
+        for item in response.results:
+            categories = {}
+            category_scores = {}
+            
+            if hasattr(item, 'categories'):
+                for cat_name in dir(item.categories):
+                    if not cat_name.startswith('_'):
+                        categories[cat_name] = getattr(item.categories, cat_name, False)
+            
+            if hasattr(item, 'category_scores'):
+                for cat_name in dir(item.category_scores):
+                    if not cat_name.startswith('_'):
+                        category_scores[cat_name] = getattr(item.category_scores, cat_name, 0.0)
+            
+            results.append(ModerationResult(
+                flagged=getattr(item, 'flagged', False),
+                categories=categories,
+                category_scores=category_scores,
+                model=model,
+                metadata=metadata or {},
+            ))
+    
+    return results
+
+
+async def amoderate(
+    input: Union[str, List[str]],
+    model: str = "text-moderation-latest",
+    timeout: float = 600.0,
+    api_key: Optional[str] = None,
+    api_base: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **kwargs
+) -> List[ModerationResult]:
+    """
+    Async: Check content for policy violations.
+    
+    See moderate() for full documentation.
+    """
+    import litellm
+    
+    call_kwargs = {
+        'input': input,
+        'model': model,
+        'timeout': timeout,
+    }
+    
+    if api_key:
+        call_kwargs['api_key'] = api_key
+    if api_base:
+        call_kwargs['api_base'] = api_base
+    
+    call_kwargs.update(kwargs)
+    
+    if metadata:
+        call_kwargs['metadata'] = metadata
+    
+    response = await litellm.amoderation(**call_kwargs)
+    
+    results = []
+    if hasattr(response, 'results'):
+        for item in response.results:
+            categories = {}
+            category_scores = {}
+            
+            if hasattr(item, 'categories'):
+                for cat_name in dir(item.categories):
+                    if not cat_name.startswith('_'):
+                        categories[cat_name] = getattr(item.categories, cat_name, False)
+            
+            if hasattr(item, 'category_scores'):
+                for cat_name in dir(item.category_scores):
+                    if not cat_name.startswith('_'):
+                        category_scores[cat_name] = getattr(item.category_scores, cat_name, 0.0)
+            
+            results.append(ModerationResult(
+                flagged=getattr(item, 'flagged', False),
+                categories=categories,
+                category_scores=category_scores,
+                model=model,
+                metadata=metadata or {},
+            ))
+    
+    return results

praisonai/capabilities/ocr.py

@@ -0,0 +1,217 @@
+"""
+OCR Capabilities Module
+
+Provides optical character recognition functionality via LiteLLM.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional, Union, List, Any, Dict
+
+
+@dataclass
+class OCRResult:
+    """Result from OCR processing."""
+    text: str
+    pages: Optional[List[Dict[str, Any]]] = None
+    model: Optional[str] = None
+    usage: Optional[Dict[str, int]] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+def ocr(
+    document: Union[str, Dict[str, str]],
+    model: str = "mistral/mistral-ocr-latest",
+    include_image_base64: bool = False,
+    pages: Optional[List[int]] = None,
+    image_limit: Optional[int] = None,
+    timeout: float = 600.0,
+    api_key: Optional[str] = None,
+    api_base: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **kwargs
+) -> OCRResult:
+    """
+    Extract text from documents or images using OCR.
+    
+    Args:
+        document: Document URL, file path, or dict with type and URL
+            - String: treated as URL or file path
+            - Dict: {"type": "document_url", "document_url": "https://..."} or
+                   {"type": "image_url", "image_url": "https://..."}
+        model: Model name (e.g., "mistral/mistral-ocr-latest")
+        include_image_base64: Whether to include base64 images in response
+        pages: Specific pages to process (for PDFs)
+        image_limit: Maximum number of images to process
+        timeout: Request timeout in seconds
+        api_key: Optional API key override
+        api_base: Optional API base URL override
+        metadata: Optional metadata for tracing
+        
+    Returns:
+        OCRResult with extracted text
+        
+    Example:
+        >>> result = ocr("https://example.com/document.pdf")
+        >>> print(result.text)
+        
+        >>> result = ocr({"type": "image_url", "image_url": "https://example.com/image.png"})
+    """
+    import litellm
+    
+    # Handle document input
+    doc_input = document
+    if isinstance(document, str):
+        # Check if it's a URL or file path
+        if document.startswith(('http://', 'https://')):
+            if document.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp')):
+                doc_input = {"type": "image_url", "image_url": document}
+            else:
+                doc_input = {"type": "document_url", "document_url": document}
+        else:
+            # Local file - read and encode as base64
+            import base64
+            with open(document, 'rb') as f:
+                content = base64.b64encode(f.read()).decode('utf-8')
+            if document.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp')):
+                doc_input = {"type": "image_url", "image_url": f"data:image/png;base64,{content}"}
+            else:
+                doc_input = {"type": "document_url", "document_url": f"data:application/pdf;base64,{content}"}
+    
+    call_kwargs = {
+        'model': model,
+        'document': doc_input,
+        'timeout': timeout,
+    }
+    
+    if include_image_base64:
+        call_kwargs['include_image_base64'] = include_image_base64
+    if pages:
+        call_kwargs['pages'] = pages
+    if image_limit:
+        call_kwargs['image_limit'] = image_limit
+    if api_key:
+        call_kwargs['api_key'] = api_key
+    if api_base:
+        call_kwargs['api_base'] = api_base
+    
+    call_kwargs.update(kwargs)
+    
+    if metadata:
+        call_kwargs['metadata'] = metadata
+    
+    response = litellm.ocr(**call_kwargs)
+    
+    # Extract text from pages
+    text_parts = []
+    pages_data = []
+    if hasattr(response, 'pages'):
+        for page in response.pages:
+            page_text = getattr(page, 'markdown', '') or getattr(page, 'text', '')
+            text_parts.append(page_text)
+            pages_data.append({
+                'index': getattr(page, 'index', 0),
+                'markdown': getattr(page, 'markdown', ''),
+                'images': getattr(page, 'images', []),
+            })
+    
+    usage = None
+    if hasattr(response, 'usage_info'):
+        usage = {
+            'pages_processed': getattr(response.usage_info, 'pages_processed', 0),
+            'doc_size_bytes': getattr(response.usage_info, 'doc_size_bytes', 0),
+        }
+    
+    return OCRResult(
+        text='\n\n'.join(text_parts),
+        pages=pages_data if pages_data else None,
+        model=model,
+        usage=usage,
+        metadata=metadata or {},
+    )
+
+
+async def aocr(
+    document: Union[str, Dict[str, str]],
+    model: str = "mistral/mistral-ocr-latest",
+    include_image_base64: bool = False,
+    pages: Optional[List[int]] = None,
+    image_limit: Optional[int] = None,
+    timeout: float = 600.0,
+    api_key: Optional[str] = None,
+    api_base: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **kwargs
+) -> OCRResult:
+    """
+    Async: Extract text from documents or images using OCR.
+    
+    See ocr() for full documentation.
+    """
+    import litellm
+    
+    doc_input = document
+    if isinstance(document, str):
+        if document.startswith(('http://', 'https://')):
+            if document.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp')):
+                doc_input = {"type": "image_url", "image_url": document}
+            else:
+                doc_input = {"type": "document_url", "document_url": document}
+        else:
+            import base64
+            with open(document, 'rb') as f:
+                content = base64.b64encode(f.read()).decode('utf-8')
+            if document.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp')):
+                doc_input = {"type": "image_url", "image_url": f"data:image/png;base64,{content}"}
+            else:
+                doc_input = {"type": "document_url", "document_url": f"data:application/pdf;base64,{content}"}
+    
+    call_kwargs = {
+        'model': model,
+        'document': doc_input,
+        'timeout': timeout,
+    }
+    
+    if include_image_base64:
+        call_kwargs['include_image_base64'] = include_image_base64
+    if pages:
+        call_kwargs['pages'] = pages
+    if image_limit:
+        call_kwargs['image_limit'] = image_limit
+    if api_key:
+        call_kwargs['api_key'] = api_key
+    if api_base:
+        call_kwargs['api_base'] = api_base
+    
+    call_kwargs.update(kwargs)
+    
+    if metadata:
+        call_kwargs['metadata'] = metadata
+    
+    response = await litellm.aocr(**call_kwargs)
+    
+    text_parts = []
+    pages_data = []
+    if hasattr(response, 'pages'):
+        for page in response.pages:
+            page_text = getattr(page, 'markdown', '') or getattr(page, 'text', '')
+            text_parts.append(page_text)
+            pages_data.append({
+                'index': getattr(page, 'index', 0),
+                'markdown': getattr(page, 'markdown', ''),
+                'images': getattr(page, 'images', []),
+            })
+    
+    usage = None
+    if hasattr(response, 'usage_info'):
+        usage = {
+            'pages_processed': getattr(response.usage_info, 'pages_processed', 0),
+            'doc_size_bytes': getattr(response.usage_info, 'doc_size_bytes', 0),
+        }
+    
+    return OCRResult(
+        text='\n\n'.join(text_parts),
+        pages=pages_data if pages_data else None,
+        model=model,
+        usage=usage,
+        metadata=metadata or {},
+    )