integrations/pdf_processing/BYOK_INTEGRATION.md

BYOK Integration for PDF Processing

Overview

The BYOK (Bring Your Own Key) system provides comprehensive AI provider management and cost optimization for PDF processing in the ATOM platform. This integration allows users to manage their own API keys, track usage, and optimize costs across multiple AI providers.

Features

🔑 Key Management

• Secure Storage: Encrypted API key storage with fallback to environment variables
• Multi-Environment Support: Separate keys for development, staging, and production
• Usage Tracking: Monitor API usage and costs per provider
• Key Rotation: Easy key updates and management

💰 Cost Optimization

• Provider Selection: Automatic selection of optimal providers based on task type and budget
• Cost Tracking: Real-time cost monitoring and usage analytics
• Budget Constraints: Set spending limits for different tasks
• Alternative Scenarios: Compare different provider options

📊 Usage Analytics

• Request Tracking: Monitor successful and failed requests
• Token Usage: Track token consumption across providers
• Cost Accumulation: Real-time cost calculations
• Rate Limit Monitoring: Track API rate limits and usage patterns

Integration Architecture

BYOK Manager

The central component that manages:

• Provider configurations
• API key storage and encryption
• Usage tracking and analytics
• Cost optimization algorithms

PDF Processing Integration

• PDF OCR Service: Uses BYOK for OpenAI Vision API key management
• PDF Memory Integration: Tracks usage for embedding generation and semantic search
• Route Optimization: BYOK-aware route selection for optimal performance

API Endpoints

BYOK Management Endpoints

GET /api/ai/providers

Get all available AI providers with status and capabilities.

Response:

{
  "providers": [
    {
      "provider": {
        "id": "openai",
        "name": "OpenAI",
        "description": "GPT models for general AI tasks and PDF processing",
        "api_key_env_var": "OPENAI_API_KEY",
        "cost_per_token": 0.002,
        "supported_tasks": ["chat", "code", "analysis", "pdf_ocr", "image_comprehension"],
        "max_requests_per_minute": 60,
        "is_active": true
      },
      "usage": {
        "total_requests": 150,
        "successful_requests": 145,
        "failed_requests": 5,
        "total_tokens_used": 45000,
        "cost_accumulated": 90.0
      },
      "has_api_keys": true,
      "status": "active"
    }
  ],
  "total_providers": 6,
  "active_providers": 3
}

POST /api/ai/providers/{provider_id}/keys

Store an API key for a provider.

Parameters:

• api_key: The API key to store
• key_name: Key identifier (default: "default")
• environment: Environment (development/staging/production)

GET /api/ai/providers/{provider_id}/keys/{key_name}

Get API key status (without revealing the key).

DELETE /api/ai/providers/{provider_id}/keys/{key_name}

Delete an API key.

PDF-Specific BYOK Endpoints

GET /api/ai/pdf/providers

Get AI providers specifically optimized for PDF processing.

Response:

{
  "pdf_providers": [
    {
      "provider_id": "openai",
      "name": "OpenAI",
      "supported_tasks": ["pdf_ocr", "image_comprehension"],
      "cost_per_token": 0.002
    }
  ],
  "total_pdf_providers": 4,
  "supported_tasks": ["pdf_ocr", "image_comprehension", "document_processing", "multimodal"]
}

POST /api/ai/pdf/optimize

Optimize provider selection for PDF processing.

Request:

{
  "pdf_type": "scanned",
  "needs_ocr": true,
  "needs_image_comprehension": true,
  "estimated_pages": 10,
  "budget_constraint": 0.05
}

Response:

{
  "pdf_analysis": {
    "pdf_type": "scanned",
    "needs_ocr": true,
    "needs_image_comprehension": true,
    "estimated_pages": 10,
    "estimated_tokens": 5000
  },
  "recommended_provider": {
    "provider_id": "openai",
    "name": "OpenAI",
    "task_type": "image_comprehension",
    "estimated_cost": 10.0,
    "cost_per_token": 0.002
  },
  "alternative_scenarios": {
    "high_quality": {
      "provider": "anthropic",
      "name": "Anthropic Claude",
      "estimated_cost": 40.0,
      "recommended_for": "Complex PDFs with images and diagrams"
    },
    "cost_effective": {
      "provider": "google_gemini",
      "name": "Google Gemini",
      "estimated_cost": 2.5,
      "recommended_for": "Simple OCR tasks on scanned documents"
    }
  }
}

Cost Optimization Endpoints

POST /api/ai/optimize-cost

Optimize AI cost usage for general tasks.

Request:

{
  "task_type": "pdf_ocr",
  "budget_constraint": 0.01,
  "estimated_tokens": 1000
}

POST /api/ai/usage/track

Track AI usage for cost monitoring.

GET /api/ai/usage/stats

Get usage statistics for AI providers.

PDF Processing Integration

PDF OCR Service BYOK Integration

The PDF OCR service automatically integrates with BYOK for:

1. API Key Management: Retrieves OpenAI API keys from BYOK system
2. Provider Optimization: Uses BYOK to select optimal providers for OCR tasks
3. Usage Tracking: Tracks token usage for cost monitoring
4. Fallback Strategy: Graceful degradation when BYOK is unavailable

Configuration:

from integrations.pdf_processing import PDFOCRService

# With BYOK integration (recommended)
service = PDFOCRService(use_byok=True)

# Without BYOK integration (fallback)
service = PDFOCRService(use_byok=False, openai_api_key="your-key")

PDF Processing Endpoints with BYOK

Enhanced PDF Processing

All PDF processing endpoints now support BYOK optimization:

• POST /api/v1/pdf/process - Process PDF with BYOK optimization
• POST /api/v1/pdf/process-url - Process PDF from URL with BYOK
• GET /api/v1/pdf/status - Get service status with BYOK integration info

Example PDF Processing with BYOK:

curl -X POST "http://localhost:5058/api/v1/pdf/process" \
  -H "Content-Type: multipart/form-data" \
  -F "file=@document.pdf" \
  -F "use_ocr=true" \
  -F "use_advanced_comprehension=true" \
  -F "optimize_with_byok=true" \
  -F "fallback_strategy=cascade"

Response includes BYOK optimization info:

{
  "processing_summary": {
    "used_ocr": true,
    "best_method": "openai_vision",
    "total_pages": 10,
    "total_characters": 2500
  },
  "byok_optimization": {
    "optimized": true,
    "task_type": "image_comprehension",
    "optimal_provider": "openai",
    "provider_name": "OpenAI",
    "cost_per_token": 0.002
  }
}

Configuration

Environment Variables

# BYOK Encryption Key (required for production)
BYOK_ENCRYPTION_KEY=your-secure-encryption-key-here

# Provider API Keys (fallback if not stored in BYOK)
OPENAI_API_KEY=your-openai-key
ANTHROPIC_API_KEY=your-anthropic-key
GOOGLE_API_KEY=your-google-key
AZURE_OPENAI_API_KEY=your-azure-key
DEEPSEEK_API_KEY=your-deepseek-key
HUGGINGFACE_API_KEY=your-huggingface-key

BYOK Configuration Files

BYOK stores configuration in two JSON files:

• ./data/byok_config.json - Provider configurations and usage stats
• ./data/byok_keys.json - Encrypted API keys

Example byok_config.json:

{
  "providers": {
    "openai": {
      "id": "openai",
      "name": "OpenAI",
      "description": "GPT models for general AI tasks and PDF processing",
      "api_key_env_var": "OPENAI_API_KEY",
      "cost_per_token": 0.002,
      "supported_tasks": ["chat", "code", "analysis", "pdf_ocr", "image_comprehension"],
      "max_requests_per_minute": 60,
      "is_active": true
    }
  },
  "usage_stats": {
    "openai": {
      "provider_id": "openai",
      "total_requests": 150,
      "successful_requests": 145,
      "failed_requests": 5,
      "total_tokens_used": 45000,
      "cost_accumulated": 90.0,
      "last_used": "2024-01-15T10:30:00"
    }
  }
}

Usage Examples

1. Setting Up API Keys

from core.byok_endpoints import get_byok_manager

byok_manager = get_byok_manager()

# Store OpenAI API key
key_id = byok_manager.store_api_key(
    provider_id="openai",
    api_key="your-openai-api-key",
    key_name="production",
    environment="production"
)

print(f"API key stored with ID: {key_id}")

2. PDF Processing with BYOK Optimization

from integrations.pdf_processing import PDFOCRService

# Initialize service with BYOK integration
service = PDFOCRService(use_byok=True)

# Process PDF with automatic provider optimization
with open('document.pdf', 'rb') as f:
    result = await service.process_pdf(
        pdf_data=f.read(),
        use_ocr=True,
        use_advanced_comprehension=True,
        fallback_strategy="cascade"
    )

print(f"Used provider: {result['byok_optimization']['optimal_provider']}")
print(f"Estimated cost: ${result['byok_optimization']['estimated_cost']}")

3. Cost Optimization for PDF Processing

from core.byok_endpoints import get_byok_manager

byok_manager = get_byok_manager()

# Optimize for PDF processing
optimization = await byok_manager.optimize_pdf_processing({
    "pdf_type": "scanned",
    "needs_ocr": True,
    "needs_image_comprehension": False,
    "estimated_pages": 20,
    "budget_constraint": 0.02
})

print(f"Recommended provider: {optimization['recommended_provider']['name']}")
print(f"Estimated cost: ${optimization['recommended_provider']['estimated_cost']}")

Security Considerations

Encryption

• API keys are encrypted using a configurable encryption key
• Default encryption uses XOR for development (not production-safe)
• For production, implement proper encryption (AES, Fernet, etc.)

Access Control

• BYOK endpoints should be protected with authentication
• Consider implementing role-based access control for key management
• Audit logging for all key operations

Key Rotation

• Regular key rotation is recommended
• BYOK supports multiple key versions per provider
• Seamless key updates without service interruption

Monitoring and Analytics

Health Checks

curl "http://localhost:5058/api/ai/health"

Usage Statistics

curl "http://localhost:5058/api/ai/usage/stats"

Provider Status

curl "http://localhost:5058/api/ai/providers/openai"

Troubleshooting

Common Issues

1. BYOK Not Available

   • Check if BYOK_ENCRYPTION_KEY is set
   • Verify BYOK configuration files are writable

2. Provider Inactive

   • Ensure API keys are stored in BYOK or environment variables
   • Check provider configuration in byok_config.json

3. Cost Optimization Fails

   • Verify supported tasks for each provider
   • Check budget constraints are realistic

4. PDF Processing Without BYOK

   • Service falls back to environment variables
   • Basic functionality remains available

Debug Mode

Enable debug logging for BYOK operations:

import logging
logging.getLogger('byok').setLevel(logging.DEBUG)

Migration Guide

From Environment Variables to BYOK

1. Export existing keys:

   export OPENAI_API_KEY=your-key
   export ANTHROPIC_API_KEY=your-key
   

2. Store keys in BYOK:

   from core.byok_endpoints import get_byok_manager
   byok_manager = get_byok_manager()
   
   byok_manager.store_api_key("openai", os.getenv("OPENAI_API_KEY"))
   byok_manager.store_api_key("anthropic", os.getenv("ANTHROPIC_API_KEY"))
   

3. Update services to use BYOK:

   # Old way
   service = PDFOCRService(openai_api_key=os.getenv("OPENAI_API_KEY"))
   
   # New way
   service = PDFOCRService(use_byok=True)
   

Best Practices

1. Key Management

• Use different keys for different environments
• Implement regular key rotation
• Monitor key usage and revoke unused keys

2. Cost Optimization

• Set realistic budget constraints
• Monitor usage patterns and adjust accordingly
• Use cost-effective providers for simple tasks

3. Error Handling

• Implement graceful fallbacks when BYOK is unavailable
• Monitor BYOK system health
• Log BYOK operations for debugging

4. Security

• Never log API keys
• Use secure encryption in production
• Implement proper access controls

Support

For issues with BYOK integration:

1. Check the BYOK health endpoint: /api/ai/health
2. Review application logs for BYOK-related errors
3. Verify API keys are properly stored and accessible
4. Check provider configurations in byok_config.json

Future Enhancements

• Multi-tenant key management
• Advanced cost prediction algorithms
• Real-time budget alerts
• Provider performance analytics
• Automated key rotation
• Integration with more AI providers