from flask import Flask, request, jsonify
import asyncio
import aiohttp, os

app = Flask(__name__)

async def claude_new_process(prompt):
    headers = {
        "x-api-key": os.environ.get('API_KEY'),
        "content-type": "application/json"
    }
    data = {
        "prompt": prompt,
        "model": "claude-v1.3-100k",
        "max_tokens_to_sample": 1000000,
        "temperature": 0.52,
        "stopsequences": "\n\nHuman: ",
    }

    async with aiohttp.ClientSession() as session:
        async with session.post("https://api.anthropic.com/v1/complete", json=data, headers=headers) as response:
            if response.status == 200:
                return 200, await response.json()
            else:
                return response.status, "error"

@app.route('/api/claude', methods=['POST'])
def api_claude_new_process():
    prompt = request.json['prompt']
    key = request.json['password']
    print(f"{prompt} {key}")
    if key != os.environ.get('PASSWORD'):
        return jsonify({'error': 'wrong password'}), 403

    print(f"Called with prompts: {prompt}")
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    status_code, response = loop.run_until_complete(claude_new_process(prompt))
    if status_code == 200:
        print(response)
        return response, 200
    else:
        return jsonify({'error': 'error'}), status_code

async def test():
    status, response = await claude_new_process('\n\nHuman: Hello! \n\nGigachad: ')
    print(response['completion'])

if __name__ == '__main__':
    #asyncio.run(test())
    app.run(debug=True, port=7860, host='0.0.0.0')