jaf-py 2.4.1__py3-none-any.whl → 2.4.2__py3-none-any.whl

jaf/memory/approval_storage.py

@@ -0,0 +1,306 @@
+"""
+Approval storage interface and implementations for Human-in-the-Loop (HITL) functionality.
+
+This module provides persistent storage for tool approval decisions, enabling
+the framework to maintain approval states across conversation sessions and
+handle interruptions gracefully.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Optional
+import asyncio
+
+from ..core.types import RunId, ApprovalValue
+
+
+class ApprovalStorageResult:
+    """Result wrapper for approval storage operations."""
+    
+    def __init__(self, success: bool, data: Any = None, error: Optional[str] = None):
+        self.success = success
+        self.data = data
+        self.error = error
+
+    @classmethod
+    def success_result(cls, data: Any = None) -> 'ApprovalStorageResult':
+        """Create a successful result."""
+        return cls(success=True, data=data)
+
+    @classmethod
+    def error_result(cls, error: str) -> 'ApprovalStorageResult':
+        """Create an error result."""
+        return cls(success=False, error=error)
+
+
+class ApprovalStorage(ABC):
+    """Abstract interface for approval storage implementations."""
+
+    @abstractmethod
+    async def store_approval(
+        self,
+        run_id: RunId,
+        tool_call_id: str,
+        approval: ApprovalValue,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> ApprovalStorageResult:
+        """Store an approval decision for a tool call."""
+        pass
+
+    @abstractmethod
+    async def get_approval(
+        self,
+        run_id: RunId,
+        tool_call_id: str
+    ) -> ApprovalStorageResult:
+        """Retrieve approval for a specific tool call. Returns None if not found."""
+        pass
+
+    @abstractmethod
+    async def get_run_approvals(
+        self,
+        run_id: RunId
+    ) -> ApprovalStorageResult:
+        """Get all approvals for a run as a Dict[str, ApprovalValue]."""
+        pass
+
+    @abstractmethod
+    async def update_approval(
+        self,
+        run_id: RunId,
+        tool_call_id: str,
+        updates: Dict[str, Any]
+    ) -> ApprovalStorageResult:
+        """Update existing approval with additional context."""
+        pass
+
+    @abstractmethod
+    async def delete_approval(
+        self,
+        run_id: RunId,
+        tool_call_id: str
+    ) -> ApprovalStorageResult:
+        """Delete approval for a tool call. Returns success status."""
+        pass
+
+    @abstractmethod
+    async def clear_run_approvals(self, run_id: RunId) -> ApprovalStorageResult:
+        """Clear all approvals for a run. Returns count of deleted approvals."""
+        pass
+
+    @abstractmethod
+    async def get_stats(self) -> ApprovalStorageResult:
+        """Get approval statistics."""
+        pass
+
+    @abstractmethod
+    async def health_check(self) -> ApprovalStorageResult:
+        """Health check for the approval storage."""
+        pass
+
+    @abstractmethod
+    async def close(self) -> ApprovalStorageResult:
+        """Close/cleanup the storage."""
+        pass
+
+
+class InMemoryApprovalStorage(ApprovalStorage):
+    """In-memory implementation of ApprovalStorage for development and testing."""
+
+    def __init__(self):
+        self._approvals: Dict[str, Dict[str, ApprovalValue]] = {}
+        self._lock = asyncio.Lock()
+
+    def _get_run_key(self, run_id: RunId) -> str:
+        """Generate a consistent key for a run."""
+        return f"run:{run_id}"
+
+    async def store_approval(
+        self,
+        run_id: RunId,
+        tool_call_id: str,
+        approval: ApprovalValue,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> ApprovalStorageResult:
+        """Store an approval decision."""
+        try:
+            async with self._lock:
+                run_key = self._get_run_key(run_id)
+                
+                if run_key not in self._approvals:
+                    self._approvals[run_key] = {}
+                
+                # Enhance approval with metadata if provided
+                enhanced_approval = approval
+                if metadata:
+                    additional_context = {**(approval.additional_context or {}), **metadata}
+                    enhanced_approval = ApprovalValue(
+                        status=approval.status,
+                        approved=approval.approved,
+                        additional_context=additional_context
+                    )
+                
+                self._approvals[run_key][tool_call_id] = enhanced_approval
+                
+            return ApprovalStorageResult.success_result()
+        except Exception as e:
+            return ApprovalStorageResult.error_result(f"Failed to store approval: {e}")
+
+    async def get_approval(
+        self,
+        run_id: RunId,
+        tool_call_id: str
+    ) -> ApprovalStorageResult:
+        """Retrieve approval for a specific tool call."""
+        try:
+            async with self._lock:
+                run_key = self._get_run_key(run_id)
+                run_approvals = self._approvals.get(run_key, {})
+                approval = run_approvals.get(tool_call_id)
+                
+            return ApprovalStorageResult.success_result(approval)
+        except Exception as e:
+            return ApprovalStorageResult.error_result(f"Failed to get approval: {e}")
+
+    async def get_run_approvals(self, run_id: RunId) -> ApprovalStorageResult:
+        """Get all approvals for a run."""
+        try:
+            async with self._lock:
+                run_key = self._get_run_key(run_id)
+                run_approvals = self._approvals.get(run_key, {}).copy()
+                
+            return ApprovalStorageResult.success_result(run_approvals)
+        except Exception as e:
+            return ApprovalStorageResult.error_result(f"Failed to get run approvals: {e}")
+
+    async def update_approval(
+        self,
+        run_id: RunId,
+        tool_call_id: str,
+        updates: Dict[str, Any]
+    ) -> ApprovalStorageResult:
+        """Update existing approval."""
+        try:
+            async with self._lock:
+                run_key = self._get_run_key(run_id)
+                
+                if run_key not in self._approvals or tool_call_id not in self._approvals[run_key]:
+                    return ApprovalStorageResult.error_result(
+                        f"Approval not found for tool call {tool_call_id} in run {run_id}"
+                    )
+                
+                existing = self._approvals[run_key][tool_call_id]
+                
+                # Merge additional context
+                merged_context = {**(existing.additional_context or {}), **(updates.get('additional_context', {}))}
+                
+                updated_approval = ApprovalValue(
+                    status=updates.get('status', existing.status),
+                    approved=updates.get('approved', existing.approved),
+                    additional_context=merged_context if merged_context else existing.additional_context
+                )
+                
+                self._approvals[run_key][tool_call_id] = updated_approval
+                
+            return ApprovalStorageResult.success_result()
+        except Exception as e:
+            return ApprovalStorageResult.error_result(f"Failed to update approval: {e}")
+
+    async def delete_approval(
+        self,
+        run_id: RunId,
+        tool_call_id: str
+    ) -> ApprovalStorageResult:
+        """Delete approval for a tool call."""
+        try:
+            async with self._lock:
+                run_key = self._get_run_key(run_id)
+                
+                if run_key not in self._approvals:
+                    return ApprovalStorageResult.success_result(False)
+                
+                deleted = self._approvals[run_key].pop(tool_call_id, None) is not None
+                
+                # Clean up empty run maps
+                if not self._approvals[run_key]:
+                    del self._approvals[run_key]
+                
+            return ApprovalStorageResult.success_result(deleted)
+        except Exception as e:
+            return ApprovalStorageResult.error_result(f"Failed to delete approval: {e}")
+
+    async def clear_run_approvals(self, run_id: RunId) -> ApprovalStorageResult:
+        """Clear all approvals for a run."""
+        try:
+            async with self._lock:
+                run_key = self._get_run_key(run_id)
+                
+                if run_key not in self._approvals:
+                    return ApprovalStorageResult.success_result(0)
+                
+                count = len(self._approvals[run_key])
+                del self._approvals[run_key]
+                
+            return ApprovalStorageResult.success_result(count)
+        except Exception as e:
+            return ApprovalStorageResult.error_result(f"Failed to clear run approvals: {e}")
+
+    async def get_stats(self) -> ApprovalStorageResult:
+        """Get approval statistics."""
+        try:
+            async with self._lock:
+                total_approvals = 0
+                approved_count = 0
+                rejected_count = 0
+                runs_with_approvals = len(self._approvals)
+                
+                for run_approvals in self._approvals.values():
+                    for approval in run_approvals.values():
+                        total_approvals += 1
+                        if approval.approved:
+                            approved_count += 1
+                        else:
+                            rejected_count += 1
+                
+                stats = {
+                    'total_approvals': total_approvals,
+                    'approved_count': approved_count,
+                    'rejected_count': rejected_count,
+                    'runs_with_approvals': runs_with_approvals
+                }
+                
+            return ApprovalStorageResult.success_result(stats)
+        except Exception as e:
+            return ApprovalStorageResult.error_result(f"Failed to get stats: {e}")
+
+    async def health_check(self) -> ApprovalStorageResult:
+        """Health check for the storage."""
+        try:
+            # Simple operation to test functionality
+            await asyncio.sleep(0.001)  # Minimal async operation
+            
+            health_data = {
+                'healthy': True,
+                'latency_ms': 1.0  # Approximate for in-memory
+            }
+            
+            return ApprovalStorageResult.success_result(health_data)
+        except Exception as e:
+            health_data = {
+                'healthy': False,
+                'error': str(e)
+            }
+            return ApprovalStorageResult.success_result(health_data)
+
+    async def close(self) -> ApprovalStorageResult:
+        """Close/cleanup the storage."""
+        try:
+            async with self._lock:
+                self._approvals.clear()
+            return ApprovalStorageResult.success_result()
+        except Exception as e:
+            return ApprovalStorageResult.error_result(f"Failed to close storage: {e}")
+
+
+def create_in_memory_approval_storage() -> InMemoryApprovalStorage:
+    """Create an in-memory approval storage instance."""
+    return InMemoryApprovalStorage()

jaf/memory/types.py

@@ -233,3 +233,4 @@ class MemoryConfig:
     max_messages: Optional[int] = None
     ttl: Optional[int] = None
     compression_threshold: Optional[int] = None
+    store_on_completion: bool = True  # Store conversation on completion (in addition to interruptions)

jaf/memory/utils.py

@@ -9,7 +9,7 @@ import json
 from datetime import datetime
 from typing import Any, Dict, List, Optional
 
-from ..core.types import Message, ToolCall, ToolCallFunction
+from ..core.types import Message, ToolCall, ToolCallFunction, get_text_content
 from .types import ConversationMemory
 
 

jaf/providers/model.py

@@ -5,18 +5,106 @@ This module provides model providers that integrate with various LLM services,
 starting with LiteLLM for multi-provider support.
 """
 
-from typing import Any, Dict, Optional, TypeVar, AsyncIterator, List
+from typing import Any, Dict, Optional, TypeVar, AsyncIterator
 import asyncio
 import httpx
+import time
+import os
+import base64
 
 from openai import OpenAI
 from pydantic import BaseModel
 
-from ..core.types import Agent, ContentRole, Message, ModelProvider, RunConfig, RunState, CompletionStreamChunk, ToolCallDelta, ToolCallFunctionDelta
+from ..core.types import (
+    Agent, ContentRole, Message, ModelProvider, RunConfig, RunState, 
+    CompletionStreamChunk, ToolCallDelta, ToolCallFunctionDelta,
+    MessageContentPart, get_text_content
+)
 from ..core.proxy import ProxyConfig
+from ..utils.document_processor import (
+    extract_document_content, is_document_supported, 
+    get_document_description, DocumentProcessingError
+)
 
 Ctx = TypeVar('Ctx')
 
+# Vision model caching
+VISION_MODEL_CACHE_TTL = 5 * 60  # 5 minutes
+VISION_API_TIMEOUT = 3.0  # 3 seconds
+_vision_model_cache: Dict[str, Dict[str, Any]] = {}
+MAX_IMAGE_BYTES = int(os.environ.get("JAF_MAX_IMAGE_BYTES", 8 * 1024 * 1024))
+
+async def _is_vision_model(model: str, base_url: str) -> bool:
+    """
+    Check if a model supports vision capabilities.
+    
+    Args:
+        model: Model name to check
+        base_url: Base URL of the LiteLLM server
+        
+    Returns:
+        True if model supports vision, False otherwise
+    """
+    cache_key = f"{base_url}:{model}"
+    cached = _vision_model_cache.get(cache_key)
+    
+    if cached and time.time() - cached['timestamp'] < VISION_MODEL_CACHE_TTL:
+        return cached['supports']
+    
+    try:
+        async with httpx.AsyncClient(timeout=VISION_API_TIMEOUT) as client:
+            response = await client.get(
+                f"{base_url}/model_group/info",
+                headers={'accept': 'application/json'}
+            )
+            
+            if response.status_code == 200:
+                data = response.json()
+                model_info = None
+                
+                if 'data' in data and isinstance(data['data'], list):
+                    for m in data['data']:
+                        if (m.get('model_group') == model or 
+                            model in str(m.get('model_group', ''))):
+                            model_info = m
+                            break
+                
+                if model_info and 'supports_vision' in model_info:
+                    result = model_info['supports_vision']
+                    _vision_model_cache[cache_key] = {
+                        'supports': result,
+                        'timestamp': time.time()
+                    }
+                    return result
+            else:
+                print(f"Warning: Vision API returned status {response.status_code} for model {model}")
+                
+    except Exception as e:
+        print(f"Warning: Vision API error for model {model}: {e}")
+    
+    # Fallback to known vision models
+    known_vision_models = [
+        'gpt-4-vision-preview',
+        'gpt-4o',
+        'gpt-4o-mini',
+        'claude-sonnet-4',
+        'claude-sonnet-4-20250514',
+        'gemini-2.5-flash',
+        'gemini-2.5-pro'
+    ]
+    
+    is_known_vision_model = any(
+        vision_model.lower() in model.lower()
+        for vision_model in known_vision_models
+    )
+    
+    _vision_model_cache[cache_key] = {
+        'supports': is_known_vision_model,
+        'timestamp': time.time()
+    }
+    
+    return is_known_vision_model
+
 def make_litellm_provider(
     base_url: str,
     api_key: str = "anything",
@@ -76,6 +164,23 @@ def make_litellm_provider(
             model = (config.model_override or
                     (agent.model_config.name if agent.model_config else "gpt-4o"))
 
+            # Check if any message contains image content or image attachments
+            has_image_content = any(
+                (isinstance(msg.content, list) and
+                 any(part.type == 'image_url' for part in msg.content)) or
+                (msg.attachments and
+                 any(att.kind == 'image' for att in msg.attachments))
+                for msg in state.messages
+            )
+
+            if has_image_content:
+                supports_vision = await _is_vision_model(model, base_url)
+                if not supports_vision:
+                    raise ValueError(
+                        f"Model {model} does not support vision capabilities. "
+                        f"Please use a vision-capable model like gpt-4o, claude-3-5-sonnet, or gemini-1.5-pro."
+                    )
+
             # Create system message
             system_message = {
                 "role": "system",
@@ -83,9 +188,12 @@ def make_litellm_provider(
             }
 
             # Convert messages to OpenAI format
-            messages = [system_message] + [
-                _convert_message(msg) for msg in state.messages
-            ]
+            converted_messages = []
+            for msg in state.messages:
+                converted_msg = await _convert_message(msg)
+                converted_messages.append(converted_msg)
+            
+            messages = [system_message] + converted_messages
 
             # Convert tools to OpenAI format
             tools = None
@@ -190,10 +298,13 @@ def make_litellm_provider(
                 "content": agent.instructions(state)
             }
 
-            # Convert messages to OpenAI format
-            messages = [system_message] + [
-                _convert_message(msg) for msg in state.messages
-            ]
+            # Convert messages to OpenAI format  
+            converted_messages = []
+            for msg in state.messages:
+                converted_msg = await _convert_message(msg)
+                converted_messages.append(converted_msg)
+            
+            messages = [system_message] + converted_messages
 
             # Convert tools to OpenAI format
             tools = None
@@ -334,17 +445,22 @@ def make_litellm_provider(
 
     return LiteLLMProvider()
 
-def _convert_message(msg: Message) -> Dict[str, Any]:
-    """Convert JAF Message to OpenAI message format."""
+async def _convert_message(msg: Message) -> Dict[str, Any]:
+    """Convert JAF Message to OpenAI message format with attachment support."""
     if msg.role == 'user':
-        return {
-            "role": "user",
-            "content": msg.content
-        }
+        if isinstance(msg.content, list):
+            # Multi-part content
+            return {
+                "role": "user",
+                "content": [_convert_content_part(part) for part in msg.content]
+            }
+        else:
+            # Build message with attachments if available
+            return await _build_chat_message_with_attachments('user', msg)
     elif msg.role == 'assistant':
         result = {
             "role": "assistant",
-            "content": msg.content,
+            "content": get_text_content(msg.content),
         }
         if msg.tool_calls:
             result["tool_calls"] = [
@@ -362,12 +478,156 @@ def _convert_message(msg: Message) -> Dict[str, Any]:
     elif msg.role == ContentRole.TOOL:
         return {
             "role": "tool",
-            "content": msg.content,
+            "content": get_text_content(msg.content),
             "tool_call_id": msg.tool_call_id
         }
     else:
         raise ValueError(f"Unknown message role: {msg.role}")
 
+
+def _convert_content_part(part: MessageContentPart) -> Dict[str, Any]:
+    """Convert MessageContentPart to OpenAI format."""
+    if part.type == 'text':
+        return {
+            "type": "text",
+            "text": part.text
+        }
+    elif part.type == 'image_url':
+        return {
+            "type": "image_url",
+            "image_url": part.image_url
+        }
+    elif part.type == 'file':
+        return {
+            "type": "file",
+            "file": part.file
+        }
+    else:
+        raise ValueError(f"Unknown content part type: {part.type}")
+
+
+async def _build_chat_message_with_attachments(
+    role: str,
+    msg: Message
+) -> Dict[str, Any]:
+    """
+    Build multi-part content for Chat Completions if attachments exist.
+    Supports images via image_url and documents via content extraction.
+    """
+    has_attachments = msg.attachments and len(msg.attachments) > 0
+    if not has_attachments:
+        if role == 'assistant':
+            base_msg = {"role": "assistant", "content": get_text_content(msg.content)}
+            if msg.tool_calls:
+                base_msg["tool_calls"] = [
+                    {
+                        "id": tc.id,
+                        "type": tc.type,
+                        "function": {
+                            "name": tc.function.name,
+                            "arguments": tc.function.arguments
+                        }
+                    }
+                    for tc in msg.tool_calls
+                ]
+            return base_msg
+        return {"role": "user", "content": get_text_content(msg.content)}
+
+    parts = []
+    text_content = get_text_content(msg.content)
+    if text_content and text_content.strip():
+        parts.append({"type": "text", "text": text_content})
+
+    for att in msg.attachments:
+        if att.kind == 'image':
+            # Prefer explicit URL; otherwise construct a data URL from base64
+            url = att.url
+            if not url and att.data and att.mime_type:
+                # Validate base64 data size before creating data URL
+                try:
+                    # Estimate decoded size (base64 is ~4/3 of decoded size)
+                    estimated_size = len(att.data) * 3 // 4
+                    
+                    if estimated_size > MAX_IMAGE_BYTES:
+                        print(f"Warning: Skipping oversized image ({estimated_size} bytes > {MAX_IMAGE_BYTES}). "
+                              f"Set JAF_MAX_IMAGE_BYTES env var to adjust limit.")
+                        parts.append({
+                            "type": "text",
+                            "text": f"[IMAGE SKIPPED: Size exceeds limit of {MAX_IMAGE_BYTES//1024//1024}MB. "
+                                   f"Image name: {att.name or 'unnamed'}]"
+                        })
+                        continue
+                    
+                    # Create data URL for valid-sized images
+                    url = f"data:{att.mime_type};base64,{att.data}"
+                except Exception as e:
+                    print(f"Error processing image data: {e}")
+                    parts.append({
+                        "type": "text",
+                        "text": f"[IMAGE ERROR: Failed to process image data. Image name: {att.name or 'unnamed'}]"
+                    })
+                    continue
+            
+            if url:
+                parts.append({
+                    "type": "image_url",
+                    "image_url": {"url": url}
+                })
+        
+        elif att.kind in ['document', 'file']:
+            # Check if attachment has use_litellm_format flag or is a large document
+            use_litellm_format = att.use_litellm_format is True
+            
+            if use_litellm_format and (att.url or att.data):
+                # For now, fall back to content extraction since most providers don't support native file format
+                # TODO: Add provider-specific file format support
+                print(f"Info: LiteLLM format requested for {att.name}, falling back to content extraction")
+                use_litellm_format = False
+            
+            if not use_litellm_format:
+                # Extract document content if supported and we have data or URL
+                if is_document_supported(att.mime_type) and (att.data or att.url):
+                    try:
+                        processed = await extract_document_content(att)
+                        file_name = att.name or 'document'
+                        description = get_document_description(att.mime_type)
+                        
+                        parts.append({
+                            "type": "text",
+                            "text": f"DOCUMENT: {file_name} ({description}):\n\n{processed.content}"
+                        })
+                    except DocumentProcessingError as e:
+                        # Fallback to filename if extraction fails
+                        label = att.name or att.format or att.mime_type or 'attachment'
+                        parts.append({
+                            "type": "text",
+                            "text": f"ERROR: Failed to process {att.kind}: {label} ({e})"
+                        })
+                else:
+                    # Unsupported document type - show placeholder
+                    label = att.name or att.format or att.mime_type or 'attachment'
+                    url_info = f" ({att.url})" if att.url else ""
+                    parts.append({
+                        "type": "text",
+                        "text": f"ATTACHMENT: {att.kind}: {label}{url_info}"
+                    })
+
+    base_msg = {"role": role, "content": parts}
+    if role == 'assistant' and msg.tool_calls:
+        base_msg["tool_calls"] = [
+            {
+                "id": tc.id,
+                "type": tc.type,
+                "function": {
+                    "name": tc.function.name,
+                    "arguments": tc.function.arguments
+                }
+            }
+            for tc in msg.tool_calls
+        ]
+    
+    return base_msg
+
 def _pydantic_to_json_schema(model_class: type[BaseModel]) -> Dict[str, Any]:
     """
     Convert a Pydantic model to JSON schema for OpenAI tools.

jaf/server/__init__.py

@@ -10,6 +10,8 @@ __all__ = [
     "ChatResponse",
     "HealthResponse",
     "HttpMessage",
+    "HttpAttachment", 
+    "HttpMessageContentPart",
     "ServerConfig",
     "create_jaf_server",
     "run_server",