# LLM service using OpenAI Chat Completions API
#
# https://platform.openai.com/docs/guides/text-generation
#
# Endpoints: GET /run?task=..., POST /run
#
# Supported request parameters (llm2-compatible):
#   task         - required, the user prompt / instruction
#   max_tokens   - max output tokens (also accepted as max_new_tokens), default: 500
#   temperature  - sampling temperature 0..2, default: 0.2
#   top_p        - nucleus sampling, default: 0.9
#   n            - number of completions to return, default: 1
#
# Removed (not applicable to OpenAI API):
#   num_beams, early_stopping, max_time, min_tokens
#
# Configuration via environment variables:
#   OPENAI_API_KEY     - required, your OpenAI API key
#   OPENAI_MODEL       - optional, default: o4-mini
#   OPENAI_BASE_URL    - optional, default: https://api.openai.com/v1
#   OPENAI_SYSTEM_PROMPT - optional, overrides the default system prompt
#
import os
import json
import logging
import time

from openai import OpenAI
from flask import Flask, abort, jsonify, request

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO,
)
logger = logging.getLogger(__name__)

# Configuration from environment variables
OPENAI_API_KEY  = os.environ.get("OPENAI_API_KEY", "")
OPENAI_MODEL    = os.environ.get("OPENAI_MODEL", "o4-mini")
OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")

DEFAULT_SYSTEM_PROMPT = (
    "You are a precise search assistant. "
    "Follow the instruction exactly. "
    "Return only the requested output — no explanations, no extra text, no markdown."
)
OPENAI_SYSTEM_PROMPT = os.environ.get("OPENAI_SYSTEM_PROMPT", DEFAULT_SYSTEM_PROMPT)

if not OPENAI_API_KEY:
    logger.warning("OPENAI_API_KEY is not set — all requests will fail!")

openai_client = OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL)

logger.info("OpenAI LLM service started")
logger.info("Model    : " + OPENAI_MODEL)
logger.info("Base URL : " + OPENAI_BASE_URL)


def run(data):
    if 'task' not in data or len(data['task']) == 0:
        abort(500)

    task = data['task']

    # max_tokens — accept both llm2 field names
    if 'max_tokens' in data:
        max_tokens = int(data['max_tokens'])
    elif 'max_new_tokens' in data:
        max_tokens = int(data['max_new_tokens'])
    else:
        max_tokens = 500

    temperature = float(data.get('temperature', 0.2))
    top_p       = float(data.get('top_p', 0.9))
    n           = int(data.get('n', 1)) if data.get('n', 0) else 1

    messages = [
        {"role": "system", "content": OPENAI_SYSTEM_PROMPT},
        {"role": "user",   "content": task}
    ]

    response = openai_client.chat.completions.create(
        model=OPENAI_MODEL,
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        n=n
    )

    if n > 1:
        result = [choice.message.content.strip() for choice in response.choices]
    else:
        result = response.choices[0].message.content.strip()

    logger.info(f"Generated: {result!r}")
    return result


#====================================
#http://localhost:8087/run?task=Who+are+you%3F

app = Flask(__name__)

@app.route("/run", methods=['GET'])
def run_get():
    task = request.args.get('task', '')
    if len(task) == 0:
        abort(500)
    logger.info('Task: ' + task)
    result = run({"task": task})
    return jsonify({'result': result})

@app.route("/run", methods=['POST'])
def run_post():
    data = json.loads(request.data)
    result = run(data)
    return jsonify({'result': result})

if __name__ == '__main__':
    app.run(host="0.0.0.0")
#    app.run(host="0.0.0.0", port=80, debug=False)
