app.py

from flask import Flask, request, jsonify
import requests
import json
import os
import logging
from typing import Optional, Dict, Any

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# Configuration from environment variables
IFLOW_API_KEY = os.getenv("IFLOW_API_KEY")
PROXY_AUTH_KEY = os.getenv("PROXY_AUTH_KEY", "your-proxy-auth-key")
IFLOW_API_URL = 'https://apis.iflow.cn/v1/chat/completions'
DEFAULT_MODEL = os.getenv("DEFAULT_MODEL", "glm-4.6")

@app.route('/')
def home():
    return "iFlow Proxy for Panda App is running! (Using /groq endpoint)"

@app.route('/groq', methods=['POST'])
def proxy_to_groq():
    """Main proxy endpoint - keeps /groq URL but uses iFlow API internally"""
    try:
        # Validate proxy authentication from header
        auth_header = request.headers.get('X-API-Key') or request.headers.get('Authorization')
        if auth_header:
            # Check if it's Bearer token format
            if auth_header.startswith('Bearer '):
                auth_key = auth_header.replace('Bearer ', '')
            else:
                auth_key = auth_header
            
            if auth_key != PROXY_AUTH_KEY:
                logger.warning("Unauthorized request with invalid API key")
                return jsonify({
                    'error': {
                        'message': 'Invalid authentication credentials',
                        'type': 'invalid_request_error',
                        'code': 'invalid_api_key'
                    }
                }), 401
        else:
            logger.warning("No authentication header provided")
            return jsonify({
                'error': {
                    'message': 'You must provide an API key',
                    'type': 'invalid_request_error',
                    'code': 'missing_api_key'
                }
            }), 401
        
        # Get the Gemini format payload from app
        gemini_payload = request.get_json()
        if not gemini_payload:
            logger.error("No JSON payload received")
            return jsonify({
                'error': {
                    'message': 'No JSON payload received',
                    'type': 'invalid_request_error'
                }
            }), 400
        
        logger.info(f"Received payload: {json.dumps(gemini_payload, indent=2)[:500]}...")
        
        # Check for streaming request
        stream = gemini_payload.get('stream', False)
        
        # Transform Gemini format to iFlow (OpenAI-compatible) format
        iflow_payload = {
            'model': gemini_payload.get('model', DEFAULT_MODEL),
            'messages': [],
            'stream': stream
        }
        
        # Copy common parameters if present
        common_params = ['temperature', 'max_tokens', 'top_p', 'frequency_penalty', 
                        'presence_penalty', 'stop', 'seed', 'response_format']
        
        for param in common_params:
            if param in gemini_payload:
                iflow_payload[param] = gemini_payload[param]
        
        # Convert messages from Gemini parts format to OpenAI content format
        messages = gemini_payload.get('messages', [])
        if not messages:
            logger.error("No messages in payload")
            return jsonify({
                'error': {
                    'message': 'No messages in payload',
                    'type': 'invalid_request_error'
                }
            }), 400
        
        for msg in messages:
            # Extract text from parts array
            parts = msg.get('parts', [])
            
            # Handle both Gemini format (parts array) and OpenAI format (content string)
            if parts:
                # Gemini format: parts array with text
                content_parts = []
                for part in parts:
                    text = part.get('text', '')
                    if text:
                        content_parts.append(text)
                
                if not content_parts:
                    logger.warning(f"Message with role '{msg.get('role')}' has no text content")
                    continue
                
                content = '\n'.join(content_parts)
            else:
                # OpenAI format: direct content field
                content = msg.get('content', '')
                if not content:
                    logger.warning(f"Message with role '{msg.get('role')}' has no content")
                    continue
            
            # Map Gemini roles to OpenAI roles
            role = msg.get('role')
            if role == 'model':
                role = 'assistant'  # OpenAI uses 'assistant', not 'model'
            
            # Handle function calls if present (for OpenAI-compatible format)
            message_dict = {
                'role': role,
                'content': content
            }
            
            # Copy any additional fields that might be present
            additional_fields = ['name', 'function_call', 'tool_calls', 'tool_call_id']
            for field in additional_fields:
                if field in msg:
                    message_dict[field] = msg[field]
            
            iflow_payload['messages'].append(message_dict)
        
        if not iflow_payload['messages']:
            logger.error("No valid messages after transformation")
            return jsonify({
                'error': {
                    'message': 'No valid messages in payload',
                    'type': 'invalid_request_error'
                }
            }), 400
        
        logger.info(f"Transformed payload model: {iflow_payload['model']}")
        logger.info(f"Message count: {len(iflow_payload['messages'])}")
        logger.info(f"Streaming: {stream}")
        
        # Call iFlow API
        headers = {
            'Authorization': f'Bearer {IFLOW_API_KEY}',
            'Content-Type': 'application/json',
            'User-Agent': 'iFlow-Proxy/1.0'
        }
        
        if stream:
            # Handle streaming response
            logger.info("Making streaming request to iFlow API")
            response = requests.post(
                IFLOW_API_URL, 
                json=iflow_payload, 
                headers=headers, 
                timeout=120,
                stream=True
            )
            
            if response.status_code != 200:
                error_text = response.text[:500]
                logger.error(f"iFlow API error: {response.status_code} - {error_text}")
                return jsonify({
                    'error': {
                        'message': f'iFlow API error: {response.status_code}',
                        'type': 'api_error'
                    }
                }), response.status_code
            
            # Return streaming response
            def generate():
                for chunk in response.iter_lines():
                    if chunk:
                        logger.debug(f"Stream chunk: {chunk[:200]}...")
                        yield chunk + b'\n\n'
            
            return app.response_class(generate(), mimetype='text/event-stream')
        
        else:
            # Handle non-streaming response
            logger.info("Making non-streaming request to iFlow API")
            response = requests.post(
                IFLOW_API_URL, 
                json=iflow_payload, 
                headers=headers, 
                timeout=60
            )
            
            logger.info(f"iFlow response status: {response.status_code}")
            
            if response.status_code != 200:
                error_text = response.text[:500]
                logger.error(f"iFlow API error: {response.status_code} - {error_text}")
                try:
                    error_json = response.json()
                    return jsonify(error_json), response.status_code
                except:
                    return jsonify({
                        'error': {
                            'message': f'iFlow API error: {response.status_code}',
                            'type': 'api_error'
                        }
                    }), response.status_code
            
            # Parse and return iFlow response
            iflow_response = response.json()
            logger.info(f"iFlow response received successfully")
            
            # Extract content from iFlow's response
            if 'choices' in iflow_response and len(iflow_response['choices']) > 0:
                choice = iflow_response['choices'][0]
                if 'message' in choice:
                    content = choice['message'].get('content', '')
                    
                    if not content:
                        logger.warning("Empty content in iFlow response")
                        content = "No response content"
                    
                    # Return plain text for /groq endpoint (original format)
                    logger.info("Successfully returning response to app via /groq")
                    return content, 200, {'Content-Type': 'text/plain'}
                else:
                    logger.error("No message in iFlow response choice")
                    return jsonify({
                        'error': {
                            'message': 'Invalid response format from iFlow API',
                            'type': 'api_error'
                        }
                    }), 500
            else:
                logger.error("No choices in iFlow response")
                return jsonify({
                    'error': {
                        'message': 'Invalid response format from iFlow API',
                        'type': 'api_error'
                    }
                }), 500
        
    except requests.exceptions.Timeout:
        logger.error("iFlow API request timeout")
        return jsonify({
            'error': {
                'message': 'Request timeout to iFlow API',
                'type': 'timeout_error'
            }
        }), 504
    except requests.exceptions.RequestException as e:
        logger.error(f"Error calling iFlow API: {e}")
        return jsonify({
            'error': {
                'message': f'Network error: {str(e)}',
                'type': 'network_error'
            }
        }), 500
    except (json.JSONDecodeError, KeyError) as e:
        logger.error(f"Error processing request: {e}")
        return jsonify({
            'error': {
                'message': f'Invalid request format: {str(e)}',
                'type': 'invalid_request_error'
            }
        }), 400
    except Exception as e:
        logger.error(f"Unexpected error: {e}", exc_info=True)
        return jsonify({
            'error': {
                'message': f'Internal server error: {str(e)}',
                'type': 'server_error'
            }
        }), 500

@app.route('/iflow', methods=['POST'])
@app.route('/chat/completions', methods=['POST'])
def proxy_to_iflow():
    """Alternative endpoints for iFlow API (returns JSON format)"""
    try:
        # Validate proxy authentication from header
        auth_header = request.headers.get('X-API-Key') or request.headers.get('Authorization')
        if auth_header:
            # Check if it's Bearer token format
            if auth_header.startswith('Bearer '):
                auth_key = auth_header.replace('Bearer ', '')
            else:
                auth_key = auth_header
            
            if auth_key != PROXY_AUTH_KEY:
                logger.warning("Unauthorized request with invalid API key")
                return jsonify({
                    'error': {
                        'message': 'Invalid authentication credentials',
                        'type': 'invalid_request_error',
                        'code': 'invalid_api_key'
                    }
                }), 401
        else:
            logger.warning("No authentication header provided")
            return jsonify({
                'error': {
                    'message': 'You must provide an API key',
                    'type': 'invalid_request_error',
                    'code': 'missing_api_key'
                }
            }), 401
        
        # Get the Gemini format payload from app
        gemini_payload = request.get_json()
        if not gemini_payload:
            logger.error("No JSON payload received")
            return jsonify({
                'error': {
                    'message': 'No JSON payload received',
                    'type': 'invalid_request_error'
                }
            }), 400
        
        # Check for streaming request
        stream = gemini_payload.get('stream', False)
        
        # Transform Gemini format to iFlow (OpenAI-compatible) format
        iflow_payload = {
            'model': gemini_payload.get('model', DEFAULT_MODEL),
            'messages': [],
            'stream': stream
        }
        
        # Copy common parameters if present
        common_params = ['temperature', 'max_tokens', 'top_p', 'frequency_penalty', 
                        'presence_penalty', 'stop', 'seed', 'response_format']
        
        for param in common_params:
            if param in gemini_payload:
                iflow_payload[param] = gemini_payload[param]
        
        # Convert messages from Gemini parts format to OpenAI content format
        messages = gemini_payload.get('messages', [])
        if not messages:
            logger.error("No messages in payload")
            return jsonify({
                'error': {
                    'message': 'No messages in payload',
                    'type': 'invalid_request_error'
                }
            }), 400
        
        for msg in messages:
            # Extract text from parts array
            parts = msg.get('parts', [])
            
            # Handle both Gemini format (parts array) and OpenAI format (content string)
            if parts:
                # Gemini format: parts array with text
                content_parts = []
                for part in parts:
                    text = part.get('text', '')
                    if text:
                        content_parts.append(text)
                
                if not content_parts:
                    logger.warning(f"Message with role '{msg.get('role')}' has no text content")
                    continue
                
                content = '\n'.join(content_parts)
            else:
                # OpenAI format: direct content field
                content = msg.get('content', '')
                if not content:
                    logger.warning(f"Message with role '{msg.get('role')}' has no content")
                    continue
            
            # Map Gemini roles to OpenAI roles
            role = msg.get('role')
            if role == 'model':
                role = 'assistant'  # OpenAI uses 'assistant', not 'model'
            
            # Handle function calls if present (for OpenAI-compatible format)
            message_dict = {
                'role': role,
                'content': content
            }
            
            # Copy any additional fields that might be present
            additional_fields = ['name', 'function_call', 'tool_calls', 'tool_call_id']
            for field in additional_fields:
                if field in msg:
                    message_dict[field] = msg[field]
            
            iflow_payload['messages'].append(message_dict)
        
        if not iflow_payload['messages']:
            logger.error("No valid messages after transformation")
            return jsonify({
                'error': {
                    'message': 'No valid messages in payload',
                    'type': 'invalid_request_error'
                }
            }), 400
        
        # Call iFlow API
        headers = {
            'Authorization': f'Bearer {IFLOW_API_KEY}',
            'Content-Type': 'application/json',
            'User-Agent': 'iFlow-Proxy/1.0'
        }
        
        response = requests.post(
            IFLOW_API_URL, 
            json=iflow_payload, 
            headers=headers, 
            timeout=60
        )
        
        if response.status_code != 200:
            error_text = response.text[:500]
            logger.error(f"iFlow API error: {response.status_code} - {error_text}")
            try:
                error_json = response.json()
                return jsonify(error_json), response.status_code
            except:
                return jsonify({
                    'error': {
                        'message': f'iFlow API error: {response.status_code}',
                        'type': 'api_error'
                    }
                }), response.status_code
        
        # Return the iFlow response directly (OpenAI-compatible JSON)
        return jsonify(response.json()), 200
        
    except Exception as e:
        logger.error(f"Error in iflow endpoint: {e}")
        return jsonify({
            'error': {
                'message': f'Internal server error: {str(e)}',
                'type': 'server_error'
            }
        }), 500

@app.route('/v1/models', methods=['GET'])
def list_models():
    """List available models (OpenAI-compatible endpoint)"""
    try:
        # Validate authentication
        auth_header = request.headers.get('Authorization')
        if auth_header and auth_header.startswith('Bearer '):
            auth_key = auth_header.replace('Bearer ', '')
            if auth_key != PROXY_AUTH_KEY:
                return jsonify({'error': 'Invalid API key'}), 401
        
        # Call iFlow models endpoint
        headers = {
            'Authorization': f'Bearer {IFLOW_API_KEY}',
            'Content-Type': 'application/json'
        }
        
        response = requests.get(
            'https://apis.iflow.cn/v1/models',
            headers=headers,
            timeout=10
        )
        
        if response.status_code == 200:
            return jsonify(response.json()), 200
        else:
            # Return a default list if iFlow API fails
            return jsonify({
                'object': 'list',
                'data': [
                    {
                        'id': 'glm-4.6',
                        'object': 'model',
                        'created': 1698100000,
                        'owned_by': 'iFlow'
                    },
                    {
                        'id': 'TBStars2-200B-A13B',
                        'object': 'model',
                        'created': 1698200000,
                        'owned_by': 'iFlow'
                    }
                ]
            }), 200
            
    except Exception as e:
        logger.error(f"Error listing models: {e}")
        return jsonify({
            'object': 'list',
            'data': [
                {
                    'id': DEFAULT_MODEL,
                    'object': 'model',
                    'created': 1698100000,
                    'owned_by': 'iFlow'
                }
            ]
        }), 200

@app.route('/test-transform', methods=['POST'])
def test_transform():
    """Test endpoint to see the transformation without calling iFlow"""
    try:
        gemini_payload = request.get_json()
        if not gemini_payload:
            return jsonify({'error': 'No payload received'}), 400
        
        # Transform as in main endpoint
        iflow_payload = {
            'model': gemini_payload.get('model', DEFAULT_MODEL),
            'messages': []
        }
        
        messages = gemini_payload.get('messages', [])
        for msg in messages:
            parts = msg.get('parts', [])
            if parts:
                content = '\n'.join(part.get('text', '') for part in parts)
            else:
                content = msg.get('content', '')
            
            role = msg.get('role')
            if role == 'model':
                role = 'assistant'
                
            iflow_payload['messages'].append({
                'role': role,
                'content': content
            })
        
        return jsonify({
            'transformed_payload': iflow_payload,
            'message_count': len(iflow_payload['messages']),
            'model': iflow_payload['model']
        }), 200
        
    except Exception as e:
        return jsonify({'error': f'Error: {str(e)}'}), 500

@app.route('/health')
def health_check():
    """Health check endpoint"""
    health_status = {
        "status": "healthy",
        "service": "iflow-proxy",
        "default_model": DEFAULT_MODEL,
        "iflow_api_key_set": bool(IFLOW_API_KEY),
        "proxy_auth_key_set": bool(PROXY_AUTH_KEY),
        "endpoints": {
            "groq_compat": "/groq (for existing apps)",
            "iflow": "/iflow",
            "chat_completions": "/chat/completions",
            "list_models": "/v1/models",
            "test_transform": "/test-transform",
            "health": "/health"
        }
    }
    
    # Test iFlow API connectivity if key is set
    if IFLOW_API_KEY:
        try:
            headers = {
                'Authorization': f'Bearer {IFLOW_API_KEY}',
                'Content-Type': 'application/json'
            }
            test_payload = {
                'model': DEFAULT_MODEL,
                'messages': [{'role': 'user', 'content': 'ping'}],
                'max_tokens': 5
            }
            test_response = requests.post(
                'https://apis.iflow.cn/v1/chat/completions',
                json=test_payload,
                headers=headers,
                timeout=15
            )
            health_status["iflow_api_accessible"] = test_response.status_code == 200
            if test_response.status_code == 200:
                health_status["iflow_api_test"] = "Connection successful"
        except Exception as e:
            health_status["iflow_api_accessible"] = False
            health_status["iflow_api_error"] = str(e)[:100]
    
    return jsonify(health_status), 200

@app.errorhandler(404)
def not_found(error):
    return jsonify({
        'error': {
            'message': 'Endpoint not found',
            'type': 'invalid_request_error'
        }
    }), 404

@app.errorhandler(405)
def method_not_allowed(error):
    return jsonify({
        'error': {
            'message': 'Method not allowed',
            'type': 'invalid_request_error'
        }
    }), 405

if __name__ == '__main__':
    logger.info("=" * 50)
    logger.info("Starting iFlow Proxy Server (with /groq compatibility)")
    logger.info(f"Default Model: {DEFAULT_MODEL}")
    logger.info(f"IFLOW_API_KEY set: {bool(IFLOW_API_KEY)}")
    logger.info(f"PROXY_AUTH_KEY set: {bool(PROXY_AUTH_KEY)}")
    logger.info(f"API URL: {IFLOW_API_URL}")
    logger.info("=" * 50)
    app.run(host='0.0.0.0', port=7860, debug=False)