import gradio as gr



# 读取图片
import base64
# Function to encode the image
def encode_image(image_path):
  with open(image_path, "rb") as image_file:
    return base64.b64encode(image_file.read()).decode('utf-8')

import requests
import os

openai_api_key = os.environ.get('openai_api_key')

def ask_image(image,text,api_token=openai_api_key):
  base64_image = encode_image(image)
  messages=[
      {
        "role": "user",
        "content": [
          {"type": "text", "text": text},
          {
            "type": "image_url",
            "image_url": {
                "url":f"data:image/jpeg;base64,{base64_image}"
  #               "url": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg",
            },
          },
        ],
      }
    ]

  # 请求头部信息
  headers = {
      'Authorization': f'Bearer {api_token}'
  }

  # 请求体信息
  data = {
      'model': 'gpt-4o',  # 可以根据需要更换其他模型
      'messages': messages,
      'temperature': 0.7  # 可以根据需要调整
  }


  # 设定最大重试次数
  max_retry = 3

  for i in range(max_retry):
      try:
          # 发送请求
          response = requests.post('https://burn.hair/v1/chat/completions', headers=headers, json=data)

          # 解析响应内容
          response_data = response.json()
          response_content = response_data['choices'][0]['message']['content']
          usage = response_data['usage']

          return response_content

      except Exception as e:
          # 如果已经达到最大重试次数，那么返回空值
          if i == max_retry - 1:
              print(f'重试次数已达上限，仍未能成功获取数据，错误信息：{e}')
              response_content = ''
              usage = {}
              return response_content
          else:
              # 如果未达到最大重试次数，打印错误信息，并继续下一次循环
              print(f'第{i+1}次请求失败，错误信息：{e}，准备进行第{i+2}次尝试')


# gradio demo

title = "Ask Image"
description = "Ask anything about your Image"

demo = gr.Interface(
    fn=ask_image,
    inputs=[gr.Image(label='',type='filepath'), gr.Text(label="Question")],
    outputs=[gr.Textbox(label="Answer",lines=3)],
    title = title,
    description = description
)
demo.queue(max_size = 20)

demo.launch(share = True)