llm_vector_server/applications/api/embedding.py

from applications.config import LOCAL_MODEL_CONFIG, VLLM_SERVER_URL, DEV_VLLM_SERVER_URL
from applications.utils import AsyncHttpClient


async def get_basic_embedding(text: str, model: str, dev=False):
    """
    embedding text into vectors
    :param dev: 是否测试环境
    :param text:
    :param model:
    :return:tong
    """
    model_path = LOCAL_MODEL_CONFIG[model]
    async with AsyncHttpClient(timeout=20) as client:
        response = await client.post(
            url=DEV_VLLM_SERVER_URL if dev else VLLM_SERVER_URL,
            json={"input": text, "model": model_path},
            headers={"Content-Type": "application/json"},
        )
        return response["data"][0]["embedding"]


__all__ = ["get_basic_embedding"]