# Embedding & Rerank 服务_HTTP协议

import openai
import httpx

# 必填：从服务管控页面获取对应服务的APIKey和API Base
api_key = "<YOUR_API_KEY>"
base_url = "https://maas-api.cn-huabei-1.xf-yun.com/v1"

client = openai.OpenAI(api_key=api_key, base_url=base_url)

def get_embedding(text, model="sde0a5839"):
    """
    调用 Embedding API。
    :param text: 需要进行向量化的一段或多段文本。
    :param model: 使用的 embedding 模型 ID。
    """
    try:
        response = client.embeddings.create(
            model=model,
            input=text,
        )
        return response.data
    except Exception as e:
        print(f"Embedding 请求出错: {e}")
        return None

def get_rerank(query, documents, model="s125c8e0e"):
    """
    调用 Rerank API。
    :param query: 查询语句。
    :param documents: 需要重排序的文档列表。
    :param model: 使用的 rerank 模型 ID。
    """
    try:
        # 注意：OpenAI SDK 没有原生的 rerank 方法，需要通过 client.post 发起原始请求
        response = client.post(
            "/rerank",
            body={
                "model": model,
                "query": query,
                "documents": documents,
            },
            cast_to=httpx.Response,
        )
        return response.json()
    except Exception as e:
        print(f"Rerank 请求出错: {e}")
        return None

if __name__ == "__main__":
    # --- Embedding Demo ---
    print("********* 1. Embedding API 调用 *********")
    embedding_text = "夏天最好吃的水果是西瓜"
    embeddings = get_embedding(embedding_text)
    if embeddings:
        for i, embedding_data in enumerate(embeddings):
            print(f"文本片段 {i} 的向量 (前5维): {embedding_data.embedding[:5]}...")
            print(f"向量总维度: {len(embedding_data.embedding)}")
    
    # --- Rerank Demo ---
    print("********* 2. Rerank API 调用 *********")
    query = "夏天最好吃的水果是？"
    documents = ["西瓜", "草莓", "荔枝"]
    rerank_results = get_rerank(query, documents)
    if rerank_results:
        print(f"查询: {query}")
        print("重排序结果:")
        # 注意：rerank 返回的 results 列表并未按分数排序，需要自行处理
        for result in rerank_results.get("results", []):
            doc_index = result.get('index')
            score = result.get('relevance_score')
            print(f"  - 文档: '{documents[doc_index]}' (相关性得分: {score})")