import requests
import json

BASE_URL = "https://"
API_KEY = "sk-"

headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {API_KEY}"
}

# ============================================================
# 使用 requests 调用 embedding 接口
# ============================================================

content = "What is the meaning of life?"

# 尝试多种端点和请求格式
endpoints_and_payloads = [
    # 格式1: OpenAI 兼容格式 /v1/embeddings
    (
        f"{BASE_URL}/v1/embeddings",
        {
            "model": "gemini-embedding-001",
            "input": content
        }
    ),
    # 格式2: 单字符串 input
    (
        f"{BASE_URL}/v1/embeddings",
        {
            "model": "gemini-embedding-001",
            "input": [content]
        }
    ),
    # 格式3: Google API 风格 embed_content
    (
        f"{BASE_URL}/v1/models/gemini-embedding-001:embedContent",
        {
            "content": {
                "parts": [{"text": content}]
            }
        }
    ),
]

result = None

for endpoint, payload in endpoints_and_payloads:
    print(f"尝试端点: {endpoint}")
    print(f"请求体: {json.dumps(payload, indent=2)}")

    try:
        resp = requests.post(endpoint, headers=headers, json=payload, timeout=30)
        print(f"  状态码: {resp.status_code}")
        print(f"  响应: {resp.text[:800]}")

        if resp.status_code == 200:
            result = resp.json()
            print(f"  ✓ 请求成功!")
            break
        else:
            print(f"  ✗ 失败，尝试下一个格式\n")
    except Exception as e:
        print(f"  异常: {e}\n")

if result is None:
    print("\n❌ 所有端点均失败")
    exit(1)

# ============================================================
# 提取 embedding 向量
# ============================================================
print("\n" + "=" * 60)
print("Embedding 结果:")
print("=" * 60)

embeddings = None

# OpenAI 格式: data[0].embedding
data_list = result.get("data", [])
if data_list and isinstance(data_list, list):
    item = data_list[0]
    if isinstance(item, dict) and "embedding" in item:
        embeddings = item["embedding"]

# Google 格式: embedding.values
if embeddings is None:
    emb_obj = result.get("embedding", {})
    if isinstance(emb_obj, dict) and "values" in emb_obj:
        embeddings = emb_obj["values"]

# Google 批量格式: embeddings[0].values
if embeddings is None:
    emb_list = result.get("embeddings", [])
    if isinstance(emb_list, list) and len(emb_list) > 0:
        item = emb_list[0]
        if isinstance(item, dict) and "values" in item:
            embeddings = item["values"]
        elif isinstance(item, list):
            embeddings = item

# 直接是列表
if embeddings is None and isinstance(result.get("embedding"), list):
    embeddings = result["embedding"]

if embeddings:
    print(f"向量维度: {len(embeddings)}")
    print(f"前10维: {embeddings[:10]}")
    print(f"完整向量: {embeddings}")
else:
    print("无法提取 embedding，完整响应:")
    print(json.dumps(result, indent=2))