lizhengwei / aigc-embedding-service

Go to a project
Toggle navigation Toggle navigation pinning
wangdongxu
f2b6aa51 1 parent ab03fe47

jira:NYJ-1460 desc:init

Inline Side-by-side
Showing 22 changed files with 655 additions and 167 deletions
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# poetry
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
#poetry.lock
# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control. Use .pdm.toml for personal settings and store shared settings in pyproject.toml.
.pdm.toml
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
environment.yml
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/
# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file.
# For a more comprehensive ignore file, see:
# https://github.com/github/gitignore/blob/main/Python.gitignore
# VS Code
/live-stream-ai-task-workline/modules/wtt_highlight/model_weights/*
*mov
/.externalToolBuilders
/.settings
.svn
/.project
/target
./target
/.classpath
.idea/
*.iml
*.iws
*.ipr
*.ids
*.orig
classes/
/bin/
*/.classpath
*/.project
*/.settings/*
######################
# Project Specific
######################
/target/www/**
/src/test/javascript/coverage/
/src/test/javascript/PhantomJS*/
######################
# Node
######################
/node/
node_tmp/
node_modules/
npm-debug.log.*
/.awcache/*
######################
# SASS
######################
.sass-cache/
######################
# Eclipse
######################
*.pydevproject
.project
.metadata
tmp/
tmp/**/*
*.tmp
*.bak
*.swp
*~.nib
local.properties
.classpath
.settings/
.loadpath
.factorypath
/src/main/resources/rebel.xml
# External tool builders
.externalToolBuilders/**
# Locally stored "Eclipse launch configurations"
*.launch
# CDT-specific
.cproject
# PDT-specific
.buildpath
######################
# Intellij
######################
.idea/
*.iml
*.iws
*.ipr
*.ids
*.orig
classes/
######################
# Visual Studio Code
######################
.vscode/
######################
# Maven
######################
/log/
/target/
######################
# Gradle
######################
.gradle/
/build/
######################
# Package Files
######################
*.jar
*.war
*.ear
*.db
######################
# Windows
######################
# Windows image file caches
Thumbs.db
# Folder com.migu.oes.power.config file
Desktop.ini
# macOS
######################
# Mac OSX
######################
.DS_Store
.AppleDouble
.LSOverride
# Linux
.svn
# Thumbnails
._*
# Files that might appear on external disk
.Spotlight-V100
.Trashes
######################
# Directories
######################
/bin/
/deploy/
######################
# Logs
######################
*.log*
######################
# Others
######################
*.class
*.*~
*~
.nfs*
\ No newline at end of file
.merge_file*
######################
# Gradle Wrapper
######################
!gradle/wrapper/gradle-wrapper.jar
######################
# Maven Wrapper
######################
!.mvn/wrapper/maven-wrapper.jar
######################
# ESLint
######################
.eslintcache
*.lst
oes-nryy-vo/target/maven-archiver/pom.properties
oes-nryy-dataservice/target/maven-archiver/pom.properties
oes-nryy-base/target/maven-archiver/pom.properties
oes-nryy-copyright/bin/.gitkeep
oes-nryy-copyright/bin/cbak
*.original
*.gz
oes-nryy-tag/target/oes-nryy-tag-1.0.0/label/bin/service.sh
oes-nryy-tag/target/maven-archiver/pom.properties
oes-nryy-data-trans/target/maven-archiver/pom.properties
oes-nryy-tag/target/oes-nryy-tag-1.0.0/label/bin/cbak
######################
# pp-scene-det
######################
__pycache__/
######################
# live-stream-ai-task-workline
######################
modules/basketball/weights
modules/singer/weights
weights/
weights1/
files/
*.log
*.log.*
cache_data/
data/
test/
# Media files
*.jpg
*.jpeg
*.png
*.gif
*.mp4
*.avi
*.mov
*.mkv
# OS
.DS_Store
Thumbs.db
... ...
from pydantic import BaseModel, Field
from typing import Generic, TypeVar, Optional
T = TypeVar('T')
class RespBean(BaseModel, Generic[T]):
code: str = Field(default="success", description="code")
message: str = Field(default="success", description='message')
data: Optional[T] = Field(description="数据", default=None)
def error(message):
return RespBean(code="error", message=message, data=None)
def success(data=None):
return RespBean(data=data)
... ...
"""API路由聚合模块."""
from fastapi import APIRouter
# 创建主路由
api_router = APIRouter(prefix="/api/v1")
# 注册各场景路由
def register_router(router: APIRouter, prefix: str):
api_router.include_router(router, prefix=prefix)
... ...
app_name: "Embedding Service"
app_version: "1.0.0"
port: 8000
debug: true
db_es_enable: true
db_es_url: "http://localhost:9200"
db_postgres_enable: false
db_postgres_url: "postgresql://postgres:postgres@localhost:5432/postgres"
... ...
from aabd.base.cfg_loader import load_yaml_by_name_with_env
from aabd.base.enhance_dict import EnhanceDict, read_prefixed_env_vars
settings = EnhanceDict(load_yaml_by_name_with_env('config'))
settings.update(read_prefixed_env_vars('APP_'))
... ...
"""自定义异常模块."""
class EmbeddingServiceError(Exception):
"""Embedding服务基础异常."""
def __init__(self, message: str, code: str = "EMBEDDING_ERROR"):
self.message = message
self.code = code
super().__init__(self.message)
class EmbeddingAPIError(EmbeddingServiceError):
"""Embedding API调用异常."""
def __init__(self, message: str, status_code: int = 500):
super().__init__(message, code="EMBEDDING_API_ERROR")
self.status_code = status_code
class DatabaseError(EmbeddingServiceError):
"""数据库操作异常."""
def __init__(self, message: str):
super().__init__(message, code="DATABASE_ERROR")
class NotFoundError(EmbeddingServiceError):
"""资源不存在异常."""
def __init__(self, message: str):
super().__init__(message, code="NOT_FOUND")
class ValidationError(EmbeddingServiceError):
"""参数校验异常."""
def __init__(self, message: str):
super().__init__(message, code="VALIDATION_ERROR")
... ...
"""数据库模块，提供统一的向量数据库抽象层.
支持Elasticsearch和PostgreSQL两种后端，通过配置自动切换。
使用示例:
from app.db import db_client
# 连接数据库
await db_client.connect()
# 创建集合
await db_client.create_collection("face_vectors", vector_dim=512)
# 插入数据
await db_client.insert("face_vectors", "user_001", vector, {"name": "张三"})
# 向量搜索
results = await db_client.search("face_vectors", query_vector, top_k=10)
# 断开连接
await db_client.disconnect()
"""
from typing import Optional
from app.core.config import settings
es_client = None
pg_client = None
async def connect():
if settings.db_established:
from .es_client import get_es_client
global es_client
es_client = get_es_client()
await es_client.connect()
if settings.db_postgres_enable:
from .postgres_client import get_pg_client
global pg_client
pg_client = get_pg_client()
await pg_client.connect()
async def disconnect():
global es_client, pg_client
if es_client is not None:
await es_client.disconnect()
es_client = None
if pg_client is not None:
await pg_client.disconnect()
... ...
"""数据库抽象基类，定义统一的CRUD接口."""
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional
class VectorDBClient(ABC):
"""向量数据库客户端抽象基类.
定义了统一的向量数据CRUD接口，支持Elasticsearch和PostgreSQL等后端。
"""
@abstractmethod
async def connect(self) -> None:
"""建立数据库连接."""
pass
@abstractmethod
async def disconnect(self) -> None:
"""断开数据库连接."""
pass
@abstractmethod
async def health_check(self) -> bool:
"""检查数据库连接健康状态.
Returns:
bool: 连接正常返回True，否则返回False
"""
pass
... ...
"""Elasticsearch向量数据库客户端实现."""
from typing import Any, Dict, List, Optional
from app.core.config import settings
from app.db.base import VectorDBClient
class ElasticsearchClient(VectorDBClient):
"""Elasticsearch向量数据库客户端.
使用Elasticsearch的dense_vector类型存储向量，支持向量相似度搜索。
"""
def __init__(self, es_url: Optional[str] = None):
"""初始化ES客户端.
Args:
es_url: Elasticsearch连接URL，默认从配置读取
"""
self.es_url = es_url or settings.get("db_es_url", "http://localhost:9200")
self._client: Optional[Any] = None
async def connect(self) -> None:
"""建立ES连接."""
from elasticsearch import AsyncElasticsearch
self._client = AsyncElasticsearch([self.es_url])
# 验证连接
await self._client.ping()
async def disconnect(self) -> None:
"""断开ES连接."""
if self._client:
await self._client.close()
self._client = None
async def health_check(self) -> bool:
"""检查ES连接健康状态."""
if not self._client:
return False
try:
return await self._client.ping()
except Exception:
return False
_es_client: Optional[ElasticsearchClient] = None
def get_es_client() -> ElasticsearchClient:
"""获取全局ES客户端实例（单例模式）.
Returns:
ElasticsearchClient: ES客户端实例
"""
global _es_client
if _es_client is None:
_es_client = ElasticsearchClient()
return _es_client
... ...
"""PostgreSQL向量数据库客户端实现."""
import json
from typing import Any, Dict, List, Optional
from app.db.base import VectorDBClient
class PostgresClient(VectorDBClient):
"""PostgreSQL向量数据库客户端.
使用pgvector扩展存储向量，支持向量相似度搜索。
需要安装pgvector扩展: CREATE EXTENSION IF NOT EXISTS vector;
"""
def __init__(self, pg_url: Optional[str] = None):
"""初始化PG客户端.
Args:
pg_url: PostgreSQL连接URL，默认从配置读取
"""
self.pg_url = pg_url
self._pool: Optional[Any] = None
async def connect(self) -> None:
"""建立PG连接池."""
import psycopg
from psycopg_pool import AsyncConnectionPool
self._pool = AsyncConnectionPool(
conninfo=self.pg_url,
min_size=1,
max_size=10,
)
# 验证连接并启用pgvector
async with self._pool.connection() as conn:
await conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
await conn.commit()
async def disconnect(self) -> None:
"""断开PG连接."""
if self._pool:
await self._pool.close()
self._pool = None
async def health_check(self) -> bool:
"""检查PG连接健康状态."""
if not self._pool:
return False
try:
async with self._pool.connection() as conn:
cursor = await conn.execute("SELECT 1")
result = await cursor.fetchone()
return result is not None and result[0] == 1
except Exception:
return False
# 全局PG客户端实例
_pg_client: Optional[PostgresClient] = None
def get_pg_client() -> PostgresClient:
"""获取全局PG客户端实例（单例模式）.
Returns:
PostgresClient: PG客户端实例
"""
global _pg_client
if _pg_client is None:
_pg_client = PostgresClient()
return _pg_client
... ...
"""FastAPI应用入口."""
from contextlib import asynccontextmanager
from typing import AsyncGenerator
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from app.api.router import api_router
from app.core.exceptions import EmbeddingServiceError
from app.core.config import settings
from app.db import connect as db_client_connect, disconnect as db_client_disconnect
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
"""应用生命周期管理."""
# 启动时建立数据库连接
try:
await db_client_connect()
db_type = settings.get("db_type", "es")
print(f"✓ {db_type.upper()} database connected")
except Exception as e:
print(f"✗ Failed to connect to database: {e}")
raise
yield
# 关闭时断开数据库连接
try:
await db_client_disconnect()
db_type = settings.get("db_type", "es")
print(f"✓ {db_type.upper()} database disconnected")
except Exception as e:
print(f"✗ Error disconnecting from database: {e}")
def create_app() -> FastAPI:
"""创建FastAPI应用."""
app = FastAPI(
title=settings.app_name,
version=settings.app_version,
description="基于向量的CRUD服务，支持多种Embedding场景",
lifespan=lifespan,
docs_url="/docs",
redoc_url="/redoc",
)
# 配置CORS
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# 注册路由
app.include_router(api_router)
# 异常处理
@app.exception_handler(EmbeddingServiceError)
async def embedding_service_exception_handler(request, exc):
from fastapi.responses import JSONResponse
return JSONResponse(
status_code=500,
content={"code": exc.code, "message": exc.message},
)
# @app.get("/health", tags=["健康检查"])
# async def health_check():
# """健康检查端点."""
# db_healthy = await db_client.health_check()
# db_type = settings.get("db_type", "es")
# return {
# "status": "healthy" if db_healthy else "unhealthy",
# "database": {
# "type": db_type,
# "status": "connected" if db_healthy else "disconnected",
# },
# }
@app.get("/", tags=["根路径"])
async def root():
"""根路径."""
return {
"name": settings.app_name,
"version": settings.app_version,
"docs": "/docs",
}
return app
# 创建应用实例
app = create_app()
if __name__ == "__main__":
import uvicorn
uvicorn.run(
"app.main:app",
host='0.0.0.0',
port=settings.port,
reload=settings.debug,
)
... ...
from app.db import es_client
... ...
[project]
name = "aigc-embedding-service"
version = "0.1.0"
description = "Add your description here"
description = "基于向量的CRUD服务，支持多种Embedding场景"
requires-python = ">=3.12"
dependencies = []
dependencies = [
"fastapi>=0.115.0",
"uvicorn[standard]>=0.32.0",
"pydantic>=2.9.0",
"pydantic-settings>=2.6.0",
"httpx>=0.27.0",
"python-multipart>=0.0.12",
"aabd>=0.4.1",
"path>=17.1.1",
"omegaconf>=2.3.0",
]
[project.optional-dependencies]
dev = [
"pytest>=8.3.0",
"pytest-asyncio>=0.24.0",
"black>=24.10.0",
"ruff>=0.7.0",
"mypy>=1.13.0",
]
es = [
"elasticsearch>=8.15.0",
]
postgres = [
"psycopg[binary,pool]>=3.2.0",
"psycopg-pool>=3.3.0",
]
[tool.black]
line-length = 100
target-version = ['py312']
[tool.ruff]
line-length = 100
select = ["E", "F", "I", "N", "W", "UP", "B", "C4", "SIM"]
ignore = ["E501"]
[tool.mypy]
python_version = "3.12"
warn_return_any = true
warn_unused_configs = true
disallow_untyped_defs = true
[tool.pytest.ini_options]
asyncio_mode = "auto"
testpaths = ["tests"]
python_files = ["test_*.py"]
python_classes = ["Test*"]
python_functions = ["test_*"]
... ...
This diff could not be displayed because it is too large.