from flask import Flask, request, jsonify
from chat import generate_response, tokenizer
import requests

app = Flask(__name__)

def get_public_ip():
    try:
        response = requests.get("https://api.ipify.org?format=json")
        response.raise_for_status()
        data = response.json()
        return data["ip"]
    except requests.exceptions.RequestException as e:
        print("Error:", e)
        return None

@app.route('/openai/v1/chat/completions', methods=['POST'])
def chat_completions():
    data = request.get_json()
    message = data['messages'][0]['content']
    history = data.get('history', [])
    response = generate_response(message, history)
    
    public_ip = get_public_ip()
    
    return jsonify({
        'id': 'chat_completion_id',
        'object': 'chat.completion',
        'created': 1234567890,
        'model': 'Samanthav3-MIXTRALDOLPHIN-LORA',
        'choices': [
            {
                'index': 0,
                'message': {
                    'role': 'assistant',
                    'content': response
                },
                'finish_reason': 'stop'
            }
        ],
        'usage': {
            'prompt_tokens': len(tokenizer.encode(message)),
            'completion_tokens': len(tokenizer.encode(response)),
            'total_tokens': len(tokenizer.encode(message)) + len(tokenizer.encode(response))
        },
        'public_ip': public_ip
    })

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=7860)