from openai import AsyncOpenAI

from env import LlmBaseConfig, EmbBaseConfig

llm_client = AsyncOpenAI(
    api_key=LlmBaseConfig.api_key,
    base_url=LlmBaseConfig.base_url,
)

emb_client = AsyncOpenAI(
    api_key=EmbBaseConfig.api_key,
    base_url=EmbBaseConfig.base_url,
)

async def generation_rule(prompt):
  response = await llm_client.chat.completions.create(
    model=LlmBaseConfig.model_name,
    messages=prompt,
    n = 1,
    stream = False,
    temperature=0.0,
    max_tokens=600,
    top_p = 1.0,
    frequency_penalty=0.0,
    presence_penalty=0.0,
    # stop = ["Q:"]
  )
  return response.choices[0].message.content

async def generation_vector(text):
  response = await emb_client.embeddings.create(
    model=EmbBaseConfig.model_name,  # 替换为实际的向量模型名称
    input=text,
    encoding_format="float"
  )
  return response.data[0].embedding

async def rerank_documents(query, documents, top_n=None):
    """
    调用 rerank 模型对文档进行重排序
    :param query: 查询语句
    :param documents: 文档列表
    :param top_n: 返回前n个结果，默认返回所有
    :return: 重排序后的结果
    """
    if top_n is None:
      top_n = min(4, len(documents))

    documents_str = "\n".join([f"{i+1}. {doc}" for i, doc in enumerate(documents)])
    messages = [
      {
        "role": "system",
        "content": "你是一个文档相关性排序助手。请根据查询语句与文档的相关性对文档进行排序，只返回文档序号的排序结果，如：[2, 1, 3]"
      },
      {
        "role": "user",
        "content": f"查询：{query}\n\n候选文档：\n{documents_str}\n\n请按相关性从高到低排序，只返回序号列表："
      }
    ]

    response = await llm_client.chat.completions.create(
      model=LlmBaseConfig.model_name,  # 使用已知可用的模型
      messages=messages,
      temperature=0.0,
      max_tokens=100
    )

    sorted_results = response.choices[0].message.content

    return sorted_results

# 将顶层await包装在异步函数中
async def main():
  # prompt = [
  #   {"role": "system", "content": "你是一个助手"},
  #   {"role": "user", "content": "你是谁"}
  # ]
  # ans = await generation_rule(prompt)

  # ans = await generation_vector("hello world")

  documents = [
    "人工智能是计算机科学的一个分支，它企图了解智能的实质，并生产出一种新的能以人类智能相似的方式做出反应的智能机器。",
    "机器学习是人工智能的一个分支，主要研究计算机如何从数据中学习规律，并利用这些规律对未知数据进行预测。",
    "深度学习是机器学习的一个子集，它模仿人脑的工作方式，通过多层神经网络来学习数据的特征。"
  ]
  ans = await rerank_documents(query="机器学习", documents=documents, top_n=2)

  print(ans)
  return ans

# 使用asyncio运行异步函数
import asyncio
if __name__ == "__main__":
    asyncio.run(main())