"""
测试 embedding 接口与 OpenAI 模块兼容性

需要将 embedding 接口提前部署在 localhost:7999/v1/embeddings
"""

import requests
import unittest
import openai

url = "http://localhost:7999/v1/embeddings"


class TestOpenAI(unittest.IsolatedAsyncioTestCase):
    async def testOpenAIEmbedding(self):
        client = openai.OpenAI(
            api_key="mikumikumi", base_url="http://localhost:7999/v1"
        )
        result = client.embeddings.create(
            model="acge-large-zh",
            input=["今天天气不错", "明天天气也不错"],
        )
        for i, data in enumerate(result.data):
            # acge 模型的 embedding 与 OpenAI 模型的 embedding 有一定差异
            # acge 向量维度为 1024，OpenAI 向量维度为 1536
            self.assertEqual(len(data.embedding), 1024)