lizhengwei / aigc-embedding-service

Go to a project
Toggle navigation Toggle navigation pinning
maxiaohui
3a30d4b1 1 parent d91dc7ad

jira:NYJ-1447 desc:router -> service -> dao

Inline Side-by-side
Showing 6 changed files with 459 additions and 18 deletions
from typing import List, Optional, Literal, Union, Dict, Any, Annotated
from pydantic import BaseModel, Field
from fastapi import APIRouter
from fastapi import APIRouter, Depends
import numpy
from pydantic import BaseModel, Field, TypeAdapter
from api.resp_bean import RespBean, success
router = APIRouter()
... ... @@ -34,7 +37,7 @@ class PutDataRequest(BaseModel):
embedding: List[float] = Field(..., description="向量")
embedding_version: str = Field(..., description="向量版本 小写数字下划线组成")
kwargs: Optional[Dict[str, Any]] = Field(None, description="扩展字段，根据type不同字段不同")
def get_typed_kwargs(self, type: str) -> Optional[KwargsType]:
"""根据type自动转换kwargs为具体类型"""
if self.kwargs is None:
... ... @@ -52,6 +55,7 @@ class PutDataRequest(BaseModel):
class DelByIdRequest(BaseModel):
"""删除数据请求体"""
ids: List[str] = Field(..., description="id 集合")
embedding_version: str = Field(..., description="向量版本 小写数字下划线组成")
class FilterItem(BaseModel):
... ... @@ -85,35 +89,57 @@ class SearchResponseData(BaseModel):
# ============== 接口定义 ==============
from service.data_service import DataService
from service import get_data_service
@router.put("/{type}/put", response_model=RespBean, tags=["数据服务"], summary="写入数据")
@router.post("/{type}/put", response_model=RespBean, tags=["数据服务"], summary="写入数据")
async def put_data(type: Literal["face", "sport_shot"], request: PutDataRequest):
async def put_data(
type: Literal["face", "sport_shot"],
request: PutDataRequest,
data_service: DataService = Depends(get_data_service)
):
"""
写入数据
"""
request.embedding = numpy.array(request.embedding)
pass
return success()
embedding_version = request.embedding_version
tb_name = f'{type}_{embedding_version}'
upserted_id = await data_service.upsert(tb_name, request.id, request.embedding,
request.get_typed_kwargs(type).model_dump())
return success(data={'id': upserted_id})
@router.delete("/{type}/del_by_id", response_model=RespBean, tags=["数据服务"], summary="删除数据")
@router.post("/{type}/del_by_id", response_model=RespBean, tags=["数据服务"], summary="删除数据")
async def del_by_id(type: Literal["face", "sport_shot"], request: DelByIdRequest):
async def del_by_id(
type: Literal["face", "sport_shot"],
request: DelByIdRequest,
data_service: DataService = Depends(get_data_service)
):
"""
删除数据
- **type**: 数据类型 (face-人脸, sport_shot-体育镜头)
- **ids**: id集合
"""
# TODO: 实现数据删除逻辑
pass
return success()
embedding_version = request.embedding_version
tb_name = f'{type}_{embedding_version}'
await data_service.delete_by_pks(tb_name, request.ids)
return success()
@router.post("/{type}/search", response_model=RespBean[SearchResponseData], tags=["数据服务"], summary="检索数据")
async def search_data(type: Literal["face", "sport_shot"], request: SearchRequest):
async def search_data(
type: Literal["face", "sport_shot"],
request: SearchRequest,
data_service: DataService = Depends(get_data_service),
adapter=TypeAdapter(List[SearchResultItem])
):
"""
检索数据
... ... @@ -123,6 +149,15 @@ async def search_data(type: Literal["face", "sport_shot"], request: SearchReques
- **topk**: 返回top k条结果
- **filters**: 过滤条件列表
"""
# TODO: 实现数据检索逻辑
pass
return success()
#
embedding_version = request.embedding_version
tb_name = f'{type}_{embedding_version}'
embedding = request.embedding
filters = request.filters
from_, size = 0, request.topk
result_data = await data_service.search(tb_name, embedding, filters, from_, size)
search_results: List[SearchResultItem] = adapter.validate_python(result_data)
return success(data=SearchResponseData(type=type, data_list=search_results))
... ...
from functools import lru_cache
from elasticsearch import AsyncElasticsearch
from config import settings
from .data_dao import DataDao
es_client = None
pg_client = None
async def connect():
if settings.db_established:
from .es_client import get_es_client
if settings.db_es_enable:
global es_client
es_client = get_es_client()
await es_client.connect()
es_client = AsyncElasticsearch([settings.db_es_url])
await es_client.ping()
if settings.db_postgres_enable:
from .postgres_client import get_pg_client
global pg_client
... ... @@ -24,3 +28,8 @@ async def disconnect():
es_client = None
if pg_client is not None:
await pg_client.disconnect()
@lru_cache()
def get_data_dao():
return DataDao(es_client)
... ...
import json
import logging
import os
import threading
import typing
from elasticsearch import AsyncElasticsearch
import numpy
logger = logging.getLogger(__name__)
class DataDao:
def __init__(self, es_client: AsyncElasticsearch):
self.es_client = es_client
self._lock = threading.RLock()
async def search(self, tb_name: str,
embedding=typing.Optional[typing.Union[numpy.ndarray, list]],
filters=None,
from_=typing.Optional[int],
size=typing.Optional[int],
):
search_body = dict()
from_ = from_ if isinstance(from_, int) and from_ >= 0 else 0
size = size if isinstance(size, int) and size > 0 else 10
search_body.update({
'from': from_,
'size': size,
})
return_embedding = os.getenv('RETURN_EMBEDDING', 'false').lower().strip()
if return_embedding != 'true':
search_body.update({
"_source": {
"excludes": ["embedding"]
}
})
condition_query = {
"match_all": {}
}
if isinstance(filters, typing.Collection) and len(filters) > 0:
tb_mapping = await self._get_mapping(tb_name)
exist_properties = tb_mapping.get('mappings', {}).get('properties', {})
clauses = []
for filter in filters:
k = filter.name
v = filter.value
opt = filter.opt
if k not in exist_properties:
continue
clause = get_filter_clause(k, v, opt, exist_properties[k])
if clause:
clauses.append(clause)
condition_query = {
"bool": {
"filter": clauses
}
}
if embedding is not None:
use_knn = os.getenv('USE_KNN', 'false').lower().strip()
if use_knn != 'true':
query_body = {
"script_score": {
"query": condition_query,
"script": {
"source": """
// 1. 计算原始分数 (范围 [0, 2])
double rawScore = cosineSimilarity(params.query_vector, 'embedding') + 1.0;
// 2. 归一化到 [0, 1]
double normalizedScore = rawScore / 2.0;
// 3. 强制截断：确保最小值为 0.0，最大值为 1.0
// Math.max 防止出现负数，Math.min 防止出现 > 1.0 的数
double clampedScore = Math.max(0.0, Math.min(1.0, normalizedScore));
// 4. 四舍五入到 4 位小数
// 原理：乘以 10000 -> 四舍五入取整 -> 除以 10000
// clampedScore = Math.round(clampedScore * 10000.0) / 10000.0;
return clampedScore;
""",
"params": {
"query_vector": embedding.tolist() if isinstance(embedding, numpy.ndarray) else list(embedding)
}
}
}
}
search_body.update({
"query": query_body
})
else:
search_body.update({
"knn": {
"field": "embedding",
"query_vector": embedding.tolist() if isinstance(embedding, numpy.ndarray) else list(embedding),
"k": from_ + size,
"num_candidates": (from_ + size) * 20,
"filter": condition_query,
}
})
else:
query_body = condition_query
search_body.update({
"query": query_body
})
logger.info(f'index {tb_name} search body: {json.dumps(search_body, ensure_ascii=False, indent=4)}')
resp = await self.es_client.search(index=tb_name, body=search_body)
logger.info(f'index {tb_name} search response: {resp}')
body = resp.body
hits = body.get('hits', {}).get('hits', [])
result_data = [
{
'id': hit['_id'],
'score': round(hit['_score'], 4),
'kwargs': hit['_source'],
}
for hit in hits
]
return result_data
async def delete_by_pks(self, tb_name: str, pks: typing.List[str]):
query_body = {
"query": {
"terms": {
"_id": pks
}
}
}
resp = await self.es_client.delete_by_query(index=tb_name, body=query_body)
logger.info(f'index {tb_name} delete_by_query response: {resp}')
body = resp.body
# return body.get('deleted')
async def _get_mapping(self, tb_name: str):
resp = await self.es_client.indices.get_mapping(index=tb_name)
logger.info(f'index {tb_name} get_mapping response: {resp}')
return resp.body.get(tb_name, {})
async def _tb_exist(self, tb_name, mapping=None):
resp = await self.es_client.indices.exists(index=tb_name)
result = resp.body
if result and mapping is not None:
tb_mapping = await self._get_mapping(tb_name)
exist_properties = tb_mapping.get('mappings', {}).get('properties', {})
expect_properties = mapping.get('mappings', {}).get('properties', {})
if exist_properties and expect_properties:
shared_keys = set(exist_properties.keys()).intersection(set(expect_properties.keys()))
for k in shared_keys or []:
if exist_properties[k]['type'] != expect_properties[k]['type']:
raise Exception(
f'index {tb_name} `{k}` type not match, expect: {expect_properties[k]["type"]}, actual: {exist_properties[k]["type"]}')
else:
if exist_properties[k]['type'] == 'dense_vector':
if exist_properties[k].get('dims') != expect_properties[k].get('dims'):
raise Exception(
f'index {tb_name} dims not match, expect: {expect_properties[k].get("dims")}, actual: {exist_properties[k].get("dims")}')
return result
async def _create_index(self, tb_name, mapping=typing.Optional[dict]) -> bool:
result = await self._tb_exist(tb_name=tb_name, mapping=mapping)
if result:
return True
with self._lock:
result = await self._tb_exist(tb_name=tb_name, mapping=mapping)
if result:
return True
resp = await self.es_client.indices.create(index=tb_name, body=mapping)
logger.info(f'index {tb_name} create response: {resp}')
return resp.body.get('acknowledged') is True
async def upsert(self, tb_name, id, embedding, params: dict):
doc = {
'embedding': embedding,
**params
}
auto_create_index = os.getenv('AUTO_CREATE_INDEX', 'true').lower().strip()
if auto_create_index == 'true':
mapping = generate_mapping(doc)
# logger.info(f'Possible mapping: {json.dumps(mapping, indent=4)}')
result = await self._create_index(tb_name=tb_name, mapping=mapping)
if isinstance(embedding, numpy.ndarray):
embedding = embedding.tolist()
# elif isinstance(embedding, typing.Iterable):
# if isinstance(embedding, typing.Mapping):
# pass
# else:
# embedding = list(embedding)
doc['embedding'] = embedding
resp = await self.es_client.index(index=tb_name, id=id, body=doc)
logger.info(f'index {tb_name} index response: {resp}')
body = resp.body
assert body.get('result') in ['created', 'updated'], f'数据插入失败, id: {id}'
return body.get('_id')
def get_filter_clause(k, v, opt, type_property):
property_type = type_property.get('type')
if property_type == 'text':
property_fields = type_property.get('fields') or {}
if 'keyword' in property_fields:
k = f'{k}.keyword'
opt = opt.lower()
# "eq", "neq", "lt", "gt", "lte", "gte", "like", "in"
if opt in ['eq', 'in']:
v = list(v) if isinstance(v, typing.Collection) else [v]
should_list = []
if None in v or len(v) == 0:
should_list.append({
"bool": {
"must_not": [
{
"exists": {
"field": k
}
}
]
}
})
no_null_v = [item for item in v if item is not None]
if len(no_null_v) > 0:
should_list.append({
"terms": {
k: no_null_v
}
})
return should_list[0] if len(should_list) == 1 else {
"bool": {
"should": should_list,
"minimum_should_match": 1
}
}
elif opt == 'neq':
must_list = []
v = list(v) if isinstance(v, typing.Collection) else [v]
if None in v or len(v) == 0:
must_list.append({
"bool": {
"must": [
{
"exists": {
"field": k
}
}
]
}
})
no_null_v = [item for item in v if item is not None]
if len(no_null_v) > 0:
must_list.append({
"bool": {
"must_not": {
"terms": {
k: no_null_v
}
}
}
})
return must_list[0] if len(must_list) == 1 else {
"bool": {
"must": must_list
}
}
elif opt in ['lt', 'lte', 'gt', 'gte']:
v = list(v) if isinstance(v, typing.Collection) else [v]
if len(v) > 0 and v[0] is not None:
return {
"range": {
k: {opt: list(v)[0]}
}
}
else:
return {
"range": {
k: {}
}
}
elif opt == 'like':
v = list(v) if isinstance(v, typing.Collection) else [v]
if len(v) > 0 and v[0] is not None:
return {
"wildcard": {
k: f"*{list(v)[0]}*"
}
}
else:
return {
"bool": {
"must_not": [
{
"exists": {
"field": k
}
}
]
}
}
else:
raise Exception(f'opt {opt} not support')
def generate_mapping(doc):
def _inner(v):
if isinstance(v, numpy.ndarray):
return 'dense_vector'
elif isinstance(v, str):
return 'text'
elif isinstance(v, int):
return 'long'
elif isinstance(v, float):
return 'float'
elif isinstance(v, typing.Mapping):
return None
elif isinstance(v, typing.Collection):
if len(v) == 0:
return None
else:
return _inner(list(v)[0])
else:
return None
properties = dict()
for key, value in (doc or {}).items():
key_type = _inner(value)
if key_type:
properties[key] = {
'type': key_type,
}
if key_type == 'dense_vector':
properties[key]['dims'] = len(value)
properties[key].update({
'dims': len(value),
'index': True,
'similarity': 'cosine'
})
elif key_type == 'text':
properties[key]['fields'] = {
"keyword": {
"type": "keyword",
# "ignore_above": 256
}
}
es_mapping = {
"settings": {
"number_of_replicas": int(os.getenv('ES_NUMBER_OF_REPLICAS', 0)),
},
"mappings": {
"properties": properties
}
}
return es_mapping
... ...
... ... @@ -12,6 +12,7 @@ dependencies = [
"omegaconf==2.3.0",
"elasticsearch>=9.3.0",
"packaging==26.0",
"aiohttp==3.13.5",
]
[project.optional-dependencies]
... ...
from functools import lru_cache
from db import get_data_dao
from .data_service import DataService
@lru_cache()
def get_data_service():
data_dao = get_data_dao()
return DataService(data_dao)
... ...
import logging
from db import DataDao
logger = logging.getLogger(__name__)
class DataService:
def __init__(self, data_dao: DataDao):
self.data_dao = data_dao
async def upsert(self, tb_name, id, embedding, params):
return await self.data_dao.upsert(tb_name, id, embedding, params)
async def delete_by_pks(self, tb_name, ids):
return await self.data_dao.delete_by_pks(tb_name, ids)
async def search(self, tb_name, embedding, filters, from_, size):
return await self.data_dao.search(tb_name, embedding, filters, from_, size)
... ...