server.py

from __future__ import annotations

import base64
import glob
import json
import logging
import os
import subprocess
from typing import Optional

import ollama
import requests
from dotenv import load_dotenv
from flask import Flask, abort, jsonify, request
from opentelemetry import metrics, trace
from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.flask import FlaskInstrumentor
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor

from .redis_helper import REDIS_CONNECTION

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

# Load environment variables from .env file
load_dotenv()

# OpenTelemetry endpoint configuration
OTEL_EXPORTER_OTLP_ENDPOINT = os.getenv('OTEL_EXPORTER_OTLP_ENDPOINT', 'http://localhost:4317')

# Configure OpenTelemetry - Shared Resource (identifies this service)
resource = Resource.create({"service.name": "ai-server"})

# ========== TRACES CONFIGURATION ==========
# TracerProvider: Factory for creating tracers (for distributed tracing)
tracer_provider = TracerProvider(resource=resource)

# OTLP Trace Exporter: Sends traces to collector
otlp_trace_exporter = OTLPSpanExporter(endpoint=OTEL_EXPORTER_OTLP_ENDPOINT, insecure=True)
span_processor = BatchSpanProcessor(otlp_trace_exporter)
tracer_provider.add_span_processor(span_processor)

# Set the global tracer provider (FlaskInstrumentor will use this)
trace.set_tracer_provider(tracer_provider)
tracer = trace.get_tracer("ai-server.tracer")

# ========== METRICS CONFIGURATION ==========
# OTLP Metric Exporter: Sends metrics to collector
otlp_metric_exporter = OTLPMetricExporter(endpoint=OTEL_EXPORTER_OTLP_ENDPOINT, insecure=True)

# PeriodicExportingMetricReader: Collects and exports metrics every 10 seconds
metric_reader = PeriodicExportingMetricReader(
    exporter=otlp_metric_exporter, export_interval_millis=10000  # Export every 10 seconds
)

# MeterProvider: Factory for creating meters (for metrics collection)
meter_provider = MeterProvider(resource=resource, metric_readers=[metric_reader])

# Set the global meter provider (FlaskInstrumentor will use this for HTTP metrics)
metrics.set_meter_provider(meter_provider)

app = Flask('AI server')
FlaskInstrumentor().instrument_app(app)

# Configuration from environment variables
DEFAULT_MODEL = os.getenv('DEFAULT_MODEL', 'deepseek-coder-v2:latest')

# Llama.cpp configuration
LLAMA_CPP_CLI = os.getenv('LLAMA_CPP_CLI', '/data1/llama.cpp/bin/llama-cli')
GGUF_DIR = os.getenv('GGUF_DIR', '/data1/GGUF')

# Llama server configuration
# e.g., http://localhost:8080 or localhost:8080
_llama_server_url = os.getenv('LLAMA_SERVER_URL')
LLAMA_SERVER_URL = (
    f"http://{_llama_server_url}"
    if _llama_server_url and not _llama_server_url.startswith(('http://', 'https://'))
    else _llama_server_url
)
SCHEMA_KEY = "schema"


def _build_messages(content: str, system_prompt: Optional[str] = None, image_files: Optional[list] = None) -> list:
    """Build messages list with optional system prompt."""
    messages = []
    if system_prompt:
        messages.append({'role': 'system', 'content': system_prompt})
    messages.append({'role': 'user', 'content': content})

    if image_files:
        messages[-1]["images"] = [base64.b64encode(image_file.read()).decode("utf-8") for image_file in image_files]
    return messages


def chat_with_llama_server_http(
    model: str,
    content: str,
    system_prompt: Optional[str] = None,
    timeout: int = 300,
    image_files: Optional[list] = None,
    json_schema: Optional[dict] = None,
    model_options: Optional[dict] = None,
) -> str:
    """Handle chat using llama-server HTTP API."""
    if not LLAMA_SERVER_URL:
        raise Exception("LLAMA_SERVER_URL environment variable not set")

    try:
        messages = _build_messages(content, system_prompt, image_files=[])  # TODO: Pass image files

        if not model_options:
            model_options = {}

        payload = {'model': model, 'messages': messages, **model_options}
        if json_schema:
            payload['json_schema'] = json_schema[SCHEMA_KEY]
        if 'stream' not in payload:
            payload['stream'] = False

        response = requests.post(
            f'{LLAMA_SERVER_URL}/v1/chat/completions',
            json=payload,
            headers={'Content-Type': 'application/json'},
            timeout=timeout,
        )

        if response.status_code == 200:
            data = response.json()
            if 'choices' in data and len(data['choices']) > 0:
                return data['choices'][0]['message']['content']
            else:
                raise Exception("Invalid response format from llama-server")
        else:
            raise Exception(f"Llama-server HTTP error")

    except requests.Timeout:
        raise Exception(f"Llama-server request timed out for model {model}")
    except requests.RequestException as e:
        raise Exception(f"Llama-server request failed: {str(e)}")


def resolve_model_path(model: str) -> Optional[str]:
    """Resolve model name to full GGUF file path using glob pattern."""
    pattern = os.path.join(GGUF_DIR, model, "*.gguf")
    matches = glob.glob(pattern)
    return matches[0] if matches else None


def is_llamacpp_available(model: str) -> bool:
    """Check if model is available in llama.cpp."""
    return resolve_model_path(model) is not None


def chat_with_ollama(
    model: str,
    content: str,
    system_prompt: Optional[str] = None,
    image_files: Optional[list] = None,
    json_schema: Optional[dict] = None,
    model_options: Optional[dict] = None,
) -> str:
    """Handle chat using ollama."""
    messages = _build_messages(content, system_prompt, image_files)

    response = ollama.chat(
        model=model,
        messages=messages,
        stream=False,
        format=json_schema[SCHEMA_KEY] if json_schema else None,
        options=model_options,
    )
    return response.message.content


def chat_with_llamacpp(
    model: str,
    content: str,
    system_prompt: Optional[str] = None,
    timeout: int = 300,
    image_files: Optional[list] = None,
    model_options: Optional[dict] = None,
    json_schema: Optional[dict] = None,
) -> str:
    """Handle chat using llama.cpp CLI."""
    model_path = resolve_model_path(model)

    if not model_path:
        raise ValueError(f"Model not found: {model}")

    cmd = [LLAMA_CPP_CLI, '-m', model_path, '--n-gpu-layers', '40', '-p', content, '-n', '512', '--single-turn']
    if json_schema:
        raw_schema = json_schema[SCHEMA_KEY] if SCHEMA_KEY in json_schema else json_schema
        cmd += ["--json-schema", json.dumps(raw_schema)]

    # Add system prompt if provided
    if system_prompt:
        cmd.extend(['--system-prompt', system_prompt])

    if model_options:
        for key, value in model_options.items():
            cmd.extend(['--model-option', key, value])

    if image_files:
        pass  # TODO: pass image files

    try:
        result = subprocess.run(cmd, capture_output=True, text=False, timeout=timeout, check=True)

        stdout_text = result.stdout.decode('utf-8', errors='replace')

        # Strip whitespace and return the response
        response = stdout_text.strip()
        return response if response else "No response generated."

    except subprocess.TimeoutExpired:
        raise Exception(f"Llama.cpp request timed out for model {model}")
    except subprocess.CalledProcessError as e:
        stderr_text = ""
        if e.stderr:
            stderr_text = e.stderr.decode('utf-8', errors='replace')
        raise Exception(f"Llama.cpp failed for {model}: {stderr_text.strip() if stderr_text else 'Unknown error'}")
    except FileNotFoundError:
        raise Exception("Llama.cpp CLI not found")


def chat_with_model(
    model: str,
    content: str,
    llama_mode: str = "cli",
    system_prompt: Optional[str] = None,
    image_files: Optional[list] = None,
    model_options: Optional[dict] = None,
    json_schema: Optional[dict] = None,
) -> str:
    """Route chat request based on llama_mode: server (external), cli, or ollama fallback; and with optional system prompt."""
    if is_llamacpp_available(model):
        if llama_mode == "server":
            if not LLAMA_SERVER_URL:
                raise Exception("LLAMA_SERVER_URL environment variable not set for server mode")
            return chat_with_llama_server_http(
                model,
                content,
                system_prompt=system_prompt,
                image_files=image_files,
                json_schema=json_schema,
                model_options=model_options,
            )
        elif llama_mode == "cli":
            return chat_with_llamacpp(
                model,
                content,
                system_prompt=system_prompt,
                image_files=image_files,
                json_schema=json_schema,
                model_options=model_options,
            )
        else:
            raise ValueError(f"Invalid llama_mode: '{llama_mode}'. Valid options are 'server' or 'cli'.")
    else:
        # Model not available in llama.cpp, use ollama
        return chat_with_ollama(
            model,
            content,
            system_prompt=system_prompt,
            image_files=image_files,
            json_schema=json_schema,
            model_options=model_options,
        )


def authenticate() -> str:
    """Authenticate the given request using an API key."""
    api_key = request.headers.get('X-API-KEY')
    client_ip = request.remote_addr
    endpoint = request.path
    if not api_key:
        logger.warning(f"Missing API key from {client_ip} at {endpoint}")
        abort(401, description="Missing API key")

    user = REDIS_CONNECTION.get(f"api-key:{api_key}")
    if not user:
        logger.warning(f"Invalid API key attempt from {client_ip} at {endpoint}")
        abort(401, description="Invalid API key")

    return user


@app.route('/chat', methods=['POST'])
def chat():
    """Handle chat request with optional llama_mode and system prompt parameters."""
    authenticate()
    model = request.form.get('model', DEFAULT_MODEL)
    content = request.form.get('content', '')
    llama_mode = request.form.get('llama_mode', 'cli')
    system_prompt = request.form.get('system_prompt')
    image_files = list(request.files.values())
    model_options = request.form.get('model_options')
    json_schema = request.form.get('json_schema')
    if json_schema:
        json_schema = json.loads(json_schema)

    if not content.strip():
        abort(400, description='Missing prompt content')

    response_content = chat_with_model(
        model, content, llama_mode, system_prompt, image_files, model_options=model_options, json_schema=json_schema
    )
    return jsonify(response_content)


@app.errorhandler(Exception)
def internal_error(error):
    return jsonify({"error": str(error)}), 500