rag_server/applications/api/embedding.py

from typing import List
from applications.config import LOCAL_MODEL_CONFIG, VLLM_SERVER_URL, DEV_VLLM_SERVER_URL
from applications.utils.http 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"]


async def get_img_embedding(url_list: List[str], dev=False):
    url = "http://117.50.199.192:8011/api/embed_image"
    async with AsyncHttpClient(timeout=20) as client:
        response = await client.post(
            url=url,
            json={"url_list": url_list},
            headers={"Content-Type": "application/json"},
        )
        return response


__all__ = ["get_basic_embedding", "get_img_embedding"]