Showing
33 changed files
with
2313 additions
and
2313 deletions
| 1 | -# Python | ||
| 2 | -__pycache__/ | ||
| 3 | -*.py[cod] | ||
| 4 | -*$py.class | ||
| 5 | -*.so | ||
| 6 | -.Python | ||
| 7 | -*.egg-info/ | ||
| 8 | -dist/ | ||
| 9 | -build/ | ||
| 10 | - | ||
| 11 | -# Virtual environments | ||
| 12 | -.venv/ | ||
| 13 | -venv/ | ||
| 14 | - | ||
| 15 | -# IDE | ||
| 16 | -.idea/ | ||
| 17 | -.vscode/ | ||
| 18 | -*.swp | ||
| 19 | -*.swo | ||
| 20 | - | ||
| 21 | -# Git | ||
| 22 | -.git/ | ||
| 23 | -.gitignore | ||
| 24 | - | ||
| 25 | -# Docker | ||
| 26 | -Dockerfile | ||
| 27 | -.dockerignore | ||
| 28 | - | ||
| 29 | -# Logs | ||
| 30 | -*.log | ||
| 31 | - | ||
| 32 | -# Local configs | ||
| 33 | -cfg/config.local.yaml | 1 | +# Python |
| 2 | +__pycache__/ | ||
| 3 | +*.py[cod] | ||
| 4 | +*$py.class | ||
| 5 | +*.so | ||
| 6 | +.Python | ||
| 7 | +*.egg-info/ | ||
| 8 | +dist/ | ||
| 9 | +build/ | ||
| 10 | + | ||
| 11 | +# Virtual environments | ||
| 12 | +.venv/ | ||
| 13 | +venv/ | ||
| 14 | + | ||
| 15 | +# IDE | ||
| 16 | +.idea/ | ||
| 17 | +.vscode/ | ||
| 18 | +*.swp | ||
| 19 | +*.swo | ||
| 20 | + | ||
| 21 | +# Git | ||
| 22 | +.git/ | ||
| 23 | +.gitignore | ||
| 24 | + | ||
| 25 | +# Docker | ||
| 26 | +Dockerfile | ||
| 27 | +.dockerignore | ||
| 28 | + | ||
| 29 | +# Logs | ||
| 30 | +*.log | ||
| 31 | + | ||
| 32 | +# Local configs | ||
| 33 | +cfg/config.local.yaml |
| 1 | -/live-stream-ai-task-workline/modules/wtt_highlight/model_weights/* | ||
| 2 | -*mov | ||
| 3 | -/.externalToolBuilders | ||
| 4 | -/.settings | ||
| 5 | -.svn | ||
| 6 | -/.project | ||
| 7 | -/target | ||
| 8 | -./target | ||
| 9 | -/.classpath | ||
| 10 | -.idea/ | ||
| 11 | -*.iml | ||
| 12 | -*.iws | ||
| 13 | -*.ipr | ||
| 14 | -*.ids | ||
| 15 | -*.orig | ||
| 16 | -classes/ | ||
| 17 | -/bin/ | ||
| 18 | -*/.classpath | ||
| 19 | -*/.project | ||
| 20 | -*/.settings/* | ||
| 21 | -###################### | ||
| 22 | -# Project Specific | ||
| 23 | -###################### | ||
| 24 | -/target/www/** | ||
| 25 | -/src/test/javascript/coverage/ | ||
| 26 | -/src/test/javascript/PhantomJS*/ | ||
| 27 | - | ||
| 28 | -###################### | ||
| 29 | -# Node | ||
| 30 | -###################### | ||
| 31 | -/node/ | ||
| 32 | -node_tmp/ | ||
| 33 | -node_modules/ | ||
| 34 | -npm-debug.log.* | ||
| 35 | -/.awcache/* | ||
| 36 | - | ||
| 37 | -###################### | ||
| 38 | -# SASS | ||
| 39 | -###################### | ||
| 40 | -.sass-cache/ | ||
| 41 | - | ||
| 42 | -###################### | ||
| 43 | -# Eclipse | ||
| 44 | -###################### | ||
| 45 | -*.pydevproject | ||
| 46 | -.project | ||
| 47 | -.metadata | ||
| 48 | -tmp/ | ||
| 49 | -tmp/**/* | ||
| 50 | -*.tmp | ||
| 51 | -*.bak | ||
| 52 | -*.swp | ||
| 53 | -*~.nib | ||
| 54 | -local.properties | ||
| 55 | -.classpath | ||
| 56 | -.settings/ | ||
| 57 | -.loadpath | ||
| 58 | -.factorypath | ||
| 59 | -/src/main/resources/rebel.xml | ||
| 60 | - | ||
| 61 | -# External tool builders | ||
| 62 | -.externalToolBuilders/** | ||
| 63 | - | ||
| 64 | -# Locally stored "Eclipse launch configurations" | ||
| 65 | -*.launch | ||
| 66 | - | ||
| 67 | -# CDT-specific | ||
| 68 | -.cproject | ||
| 69 | - | ||
| 70 | -# PDT-specific | ||
| 71 | -.buildpath | ||
| 72 | - | ||
| 73 | -###################### | ||
| 74 | -# Intellij | ||
| 75 | -###################### | ||
| 76 | -.idea/ | ||
| 77 | -*.iml | ||
| 78 | -*.iws | ||
| 79 | -*.ipr | ||
| 80 | -*.ids | ||
| 81 | -*.orig | ||
| 82 | -classes/ | ||
| 83 | - | ||
| 84 | -###################### | ||
| 85 | -# Visual Studio Code | ||
| 86 | -###################### | ||
| 87 | -.vscode/ | ||
| 88 | - | ||
| 89 | -###################### | ||
| 90 | -# Maven | ||
| 91 | -###################### | ||
| 92 | -/log/ | ||
| 93 | -/target/ | ||
| 94 | - | ||
| 95 | -###################### | ||
| 96 | -# Gradle | ||
| 97 | -###################### | ||
| 98 | -.gradle/ | ||
| 99 | -/build/ | ||
| 100 | - | ||
| 101 | -###################### | ||
| 102 | -# Package Files | ||
| 103 | -###################### | ||
| 104 | -*.jar | ||
| 105 | -*.war | ||
| 106 | -*.ear | ||
| 107 | -*.db | ||
| 108 | - | ||
| 109 | -###################### | ||
| 110 | -# Windows | ||
| 111 | -###################### | ||
| 112 | -# Windows image file caches | ||
| 113 | -Thumbs.db | ||
| 114 | - | ||
| 115 | -# Folder com.migu.oes.power.config file | ||
| 116 | -Desktop.ini | ||
| 117 | - | ||
| 118 | -###################### | ||
| 119 | -# Mac OSX | ||
| 120 | -###################### | ||
| 121 | -.DS_Store | ||
| 122 | -.svn | ||
| 123 | - | ||
| 124 | -# Thumbnails | ||
| 125 | -._* | ||
| 126 | - | ||
| 127 | -# Files that might appear on external disk | ||
| 128 | -.Spotlight-V100 | ||
| 129 | -.Trashes | ||
| 130 | - | ||
| 131 | -###################### | ||
| 132 | -# Directories | ||
| 133 | -###################### | ||
| 134 | -/bin/ | ||
| 135 | -/deploy/ | ||
| 136 | - | ||
| 137 | -###################### | ||
| 138 | -# Logs | ||
| 139 | -###################### | ||
| 140 | -*.log* | ||
| 141 | - | ||
| 142 | -###################### | ||
| 143 | -# Others | ||
| 144 | -###################### | ||
| 145 | -*.class | ||
| 146 | -*.*~ | ||
| 147 | -*~ | ||
| 148 | -.merge_file* | ||
| 149 | - | ||
| 150 | -###################### | ||
| 151 | -# Gradle Wrapper | ||
| 152 | -###################### | ||
| 153 | -!gradle/wrapper/gradle-wrapper.jar | ||
| 154 | - | ||
| 155 | -###################### | ||
| 156 | -# Maven Wrapper | ||
| 157 | -###################### | ||
| 158 | -!.mvn/wrapper/maven-wrapper.jar | ||
| 159 | - | ||
| 160 | -###################### | ||
| 161 | -# ESLint | ||
| 162 | -###################### | ||
| 163 | -.eslintcache | ||
| 164 | -*.lst | ||
| 165 | -oes-nryy-vo/target/maven-archiver/pom.properties | ||
| 166 | -oes-nryy-dataservice/target/maven-archiver/pom.properties | ||
| 167 | -oes-nryy-base/target/maven-archiver/pom.properties | ||
| 168 | -oes-nryy-copyright/bin/.gitkeep | ||
| 169 | -oes-nryy-copyright/bin/cbak | ||
| 170 | -*.original | ||
| 171 | -*.gz | ||
| 172 | -oes-nryy-tag/target/oes-nryy-tag-1.0.0/label/bin/service.sh | ||
| 173 | -oes-nryy-tag/target/maven-archiver/pom.properties | ||
| 174 | -oes-nryy-data-trans/target/maven-archiver/pom.properties | ||
| 175 | -oes-nryy-tag/target/oes-nryy-tag-1.0.0/label/bin/cbak | ||
| 176 | - | ||
| 177 | - | ||
| 178 | -###################### | ||
| 179 | -# pp-scene-det | ||
| 180 | -###################### | ||
| 181 | -__pycache__/ | ||
| 182 | - | ||
| 183 | -###################### | ||
| 184 | -# live-stream-ai-task-workline | ||
| 185 | -###################### | ||
| 186 | -modules/basketball/weights | ||
| 187 | -modules/singer/weights | ||
| 188 | -weights/ | ||
| 189 | -weights1/ | ||
| 190 | -files/ | ||
| 191 | -*.log | ||
| 192 | -*.log.* | ||
| 193 | -cache_data/ | ||
| 194 | -data/ | ||
| 195 | -test/ | ||
| 196 | -tests/ | ||
| 197 | -# Media files | ||
| 198 | -*.jpg | ||
| 199 | -*.jpeg | ||
| 200 | -*.png | ||
| 201 | -*.gif | ||
| 202 | -*.mp4 | ||
| 203 | -*.avi | ||
| 204 | -*.mov | ||
| 205 | -*.mkv | ||
| 206 | - | ||
| 207 | -# OS | ||
| 208 | -.DS_Store | ||
| 209 | -Thumbs.db | ||
| 210 | -logs/ | 1 | +/live-stream-ai-task-workline/modules/wtt_highlight/model_weights/* |
| 2 | +*mov | ||
| 3 | +/.externalToolBuilders | ||
| 4 | +/.settings | ||
| 5 | +.svn | ||
| 6 | +/.project | ||
| 7 | +/target | ||
| 8 | +./target | ||
| 9 | +/.classpath | ||
| 10 | +.idea/ | ||
| 11 | +*.iml | ||
| 12 | +*.iws | ||
| 13 | +*.ipr | ||
| 14 | +*.ids | ||
| 15 | +*.orig | ||
| 16 | +classes/ | ||
| 17 | +/bin/ | ||
| 18 | +*/.classpath | ||
| 19 | +*/.project | ||
| 20 | +*/.settings/* | ||
| 21 | +###################### | ||
| 22 | +# Project Specific | ||
| 23 | +###################### | ||
| 24 | +/target/www/** | ||
| 25 | +/src/test/javascript/coverage/ | ||
| 26 | +/src/test/javascript/PhantomJS*/ | ||
| 27 | + | ||
| 28 | +###################### | ||
| 29 | +# Node | ||
| 30 | +###################### | ||
| 31 | +/node/ | ||
| 32 | +node_tmp/ | ||
| 33 | +node_modules/ | ||
| 34 | +npm-debug.log.* | ||
| 35 | +/.awcache/* | ||
| 36 | + | ||
| 37 | +###################### | ||
| 38 | +# SASS | ||
| 39 | +###################### | ||
| 40 | +.sass-cache/ | ||
| 41 | + | ||
| 42 | +###################### | ||
| 43 | +# Eclipse | ||
| 44 | +###################### | ||
| 45 | +*.pydevproject | ||
| 46 | +.project | ||
| 47 | +.metadata | ||
| 48 | +tmp/ | ||
| 49 | +tmp/**/* | ||
| 50 | +*.tmp | ||
| 51 | +*.bak | ||
| 52 | +*.swp | ||
| 53 | +*~.nib | ||
| 54 | +local.properties | ||
| 55 | +.classpath | ||
| 56 | +.settings/ | ||
| 57 | +.loadpath | ||
| 58 | +.factorypath | ||
| 59 | +/src/main/resources/rebel.xml | ||
| 60 | + | ||
| 61 | +# External tool builders | ||
| 62 | +.externalToolBuilders/** | ||
| 63 | + | ||
| 64 | +# Locally stored "Eclipse launch configurations" | ||
| 65 | +*.launch | ||
| 66 | + | ||
| 67 | +# CDT-specific | ||
| 68 | +.cproject | ||
| 69 | + | ||
| 70 | +# PDT-specific | ||
| 71 | +.buildpath | ||
| 72 | + | ||
| 73 | +###################### | ||
| 74 | +# Intellij | ||
| 75 | +###################### | ||
| 76 | +.idea/ | ||
| 77 | +*.iml | ||
| 78 | +*.iws | ||
| 79 | +*.ipr | ||
| 80 | +*.ids | ||
| 81 | +*.orig | ||
| 82 | +classes/ | ||
| 83 | + | ||
| 84 | +###################### | ||
| 85 | +# Visual Studio Code | ||
| 86 | +###################### | ||
| 87 | +.vscode/ | ||
| 88 | + | ||
| 89 | +###################### | ||
| 90 | +# Maven | ||
| 91 | +###################### | ||
| 92 | +/log/ | ||
| 93 | +/target/ | ||
| 94 | + | ||
| 95 | +###################### | ||
| 96 | +# Gradle | ||
| 97 | +###################### | ||
| 98 | +.gradle/ | ||
| 99 | +/build/ | ||
| 100 | + | ||
| 101 | +###################### | ||
| 102 | +# Package Files | ||
| 103 | +###################### | ||
| 104 | +*.jar | ||
| 105 | +*.war | ||
| 106 | +*.ear | ||
| 107 | +*.db | ||
| 108 | + | ||
| 109 | +###################### | ||
| 110 | +# Windows | ||
| 111 | +###################### | ||
| 112 | +# Windows image file caches | ||
| 113 | +Thumbs.db | ||
| 114 | + | ||
| 115 | +# Folder com.migu.oes.power.config file | ||
| 116 | +Desktop.ini | ||
| 117 | + | ||
| 118 | +###################### | ||
| 119 | +# Mac OSX | ||
| 120 | +###################### | ||
| 121 | +.DS_Store | ||
| 122 | +.svn | ||
| 123 | + | ||
| 124 | +# Thumbnails | ||
| 125 | +._* | ||
| 126 | + | ||
| 127 | +# Files that might appear on external disk | ||
| 128 | +.Spotlight-V100 | ||
| 129 | +.Trashes | ||
| 130 | + | ||
| 131 | +###################### | ||
| 132 | +# Directories | ||
| 133 | +###################### | ||
| 134 | +/bin/ | ||
| 135 | +/deploy/ | ||
| 136 | + | ||
| 137 | +###################### | ||
| 138 | +# Logs | ||
| 139 | +###################### | ||
| 140 | +*.log* | ||
| 141 | + | ||
| 142 | +###################### | ||
| 143 | +# Others | ||
| 144 | +###################### | ||
| 145 | +*.class | ||
| 146 | +*.*~ | ||
| 147 | +*~ | ||
| 148 | +.merge_file* | ||
| 149 | + | ||
| 150 | +###################### | ||
| 151 | +# Gradle Wrapper | ||
| 152 | +###################### | ||
| 153 | +!gradle/wrapper/gradle-wrapper.jar | ||
| 154 | + | ||
| 155 | +###################### | ||
| 156 | +# Maven Wrapper | ||
| 157 | +###################### | ||
| 158 | +!.mvn/wrapper/maven-wrapper.jar | ||
| 159 | + | ||
| 160 | +###################### | ||
| 161 | +# ESLint | ||
| 162 | +###################### | ||
| 163 | +.eslintcache | ||
| 164 | +*.lst | ||
| 165 | +oes-nryy-vo/target/maven-archiver/pom.properties | ||
| 166 | +oes-nryy-dataservice/target/maven-archiver/pom.properties | ||
| 167 | +oes-nryy-base/target/maven-archiver/pom.properties | ||
| 168 | +oes-nryy-copyright/bin/.gitkeep | ||
| 169 | +oes-nryy-copyright/bin/cbak | ||
| 170 | +*.original | ||
| 171 | +*.gz | ||
| 172 | +oes-nryy-tag/target/oes-nryy-tag-1.0.0/label/bin/service.sh | ||
| 173 | +oes-nryy-tag/target/maven-archiver/pom.properties | ||
| 174 | +oes-nryy-data-trans/target/maven-archiver/pom.properties | ||
| 175 | +oes-nryy-tag/target/oes-nryy-tag-1.0.0/label/bin/cbak | ||
| 176 | + | ||
| 177 | + | ||
| 178 | +###################### | ||
| 179 | +# pp-scene-det | ||
| 180 | +###################### | ||
| 181 | +__pycache__/ | ||
| 182 | + | ||
| 183 | +###################### | ||
| 184 | +# live-stream-ai-task-workline | ||
| 185 | +###################### | ||
| 186 | +modules/basketball/weights | ||
| 187 | +modules/singer/weights | ||
| 188 | +weights/ | ||
| 189 | +weights1/ | ||
| 190 | +files/ | ||
| 191 | +*.log | ||
| 192 | +*.log.* | ||
| 193 | +cache_data/ | ||
| 194 | +data/ | ||
| 195 | +test/ | ||
| 196 | +tests/ | ||
| 197 | +# Media files | ||
| 198 | +*.jpg | ||
| 199 | +*.jpeg | ||
| 200 | +*.png | ||
| 201 | +*.gif | ||
| 202 | +*.mp4 | ||
| 203 | +*.avi | ||
| 204 | +*.mov | ||
| 205 | +*.mkv | ||
| 206 | + | ||
| 207 | +# OS | ||
| 208 | +.DS_Store | ||
| 209 | +Thumbs.db | ||
| 210 | +logs/ | ||
| 211 | cache/ | 211 | cache/ |
| 1 | -# ============================================================================ | ||
| 2 | -# embedding-server Dockerfile | ||
| 3 | -# ============================================================================ | ||
| 4 | -# 构建命令: | ||
| 5 | -# docker build -f Dockerfile.embedding -t embedding-server:latest . | ||
| 6 | -# | ||
| 7 | -# 运行命令: | ||
| 8 | -# docker run -d -p 8000:8000 --name embedding-server embedding-server:latest | ||
| 9 | -# ============================================================================ | ||
| 10 | - | ||
| 11 | -ARG PYTHON_VERSION=3.12 | ||
| 12 | -#ARG UV_VERSION=0.6.0 | ||
| 13 | - | ||
| 14 | -# ---------------------------------------------------------------------------- | ||
| 15 | -# 构建阶段 | ||
| 16 | -# ---------------------------------------------------------------------------- | ||
| 17 | -FROM python:${PYTHON_VERSION}-slim AS builder | ||
| 18 | -# 使用 Debian 的阿里云镜像 | ||
| 19 | -RUN echo "deb http://mirrors.aliyun.com/debian/ bookworm main contrib non-free" > /etc/apt/sources.list && \ | ||
| 20 | - echo "deb http://mirrors.aliyun.com/debian/ bookworm-updates main contrib non-free" >> /etc/apt/sources.list && \ | ||
| 21 | - echo "deb http://mirrors.aliyun.com/debian-security/ bookworm-security main contrib non-free" >> /etc/apt/sources.list | ||
| 22 | - | ||
| 23 | -# 安装编译依赖和 uv | ||
| 24 | -RUN apt-get update \ | ||
| 25 | - && apt-get install -y --no-install-recommends gcc curl vim \ | ||
| 26 | - && rm -rf /var/lib/apt/lists/* | ||
| 27 | - | ||
| 28 | -RUN pip install --no-cache-dir uv | ||
| 29 | - | ||
| 30 | -WORKDIR /app | ||
| 31 | - | ||
| 32 | -# 复制项目文件（仅 embedding-server 所需） | ||
| 33 | -COPY pyproject.toml uv.lock ./ | ||
| 34 | -COPY src/embedding_sever/ ./src/embedding_sever/ | ||
| 35 | - | ||
| 36 | -# 安装依赖：核心依赖 + embedding 专用依赖 | ||
| 37 | -RUN uv sync --no-dev --no-editable --extra embedding | ||
| 38 | - | ||
| 39 | -# ---------------------------------------------------------------------------- | ||
| 40 | -# 运行阶段 | ||
| 41 | -# ---------------------------------------------------------------------------- | ||
| 42 | -FROM python:${PYTHON_VERSION}-slim | ||
| 43 | - | ||
| 44 | -WORKDIR /app | ||
| 45 | - | ||
| 46 | -# 从构建阶段复制虚拟环境和源码 | ||
| 47 | -COPY --from=builder /app/.venv /app/.venv | ||
| 48 | -COPY --from=builder /app/src/embedding_sever /app/src/embedding_sever | ||
| 49 | - | ||
| 50 | -# 环境变量 | ||
| 51 | -ENV PATH="/app/.venv/bin:$PATH" | ||
| 52 | -ENV PYTHONPATH="/app/src" | ||
| 53 | -ENV PROJECT_ROOT="/app/src" | ||
| 54 | -ENV PYTHONUNBUFFERED=1 | ||
| 55 | -ENV ENV=prod | ||
| 56 | -ENV APP_LOG_TYPE=file,console | ||
| 57 | - | ||
| 58 | -EXPOSE 8000 | ||
| 59 | - | ||
| 60 | -HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \ | ||
| 61 | - CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/')" || exit 1 | ||
| 62 | - | ||
| 63 | -CMD ["embedding-server"] | 1 | +# ============================================================================ |
| 2 | +# embedding-server Dockerfile | ||
| 3 | +# ============================================================================ | ||
| 4 | +# 构建命令: | ||
| 5 | +# docker build -f Dockerfile.embedding -t embedding-server:latest . | ||
| 6 | +# | ||
| 7 | +# 运行命令: | ||
| 8 | +# docker run -d -p 8000:8000 --name embedding-server embedding-server:latest | ||
| 9 | +# ============================================================================ | ||
| 10 | + | ||
| 11 | +ARG PYTHON_VERSION=3.12 | ||
| 12 | +#ARG UV_VERSION=0.6.0 | ||
| 13 | + | ||
| 14 | +# ---------------------------------------------------------------------------- | ||
| 15 | +# 构建阶段 | ||
| 16 | +# ---------------------------------------------------------------------------- | ||
| 17 | +FROM python:${PYTHON_VERSION}-slim AS builder | ||
| 18 | +# 使用 Debian 的阿里云镜像 | ||
| 19 | +RUN echo "deb http://mirrors.aliyun.com/debian/ bookworm main contrib non-free" > /etc/apt/sources.list && \ | ||
| 20 | + echo "deb http://mirrors.aliyun.com/debian/ bookworm-updates main contrib non-free" >> /etc/apt/sources.list && \ | ||
| 21 | + echo "deb http://mirrors.aliyun.com/debian-security/ bookworm-security main contrib non-free" >> /etc/apt/sources.list | ||
| 22 | + | ||
| 23 | +# 安装编译依赖和 uv | ||
| 24 | +RUN apt-get update \ | ||
| 25 | + && apt-get install -y --no-install-recommends gcc curl vim \ | ||
| 26 | + && rm -rf /var/lib/apt/lists/* | ||
| 27 | + | ||
| 28 | +RUN pip install --no-cache-dir uv | ||
| 29 | + | ||
| 30 | +WORKDIR /app | ||
| 31 | + | ||
| 32 | +# 复制项目文件（仅 embedding-server 所需） | ||
| 33 | +COPY pyproject.toml uv.lock ./ | ||
| 34 | +COPY src/embedding_sever/ ./src/embedding_sever/ | ||
| 35 | + | ||
| 36 | +# 安装依赖：核心依赖 + embedding 专用依赖 | ||
| 37 | +RUN uv sync --no-dev --no-editable --extra embedding | ||
| 38 | + | ||
| 39 | +# ---------------------------------------------------------------------------- | ||
| 40 | +# 运行阶段 | ||
| 41 | +# ---------------------------------------------------------------------------- | ||
| 42 | +FROM python:${PYTHON_VERSION}-slim | ||
| 43 | + | ||
| 44 | +WORKDIR /app | ||
| 45 | + | ||
| 46 | +# 从构建阶段复制虚拟环境和源码 | ||
| 47 | +COPY --from=builder /app/.venv /app/.venv | ||
| 48 | +COPY --from=builder /app/src/embedding_sever /app/src/embedding_sever | ||
| 49 | + | ||
| 50 | +# 环境变量 | ||
| 51 | +ENV PATH="/app/.venv/bin:$PATH" | ||
| 52 | +ENV PYTHONPATH="/app/src" | ||
| 53 | +ENV PROJECT_ROOT="/app/src" | ||
| 54 | +ENV PYTHONUNBUFFERED=1 | ||
| 55 | +ENV ENV=prod | ||
| 56 | +ENV APP_LOG_TYPE=file,console | ||
| 57 | + | ||
| 58 | +EXPOSE 8000 | ||
| 59 | + | ||
| 60 | +HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \ | ||
| 61 | + CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/')" || exit 1 | ||
| 62 | + | ||
| 63 | +CMD ["embedding-server"] |
| 1 | -### 一些启动命令 | ||
| 2 | -1. 启动 embedding-server | ||
| 3 | - | ||
| 4 | -默认 | ||
| 5 | -```bash | ||
| 6 | -uv run embedding-server | ||
| 7 | -``` | ||
| 8 | -window 开发 | ||
| 9 | -```bash | ||
| 10 | -$env:ENV="dev";uv run embedding-server | ||
| 11 | -$env:ENV="dev";uv run football-repaly-match-server | ||
| 12 | -``` | ||
| 13 | - | ||
| 14 | -安装依赖 | ||
| 15 | -```shell | ||
| 16 | -uv sync --no-dev --no-editable --extra football-replay-match --extra embedding | 1 | +### 一些启动命令 |
| 2 | +1. 启动 embedding-server | ||
| 3 | + | ||
| 4 | +默认 | ||
| 5 | +```bash | ||
| 6 | +uv run embedding-server | ||
| 7 | +``` | ||
| 8 | +window 开发 | ||
| 9 | +```bash | ||
| 10 | +$env:ENV="dev";uv run embedding-server | ||
| 11 | +$env:ENV="dev";uv run football-repaly-match-server | ||
| 12 | +``` | ||
| 13 | + | ||
| 14 | +安装依赖 | ||
| 15 | +```shell | ||
| 16 | +uv sync --no-dev --no-editable --extra football-replay-match --extra embedding | ||
| 17 | ``` | 17 | ``` |
| 1 | -[build-system] | ||
| 2 | -requires = ["hatchling"] | ||
| 3 | -build-backend = "hatchling.build" | ||
| 4 | - | ||
| 5 | -[project] | ||
| 6 | -name = "ai_server" | ||
| 7 | -version = "0.1.0" | ||
| 8 | -description = "Add your description here" | ||
| 9 | -requires-python = ">=3.12" | ||
| 10 | -dependencies = [ | ||
| 11 | - "fastapi==0.136.1", | ||
| 12 | - "uvicorn[standard]==0.46.0", | ||
| 13 | - "pydantic==2.13.4", | ||
| 14 | - "numpy>=1.24.0", | ||
| 15 | - "aiohttp>=3.8.0", | ||
| 16 | - "aabd==0.4.8", | ||
| 17 | - "path==17.1.1", | ||
| 18 | - "omegaconf==2.3.0", | ||
| 19 | - "packaging==26.2", | ||
| 20 | -] | ||
| 21 | - | ||
| 22 | -[project.optional-dependencies] | ||
| 23 | -# embedding-server 专用依赖 | ||
| 24 | -embedding = [ | ||
| 25 | - "elasticsearch==9.3.0", | ||
| 26 | -] | ||
| 27 | -football-replay-match = [ | ||
| 28 | - "Pillow==12.2.0", | ||
| 29 | - "av==13.1.0", | ||
| 30 | - "confluent_kafka==2.6.1", | ||
| 31 | - "langchain-openai>=1.1.15", | ||
| 32 | - "opencv-contrib-python==4.13.0.92" | ||
| 33 | -] | ||
| 34 | -# task-server 专用依赖 | ||
| 35 | -task = [ | ||
| 36 | - | ||
| 37 | -] | ||
| 38 | - | ||
| 39 | -[tool.uv] | ||
| 40 | -package = true | ||
| 41 | - | ||
| 42 | -[tool.hatch.build.targets.wheel] | ||
| 43 | -packages = ["src/embedding_sever"] | ||
| 44 | - | ||
| 45 | -[project.scripts] | ||
| 46 | -embedding-server = "embedding_sever.main:main" | ||
| 47 | -#task-server = "task_server.main:main" | ||
| 48 | -football-repaly-match-server = "football_replay_match.main:main" | 1 | +[build-system] |
| 2 | +requires = ["hatchling"] | ||
| 3 | +build-backend = "hatchling.build" | ||
| 4 | + | ||
| 5 | +[project] | ||
| 6 | +name = "ai_server" | ||
| 7 | +version = "0.1.0" | ||
| 8 | +description = "Add your description here" | ||
| 9 | +requires-python = ">=3.12" | ||
| 10 | +dependencies = [ | ||
| 11 | + "fastapi==0.136.1", | ||
| 12 | + "uvicorn[standard]==0.46.0", | ||
| 13 | + "pydantic==2.13.4", | ||
| 14 | + "numpy>=1.24.0", | ||
| 15 | + "aiohttp>=3.8.0", | ||
| 16 | + "aabd==0.4.8", | ||
| 17 | + "path==17.1.1", | ||
| 18 | + "omegaconf==2.3.0", | ||
| 19 | + "packaging==26.2", | ||
| 20 | +] | ||
| 21 | + | ||
| 22 | +[project.optional-dependencies] | ||
| 23 | +# embedding-server 专用依赖 | ||
| 24 | +embedding = [ | ||
| 25 | + "elasticsearch==9.3.0", | ||
| 26 | +] | ||
| 27 | +football-replay-match = [ | ||
| 28 | + "Pillow==12.2.0", | ||
| 29 | + "av==13.1.0", | ||
| 30 | + "confluent_kafka==2.6.1", | ||
| 31 | + "langchain-openai>=1.1.15", | ||
| 32 | + "opencv-contrib-python==4.13.0.92" | ||
| 33 | +] | ||
| 34 | +# task-server 专用依赖 | ||
| 35 | +task = [ | ||
| 36 | + | ||
| 37 | +] | ||
| 38 | + | ||
| 39 | +[tool.uv] | ||
| 40 | +package = true | ||
| 41 | + | ||
| 42 | +[tool.hatch.build.targets.wheel] | ||
| 43 | +packages = ["src/embedding_sever"] | ||
| 44 | + | ||
| 45 | +[project.scripts] | ||
| 46 | +embedding-server = "embedding_sever.main:main" | ||
| 47 | +#task-server = "task_server.main:main" | ||
| 48 | +football-repaly-match-server = "football_replay_match.main:main" |
| 1 | -from typing import List, Optional, Literal, Union, Dict, Any | ||
| 2 | - | ||
| 3 | -from fastapi import APIRouter, Depends | ||
| 4 | -import numpy | ||
| 5 | -from pydantic import BaseModel, Field, TypeAdapter | ||
| 6 | - | ||
| 7 | -from embedding_sever.api.resp_bean import RespBean, success | ||
| 8 | - | ||
| 9 | -router = APIRouter() | ||
| 10 | - | ||
| 11 | - | ||
| 12 | -# ============== 请求/响应模型定义 ============== | ||
| 13 | - | ||
| 14 | -class FaceKwargs(BaseModel): | ||
| 15 | - """人脸类型扩展字段""" | ||
| 16 | - person_id: Optional[str] = Field(None, description="人物id") | ||
| 17 | - person_name: Optional[str] = Field(None, description="人物姓名") | ||
| 18 | - | ||
| 19 | - | ||
| 20 | -class SportShotKwargs(BaseModel): | ||
| 21 | - """体育镜头类型扩展字段""" | ||
| 22 | - match_name: Optional[str] = Field(None, description="比赛名称") | ||
| 23 | - person_name: Optional[str] = Field(None, description="人名") | ||
| 24 | - video_desc: Optional[str] = Field(None, description="视频描述") | ||
| 25 | - shot_cls: Optional[str] = Field(None, description="镜头分类") | ||
| 26 | - match_id: Optional[str] = Field(None, description="比赛id") | ||
| 27 | - program_id: Optional[str] = Field(None, description="节目id") | ||
| 28 | - | ||
| 29 | - | ||
| 30 | -# 定义类型映射，用于根据type自动转换kwargs | ||
| 31 | -KwargsType = Union[FaceKwargs, SportShotKwargs] | ||
| 32 | - | ||
| 33 | - | ||
| 34 | -class PutDataRequest(BaseModel): | ||
| 35 | - """写入数据请求体""" | ||
| 36 | - id: str = Field(..., description="唯一键id") | ||
| 37 | - embedding: List[float] = Field(..., description="向量") | ||
| 38 | - embedding_version: str = Field(..., description="向量版本 小写数字下划线组成") | ||
| 39 | - kwargs: Optional[Dict[str, Any]] = Field(None, description="扩展字段，根据type不同字段不同") | ||
| 40 | - | ||
| 41 | - def get_typed_kwargs(self, type: str) -> Optional[KwargsType]: | ||
| 42 | - """根据type自动转换kwargs为具体类型""" | ||
| 43 | - if self.kwargs is None: | ||
| 44 | - return None | ||
| 45 | - type_map = { | ||
| 46 | - 'face': FaceKwargs, | ||
| 47 | - 'sport_shot': SportShotKwargs | ||
| 48 | - } | ||
| 49 | - kwargs_class = type_map.get(type) | ||
| 50 | - if kwargs_class: | ||
| 51 | - return kwargs_class(**self.kwargs) | ||
| 52 | - return None | ||
| 53 | - | ||
| 54 | - | ||
| 55 | -class DelByIdRequest(BaseModel): | ||
| 56 | - """删除数据请求体""" | ||
| 57 | - ids: List[str] = Field(..., description="id 集合") | ||
| 58 | - embedding_version: str = Field(..., description="向量版本 小写数字下划线组成") | ||
| 59 | - | ||
| 60 | - | ||
| 61 | -class FilterItem(BaseModel): | ||
| 62 | - """过滤条件项""" | ||
| 63 | - name: str = Field(..., description="字段名") | ||
| 64 | - value: List = Field(..., description="字段值") | ||
| 65 | - opt: Optional[Literal["eq", "neq", "lt", "gt", "lte", "gte", "like", "in"]] = Field( | ||
| 66 | - None, description="操作类型: eq相等, neq不等, lt小于, gt大于, lte小于等于, gte大于等于, like模糊匹配, in内" | ||
| 67 | - ) | ||
| 68 | - | ||
| 69 | - | ||
| 70 | -class SearchRequest(BaseModel): | ||
| 71 | - """检索数据请求体""" | ||
| 72 | - embedding: List[float] = Field(..., description="向量") | ||
| 73 | - embedding_version: str = Field(..., description="向量版本 小写数字下划线组成") | ||
| 74 | - topk: int = Field(..., description="top k") | ||
| 75 | - filters: Optional[List[FilterItem]] = Field(None, description="过滤字段") | ||
| 76 | - | ||
| 77 | - | ||
| 78 | -class SearchResultItem(BaseModel): | ||
| 79 | - """检索结果单项""" | ||
| 80 | - id: str = Field(..., description="id") | ||
| 81 | - kwargs: dict = Field(..., description="扩展字段") | ||
| 82 | - score: float = Field(..., description="分值") | ||
| 83 | - | ||
| 84 | - | ||
| 85 | -class SearchResponseData(BaseModel): | ||
| 86 | - """检索响应数据""" | ||
| 87 | - type: str = Field(..., description="type") | ||
| 88 | - data_list: List[SearchResultItem] = Field(..., description="topk结果列表") | ||
| 89 | - | ||
| 90 | - | ||
| 91 | -# ============== 接口定义 ============== | ||
| 92 | -from embedding_sever.service.data_service import DataService | ||
| 93 | -from embedding_sever.service import get_data_service | ||
| 94 | - | ||
| 95 | - | ||
| 96 | -@router.put("/{type}/put", response_model=RespBean, tags=["数据服务"], summary="写入数据") | ||
| 97 | -@router.post("/{type}/put", response_model=RespBean, tags=["数据服务"], summary="写入数据") | ||
| 98 | -async def put_data( | ||
| 99 | - type: Literal["face", "sport_shot"], | ||
| 100 | - request: PutDataRequest, | ||
| 101 | - data_service: DataService = Depends(get_data_service) | ||
| 102 | -): | ||
| 103 | - """ | ||
| 104 | - 写入数据 | ||
| 105 | - """ | ||
| 106 | - request.embedding = numpy.array(request.embedding) | ||
| 107 | - | ||
| 108 | - embedding_version = request.embedding_version | ||
| 109 | - tb_name = f'{type}_{embedding_version}' | ||
| 110 | - upserted_id = await data_service.upsert(tb_name, request.id, request.embedding, | ||
| 111 | - request.get_typed_kwargs(type).model_dump()) | ||
| 112 | - | ||
| 113 | - return success(data={'id': upserted_id}) | ||
| 114 | - | ||
| 115 | - | ||
| 116 | -@router.delete("/{type}/del_by_id", response_model=RespBean, tags=["数据服务"], summary="删除数据") | ||
| 117 | -@router.post("/{type}/del_by_id", response_model=RespBean, tags=["数据服务"], summary="删除数据") | ||
| 118 | -async def del_by_id( | ||
| 119 | - type: Literal["face", "sport_shot"], | ||
| 120 | - request: DelByIdRequest, | ||
| 121 | - data_service: DataService = Depends(get_data_service) | ||
| 122 | -): | ||
| 123 | - """ | ||
| 124 | - 删除数据 | ||
| 125 | - | ||
| 126 | - - **type**: 数据类型 (face-人脸, sport_shot-体育镜头) | ||
| 127 | - - **ids**: id集合 | ||
| 128 | - """ | ||
| 129 | - embedding_version = request.embedding_version | ||
| 130 | - tb_name = f'{type}_{embedding_version}' | ||
| 131 | - await data_service.delete_by_pks(tb_name, request.ids) | ||
| 132 | - | ||
| 133 | - return success() | ||
| 134 | - | ||
| 135 | - | ||
| 136 | -@router.post("/{type}/search", response_model=RespBean[SearchResponseData], tags=["数据服务"], summary="检索数据") | ||
| 137 | -async def search_data( | ||
| 138 | - type: Literal["face", "sport_shot"], | ||
| 139 | - request: SearchRequest, | ||
| 140 | - data_service: DataService = Depends(get_data_service), | ||
| 141 | - adapter=TypeAdapter(List[SearchResultItem]) | ||
| 142 | -): | ||
| 143 | - """ | ||
| 144 | - 检索数据 | ||
| 145 | - | ||
| 146 | - - **type**: 数据类型 (face-人脸, sport_shot-体育镜头) | ||
| 147 | - - **embedding**: 查询向量 | ||
| 148 | - - **embedding_version**: 向量版本 | ||
| 149 | - - **topk**: 返回top k条结果 | ||
| 150 | - - **filters**: 过滤条件列表 | ||
| 151 | - """ | ||
| 152 | - # | ||
| 153 | - embedding_version = request.embedding_version | ||
| 154 | - tb_name = f'{type}_{embedding_version}' | ||
| 155 | - | ||
| 156 | - embedding = request.embedding | ||
| 157 | - filters = request.filters | ||
| 158 | - from_, size = 0, request.topk | ||
| 159 | - | ||
| 160 | - result_data = await data_service.search(tb_name, embedding, filters, from_, size) | ||
| 161 | - search_results: List[SearchResultItem] = adapter.validate_python(result_data) | ||
| 162 | - | ||
| 163 | - return success(data=SearchResponseData(type=type, data_list=search_results)) | 1 | +from typing import List, Optional, Literal, Union, Dict, Any |
| 2 | + | ||
| 3 | +from fastapi import APIRouter, Depends | ||
| 4 | +import numpy | ||
| 5 | +from pydantic import BaseModel, Field, TypeAdapter | ||
| 6 | + | ||
| 7 | +from embedding_sever.api.resp_bean import RespBean, success | ||
| 8 | + | ||
| 9 | +router = APIRouter() | ||
| 10 | + | ||
| 11 | + | ||
| 12 | +# ============== 请求/响应模型定义 ============== | ||
| 13 | + | ||
| 14 | +class FaceKwargs(BaseModel): | ||
| 15 | + """人脸类型扩展字段""" | ||
| 16 | + person_id: Optional[str] = Field(None, description="人物id") | ||
| 17 | + person_name: Optional[str] = Field(None, description="人物姓名") | ||
| 18 | + | ||
| 19 | + | ||
| 20 | +class SportShotKwargs(BaseModel): | ||
| 21 | + """体育镜头类型扩展字段""" | ||
| 22 | + match_name: Optional[str] = Field(None, description="比赛名称") | ||
| 23 | + person_name: Optional[str] = Field(None, description="人名") | ||
| 24 | + video_desc: Optional[str] = Field(None, description="视频描述") | ||
| 25 | + shot_cls: Optional[str] = Field(None, description="镜头分类") | ||
| 26 | + match_id: Optional[str] = Field(None, description="比赛id") | ||
| 27 | + program_id: Optional[str] = Field(None, description="节目id") | ||
| 28 | + | ||
| 29 | + | ||
| 30 | +# 定义类型映射，用于根据type自动转换kwargs | ||
| 31 | +KwargsType = Union[FaceKwargs, SportShotKwargs] | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +class PutDataRequest(BaseModel): | ||
| 35 | + """写入数据请求体""" | ||
| 36 | + id: str = Field(..., description="唯一键id") | ||
| 37 | + embedding: List[float] = Field(..., description="向量") | ||
| 38 | + embedding_version: str = Field(..., description="向量版本 小写数字下划线组成") | ||
| 39 | + kwargs: Optional[Dict[str, Any]] = Field(None, description="扩展字段，根据type不同字段不同") | ||
| 40 | + | ||
| 41 | + def get_typed_kwargs(self, type: str) -> Optional[KwargsType]: | ||
| 42 | + """根据type自动转换kwargs为具体类型""" | ||
| 43 | + if self.kwargs is None: | ||
| 44 | + return None | ||
| 45 | + type_map = { | ||
| 46 | + 'face': FaceKwargs, | ||
| 47 | + 'sport_shot': SportShotKwargs | ||
| 48 | + } | ||
| 49 | + kwargs_class = type_map.get(type) | ||
| 50 | + if kwargs_class: | ||
| 51 | + return kwargs_class(**self.kwargs) | ||
| 52 | + return None | ||
| 53 | + | ||
| 54 | + | ||
| 55 | +class DelByIdRequest(BaseModel): | ||
| 56 | + """删除数据请求体""" | ||
| 57 | + ids: List[str] = Field(..., description="id 集合") | ||
| 58 | + embedding_version: str = Field(..., description="向量版本 小写数字下划线组成") | ||
| 59 | + | ||
| 60 | + | ||
| 61 | +class FilterItem(BaseModel): | ||
| 62 | + """过滤条件项""" | ||
| 63 | + name: str = Field(..., description="字段名") | ||
| 64 | + value: List = Field(..., description="字段值") | ||
| 65 | + opt: Optional[Literal["eq", "neq", "lt", "gt", "lte", "gte", "like", "in"]] = Field( | ||
| 66 | + None, description="操作类型: eq相等, neq不等, lt小于, gt大于, lte小于等于, gte大于等于, like模糊匹配, in内" | ||
| 67 | + ) | ||
| 68 | + | ||
| 69 | + | ||
| 70 | +class SearchRequest(BaseModel): | ||
| 71 | + """检索数据请求体""" | ||
| 72 | + embedding: List[float] = Field(..., description="向量") | ||
| 73 | + embedding_version: str = Field(..., description="向量版本 小写数字下划线组成") | ||
| 74 | + topk: int = Field(..., description="top k") | ||
| 75 | + filters: Optional[List[FilterItem]] = Field(None, description="过滤字段") | ||
| 76 | + | ||
| 77 | + | ||
| 78 | +class SearchResultItem(BaseModel): | ||
| 79 | + """检索结果单项""" | ||
| 80 | + id: str = Field(..., description="id") | ||
| 81 | + kwargs: dict = Field(..., description="扩展字段") | ||
| 82 | + score: float = Field(..., description="分值") | ||
| 83 | + | ||
| 84 | + | ||
| 85 | +class SearchResponseData(BaseModel): | ||
| 86 | + """检索响应数据""" | ||
| 87 | + type: str = Field(..., description="type") | ||
| 88 | + data_list: List[SearchResultItem] = Field(..., description="topk结果列表") | ||
| 89 | + | ||
| 90 | + | ||
| 91 | +# ============== 接口定义 ============== | ||
| 92 | +from embedding_sever.service.data_service import DataService | ||
| 93 | +from embedding_sever.service import get_data_service | ||
| 94 | + | ||
| 95 | + | ||
| 96 | +@router.put("/{type}/put", response_model=RespBean, tags=["数据服务"], summary="写入数据") | ||
| 97 | +@router.post("/{type}/put", response_model=RespBean, tags=["数据服务"], summary="写入数据") | ||
| 98 | +async def put_data( | ||
| 99 | + type: Literal["face", "sport_shot"], | ||
| 100 | + request: PutDataRequest, | ||
| 101 | + data_service: DataService = Depends(get_data_service) | ||
| 102 | +): | ||
| 103 | + """ | ||
| 104 | + 写入数据 | ||
| 105 | + """ | ||
| 106 | + request.embedding = numpy.array(request.embedding) | ||
| 107 | + | ||
| 108 | + embedding_version = request.embedding_version | ||
| 109 | + tb_name = f'{type}_{embedding_version}' | ||
| 110 | + upserted_id = await data_service.upsert(tb_name, request.id, request.embedding, | ||
| 111 | + request.get_typed_kwargs(type).model_dump()) | ||
| 112 | + | ||
| 113 | + return success(data={'id': upserted_id}) | ||
| 114 | + | ||
| 115 | + | ||
| 116 | +@router.delete("/{type}/del_by_id", response_model=RespBean, tags=["数据服务"], summary="删除数据") | ||
| 117 | +@router.post("/{type}/del_by_id", response_model=RespBean, tags=["数据服务"], summary="删除数据") | ||
| 118 | +async def del_by_id( | ||
| 119 | + type: Literal["face", "sport_shot"], | ||
| 120 | + request: DelByIdRequest, | ||
| 121 | + data_service: DataService = Depends(get_data_service) | ||
| 122 | +): | ||
| 123 | + """ | ||
| 124 | + 删除数据 | ||
| 125 | + | ||
| 126 | + - **type**: 数据类型 (face-人脸, sport_shot-体育镜头) | ||
| 127 | + - **ids**: id集合 | ||
| 128 | + """ | ||
| 129 | + embedding_version = request.embedding_version | ||
| 130 | + tb_name = f'{type}_{embedding_version}' | ||
| 131 | + await data_service.delete_by_pks(tb_name, request.ids) | ||
| 132 | + | ||
| 133 | + return success() | ||
| 134 | + | ||
| 135 | + | ||
| 136 | +@router.post("/{type}/search", response_model=RespBean[SearchResponseData], tags=["数据服务"], summary="检索数据") | ||
| 137 | +async def search_data( | ||
| 138 | + type: Literal["face", "sport_shot"], | ||
| 139 | + request: SearchRequest, | ||
| 140 | + data_service: DataService = Depends(get_data_service), | ||
| 141 | + adapter=TypeAdapter(List[SearchResultItem]) | ||
| 142 | +): | ||
| 143 | + """ | ||
| 144 | + 检索数据 | ||
| 145 | + | ||
| 146 | + - **type**: 数据类型 (face-人脸, sport_shot-体育镜头) | ||
| 147 | + - **embedding**: 查询向量 | ||
| 148 | + - **embedding_version**: 向量版本 | ||
| 149 | + - **topk**: 返回top k条结果 | ||
| 150 | + - **filters**: 过滤条件列表 | ||
| 151 | + """ | ||
| 152 | + # | ||
| 153 | + embedding_version = request.embedding_version | ||
| 154 | + tb_name = f'{type}_{embedding_version}' | ||
| 155 | + | ||
| 156 | + embedding = request.embedding | ||
| 157 | + filters = request.filters | ||
| 158 | + from_, size = 0, request.topk | ||
| 159 | + | ||
| 160 | + result_data = await data_service.search(tb_name, embedding, filters, from_, size) | ||
| 161 | + search_results: List[SearchResultItem] = adapter.validate_python(result_data) | ||
| 162 | + | ||
| 163 | + return success(data=SearchResponseData(type=type, data_list=search_results)) |
| 1 | -from fastapi.middleware.cors import CORSMiddleware | ||
| 2 | -from starlette.responses import JSONResponse | ||
| 3 | -from embedding_sever.api.http_log import LoggingMiddleware | ||
| 4 | -from fastapi import FastAPI | ||
| 5 | -from embedding_sever.api.router import api_router | ||
| 6 | -from embedding_sever.config import settings | ||
| 7 | - | ||
| 8 | - | ||
| 9 | -def config_fastapi(fastapi_app: FastAPI): | ||
| 10 | - # 配置请求/响应日志中间件（需在CORS之前添加） | ||
| 11 | - fastapi_app.add_middleware(LoggingMiddleware) | ||
| 12 | - | ||
| 13 | - # 配置CORS | ||
| 14 | - fastapi_app.add_middleware( | ||
| 15 | - CORSMiddleware, | ||
| 16 | - allow_origins=["*"], | ||
| 17 | - allow_credentials=True, | ||
| 18 | - allow_methods=["*"], | ||
| 19 | - allow_headers=["*"], | ||
| 20 | - ) | ||
| 21 | - | ||
| 22 | - # 注册路由 | ||
| 23 | - fastapi_app.include_router(api_router) | ||
| 24 | - | ||
| 25 | - @fastapi_app.exception_handler(Exception) | ||
| 26 | - async def generic_exception_handler(request, exc: Exception): | ||
| 27 | - # 返回自定义的错误响应 | ||
| 28 | - error_detail = str(exc) | ||
| 29 | - return JSONResponse( | ||
| 30 | - status_code=200, | ||
| 31 | - content={"code": "error", "message": error_detail}, | ||
| 32 | - ) | ||
| 33 | - | ||
| 34 | - @fastapi_app.get("/", tags=["根路径"]) | ||
| 35 | - async def root(): | ||
| 36 | - """根路径.""" | ||
| 37 | - return { | ||
| 38 | - "name": settings.app_name, | ||
| 39 | - "version": settings.app_version, | ||
| 40 | - "docs": "/docs", | ||
| 41 | - } | ||
| 42 | - | ||
| 43 | - return fastapi_app | 1 | +from fastapi.middleware.cors import CORSMiddleware |
| 2 | +from starlette.responses import JSONResponse | ||
| 3 | +from embedding_sever.api.http_log import LoggingMiddleware | ||
| 4 | +from fastapi import FastAPI | ||
| 5 | +from embedding_sever.api.router import api_router | ||
| 6 | +from embedding_sever.config import settings | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +def config_fastapi(fastapi_app: FastAPI): | ||
| 10 | + # 配置请求/响应日志中间件（需在CORS之前添加） | ||
| 11 | + fastapi_app.add_middleware(LoggingMiddleware) | ||
| 12 | + | ||
| 13 | + # 配置CORS | ||
| 14 | + fastapi_app.add_middleware( | ||
| 15 | + CORSMiddleware, | ||
| 16 | + allow_origins=["*"], | ||
| 17 | + allow_credentials=True, | ||
| 18 | + allow_methods=["*"], | ||
| 19 | + allow_headers=["*"], | ||
| 20 | + ) | ||
| 21 | + | ||
| 22 | + # 注册路由 | ||
| 23 | + fastapi_app.include_router(api_router) | ||
| 24 | + | ||
| 25 | + @fastapi_app.exception_handler(Exception) | ||
| 26 | + async def generic_exception_handler(request, exc: Exception): | ||
| 27 | + # 返回自定义的错误响应 | ||
| 28 | + error_detail = str(exc) | ||
| 29 | + return JSONResponse( | ||
| 30 | + status_code=200, | ||
| 31 | + content={"code": "error", "message": error_detail}, | ||
| 32 | + ) | ||
| 33 | + | ||
| 34 | + @fastapi_app.get("/", tags=["根路径"]) | ||
| 35 | + async def root(): | ||
| 36 | + """根路径.""" | ||
| 37 | + return { | ||
| 38 | + "name": settings.app_name, | ||
| 39 | + "version": settings.app_version, | ||
| 40 | + "docs": "/docs", | ||
| 41 | + } | ||
| 42 | + | ||
| 43 | + return fastapi_app |
| 1 | -import json | ||
| 2 | -import time | ||
| 3 | -import logging | ||
| 4 | - | ||
| 5 | -logger = logging.getLogger(__name__) | ||
| 6 | -from fastapi import Request | ||
| 7 | -from starlette.middleware.base import BaseHTTPMiddleware | ||
| 8 | -from starlette.responses import StreamingResponse | ||
| 9 | -from embedding_sever.config import settings | ||
| 10 | - | ||
| 11 | - | ||
| 12 | -class LoggingMiddleware(BaseHTTPMiddleware): | ||
| 13 | - | ||
| 14 | - async def dispatch(self, request: Request, call_next): | ||
| 15 | - # 获取请求信息 | ||
| 16 | - method = request.method | ||
| 17 | - url = str(request.url) | ||
| 18 | - client_host = request.client.host if request.client else "unknown" | ||
| 19 | - | ||
| 20 | - # 判断是否需要记录日志 | ||
| 21 | - should_log = request.url.path.startswith(settings.url_prefix) | ||
| 22 | - | ||
| 23 | - if should_log: | ||
| 24 | - # 记录请求开始时间 | ||
| 25 | - start_time = time.time() | ||
| 26 | - | ||
| 27 | - # 读取请求体 | ||
| 28 | - request_body = None | ||
| 29 | - if method in ("POST", "PUT", "PATCH"): | ||
| 30 | - try: | ||
| 31 | - body = await request.body() | ||
| 32 | - if body: | ||
| 33 | - request_body = body.decode("utf-8", errors="replace") | ||
| 34 | - # 尝试格式化为JSON | ||
| 35 | - try: | ||
| 36 | - request_body = json.loads(request_body) | ||
| 37 | - except json.JSONDecodeError: | ||
| 38 | - pass # 保持原始字符串 | ||
| 39 | - except Exception: | ||
| 40 | - pass | ||
| 41 | - | ||
| 42 | - # 处理请求 | ||
| 43 | - response = await call_next(request) | ||
| 44 | - | ||
| 45 | - # 计算处理时间 | ||
| 46 | - process_time = time.time() - start_time | ||
| 47 | - | ||
| 48 | - # 读取响应体（仅处理非流式响应） | ||
| 49 | - response_body = None | ||
| 50 | - if not isinstance(response, StreamingResponse): | ||
| 51 | - try: | ||
| 52 | - response_body_bytes = b"" | ||
| 53 | - async for chunk in response.body_iterator: | ||
| 54 | - response_body_bytes += chunk | ||
| 55 | - | ||
| 56 | - # 重新构建响应 | ||
| 57 | - response_body = response_body_bytes.decode("utf-8", errors="replace") | ||
| 58 | - # 尝试格式化为JSON | ||
| 59 | - try: | ||
| 60 | - response_body = json.loads(response_body) | ||
| 61 | - except json.JSONDecodeError: | ||
| 62 | - pass # 保持原始字符串 | ||
| 63 | - | ||
| 64 | - # 重建响应对象 | ||
| 65 | - response = StreamingResponse( | ||
| 66 | - iter([response_body_bytes]), | ||
| 67 | - status_code=response.status_code, | ||
| 68 | - headers=dict(response.headers), | ||
| 69 | - media_type=response.media_type, | ||
| 70 | - ) | ||
| 71 | - except Exception: | ||
| 72 | - pass | ||
| 73 | - | ||
| 74 | - # 构建单条日志记录 | ||
| 75 | - status_code = response.status_code | ||
| 76 | - log_data = { | ||
| 77 | - "method": method, | ||
| 78 | - "url": url, | ||
| 79 | - "client": client_host, | ||
| 80 | - "status": status_code, | ||
| 81 | - "time": f"{process_time:.3f}s" | ||
| 82 | - } | ||
| 83 | - | ||
| 84 | - if request_body: | ||
| 85 | - log_data["request_body"] = request_body | ||
| 86 | - if response_body: | ||
| 87 | - log_data["response_body"] = response_body | ||
| 88 | - | ||
| 89 | - logger.info(json.dumps(log_data, ensure_ascii=False)) | ||
| 90 | - | ||
| 91 | - return response | ||
| 92 | - else: | ||
| 93 | - # 不需要记录日志的请求，直接处理 | ||
| 94 | - return await call_next(request) | 1 | +import json |
| 2 | +import time | ||
| 3 | +import logging | ||
| 4 | + | ||
| 5 | +logger = logging.getLogger(__name__) | ||
| 6 | +from fastapi import Request | ||
| 7 | +from starlette.middleware.base import BaseHTTPMiddleware | ||
| 8 | +from starlette.responses import StreamingResponse | ||
| 9 | +from embedding_sever.config import settings | ||
| 10 | + | ||
| 11 | + | ||
| 12 | +class LoggingMiddleware(BaseHTTPMiddleware): | ||
| 13 | + | ||
| 14 | + async def dispatch(self, request: Request, call_next): | ||
| 15 | + # 获取请求信息 | ||
| 16 | + method = request.method | ||
| 17 | + url = str(request.url) | ||
| 18 | + client_host = request.client.host if request.client else "unknown" | ||
| 19 | + | ||
| 20 | + # 判断是否需要记录日志 | ||
| 21 | + should_log = request.url.path.startswith(settings.url_prefix) | ||
| 22 | + | ||
| 23 | + if should_log: | ||
| 24 | + # 记录请求开始时间 | ||
| 25 | + start_time = time.time() | ||
| 26 | + | ||
| 27 | + # 读取请求体 | ||
| 28 | + request_body = None | ||
| 29 | + if method in ("POST", "PUT", "PATCH"): | ||
| 30 | + try: | ||
| 31 | + body = await request.body() | ||
| 32 | + if body: | ||
| 33 | + request_body = body.decode("utf-8", errors="replace") | ||
| 34 | + # 尝试格式化为JSON | ||
| 35 | + try: | ||
| 36 | + request_body = json.loads(request_body) | ||
| 37 | + except json.JSONDecodeError: | ||
| 38 | + pass # 保持原始字符串 | ||
| 39 | + except Exception: | ||
| 40 | + pass | ||
| 41 | + | ||
| 42 | + # 处理请求 | ||
| 43 | + response = await call_next(request) | ||
| 44 | + | ||
| 45 | + # 计算处理时间 | ||
| 46 | + process_time = time.time() - start_time | ||
| 47 | + | ||
| 48 | + # 读取响应体（仅处理非流式响应） | ||
| 49 | + response_body = None | ||
| 50 | + if not isinstance(response, StreamingResponse): | ||
| 51 | + try: | ||
| 52 | + response_body_bytes = b"" | ||
| 53 | + async for chunk in response.body_iterator: | ||
| 54 | + response_body_bytes += chunk | ||
| 55 | + | ||
| 56 | + # 重新构建响应 | ||
| 57 | + response_body = response_body_bytes.decode("utf-8", errors="replace") | ||
| 58 | + # 尝试格式化为JSON | ||
| 59 | + try: | ||
| 60 | + response_body = json.loads(response_body) | ||
| 61 | + except json.JSONDecodeError: | ||
| 62 | + pass # 保持原始字符串 | ||
| 63 | + | ||
| 64 | + # 重建响应对象 | ||
| 65 | + response = StreamingResponse( | ||
| 66 | + iter([response_body_bytes]), | ||
| 67 | + status_code=response.status_code, | ||
| 68 | + headers=dict(response.headers), | ||
| 69 | + media_type=response.media_type, | ||
| 70 | + ) | ||
| 71 | + except Exception: | ||
| 72 | + pass | ||
| 73 | + | ||
| 74 | + # 构建单条日志记录 | ||
| 75 | + status_code = response.status_code | ||
| 76 | + log_data = { | ||
| 77 | + "method": method, | ||
| 78 | + "url": url, | ||
| 79 | + "client": client_host, | ||
| 80 | + "status": status_code, | ||
| 81 | + "time": f"{process_time:.3f}s" | ||
| 82 | + } | ||
| 83 | + | ||
| 84 | + if request_body: | ||
| 85 | + log_data["request_body"] = request_body | ||
| 86 | + if response_body: | ||
| 87 | + log_data["response_body"] = response_body | ||
| 88 | + | ||
| 89 | + logger.info(json.dumps(log_data, ensure_ascii=False)) | ||
| 90 | + | ||
| 91 | + return response | ||
| 92 | + else: | ||
| 93 | + # 不需要记录日志的请求，直接处理 | ||
| 94 | + return await call_next(request) |
| 1 | -from pydantic import BaseModel, Field | ||
| 2 | -from typing import Generic, TypeVar, Optional | ||
| 3 | - | ||
| 4 | -T = TypeVar('T') | ||
| 5 | - | ||
| 6 | - | ||
| 7 | -class RespBean(BaseModel, Generic[T]): | ||
| 8 | - code: str = Field(default="success", description="code") | ||
| 9 | - message: str = Field(default="success", description='message') | ||
| 10 | - data: Optional[T] = Field(description="数据", default=None) | ||
| 11 | - | ||
| 12 | - | ||
| 13 | -def error(message): | ||
| 14 | - return RespBean(code="error", message=message, data=None) | ||
| 15 | - | ||
| 16 | - | ||
| 17 | -def success(data=None): | ||
| 18 | - return RespBean(data=data) | 1 | +from pydantic import BaseModel, Field |
| 2 | +from typing import Generic, TypeVar, Optional | ||
| 3 | + | ||
| 4 | +T = TypeVar('T') | ||
| 5 | + | ||
| 6 | + | ||
| 7 | +class RespBean(BaseModel, Generic[T]): | ||
| 8 | + code: str = Field(default="success", description="code") | ||
| 9 | + message: str = Field(default="success", description='message') | ||
| 10 | + data: Optional[T] = Field(description="数据", default=None) | ||
| 11 | + | ||
| 12 | + | ||
| 13 | +def error(message): | ||
| 14 | + return RespBean(code="error", message=message, data=None) | ||
| 15 | + | ||
| 16 | + | ||
| 17 | +def success(data=None): | ||
| 18 | + return RespBean(data=data) |
| 1 | -from fastapi import APIRouter | ||
| 2 | -from embedding_sever.config import settings | ||
| 3 | -from embedding_sever.api.data_router import router as data_router | ||
| 4 | - | ||
| 5 | -api_router = APIRouter(prefix=settings.url_prefix) | ||
| 6 | - | ||
| 7 | - | ||
| 8 | -def register_router(router: APIRouter, prefix: str): | ||
| 9 | - api_router.include_router(router, prefix=prefix) | ||
| 10 | - | ||
| 11 | - | ||
| 12 | -# 注册数据路由 | ||
| 13 | -register_router(data_router, prefix="/data") | 1 | +from fastapi import APIRouter |
| 2 | +from embedding_sever.config import settings | ||
| 3 | +from embedding_sever.api.data_router import router as data_router | ||
| 4 | + | ||
| 5 | +api_router = APIRouter(prefix=settings.url_prefix) | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +def register_router(router: APIRouter, prefix: str): | ||
| 9 | + api_router.include_router(router, prefix=prefix) | ||
| 10 | + | ||
| 11 | + | ||
| 12 | +# 注册数据路由 | ||
| 13 | +register_router(data_router, prefix="/data") |
| 1 | -app_name: "Embedding Service" | ||
| 2 | -app_version: "1.0.0" | ||
| 3 | -port: 8000 | ||
| 4 | -debug: false | ||
| 5 | -url_prefix: "/aigc-embedding/api" | ||
| 6 | - | ||
| 7 | -db_es_enable: true | ||
| 8 | -db_es_url: "http://192.168.0.14:9200" | ||
| 9 | - | ||
| 10 | -#db_postgres_enable: false | 1 | +app_name: "Embedding Service" |
| 2 | +app_version: "1.0.0" | ||
| 3 | +port: 8000 | ||
| 4 | +debug: false | ||
| 5 | +url_prefix: "/aigc-embedding/api" | ||
| 6 | + | ||
| 7 | +db_es_enable: true | ||
| 8 | +db_es_url: "http://192.168.0.14:9200" | ||
| 9 | + | ||
| 10 | +#db_postgres_enable: false | ||
| 11 | #db_postgres_url: "postgresql://postgres:postgres@localhost:5432/postgres" | 11 | #db_postgres_url: "postgresql://postgres:postgres@localhost:5432/postgres" |
| 1 | -app_name: "Embedding Service" | ||
| 2 | -app_version: "1.0.0" | ||
| 3 | -port: 8000 | ||
| 4 | -debug: false | ||
| 5 | -url_prefix: "/aigc-embedding/api" | ||
| 6 | - | ||
| 7 | -db_es_enable: true | ||
| 8 | -db_es_url: "http://localhost:9200" | ||
| 9 | - | ||
| 10 | -#db_postgres_enable: false | 1 | +app_name: "Embedding Service" |
| 2 | +app_version: "1.0.0" | ||
| 3 | +port: 8000 | ||
| 4 | +debug: false | ||
| 5 | +url_prefix: "/aigc-embedding/api" | ||
| 6 | + | ||
| 7 | +db_es_enable: true | ||
| 8 | +db_es_url: "http://localhost:9200" | ||
| 9 | + | ||
| 10 | +#db_postgres_enable: false | ||
| 11 | #db_postgres_url: "postgresql://postgres:postgres@localhost:5432/postgres" | 11 | #db_postgres_url: "postgresql://postgres:postgres@localhost:5432/postgres" |
| 1 | -app_name: "Embedding Service" | ||
| 2 | -app_version: "1.0.0" | ||
| 3 | -port: 8000 | ||
| 4 | -debug: false | ||
| 5 | -url_prefix: "/aigc-embedding/api" | ||
| 6 | - | ||
| 7 | -db_es_enable: true | ||
| 8 | -db_es_url: "http://localhost:9200" | ||
| 9 | - | ||
| 10 | -#db_postgres_enable: false | 1 | +app_name: "Embedding Service" |
| 2 | +app_version: "1.0.0" | ||
| 3 | +port: 8000 | ||
| 4 | +debug: false | ||
| 5 | +url_prefix: "/aigc-embedding/api" | ||
| 6 | + | ||
| 7 | +db_es_enable: true | ||
| 8 | +db_es_url: "http://localhost:9200" | ||
| 9 | + | ||
| 10 | +#db_postgres_enable: false | ||
| 11 | #db_postgres_url: "postgresql://postgres:postgres@localhost:5432/postgres" | 11 | #db_postgres_url: "postgresql://postgres:postgres@localhost:5432/postgres" |
| 1 | -from aabd.base.cfg_loader import load_yaml_by_file_with_env | ||
| 2 | -from aabd.base.enhance_dict import EnhanceDict, read_prefixed_env_vars | ||
| 3 | -from pathlib import Path | ||
| 4 | -from logging import getLogger | ||
| 5 | -logger = getLogger(__name__) | ||
| 6 | -cfg_dir = Path(__file__).parent.absolute() / 'cfg' | ||
| 7 | -settings = EnhanceDict(load_yaml_by_file_with_env((cfg_dir / 'config.yaml').as_posix())) | ||
| 8 | -settings.update(read_prefixed_env_vars('EMB_')) | ||
| 9 | - | 1 | +from aabd.base.cfg_loader import load_yaml_by_file_with_env |
| 2 | +from aabd.base.enhance_dict import EnhanceDict, read_prefixed_env_vars | ||
| 3 | +from pathlib import Path | ||
| 4 | +from logging import getLogger | ||
| 5 | +logger = getLogger(__name__) | ||
| 6 | +cfg_dir = Path(__file__).parent.absolute() / 'cfg' | ||
| 7 | +settings = EnhanceDict(load_yaml_by_file_with_env((cfg_dir / 'config.yaml').as_posix())) | ||
| 8 | +settings.update(read_prefixed_env_vars('EMB_')) | ||
| 9 | + | ||
| 10 | logger.info(f'Settings: {settings}') | 10 | logger.info(f'Settings: {settings}') |
| 1 | -from functools import lru_cache | ||
| 2 | - | ||
| 3 | -from elasticsearch import AsyncElasticsearch | ||
| 4 | - | ||
| 5 | -from embedding_sever.config import settings | ||
| 6 | -from .data_dao import DataDao | ||
| 7 | - | ||
| 8 | -es_client = None | ||
| 9 | -pg_client = None | ||
| 10 | - | ||
| 11 | - | ||
| 12 | -async def connect(): | ||
| 13 | - if settings.db_es_enable: | ||
| 14 | - global es_client | ||
| 15 | - es_client = AsyncElasticsearch([settings.db_es_url]) | ||
| 16 | - await es_client.ping() | ||
| 17 | - if settings.db_postgres_enable: | ||
| 18 | - from .postgres_client import get_pg_client | ||
| 19 | - global pg_client | ||
| 20 | - pg_client = get_pg_client() | ||
| 21 | - await pg_client.connect() | ||
| 22 | - | ||
| 23 | - | ||
| 24 | -async def disconnect(): | ||
| 25 | - global es_client, pg_client | ||
| 26 | - if es_client is not None: | ||
| 27 | - await es_client.disconnect() | ||
| 28 | - es_client = None | ||
| 29 | - if pg_client is not None: | ||
| 30 | - await pg_client.disconnect() | ||
| 31 | - | ||
| 32 | - | ||
| 33 | -@lru_cache() | ||
| 34 | -def get_data_dao(): | ||
| 35 | - return DataDao(es_client) | 1 | +from functools import lru_cache |
| 2 | + | ||
| 3 | +from elasticsearch import AsyncElasticsearch | ||
| 4 | + | ||
| 5 | +from embedding_sever.config import settings | ||
| 6 | +from .data_dao import DataDao | ||
| 7 | + | ||
| 8 | +es_client = None | ||
| 9 | +pg_client = None | ||
| 10 | + | ||
| 11 | + | ||
| 12 | +async def connect(): | ||
| 13 | + if settings.db_es_enable: | ||
| 14 | + global es_client | ||
| 15 | + es_client = AsyncElasticsearch([settings.db_es_url]) | ||
| 16 | + await es_client.ping() | ||
| 17 | + if settings.db_postgres_enable: | ||
| 18 | + from .postgres_client import get_pg_client | ||
| 19 | + global pg_client | ||
| 20 | + pg_client = get_pg_client() | ||
| 21 | + await pg_client.connect() | ||
| 22 | + | ||
| 23 | + | ||
| 24 | +async def disconnect(): | ||
| 25 | + global es_client, pg_client | ||
| 26 | + if es_client is not None: | ||
| 27 | + await es_client.disconnect() | ||
| 28 | + es_client = None | ||
| 29 | + if pg_client is not None: | ||
| 30 | + await pg_client.disconnect() | ||
| 31 | + | ||
| 32 | + | ||
| 33 | +@lru_cache() | ||
| 34 | +def get_data_dao(): | ||
| 35 | + return DataDao(es_client) |
| 1 | -"""数据库抽象基类，定义统一的CRUD接口.""" | ||
| 2 | - | ||
| 3 | -from abc import ABC, abstractmethod | ||
| 4 | -from typing import Any, Dict, List, Optional | ||
| 5 | - | ||
| 6 | - | ||
| 7 | -class VectorDBClient(ABC): | ||
| 8 | - """向量数据库客户端抽象基类. | ||
| 9 | - | ||
| 10 | - 定义了统一的向量数据CRUD接口，支持Elasticsearch和PostgreSQL等后端。 | ||
| 11 | - """ | ||
| 12 | - | ||
| 13 | - @abstractmethod | ||
| 14 | - async def connect(self) -> None: | ||
| 15 | - """建立数据库连接.""" | ||
| 16 | - pass | ||
| 17 | - | ||
| 18 | - @abstractmethod | ||
| 19 | - async def disconnect(self) -> None: | ||
| 20 | - """断开数据库连接.""" | ||
| 21 | - pass | ||
| 22 | - | ||
| 23 | - @abstractmethod | ||
| 24 | - async def health_check(self) -> bool: | ||
| 25 | - """检查数据库连接健康状态. | ||
| 26 | - | ||
| 27 | - Returns: | ||
| 28 | - bool: 连接正常返回True，否则返回False | ||
| 29 | - """ | ||
| 30 | - pass | ||
| 31 | - | ||
| 32 | - | ||
| 33 | -class TableInit(ABC): | ||
| 34 | - | ||
| 35 | - @abstractmethod | ||
| 36 | - async def create_table(self, name, embedding_version, embedding_dim) -> None: | ||
| 37 | - pass | 1 | +"""数据库抽象基类，定义统一的CRUD接口.""" |
| 2 | + | ||
| 3 | +from abc import ABC, abstractmethod | ||
| 4 | +from typing import Any, Dict, List, Optional | ||
| 5 | + | ||
| 6 | + | ||
| 7 | +class VectorDBClient(ABC): | ||
| 8 | + """向量数据库客户端抽象基类. | ||
| 9 | + | ||
| 10 | + 定义了统一的向量数据CRUD接口，支持Elasticsearch和PostgreSQL等后端。 | ||
| 11 | + """ | ||
| 12 | + | ||
| 13 | + @abstractmethod | ||
| 14 | + async def connect(self) -> None: | ||
| 15 | + """建立数据库连接.""" | ||
| 16 | + pass | ||
| 17 | + | ||
| 18 | + @abstractmethod | ||
| 19 | + async def disconnect(self) -> None: | ||
| 20 | + """断开数据库连接.""" | ||
| 21 | + pass | ||
| 22 | + | ||
| 23 | + @abstractmethod | ||
| 24 | + async def health_check(self) -> bool: | ||
| 25 | + """检查数据库连接健康状态. | ||
| 26 | + | ||
| 27 | + Returns: | ||
| 28 | + bool: 连接正常返回True，否则返回False | ||
| 29 | + """ | ||
| 30 | + pass | ||
| 31 | + | ||
| 32 | + | ||
| 33 | +class TableInit(ABC): | ||
| 34 | + | ||
| 35 | + @abstractmethod | ||
| 36 | + async def create_table(self, name, embedding_version, embedding_dim) -> None: | ||
| 37 | + pass |
| 1 | -import json | ||
| 2 | -import logging | ||
| 3 | -import os | ||
| 4 | -import threading | ||
| 5 | -import typing | ||
| 6 | - | ||
| 7 | -from elasticsearch import AsyncElasticsearch | ||
| 8 | -import numpy | ||
| 9 | - | ||
| 10 | -logger = logging.getLogger(__name__) | ||
| 11 | - | ||
| 12 | - | ||
| 13 | -class DataDao: | ||
| 14 | - | ||
| 15 | - def __init__(self, es_client: AsyncElasticsearch): | ||
| 16 | - self.es_client = es_client | ||
| 17 | - | ||
| 18 | - self._lock = threading.RLock() | ||
| 19 | - | ||
| 20 | - async def search(self, tb_name: str, | ||
| 21 | - embedding=typing.Optional[typing.Union[numpy.ndarray, list]], | ||
| 22 | - filters=None, | ||
| 23 | - from_=typing.Optional[int], | ||
| 24 | - size=typing.Optional[int], | ||
| 25 | - ): | ||
| 26 | - search_body = dict() | ||
| 27 | - | ||
| 28 | - from_ = from_ if isinstance(from_, int) and from_ >= 0 else 0 | ||
| 29 | - size = size if isinstance(size, int) and size > 0 else 10 | ||
| 30 | - search_body.update({ | ||
| 31 | - 'from': from_, | ||
| 32 | - 'size': size, | ||
| 33 | - }) | ||
| 34 | - | ||
| 35 | - return_embedding = os.getenv('RETURN_EMBEDDING', 'false').lower().strip() | ||
| 36 | - if return_embedding != 'true': | ||
| 37 | - search_body.update({ | ||
| 38 | - "_source": { | ||
| 39 | - "excludes": ["embedding"] | ||
| 40 | - } | ||
| 41 | - }) | ||
| 42 | - | ||
| 43 | - condition_query = { | ||
| 44 | - "match_all": {} | ||
| 45 | - } | ||
| 46 | - if isinstance(filters, typing.Collection) and len(filters) > 0: | ||
| 47 | - tb_mapping = await self._get_mapping(tb_name) | ||
| 48 | - exist_properties = tb_mapping.get('mappings', {}).get('properties', {}) | ||
| 49 | - | ||
| 50 | - clauses = [] | ||
| 51 | - for filter in filters: | ||
| 52 | - k = filter.name | ||
| 53 | - v = filter.value | ||
| 54 | - opt = filter.opt | ||
| 55 | - if k not in exist_properties: | ||
| 56 | - continue | ||
| 57 | - clause = get_filter_clause(k, v, opt, exist_properties[k]) | ||
| 58 | - if clause: | ||
| 59 | - clauses.append(clause) | ||
| 60 | - condition_query = { | ||
| 61 | - "bool": { | ||
| 62 | - "filter": clauses | ||
| 63 | - } | ||
| 64 | - } | ||
| 65 | - if embedding is not None: | ||
| 66 | - use_knn = os.getenv('USE_KNN', 'false').lower().strip() | ||
| 67 | - if use_knn != 'true': | ||
| 68 | - query_body = { | ||
| 69 | - "script_score": { | ||
| 70 | - "query": condition_query, | ||
| 71 | - "script": { | ||
| 72 | - "source": """ | ||
| 73 | - // 1. 计算原始分数 (范围 [0, 2]) | ||
| 74 | - double rawScore = cosineSimilarity(params.query_vector, 'embedding') + 1.0; | ||
| 75 | - | ||
| 76 | - // 2. 归一化到 [0, 1] | ||
| 77 | - double normalizedScore = rawScore / 2.0; | ||
| 78 | - | ||
| 79 | - // 3. 强制截断：确保最小值为 0.0，最大值为 1.0 | ||
| 80 | - // Math.max 防止出现负数，Math.min 防止出现 > 1.0 的数 | ||
| 81 | - double clampedScore = Math.max(0.0, Math.min(1.0, normalizedScore)); | ||
| 82 | - | ||
| 83 | - // 4. 四舍五入到 4 位小数 | ||
| 84 | - // 原理：乘以 10000 -> 四舍五入取整 -> 除以 10000 | ||
| 85 | - // clampedScore = Math.round(clampedScore * 10000.0) / 10000.0; | ||
| 86 | - return clampedScore; | ||
| 87 | - """, | ||
| 88 | - "params": { | ||
| 89 | - "query_vector": embedding.tolist() if isinstance(embedding, numpy.ndarray) else list(embedding) | ||
| 90 | - } | ||
| 91 | - } | ||
| 92 | - } | ||
| 93 | - } | ||
| 94 | - search_body.update({ | ||
| 95 | - "query": query_body | ||
| 96 | - }) | ||
| 97 | - else: | ||
| 98 | - search_body.update({ | ||
| 99 | - "knn": { | ||
| 100 | - "field": "embedding", | ||
| 101 | - "query_vector": embedding.tolist() if isinstance(embedding, numpy.ndarray) else list(embedding), | ||
| 102 | - "k": from_ + size, | ||
| 103 | - "num_candidates": (from_ + size) * 20, | ||
| 104 | - "filter": condition_query, | ||
| 105 | - } | ||
| 106 | - }) | ||
| 107 | - else: | ||
| 108 | - query_body = condition_query | ||
| 109 | - search_body.update({ | ||
| 110 | - "query": query_body | ||
| 111 | - }) | ||
| 112 | - | ||
| 113 | - logger.info(f'index {tb_name} search body: {json.dumps(search_body, ensure_ascii=False, indent=4)}') | ||
| 114 | - resp = await self.es_client.search(index=tb_name, body=search_body) | ||
| 115 | - logger.info(f'index {tb_name} search response: {resp}') | ||
| 116 | - | ||
| 117 | - body = resp.body | ||
| 118 | - hits = body.get('hits', {}).get('hits', []) | ||
| 119 | - result_data = [ | ||
| 120 | - { | ||
| 121 | - 'id': hit['_id'], | ||
| 122 | - 'score': round(hit['_score'], 4), | ||
| 123 | - 'kwargs': hit['_source'], | ||
| 124 | - } | ||
| 125 | - for hit in hits | ||
| 126 | - ] | ||
| 127 | - return result_data | ||
| 128 | - | ||
| 129 | - async def delete_by_pks(self, tb_name: str, pks: typing.List[str]): | ||
| 130 | - query_body = { | ||
| 131 | - "query": { | ||
| 132 | - "terms": { | ||
| 133 | - "_id": pks | ||
| 134 | - } | ||
| 135 | - } | ||
| 136 | - } | ||
| 137 | - resp = await self.es_client.delete_by_query(index=tb_name, body=query_body) | ||
| 138 | - logger.info(f'index {tb_name} delete_by_query response: {resp}') | ||
| 139 | - | ||
| 140 | - body = resp.body | ||
| 141 | - # return body.get('deleted') | ||
| 142 | - | ||
| 143 | - async def _get_mapping(self, tb_name: str): | ||
| 144 | - resp = await self.es_client.indices.get_mapping(index=tb_name) | ||
| 145 | - logger.info(f'index {tb_name} get_mapping response: {resp}') | ||
| 146 | - return resp.body.get(tb_name, {}) | ||
| 147 | - | ||
| 148 | - async def _tb_exist(self, tb_name, mapping=None): | ||
| 149 | - resp = await self.es_client.indices.exists(index=tb_name) | ||
| 150 | - result = resp.body | ||
| 151 | - if result and mapping is not None: | ||
| 152 | - tb_mapping = await self._get_mapping(tb_name) | ||
| 153 | - exist_properties = tb_mapping.get('mappings', {}).get('properties', {}) | ||
| 154 | - | ||
| 155 | - expect_properties = mapping.get('mappings', {}).get('properties', {}) | ||
| 156 | - if exist_properties and expect_properties: | ||
| 157 | - shared_keys = set(exist_properties.keys()).intersection(set(expect_properties.keys())) | ||
| 158 | - for k in shared_keys or []: | ||
| 159 | - if exist_properties[k]['type'] != expect_properties[k]['type']: | ||
| 160 | - raise Exception( | ||
| 161 | - f'index {tb_name} `{k}` type not match, expect: {expect_properties[k]["type"]}, actual: {exist_properties[k]["type"]}') | ||
| 162 | - else: | ||
| 163 | - if exist_properties[k]['type'] == 'dense_vector': | ||
| 164 | - if exist_properties[k].get('dims') != expect_properties[k].get('dims'): | ||
| 165 | - raise Exception( | ||
| 166 | - f'index {tb_name} dims not match, expect: {expect_properties[k].get("dims")}, actual: {exist_properties[k].get("dims")}') | ||
| 167 | - return result | ||
| 168 | - | ||
| 169 | - async def _create_index(self, tb_name, mapping=typing.Optional[dict]) -> bool: | ||
| 170 | - result = await self._tb_exist(tb_name=tb_name, mapping=mapping) | ||
| 171 | - if result: | ||
| 172 | - return True | ||
| 173 | - with self._lock: | ||
| 174 | - result = await self._tb_exist(tb_name=tb_name, mapping=mapping) | ||
| 175 | - if result: | ||
| 176 | - return True | ||
| 177 | - resp = await self.es_client.indices.create(index=tb_name, body=mapping) | ||
| 178 | - logger.info(f'index {tb_name} create response: {resp}') | ||
| 179 | - return resp.body.get('acknowledged') is True | ||
| 180 | - | ||
| 181 | - async def upsert(self, tb_name, id, embedding, params: dict): | ||
| 182 | - doc = { | ||
| 183 | - 'embedding': embedding, | ||
| 184 | - **params | ||
| 185 | - } | ||
| 186 | - | ||
| 187 | - auto_create_index = os.getenv('AUTO_CREATE_INDEX', 'true').lower().strip() | ||
| 188 | - if auto_create_index == 'true': | ||
| 189 | - mapping = generate_mapping(doc) | ||
| 190 | - # logger.info(f'Possible mapping: {json.dumps(mapping, indent=4)}') | ||
| 191 | - result = await self._create_index(tb_name=tb_name, mapping=mapping) | ||
| 192 | - | ||
| 193 | - if isinstance(embedding, numpy.ndarray): | ||
| 194 | - embedding = embedding.tolist() | ||
| 195 | - # elif isinstance(embedding, typing.Iterable): | ||
| 196 | - # if isinstance(embedding, typing.Mapping): | ||
| 197 | - # pass | ||
| 198 | - # else: | ||
| 199 | - # embedding = list(embedding) | ||
| 200 | - doc['embedding'] = embedding | ||
| 201 | - | ||
| 202 | - resp = await self.es_client.index(index=tb_name, id=id, body=doc) | ||
| 203 | - logger.info(f'index {tb_name} index response: {resp}') | ||
| 204 | - | ||
| 205 | - body = resp.body | ||
| 206 | - assert body.get('result') in ['created', 'updated'], f'数据插入失败, id: {id}' | ||
| 207 | - return body.get('_id') | ||
| 208 | - | ||
| 209 | - | ||
| 210 | -def get_filter_clause(k, v, opt, type_property): | ||
| 211 | - property_type = type_property.get('type') | ||
| 212 | - if property_type == 'text': | ||
| 213 | - property_fields = type_property.get('fields') or {} | ||
| 214 | - if 'keyword' in property_fields: | ||
| 215 | - k = f'{k}.keyword' | ||
| 216 | - | ||
| 217 | - opt = opt.lower() | ||
| 218 | - # "eq", "neq", "lt", "gt", "lte", "gte", "like", "in" | ||
| 219 | - if opt in ['eq', 'in']: | ||
| 220 | - v = list(v) if isinstance(v, typing.Collection) else [v] | ||
| 221 | - should_list = [] | ||
| 222 | - if None in v or len(v) == 0: | ||
| 223 | - should_list.append({ | ||
| 224 | - "bool": { | ||
| 225 | - "must_not": [ | ||
| 226 | - { | ||
| 227 | - "exists": { | ||
| 228 | - "field": k | ||
| 229 | - } | ||
| 230 | - } | ||
| 231 | - ] | ||
| 232 | - } | ||
| 233 | - }) | ||
| 234 | - no_null_v = [item for item in v if item is not None] | ||
| 235 | - if len(no_null_v) > 0: | ||
| 236 | - should_list.append({ | ||
| 237 | - "terms": { | ||
| 238 | - k: no_null_v | ||
| 239 | - } | ||
| 240 | - }) | ||
| 241 | - return should_list[0] if len(should_list) == 1 else { | ||
| 242 | - "bool": { | ||
| 243 | - "should": should_list, | ||
| 244 | - "minimum_should_match": 1 | ||
| 245 | - } | ||
| 246 | - } | ||
| 247 | - elif opt == 'neq': | ||
| 248 | - must_list = [] | ||
| 249 | - v = list(v) if isinstance(v, typing.Collection) else [v] | ||
| 250 | - | ||
| 251 | - if None in v or len(v) == 0: | ||
| 252 | - must_list.append({ | ||
| 253 | - "bool": { | ||
| 254 | - "must": [ | ||
| 255 | - { | ||
| 256 | - "exists": { | ||
| 257 | - "field": k | ||
| 258 | - } | ||
| 259 | - } | ||
| 260 | - ] | ||
| 261 | - } | ||
| 262 | - }) | ||
| 263 | - no_null_v = [item for item in v if item is not None] | ||
| 264 | - if len(no_null_v) > 0: | ||
| 265 | - must_list.append({ | ||
| 266 | - "bool": { | ||
| 267 | - "must_not": { | ||
| 268 | - "terms": { | ||
| 269 | - k: no_null_v | ||
| 270 | - } | ||
| 271 | - } | ||
| 272 | - } | ||
| 273 | - }) | ||
| 274 | - return must_list[0] if len(must_list) == 1 else { | ||
| 275 | - "bool": { | ||
| 276 | - "must": must_list | ||
| 277 | - } | ||
| 278 | - } | ||
| 279 | - elif opt in ['lt', 'lte', 'gt', 'gte']: | ||
| 280 | - v = list(v) if isinstance(v, typing.Collection) else [v] | ||
| 281 | - if len(v) > 0 and v[0] is not None: | ||
| 282 | - return { | ||
| 283 | - "range": { | ||
| 284 | - k: {opt: list(v)[0]} | ||
| 285 | - } | ||
| 286 | - } | ||
| 287 | - else: | ||
| 288 | - return { | ||
| 289 | - "range": { | ||
| 290 | - k: {} | ||
| 291 | - } | ||
| 292 | - } | ||
| 293 | - elif opt == 'like': | ||
| 294 | - v = list(v) if isinstance(v, typing.Collection) else [v] | ||
| 295 | - if len(v) > 0 and v[0] is not None: | ||
| 296 | - return { | ||
| 297 | - "wildcard": { | ||
| 298 | - k: f"*{list(v)[0]}*" | ||
| 299 | - } | ||
| 300 | - } | ||
| 301 | - else: | ||
| 302 | - return { | ||
| 303 | - "bool": { | ||
| 304 | - "must_not": [ | ||
| 305 | - { | ||
| 306 | - "exists": { | ||
| 307 | - "field": k | ||
| 308 | - } | ||
| 309 | - } | ||
| 310 | - ] | ||
| 311 | - } | ||
| 312 | - } | ||
| 313 | - else: | ||
| 314 | - raise Exception(f'opt {opt} not support') | ||
| 315 | - | ||
| 316 | - | ||
| 317 | -def generate_mapping(doc): | ||
| 318 | - def _inner(v): | ||
| 319 | - if isinstance(v, numpy.ndarray): | ||
| 320 | - return 'dense_vector' | ||
| 321 | - elif isinstance(v, str): | ||
| 322 | - return 'text' | ||
| 323 | - elif isinstance(v, int): | ||
| 324 | - return 'long' | ||
| 325 | - elif isinstance(v, float): | ||
| 326 | - return 'float' | ||
| 327 | - elif isinstance(v, typing.Mapping): | ||
| 328 | - return None | ||
| 329 | - elif isinstance(v, typing.Collection): | ||
| 330 | - if len(v) == 0: | ||
| 331 | - return None | ||
| 332 | - else: | ||
| 333 | - return _inner(list(v)[0]) | ||
| 334 | - else: | ||
| 335 | - return None | ||
| 336 | - | ||
| 337 | - properties = dict() | ||
| 338 | - for key, value in (doc or {}).items(): | ||
| 339 | - key_type = _inner(value) | ||
| 340 | - if key_type: | ||
| 341 | - properties[key] = { | ||
| 342 | - 'type': key_type, | ||
| 343 | - } | ||
| 344 | - if key_type == 'dense_vector': | ||
| 345 | - properties[key]['dims'] = len(value) | ||
| 346 | - properties[key].update({ | ||
| 347 | - 'dims': len(value), | ||
| 348 | - 'index': True, | ||
| 349 | - 'similarity': 'cosine' | ||
| 350 | - }) | ||
| 351 | - elif key_type == 'text': | ||
| 352 | - properties[key]['fields'] = { | ||
| 353 | - "keyword": { | ||
| 354 | - "type": "keyword", | ||
| 355 | - # "ignore_above": 256 | ||
| 356 | - } | ||
| 357 | - } | ||
| 358 | - es_mapping = { | ||
| 359 | - "settings": { | ||
| 360 | - "number_of_replicas": int(os.getenv('ES_NUMBER_OF_REPLICAS', 0)), | ||
| 361 | - }, | ||
| 362 | - "mappings": { | ||
| 363 | - "properties": properties | ||
| 364 | - } | ||
| 365 | - } | ||
| 366 | - return es_mapping | 1 | +import json |
| 2 | +import logging | ||
| 3 | +import os | ||
| 4 | +import threading | ||
| 5 | +import typing | ||
| 6 | + | ||
| 7 | +from elasticsearch import AsyncElasticsearch | ||
| 8 | +import numpy | ||
| 9 | + | ||
| 10 | +logger = logging.getLogger(__name__) | ||
| 11 | + | ||
| 12 | + | ||
| 13 | +class DataDao: | ||
| 14 | + | ||
| 15 | + def __init__(self, es_client: AsyncElasticsearch): | ||
| 16 | + self.es_client = es_client | ||
| 17 | + | ||
| 18 | + self._lock = threading.RLock() | ||
| 19 | + | ||
| 20 | + async def search(self, tb_name: str, | ||
| 21 | + embedding=typing.Optional[typing.Union[numpy.ndarray, list]], | ||
| 22 | + filters=None, | ||
| 23 | + from_=typing.Optional[int], | ||
| 24 | + size=typing.Optional[int], | ||
| 25 | + ): | ||
| 26 | + search_body = dict() | ||
| 27 | + | ||
| 28 | + from_ = from_ if isinstance(from_, int) and from_ >= 0 else 0 | ||
| 29 | + size = size if isinstance(size, int) and size > 0 else 10 | ||
| 30 | + search_body.update({ | ||
| 31 | + 'from': from_, | ||
| 32 | + 'size': size, | ||
| 33 | + }) | ||
| 34 | + | ||
| 35 | + return_embedding = os.getenv('RETURN_EMBEDDING', 'false').lower().strip() | ||
| 36 | + if return_embedding != 'true': | ||
| 37 | + search_body.update({ | ||
| 38 | + "_source": { | ||
| 39 | + "excludes": ["embedding"] | ||
| 40 | + } | ||
| 41 | + }) | ||
| 42 | + | ||
| 43 | + condition_query = { | ||
| 44 | + "match_all": {} | ||
| 45 | + } | ||
| 46 | + if isinstance(filters, typing.Collection) and len(filters) > 0: | ||
| 47 | + tb_mapping = await self._get_mapping(tb_name) | ||
| 48 | + exist_properties = tb_mapping.get('mappings', {}).get('properties', {}) | ||
| 49 | + | ||
| 50 | + clauses = [] | ||
| 51 | + for filter in filters: | ||
| 52 | + k = filter.name | ||
| 53 | + v = filter.value | ||
| 54 | + opt = filter.opt | ||
| 55 | + if k not in exist_properties: | ||
| 56 | + continue | ||
| 57 | + clause = get_filter_clause(k, v, opt, exist_properties[k]) | ||
| 58 | + if clause: | ||
| 59 | + clauses.append(clause) | ||
| 60 | + condition_query = { | ||
| 61 | + "bool": { | ||
| 62 | + "filter": clauses | ||
| 63 | + } | ||
| 64 | + } | ||
| 65 | + if embedding is not None: | ||
| 66 | + use_knn = os.getenv('USE_KNN', 'false').lower().strip() | ||
| 67 | + if use_knn != 'true': | ||
| 68 | + query_body = { | ||
| 69 | + "script_score": { | ||
| 70 | + "query": condition_query, | ||
| 71 | + "script": { | ||
| 72 | + "source": """ | ||
| 73 | + // 1. 计算原始分数 (范围 [0, 2]) | ||
| 74 | + double rawScore = cosineSimilarity(params.query_vector, 'embedding') + 1.0; | ||
| 75 | + | ||
| 76 | + // 2. 归一化到 [0, 1] | ||
| 77 | + double normalizedScore = rawScore / 2.0; | ||
| 78 | + | ||
| 79 | + // 3. 强制截断：确保最小值为 0.0，最大值为 1.0 | ||
| 80 | + // Math.max 防止出现负数，Math.min 防止出现 > 1.0 的数 | ||
| 81 | + double clampedScore = Math.max(0.0, Math.min(1.0, normalizedScore)); | ||
| 82 | + | ||
| 83 | + // 4. 四舍五入到 4 位小数 | ||
| 84 | + // 原理：乘以 10000 -> 四舍五入取整 -> 除以 10000 | ||
| 85 | + // clampedScore = Math.round(clampedScore * 10000.0) / 10000.0; | ||
| 86 | + return clampedScore; | ||
| 87 | + """, | ||
| 88 | + "params": { | ||
| 89 | + "query_vector": embedding.tolist() if isinstance(embedding, numpy.ndarray) else list(embedding) | ||
| 90 | + } | ||
| 91 | + } | ||
| 92 | + } | ||
| 93 | + } | ||
| 94 | + search_body.update({ | ||
| 95 | + "query": query_body | ||
| 96 | + }) | ||
| 97 | + else: | ||
| 98 | + search_body.update({ | ||
| 99 | + "knn": { | ||
| 100 | + "field": "embedding", | ||
| 101 | + "query_vector": embedding.tolist() if isinstance(embedding, numpy.ndarray) else list(embedding), | ||
| 102 | + "k": from_ + size, | ||
| 103 | + "num_candidates": (from_ + size) * 20, | ||
| 104 | + "filter": condition_query, | ||
| 105 | + } | ||
| 106 | + }) | ||
| 107 | + else: | ||
| 108 | + query_body = condition_query | ||
| 109 | + search_body.update({ | ||
| 110 | + "query": query_body | ||
| 111 | + }) | ||
| 112 | + | ||
| 113 | + logger.info(f'index {tb_name} search body: {json.dumps(search_body, ensure_ascii=False, indent=4)}') | ||
| 114 | + resp = await self.es_client.search(index=tb_name, body=search_body) | ||
| 115 | + logger.info(f'index {tb_name} search response: {resp}') | ||
| 116 | + | ||
| 117 | + body = resp.body | ||
| 118 | + hits = body.get('hits', {}).get('hits', []) | ||
| 119 | + result_data = [ | ||
| 120 | + { | ||
| 121 | + 'id': hit['_id'], | ||
| 122 | + 'score': round(hit['_score'], 4), | ||
| 123 | + 'kwargs': hit['_source'], | ||
| 124 | + } | ||
| 125 | + for hit in hits | ||
| 126 | + ] | ||
| 127 | + return result_data | ||
| 128 | + | ||
| 129 | + async def delete_by_pks(self, tb_name: str, pks: typing.List[str]): | ||
| 130 | + query_body = { | ||
| 131 | + "query": { | ||
| 132 | + "terms": { | ||
| 133 | + "_id": pks | ||
| 134 | + } | ||
| 135 | + } | ||
| 136 | + } | ||
| 137 | + resp = await self.es_client.delete_by_query(index=tb_name, body=query_body) | ||
| 138 | + logger.info(f'index {tb_name} delete_by_query response: {resp}') | ||
| 139 | + | ||
| 140 | + body = resp.body | ||
| 141 | + # return body.get('deleted') | ||
| 142 | + | ||
| 143 | + async def _get_mapping(self, tb_name: str): | ||
| 144 | + resp = await self.es_client.indices.get_mapping(index=tb_name) | ||
| 145 | + logger.info(f'index {tb_name} get_mapping response: {resp}') | ||
| 146 | + return resp.body.get(tb_name, {}) | ||
| 147 | + | ||
| 148 | + async def _tb_exist(self, tb_name, mapping=None): | ||
| 149 | + resp = await self.es_client.indices.exists(index=tb_name) | ||
| 150 | + result = resp.body | ||
| 151 | + if result and mapping is not None: | ||
| 152 | + tb_mapping = await self._get_mapping(tb_name) | ||
| 153 | + exist_properties = tb_mapping.get('mappings', {}).get('properties', {}) | ||
| 154 | + | ||
| 155 | + expect_properties = mapping.get('mappings', {}).get('properties', {}) | ||
| 156 | + if exist_properties and expect_properties: | ||
| 157 | + shared_keys = set(exist_properties.keys()).intersection(set(expect_properties.keys())) | ||
| 158 | + for k in shared_keys or []: | ||
| 159 | + if exist_properties[k]['type'] != expect_properties[k]['type']: | ||
| 160 | + raise Exception( | ||
| 161 | + f'index {tb_name} `{k}` type not match, expect: {expect_properties[k]["type"]}, actual: {exist_properties[k]["type"]}') | ||
| 162 | + else: | ||
| 163 | + if exist_properties[k]['type'] == 'dense_vector': | ||
| 164 | + if exist_properties[k].get('dims') != expect_properties[k].get('dims'): | ||
| 165 | + raise Exception( | ||
| 166 | + f'index {tb_name} dims not match, expect: {expect_properties[k].get("dims")}, actual: {exist_properties[k].get("dims")}') | ||
| 167 | + return result | ||
| 168 | + | ||
| 169 | + async def _create_index(self, tb_name, mapping=typing.Optional[dict]) -> bool: | ||
| 170 | + result = await self._tb_exist(tb_name=tb_name, mapping=mapping) | ||
| 171 | + if result: | ||
| 172 | + return True | ||
| 173 | + with self._lock: | ||
| 174 | + result = await self._tb_exist(tb_name=tb_name, mapping=mapping) | ||
| 175 | + if result: | ||
| 176 | + return True | ||
| 177 | + resp = await self.es_client.indices.create(index=tb_name, body=mapping) | ||
| 178 | + logger.info(f'index {tb_name} create response: {resp}') | ||
| 179 | + return resp.body.get('acknowledged') is True | ||
| 180 | + | ||
| 181 | + async def upsert(self, tb_name, id, embedding, params: dict): | ||
| 182 | + doc = { | ||
| 183 | + 'embedding': embedding, | ||
| 184 | + **params | ||
| 185 | + } | ||
| 186 | + | ||
| 187 | + auto_create_index = os.getenv('AUTO_CREATE_INDEX', 'true').lower().strip() | ||
| 188 | + if auto_create_index == 'true': | ||
| 189 | + mapping = generate_mapping(doc) | ||
| 190 | + # logger.info(f'Possible mapping: {json.dumps(mapping, indent=4)}') | ||
| 191 | + result = await self._create_index(tb_name=tb_name, mapping=mapping) | ||
| 192 | + | ||
| 193 | + if isinstance(embedding, numpy.ndarray): | ||
| 194 | + embedding = embedding.tolist() | ||
| 195 | + # elif isinstance(embedding, typing.Iterable): | ||
| 196 | + # if isinstance(embedding, typing.Mapping): | ||
| 197 | + # pass | ||
| 198 | + # else: | ||
| 199 | + # embedding = list(embedding) | ||
| 200 | + doc['embedding'] = embedding | ||
| 201 | + | ||
| 202 | + resp = await self.es_client.index(index=tb_name, id=id, body=doc) | ||
| 203 | + logger.info(f'index {tb_name} index response: {resp}') | ||
| 204 | + | ||
| 205 | + body = resp.body | ||
| 206 | + assert body.get('result') in ['created', 'updated'], f'数据插入失败, id: {id}' | ||
| 207 | + return body.get('_id') | ||
| 208 | + | ||
| 209 | + | ||
| 210 | +def get_filter_clause(k, v, opt, type_property): | ||
| 211 | + property_type = type_property.get('type') | ||
| 212 | + if property_type == 'text': | ||
| 213 | + property_fields = type_property.get('fields') or {} | ||
| 214 | + if 'keyword' in property_fields: | ||
| 215 | + k = f'{k}.keyword' | ||
| 216 | + | ||
| 217 | + opt = opt.lower() | ||
| 218 | + # "eq", "neq", "lt", "gt", "lte", "gte", "like", "in" | ||
| 219 | + if opt in ['eq', 'in']: | ||
| 220 | + v = list(v) if isinstance(v, typing.Collection) else [v] | ||
| 221 | + should_list = [] | ||
| 222 | + if None in v or len(v) == 0: | ||
| 223 | + should_list.append({ | ||
| 224 | + "bool": { | ||
| 225 | + "must_not": [ | ||
| 226 | + { | ||
| 227 | + "exists": { | ||
| 228 | + "field": k | ||
| 229 | + } | ||
| 230 | + } | ||
| 231 | + ] | ||
| 232 | + } | ||
| 233 | + }) | ||
| 234 | + no_null_v = [item for item in v if item is not None] | ||
| 235 | + if len(no_null_v) > 0: | ||
| 236 | + should_list.append({ | ||
| 237 | + "terms": { | ||
| 238 | + k: no_null_v | ||
| 239 | + } | ||
| 240 | + }) | ||
| 241 | + return should_list[0] if len(should_list) == 1 else { | ||
| 242 | + "bool": { | ||
| 243 | + "should": should_list, | ||
| 244 | + "minimum_should_match": 1 | ||
| 245 | + } | ||
| 246 | + } | ||
| 247 | + elif opt == 'neq': | ||
| 248 | + must_list = [] | ||
| 249 | + v = list(v) if isinstance(v, typing.Collection) else [v] | ||
| 250 | + | ||
| 251 | + if None in v or len(v) == 0: | ||
| 252 | + must_list.append({ | ||
| 253 | + "bool": { | ||
| 254 | + "must": [ | ||
| 255 | + { | ||
| 256 | + "exists": { | ||
| 257 | + "field": k | ||
| 258 | + } | ||
| 259 | + } | ||
| 260 | + ] | ||
| 261 | + } | ||
| 262 | + }) | ||
| 263 | + no_null_v = [item for item in v if item is not None] | ||
| 264 | + if len(no_null_v) > 0: | ||
| 265 | + must_list.append({ | ||
| 266 | + "bool": { | ||
| 267 | + "must_not": { | ||
| 268 | + "terms": { | ||
| 269 | + k: no_null_v | ||
| 270 | + } | ||
| 271 | + } | ||
| 272 | + } | ||
| 273 | + }) | ||
| 274 | + return must_list[0] if len(must_list) == 1 else { | ||
| 275 | + "bool": { | ||
| 276 | + "must": must_list | ||
| 277 | + } | ||
| 278 | + } | ||
| 279 | + elif opt in ['lt', 'lte', 'gt', 'gte']: | ||
| 280 | + v = list(v) if isinstance(v, typing.Collection) else [v] | ||
| 281 | + if len(v) > 0 and v[0] is not None: | ||
| 282 | + return { | ||
| 283 | + "range": { | ||
| 284 | + k: {opt: list(v)[0]} | ||
| 285 | + } | ||
| 286 | + } | ||
| 287 | + else: | ||
| 288 | + return { | ||
| 289 | + "range": { | ||
| 290 | + k: {} | ||
| 291 | + } | ||
| 292 | + } | ||
| 293 | + elif opt == 'like': | ||
| 294 | + v = list(v) if isinstance(v, typing.Collection) else [v] | ||
| 295 | + if len(v) > 0 and v[0] is not None: | ||
| 296 | + return { | ||
| 297 | + "wildcard": { | ||
| 298 | + k: f"*{list(v)[0]}*" | ||
| 299 | + } | ||
| 300 | + } | ||
| 301 | + else: | ||
| 302 | + return { | ||
| 303 | + "bool": { | ||
| 304 | + "must_not": [ | ||
| 305 | + { | ||
| 306 | + "exists": { | ||
| 307 | + "field": k | ||
| 308 | + } | ||
| 309 | + } | ||
| 310 | + ] | ||
| 311 | + } | ||
| 312 | + } | ||
| 313 | + else: | ||
| 314 | + raise Exception(f'opt {opt} not support') | ||
| 315 | + | ||
| 316 | + | ||
| 317 | +def generate_mapping(doc): | ||
| 318 | + def _inner(v): | ||
| 319 | + if isinstance(v, numpy.ndarray): | ||
| 320 | + return 'dense_vector' | ||
| 321 | + elif isinstance(v, str): | ||
| 322 | + return 'text' | ||
| 323 | + elif isinstance(v, int): | ||
| 324 | + return 'long' | ||
| 325 | + elif isinstance(v, float): | ||
| 326 | + return 'float' | ||
| 327 | + elif isinstance(v, typing.Mapping): | ||
| 328 | + return None | ||
| 329 | + elif isinstance(v, typing.Collection): | ||
| 330 | + if len(v) == 0: | ||
| 331 | + return None | ||
| 332 | + else: | ||
| 333 | + return _inner(list(v)[0]) | ||
| 334 | + else: | ||
| 335 | + return None | ||
| 336 | + | ||
| 337 | + properties = dict() | ||
| 338 | + for key, value in (doc or {}).items(): | ||
| 339 | + key_type = _inner(value) | ||
| 340 | + if key_type: | ||
| 341 | + properties[key] = { | ||
| 342 | + 'type': key_type, | ||
| 343 | + } | ||
| 344 | + if key_type == 'dense_vector': | ||
| 345 | + properties[key]['dims'] = len(value) | ||
| 346 | + properties[key].update({ | ||
| 347 | + 'dims': len(value), | ||
| 348 | + 'index': True, | ||
| 349 | + 'similarity': 'cosine' | ||
| 350 | + }) | ||
| 351 | + elif key_type == 'text': | ||
| 352 | + properties[key]['fields'] = { | ||
| 353 | + "keyword": { | ||
| 354 | + "type": "keyword", | ||
| 355 | + # "ignore_above": 256 | ||
| 356 | + } | ||
| 357 | + } | ||
| 358 | + es_mapping = { | ||
| 359 | + "settings": { | ||
| 360 | + "number_of_replicas": int(os.getenv('ES_NUMBER_OF_REPLICAS', 0)), | ||
| 361 | + }, | ||
| 362 | + "mappings": { | ||
| 363 | + "properties": properties | ||
| 364 | + } | ||
| 365 | + } | ||
| 366 | + return es_mapping |
| 1 | -"""Elasticsearch向量数据库客户端实现.""" | ||
| 2 | - | ||
| 3 | -from typing import Any, Optional | ||
| 4 | - | ||
| 5 | -from embedding_sever.db.base import VectorDBClient | ||
| 6 | - | ||
| 7 | - | ||
| 8 | -class ElasticsearchClient(VectorDBClient): | ||
| 9 | - """Elasticsearch向量数据库客户端. | ||
| 10 | - | ||
| 11 | - 使用Elasticsearch的dense_vector类型存储向量，支持向量相似度搜索。 | ||
| 12 | - """ | ||
| 13 | - | ||
| 14 | - def __init__(self, es_url: Optional[str] = None): | ||
| 15 | - """初始化ES客户端. | ||
| 16 | - | ||
| 17 | - Args: | ||
| 18 | - es_url: Elasticsearch连接URL，默认从配置读取 | ||
| 19 | - """ | ||
| 20 | - self.es_url = es_url | ||
| 21 | - self._client: Optional[Any] = None | ||
| 22 | - | ||
| 23 | - async def connect(self) -> None: | ||
| 24 | - """建立ES连接.""" | ||
| 25 | - from elasticsearch import AsyncElasticsearch | ||
| 26 | - | ||
| 27 | - self._client = AsyncElasticsearch([self.es_url]) | ||
| 28 | - # 验证连接 | ||
| 29 | - await self._client.ping() | ||
| 30 | - | ||
| 31 | - async def disconnect(self) -> None: | ||
| 32 | - """断开ES连接.""" | ||
| 33 | - if self._client: | ||
| 34 | - await self._client.close() | ||
| 35 | - self._client = None | ||
| 36 | - | ||
| 37 | - async def health_check(self) -> bool: | ||
| 38 | - """检查ES连接健康状态.""" | ||
| 39 | - if not self._client: | ||
| 40 | - return False | ||
| 41 | - try: | ||
| 42 | - return await self._client.ping() | ||
| 43 | - except Exception: | ||
| 44 | - return False | ||
| 45 | -_es_client: Optional[ElasticsearchClient] = None | ||
| 46 | - | ||
| 47 | - | ||
| 48 | -def get_es_client() -> ElasticsearchClient: | ||
| 49 | - """获取全局ES客户端实例（单例模式）. | ||
| 50 | - | ||
| 51 | - Returns: | ||
| 52 | - ElasticsearchClient: ES客户端实例 | ||
| 53 | - """ | ||
| 54 | - global _es_client | ||
| 55 | - if _es_client is None: | ||
| 56 | - _es_client = ElasticsearchClient() | ||
| 57 | - return _es_client | 1 | +"""Elasticsearch向量数据库客户端实现.""" |
| 2 | + | ||
| 3 | +from typing import Any, Optional | ||
| 4 | + | ||
| 5 | +from embedding_sever.db.base import VectorDBClient | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +class ElasticsearchClient(VectorDBClient): | ||
| 9 | + """Elasticsearch向量数据库客户端. | ||
| 10 | + | ||
| 11 | + 使用Elasticsearch的dense_vector类型存储向量，支持向量相似度搜索。 | ||
| 12 | + """ | ||
| 13 | + | ||
| 14 | + def __init__(self, es_url: Optional[str] = None): | ||
| 15 | + """初始化ES客户端. | ||
| 16 | + | ||
| 17 | + Args: | ||
| 18 | + es_url: Elasticsearch连接URL，默认从配置读取 | ||
| 19 | + """ | ||
| 20 | + self.es_url = es_url | ||
| 21 | + self._client: Optional[Any] = None | ||
| 22 | + | ||
| 23 | + async def connect(self) -> None: | ||
| 24 | + """建立ES连接.""" | ||
| 25 | + from elasticsearch import AsyncElasticsearch | ||
| 26 | + | ||
| 27 | + self._client = AsyncElasticsearch([self.es_url]) | ||
| 28 | + # 验证连接 | ||
| 29 | + await self._client.ping() | ||
| 30 | + | ||
| 31 | + async def disconnect(self) -> None: | ||
| 32 | + """断开ES连接.""" | ||
| 33 | + if self._client: | ||
| 34 | + await self._client.close() | ||
| 35 | + self._client = None | ||
| 36 | + | ||
| 37 | + async def health_check(self) -> bool: | ||
| 38 | + """检查ES连接健康状态.""" | ||
| 39 | + if not self._client: | ||
| 40 | + return False | ||
| 41 | + try: | ||
| 42 | + return await self._client.ping() | ||
| 43 | + except Exception: | ||
| 44 | + return False | ||
| 45 | +_es_client: Optional[ElasticsearchClient] = None | ||
| 46 | + | ||
| 47 | + | ||
| 48 | +def get_es_client() -> ElasticsearchClient: | ||
| 49 | + """获取全局ES客户端实例（单例模式）. | ||
| 50 | + | ||
| 51 | + Returns: | ||
| 52 | + ElasticsearchClient: ES客户端实例 | ||
| 53 | + """ | ||
| 54 | + global _es_client | ||
| 55 | + if _es_client is None: | ||
| 56 | + _es_client = ElasticsearchClient() | ||
| 57 | + return _es_client |
| 1 | -from .base import TableInit | ||
| 2 | -from . import es_client | ||
| 3 | - | ||
| 4 | - | ||
| 5 | -class ESTableInit(TableInit): | ||
| 6 | - | ||
| 7 | - def __init__(self): | ||
| 8 | - self.es_client = es_client | ||
| 9 | - | ||
| 10 | - async def create_table(self, name, embedding_version, embedding_dim) -> None: | ||
| 11 | - pass | 1 | +from .base import TableInit |
| 2 | +from . import es_client | ||
| 3 | + | ||
| 4 | + | ||
| 5 | +class ESTableInit(TableInit): | ||
| 6 | + | ||
| 7 | + def __init__(self): | ||
| 8 | + self.es_client = es_client | ||
| 9 | + | ||
| 10 | + async def create_table(self, name, embedding_version, embedding_dim) -> None: | ||
| 11 | + pass |
| 1 | -"""PostgreSQL向量数据库客户端实现.""" | ||
| 2 | - | ||
| 3 | -from typing import Any, Optional | ||
| 4 | - | ||
| 5 | -from embedding_sever.db.base import VectorDBClient | ||
| 6 | - | ||
| 7 | - | ||
| 8 | -class PostgresClient(VectorDBClient): | ||
| 9 | - """PostgreSQL向量数据库客户端. | ||
| 10 | - | ||
| 11 | - 使用pgvector扩展存储向量，支持向量相似度搜索。 | ||
| 12 | - 需要安装pgvector扩展: CREATE EXTENSION IF NOT EXISTS vector; | ||
| 13 | - """ | ||
| 14 | - | ||
| 15 | - def __init__(self, pg_url: Optional[str] = None): | ||
| 16 | - """初始化PG客户端. | ||
| 17 | - | ||
| 18 | - Args: | ||
| 19 | - pg_url: PostgreSQL连接URL，默认从配置读取 | ||
| 20 | - """ | ||
| 21 | - self.pg_url = pg_url | ||
| 22 | - self._pool: Optional[Any] = None | ||
| 23 | - | ||
| 24 | - async def connect(self) -> None: | ||
| 25 | - """建立PG连接池.""" | ||
| 26 | - from psycopg_pool import AsyncConnectionPool | ||
| 27 | - | ||
| 28 | - self._pool = AsyncConnectionPool( | ||
| 29 | - conninfo=self.pg_url, | ||
| 30 | - min_size=1, | ||
| 31 | - max_size=10, | ||
| 32 | - ) | ||
| 33 | - | ||
| 34 | - # 验证连接并启用pgvector | ||
| 35 | - async with self._pool.connection() as conn: | ||
| 36 | - await conn.execute("CREATE EXTENSION IF NOT EXISTS vector") | ||
| 37 | - await conn.commit() | ||
| 38 | - | ||
| 39 | - async def disconnect(self) -> None: | ||
| 40 | - """断开PG连接.""" | ||
| 41 | - if self._pool: | ||
| 42 | - await self._pool.close() | ||
| 43 | - self._pool = None | ||
| 44 | - | ||
| 45 | - async def health_check(self) -> bool: | ||
| 46 | - """检查PG连接健康状态.""" | ||
| 47 | - if not self._pool: | ||
| 48 | - return False | ||
| 49 | - try: | ||
| 50 | - async with self._pool.connection() as conn: | ||
| 51 | - cursor = await conn.execute("SELECT 1") | ||
| 52 | - result = await cursor.fetchone() | ||
| 53 | - return result is not None and result[0] == 1 | ||
| 54 | - except Exception: | ||
| 55 | - return False | ||
| 56 | - | ||
| 57 | - | ||
| 58 | -# 全局PG客户端实例 | ||
| 59 | -_pg_client: Optional[PostgresClient] = None | ||
| 60 | - | ||
| 61 | - | ||
| 62 | -def get_pg_client() -> PostgresClient: | ||
| 63 | - """获取全局PG客户端实例（单例模式）. | ||
| 64 | - | ||
| 65 | - Returns: | ||
| 66 | - PostgresClient: PG客户端实例 | ||
| 67 | - """ | ||
| 68 | - global _pg_client | ||
| 69 | - if _pg_client is None: | ||
| 70 | - _pg_client = PostgresClient() | ||
| 71 | - return _pg_client | 1 | +"""PostgreSQL向量数据库客户端实现.""" |
| 2 | + | ||
| 3 | +from typing import Any, Optional | ||
| 4 | + | ||
| 5 | +from embedding_sever.db.base import VectorDBClient | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +class PostgresClient(VectorDBClient): | ||
| 9 | + """PostgreSQL向量数据库客户端. | ||
| 10 | + | ||
| 11 | + 使用pgvector扩展存储向量，支持向量相似度搜索。 | ||
| 12 | + 需要安装pgvector扩展: CREATE EXTENSION IF NOT EXISTS vector; | ||
| 13 | + """ | ||
| 14 | + | ||
| 15 | + def __init__(self, pg_url: Optional[str] = None): | ||
| 16 | + """初始化PG客户端. | ||
| 17 | + | ||
| 18 | + Args: | ||
| 19 | + pg_url: PostgreSQL连接URL，默认从配置读取 | ||
| 20 | + """ | ||
| 21 | + self.pg_url = pg_url | ||
| 22 | + self._pool: Optional[Any] = None | ||
| 23 | + | ||
| 24 | + async def connect(self) -> None: | ||
| 25 | + """建立PG连接池.""" | ||
| 26 | + from psycopg_pool import AsyncConnectionPool | ||
| 27 | + | ||
| 28 | + self._pool = AsyncConnectionPool( | ||
| 29 | + conninfo=self.pg_url, | ||
| 30 | + min_size=1, | ||
| 31 | + max_size=10, | ||
| 32 | + ) | ||
| 33 | + | ||
| 34 | + # 验证连接并启用pgvector | ||
| 35 | + async with self._pool.connection() as conn: | ||
| 36 | + await conn.execute("CREATE EXTENSION IF NOT EXISTS vector") | ||
| 37 | + await conn.commit() | ||
| 38 | + | ||
| 39 | + async def disconnect(self) -> None: | ||
| 40 | + """断开PG连接.""" | ||
| 41 | + if self._pool: | ||
| 42 | + await self._pool.close() | ||
| 43 | + self._pool = None | ||
| 44 | + | ||
| 45 | + async def health_check(self) -> bool: | ||
| 46 | + """检查PG连接健康状态.""" | ||
| 47 | + if not self._pool: | ||
| 48 | + return False | ||
| 49 | + try: | ||
| 50 | + async with self._pool.connection() as conn: | ||
| 51 | + cursor = await conn.execute("SELECT 1") | ||
| 52 | + result = await cursor.fetchone() | ||
| 53 | + return result is not None and result[0] == 1 | ||
| 54 | + except Exception: | ||
| 55 | + return False | ||
| 56 | + | ||
| 57 | + | ||
| 58 | +# 全局PG客户端实例 | ||
| 59 | +_pg_client: Optional[PostgresClient] = None | ||
| 60 | + | ||
| 61 | + | ||
| 62 | +def get_pg_client() -> PostgresClient: | ||
| 63 | + """获取全局PG客户端实例（单例模式）. | ||
| 64 | + | ||
| 65 | + Returns: | ||
| 66 | + PostgresClient: PG客户端实例 | ||
| 67 | + """ | ||
| 68 | + global _pg_client | ||
| 69 | + if _pg_client is None: | ||
| 70 | + _pg_client = PostgresClient() | ||
| 71 | + return _pg_client |
| 1 | -import os | ||
| 2 | -from pathlib import Path | ||
| 3 | - | ||
| 4 | -if not os.environ.get("PROJECT_ROOT"): | ||
| 5 | - os.environ['PROJECT_ROOT'] = Path(__file__).parent.absolute().as_posix() | ||
| 6 | -from aabd.base.patched_logging import init_logging, get_logger | ||
| 7 | - | ||
| 8 | -init_logging() | ||
| 9 | -logger = get_logger(__name__) | ||
| 10 | - | ||
| 11 | -from contextlib import asynccontextmanager | ||
| 12 | -from typing import AsyncGenerator | ||
| 13 | -from fastapi import FastAPI | ||
| 14 | -from embedding_sever.api.http_config import config_fastapi | ||
| 15 | -from embedding_sever.config import settings | ||
| 16 | -from embedding_sever.db import connect as db_client_connect, disconnect as db_client_disconnect | ||
| 17 | - | ||
| 18 | - | ||
| 19 | -@asynccontextmanager | ||
| 20 | -async def lifespan(fastapi_app: FastAPI) -> AsyncGenerator[None, None]: | ||
| 21 | - """应用生命周期管理.""" | ||
| 22 | - # 启动时建立数据库连接 | ||
| 23 | - try: | ||
| 24 | - await db_client_connect() | ||
| 25 | - logger.info(f"✓ database connected") | ||
| 26 | - except Exception as e: | ||
| 27 | - logger.error(f"✗ Failed to connect to database: {e}") | ||
| 28 | - raise | ||
| 29 | - | ||
| 30 | - yield | ||
| 31 | - | ||
| 32 | - # 关闭时断开数据库连接 | ||
| 33 | - try: | ||
| 34 | - await db_client_disconnect() | ||
| 35 | - logger.info(f"✓ database disconnected") | ||
| 36 | - except Exception as e: | ||
| 37 | - logger.error(f"✗ Error disconnecting from database: {e}") | ||
| 38 | - | ||
| 39 | - | ||
| 40 | -# 创建应用实例 | ||
| 41 | -app = FastAPI( | ||
| 42 | - title=settings.app_name, | ||
| 43 | - version=settings.app_version, | ||
| 44 | - description="基于向量的CRUD服务，支持多种Embedding场景", | ||
| 45 | - lifespan=lifespan, | ||
| 46 | - docs_url="/docs", | ||
| 47 | - redoc_url="/redoc", | ||
| 48 | -) | ||
| 49 | -config_fastapi(app) | ||
| 50 | - | ||
| 51 | -def main(): | ||
| 52 | - """应用入口函数.""" | ||
| 53 | - import uvicorn | ||
| 54 | - | ||
| 55 | - logger.info(f'docs_url: http://127.0.0.1:{settings.port}/docs') | ||
| 56 | - uvicorn.run( | ||
| 57 | - app, | ||
| 58 | - host='0.0.0.0', | ||
| 59 | - port=settings.port, | ||
| 60 | - reload=settings.debug, | ||
| 61 | - log_config=None | ||
| 62 | - ) | ||
| 63 | - | ||
| 64 | - | ||
| 65 | -if __name__ == "__main__": | ||
| 66 | - main() | 1 | +import os |
| 2 | +from pathlib import Path | ||
| 3 | + | ||
| 4 | +if not os.environ.get("PROJECT_ROOT"): | ||
| 5 | + os.environ['PROJECT_ROOT'] = Path(__file__).parent.absolute().as_posix() | ||
| 6 | +from aabd.base.patched_logging import init_logging, get_logger | ||
| 7 | + | ||
| 8 | +init_logging() | ||
| 9 | +logger = get_logger(__name__) | ||
| 10 | + | ||
| 11 | +from contextlib import asynccontextmanager | ||
| 12 | +from typing import AsyncGenerator | ||
| 13 | +from fastapi import FastAPI | ||
| 14 | +from embedding_sever.api.http_config import config_fastapi | ||
| 15 | +from embedding_sever.config import settings | ||
| 16 | +from embedding_sever.db import connect as db_client_connect, disconnect as db_client_disconnect | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +@asynccontextmanager | ||
| 20 | +async def lifespan(fastapi_app: FastAPI) -> AsyncGenerator[None, None]: | ||
| 21 | + """应用生命周期管理.""" | ||
| 22 | + # 启动时建立数据库连接 | ||
| 23 | + try: | ||
| 24 | + await db_client_connect() | ||
| 25 | + logger.info(f"✓ database connected") | ||
| 26 | + except Exception as e: | ||
| 27 | + logger.error(f"✗ Failed to connect to database: {e}") | ||
| 28 | + raise | ||
| 29 | + | ||
| 30 | + yield | ||
| 31 | + | ||
| 32 | + # 关闭时断开数据库连接 | ||
| 33 | + try: | ||
| 34 | + await db_client_disconnect() | ||
| 35 | + logger.info(f"✓ database disconnected") | ||
| 36 | + except Exception as e: | ||
| 37 | + logger.error(f"✗ Error disconnecting from database: {e}") | ||
| 38 | + | ||
| 39 | + | ||
| 40 | +# 创建应用实例 | ||
| 41 | +app = FastAPI( | ||
| 42 | + title=settings.app_name, | ||
| 43 | + version=settings.app_version, | ||
| 44 | + description="基于向量的CRUD服务，支持多种Embedding场景", | ||
| 45 | + lifespan=lifespan, | ||
| 46 | + docs_url="/docs", | ||
| 47 | + redoc_url="/redoc", | ||
| 48 | +) | ||
| 49 | +config_fastapi(app) | ||
| 50 | + | ||
| 51 | +def main(): | ||
| 52 | + """应用入口函数.""" | ||
| 53 | + import uvicorn | ||
| 54 | + | ||
| 55 | + logger.info(f'docs_url: http://127.0.0.1:{settings.port}/docs') | ||
| 56 | + uvicorn.run( | ||
| 57 | + app, | ||
| 58 | + host='0.0.0.0', | ||
| 59 | + port=settings.port, | ||
| 60 | + reload=settings.debug, | ||
| 61 | + log_config=None | ||
| 62 | + ) | ||
| 63 | + | ||
| 64 | + | ||
| 65 | +if __name__ == "__main__": | ||
| 66 | + main() |
| 1 | -from functools import lru_cache | ||
| 2 | - | ||
| 3 | -from embedding_sever.db import get_data_dao | ||
| 4 | -from .data_service import DataService | ||
| 5 | - | ||
| 6 | - | ||
| 7 | -@lru_cache() | ||
| 8 | -def get_data_service(): | ||
| 9 | - data_dao = get_data_dao() | ||
| 10 | - return DataService(data_dao) | 1 | +from functools import lru_cache |
| 2 | + | ||
| 3 | +from embedding_sever.db import get_data_dao | ||
| 4 | +from .data_service import DataService | ||
| 5 | + | ||
| 6 | + | ||
| 7 | +@lru_cache() | ||
| 8 | +def get_data_service(): | ||
| 9 | + data_dao = get_data_dao() | ||
| 10 | + return DataService(data_dao) |
| 1 | -import logging | ||
| 2 | - | ||
| 3 | -from embedding_sever.db import DataDao | ||
| 4 | - | ||
| 5 | -logger = logging.getLogger(__name__) | ||
| 6 | - | ||
| 7 | - | ||
| 8 | -class DataService: | ||
| 9 | - | ||
| 10 | - def __init__(self, data_dao: DataDao): | ||
| 11 | - self.data_dao = data_dao | ||
| 12 | - | ||
| 13 | - async def upsert(self, tb_name, id, embedding, params): | ||
| 14 | - return await self.data_dao.upsert(tb_name, id, embedding, params) | ||
| 15 | - | ||
| 16 | - async def delete_by_pks(self, tb_name, ids): | ||
| 17 | - return await self.data_dao.delete_by_pks(tb_name, ids) | ||
| 18 | - | ||
| 19 | - async def search(self, tb_name, embedding, filters, from_, size): | ||
| 20 | - return await self.data_dao.search(tb_name, embedding, filters, from_, size) | 1 | +import logging |
| 2 | + | ||
| 3 | +from embedding_sever.db import DataDao | ||
| 4 | + | ||
| 5 | +logger = logging.getLogger(__name__) | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +class DataService: | ||
| 9 | + | ||
| 10 | + def __init__(self, data_dao: DataDao): | ||
| 11 | + self.data_dao = data_dao | ||
| 12 | + | ||
| 13 | + async def upsert(self, tb_name, id, embedding, params): | ||
| 14 | + return await self.data_dao.upsert(tb_name, id, embedding, params) | ||
| 15 | + | ||
| 16 | + async def delete_by_pks(self, tb_name, ids): | ||
| 17 | + return await self.data_dao.delete_by_pks(tb_name, ids) | ||
| 18 | + | ||
| 19 | + async def search(self, tb_name, embedding, filters, from_, size): | ||
| 20 | + return await self.data_dao.search(tb_name, embedding, filters, from_, size) |
| 1 | -kafka_servers: 192.168.0.14:9092 | ||
| 2 | - | ||
| 3 | -llm: | ||
| 4 | - base_url: http://192.168.1.59:11434/v1 | ||
| 5 | - model_name: 'Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf' | ||
| 6 | - | ||
| 7 | -common: | 1 | +kafka_servers: 192.168.0.14:9092 |
| 2 | + | ||
| 3 | +llm: | ||
| 4 | + base_url: http://192.168.1.59:11434/v1 | ||
| 5 | + model_name: 'Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf' | ||
| 6 | + | ||
| 7 | +common: | ||
| 8 | cache_dir: ./cache | 8 | cache_dir: ./cache |
| 1 | -app_name: "Football Replay Service" | ||
| 2 | -app_version: "1.0.0" | ||
| 3 | - | ||
| 4 | -match_by_time_threshold: 30 # 通过时间匹配进球事件的阈值（秒） | ||
| 5 | - | ||
| 6 | -#kafka_servers: 192.168.0.14:9092 | ||
| 7 | -kafka_username: | ||
| 8 | -kafka_password: | ||
| 9 | - | ||
| 10 | -input_kafka: | ||
| 11 | - servers: ${kafka_servers} | ||
| 12 | - group_id: ai_match_service | ||
| 13 | - topic: football_replay_match_req | ||
| 14 | - username: ${kafka_username} | ||
| 15 | - password: ${kafka_password} | ||
| 16 | - | ||
| 17 | -output_kafka: | ||
| 18 | - servers: ${kafka_servers} | ||
| 19 | - topic: football_replay_match_resp | ||
| 20 | - username: ${kafka_username} | ||
| 21 | - password: ${kafka_password} | ||
| 22 | - | ||
| 23 | -llm: | ||
| 24 | - base_url: http://192.168.1.59:11434/v1 | ||
| 25 | - model_name: Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf | ||
| 26 | - temperature: 0.7 | ||
| 27 | - | ||
| 28 | -common: | 1 | +app_name: "Football Replay Service" |
| 2 | +app_version: "1.0.0" | ||
| 3 | + | ||
| 4 | +match_by_time_threshold: 30 # 通过时间匹配进球事件的阈值（秒） | ||
| 5 | + | ||
| 6 | +#kafka_servers: 192.168.0.14:9092 | ||
| 7 | +kafka_username: | ||
| 8 | +kafka_password: | ||
| 9 | + | ||
| 10 | +input_kafka: | ||
| 11 | + servers: ${kafka_servers} | ||
| 12 | + group_id: ai_match_service | ||
| 13 | + topic: football_replay_match_req | ||
| 14 | + username: ${kafka_username} | ||
| 15 | + password: ${kafka_password} | ||
| 16 | + | ||
| 17 | +output_kafka: | ||
| 18 | + servers: ${kafka_servers} | ||
| 19 | + topic: football_replay_match_resp | ||
| 20 | + username: ${kafka_username} | ||
| 21 | + password: ${kafka_password} | ||
| 22 | + | ||
| 23 | +llm: | ||
| 24 | + base_url: http://192.168.1.59:11434/v1 | ||
| 25 | + model_name: Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf | ||
| 26 | + temperature: 0.7 | ||
| 27 | + | ||
| 28 | +common: | ||
| 29 | cache_dir: ./cache | 29 | cache_dir: ./cache |
| 1 | -from aabd.base.cfg_loader import load_yaml_by_file_with_env | ||
| 2 | -from aabd.base.enhance_dict import EnhanceDict, read_prefixed_env_vars | ||
| 3 | -from pathlib import Path | ||
| 4 | -from logging import getLogger | ||
| 5 | - | ||
| 6 | -logger = getLogger(__name__) | ||
| 7 | -cfg_dir = Path(__file__).parent.absolute() / 'cfg' | ||
| 8 | -settings = EnhanceDict(load_yaml_by_file_with_env((cfg_dir / 'config.yaml').as_posix())) | ||
| 9 | -settings.update(read_prefixed_env_vars('EMB_')) | ||
| 10 | - | ||
| 11 | -logger.info(f'Settings: {settings}') | 1 | +from aabd.base.cfg_loader import load_yaml_by_file_with_env |
| 2 | +from aabd.base.enhance_dict import EnhanceDict, read_prefixed_env_vars | ||
| 3 | +from pathlib import Path | ||
| 4 | +from logging import getLogger | ||
| 5 | + | ||
| 6 | +logger = getLogger(__name__) | ||
| 7 | +cfg_dir = Path(__file__).parent.absolute() / 'cfg' | ||
| 8 | +settings = EnhanceDict(load_yaml_by_file_with_env((cfg_dir / 'config.yaml').as_posix())) | ||
| 9 | +settings.update(read_prefixed_env_vars('EMB_')) | ||
| 10 | + | ||
| 11 | +logger.info(f'Settings: {settings}') |
| 1 | -import os | ||
| 2 | -import logging | ||
| 3 | -from aabd.base.enhance_dict import value_or_default | ||
| 4 | - | ||
| 5 | -logger = logging.getLogger(__name__) | ||
| 6 | - | ||
| 7 | -from utils.football_replay_video_event_by_llm import FootballReplayVideoEvent | ||
| 8 | -from utils.football_replay_match_live import FootballReplayMatchLive | ||
| 9 | - | ||
| 10 | -class FootballReplayMatch: | ||
| 11 | - def __init__(self, settings): | ||
| 12 | - self.settings = settings | ||
| 13 | - self.match_by_time_threshold = value_or_default(settings.match_by_time_threshold, 30) * 1000 | ||
| 14 | - | ||
| 15 | - llm_base_url = value_or_default(settings.llm.base_url, None) | ||
| 16 | - model_name = value_or_default(settings.llm.model_name, None) | ||
| 17 | - temperature = value_or_default(settings.llm.temperature, 0.7) | ||
| 18 | - self.cache_dir = value_or_default(settings.common.cache_dir, None) | ||
| 19 | - save_frames_enable = value_or_default(settings.save_frames_enable, False) | ||
| 20 | - | ||
| 21 | - self.videoEventRecognition = FootballReplayVideoEvent(llm_base_url, model_name, temperature, 'no_key', self.cache_dir, save_frames_enable) | ||
| 22 | - self.videoMatchLive = FootballReplayMatchLive(llm_base_url, model_name, temperature, 'no_key', self.cache_dir, save_frames_enable) | ||
| 23 | - | ||
| 24 | - def match_by_time(self, replay, events): | ||
| 25 | - start_utc = replay.get('start_utc') | ||
| 26 | - | ||
| 27 | - events = [(start_utc - e.get('event_utc'), e) for e in events] | ||
| 28 | - events.sort(key=lambda x: x[0]) | ||
| 29 | - | ||
| 30 | - target_index = next((i for i, e in enumerate(events) if e[0] > 0), -1) # 找到第一个 event_utc 早于 start_utc 的事件 | ||
| 31 | - | ||
| 32 | - if target_index == -1: | ||
| 33 | - return None | ||
| 34 | - | ||
| 35 | - time_diff, event = events[target_index] | ||
| 36 | - if time_diff < self.match_by_time_threshold: | ||
| 37 | - return event | ||
| 38 | - else: | ||
| 39 | - return None | ||
| 40 | - | ||
| 41 | - def det_goal_replay(self, replay, task_id): | ||
| 42 | - return self.videoEventRecognition.video_event(replay, cache_dir=os.path.join(self.cache_dir, task_id, "replay")) | ||
| 43 | - | ||
| 44 | - | ||
| 45 | - def match_by_llm(self, replay, events, task_id): | ||
| 46 | - return self.videoMatchLive.match_batch(replay, events, max_parallel=2, cache_dir=os.path.join(self.cache_dir, task_id, "envents")) | ||
| 47 | - | ||
| 48 | - def replay_match_event(self, data): | ||
| 49 | - """ | ||
| 50 | - | ||
| 51 | - :param data: | ||
| 52 | - { | ||
| 53 | - id:xx, | ||
| 54 | - match_id: 比赛id， | ||
| 55 | - replay:{ | ||
| 56 | - id:, | ||
| 57 | - url:, | ||
| 58 | - start_utc:, | ||
| 59 | - end_utc:, | ||
| 60 | - } | ||
| 61 | - events:[{ | ||
| 62 | - id:xx, | ||
| 63 | - type:xx, | ||
| 64 | - url:xx, | ||
| 65 | - event_utc, | ||
| 66 | - }] | ||
| 67 | - } | ||
| 68 | - :return: | ||
| 69 | - """ | ||
| 70 | - task_id = data.get("id") | ||
| 71 | - match_id = data.get("match_id") | ||
| 72 | - | ||
| 73 | - matched_event_id = None | ||
| 74 | - | ||
| 75 | - # 该时间段附近是否有进球事件 | ||
| 76 | - replay_info = data['replay'] | ||
| 77 | - replay_id = replay_info.get('id') | ||
| 78 | - live_events = data['events'] | ||
| 79 | - goal_live_events = [e for e in live_events if str(e.get('type', '')) == '1'] | ||
| 80 | - | ||
| 81 | - matched_event = self.match_by_time(replay_info, goal_live_events) | ||
| 82 | - if matched_event is None: | ||
| 83 | - replay_det_info = self.det_goal_replay(replay_info, task_id) | ||
| 84 | - replay_event_name = replay_det_info.get('event_name', '') | ||
| 85 | - if replay_event_name == '进球': | ||
| 86 | - matched_event = self.match_by_llm(replay_info, goal_live_events, task_id) | ||
| 87 | - if matched_event is None: | ||
| 88 | - logger.info(f"{task_id}_LLM认为视频{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 是进球但是未能找到匹配的事件") | ||
| 89 | - # print(f"{task_id}_LLM认为视频{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 是进球但是未能找到匹配的事件") | ||
| 90 | - else: | ||
| 91 | - matched_event_id = matched_event.get('video_id') | ||
| 92 | - logger.info(f"{task_id}_LLM认为视频{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 对应的进球的事件视频是{matched_event_id}") | ||
| 93 | - # print(f"{task_id}_LLM认为视频{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 对应的进球的事件视频是{matched_event_id}") | ||
| 94 | - else: | ||
| 95 | - logger.info(f"{task_id}_LLM判断{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 不是进球，无需匹配") | ||
| 96 | - # print(f"{task_id}_LLM判断{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 不是进球，无需匹配") | ||
| 97 | - else: | ||
| 98 | - matched_event_id = matched_event.get('id') | ||
| 99 | - logger.info(f"{task_id}_通过时间找到匹配认为视频{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 对应的进球的事件视频是{matched_event_id}") | ||
| 100 | - # print(f"{task_id}_通过时间找到匹配认为视频{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 对应的进球的事件视频是{matched_event_id}") | ||
| 101 | - return { | ||
| 102 | - "id": task_id, | ||
| 103 | - "match_id": match_id, | ||
| 104 | - "replay_id": replay_id, | ||
| 105 | - "event_id": matched_event_id | 1 | +import os |
| 2 | +import logging | ||
| 3 | +from aabd.base.enhance_dict import value_or_default | ||
| 4 | + | ||
| 5 | +logger = logging.getLogger(__name__) | ||
| 6 | + | ||
| 7 | +from utils.football_replay_video_event_by_llm import FootballReplayVideoEvent | ||
| 8 | +from utils.football_replay_match_live import FootballReplayMatchLive | ||
| 9 | + | ||
| 10 | +class FootballReplayMatch: | ||
| 11 | + def __init__(self, settings): | ||
| 12 | + self.settings = settings | ||
| 13 | + self.match_by_time_threshold = value_or_default(settings.match_by_time_threshold, 30) * 1000 | ||
| 14 | + | ||
| 15 | + llm_base_url = value_or_default(settings.llm.base_url, None) | ||
| 16 | + model_name = value_or_default(settings.llm.model_name, None) | ||
| 17 | + temperature = value_or_default(settings.llm.temperature, 0.7) | ||
| 18 | + self.cache_dir = value_or_default(settings.common.cache_dir, None) | ||
| 19 | + save_frames_enable = value_or_default(settings.save_frames_enable, False) | ||
| 20 | + | ||
| 21 | + self.videoEventRecognition = FootballReplayVideoEvent(llm_base_url, model_name, temperature, 'no_key', self.cache_dir, save_frames_enable) | ||
| 22 | + self.videoMatchLive = FootballReplayMatchLive(llm_base_url, model_name, temperature, 'no_key', self.cache_dir, save_frames_enable) | ||
| 23 | + | ||
| 24 | + def match_by_time(self, replay, events): | ||
| 25 | + start_utc = replay.get('start_utc') | ||
| 26 | + | ||
| 27 | + events = [(start_utc - e.get('event_utc'), e) for e in events] | ||
| 28 | + events.sort(key=lambda x: x[0]) | ||
| 29 | + | ||
| 30 | + target_index = next((i for i, e in enumerate(events) if e[0] > 0), -1) # 找到第一个 event_utc 早于 start_utc 的事件 | ||
| 31 | + | ||
| 32 | + if target_index == -1: | ||
| 33 | + return None | ||
| 34 | + | ||
| 35 | + time_diff, event = events[target_index] | ||
| 36 | + if time_diff < self.match_by_time_threshold: | ||
| 37 | + return event | ||
| 38 | + else: | ||
| 39 | + return None | ||
| 40 | + | ||
| 41 | + def det_goal_replay(self, replay, task_id): | ||
| 42 | + return self.videoEventRecognition.video_event(replay, cache_dir=os.path.join(self.cache_dir, task_id, "replay")) | ||
| 43 | + | ||
| 44 | + | ||
| 45 | + def match_by_llm(self, replay, events, task_id): | ||
| 46 | + return self.videoMatchLive.match_batch(replay, events, max_parallel=2, cache_dir=os.path.join(self.cache_dir, task_id, "envents")) | ||
| 47 | + | ||
| 48 | + def replay_match_event(self, data): | ||
| 49 | + """ | ||
| 50 | + | ||
| 51 | + :param data: | ||
| 52 | + { | ||
| 53 | + id:xx, | ||
| 54 | + match_id: 比赛id， | ||
| 55 | + replay:{ | ||
| 56 | + id:, | ||
| 57 | + url:, | ||
| 58 | + start_utc:, | ||
| 59 | + end_utc:, | ||
| 60 | + } | ||
| 61 | + events:[{ | ||
| 62 | + id:xx, | ||
| 63 | + type:xx, | ||
| 64 | + url:xx, | ||
| 65 | + event_utc, | ||
| 66 | + }] | ||
| 67 | + } | ||
| 68 | + :return: | ||
| 69 | + """ | ||
| 70 | + task_id = data.get("id") | ||
| 71 | + match_id = data.get("match_id") | ||
| 72 | + | ||
| 73 | + matched_event_id = None | ||
| 74 | + | ||
| 75 | + # 该时间段附近是否有进球事件 | ||
| 76 | + replay_info = data['replay'] | ||
| 77 | + replay_id = replay_info.get('id') | ||
| 78 | + live_events = data['events'] | ||
| 79 | + goal_live_events = [e for e in live_events if str(e.get('type', '')) == '1'] | ||
| 80 | + | ||
| 81 | + matched_event = self.match_by_time(replay_info, goal_live_events) | ||
| 82 | + if matched_event is None: | ||
| 83 | + replay_det_info = self.det_goal_replay(replay_info, task_id) | ||
| 84 | + replay_event_name = replay_det_info.get('event_name', '') | ||
| 85 | + if replay_event_name == '进球': | ||
| 86 | + matched_event = self.match_by_llm(replay_info, goal_live_events, task_id) | ||
| 87 | + if matched_event is None: | ||
| 88 | + logger.info(f"{task_id}_LLM认为视频{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 是进球但是未能找到匹配的事件") | ||
| 89 | + # print(f"{task_id}_LLM认为视频{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 是进球但是未能找到匹配的事件") | ||
| 90 | + else: | ||
| 91 | + matched_event_id = matched_event.get('video_id') | ||
| 92 | + logger.info(f"{task_id}_LLM认为视频{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 对应的进球的事件视频是{matched_event_id}") | ||
| 93 | + # print(f"{task_id}_LLM认为视频{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 对应的进球的事件视频是{matched_event_id}") | ||
| 94 | + else: | ||
| 95 | + logger.info(f"{task_id}_LLM判断{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 不是进球，无需匹配") | ||
| 96 | + # print(f"{task_id}_LLM判断{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 不是进球，无需匹配") | ||
| 97 | + else: | ||
| 98 | + matched_event_id = matched_event.get('id') | ||
| 99 | + logger.info(f"{task_id}_通过时间找到匹配认为视频{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 对应的进球的事件视频是{matched_event_id}") | ||
| 100 | + # print(f"{task_id}_通过时间找到匹配认为视频{replay_info['url']}_{replay_info['start_utc']}_{replay_info['end_utc']} 对应的进球的事件视频是{matched_event_id}") | ||
| 101 | + return { | ||
| 102 | + "id": task_id, | ||
| 103 | + "match_id": match_id, | ||
| 104 | + "replay_id": replay_id, | ||
| 105 | + "event_id": matched_event_id | ||
| 106 | } | 106 | } |
| 1 | -import json | ||
| 2 | -import os.path | ||
| 3 | -from pathlib import Path | ||
| 4 | -from datetime import timedelta | ||
| 5 | - | ||
| 6 | -from langchain_core.messages import SystemMessage, HumanMessage | ||
| 7 | -from langchain_openai import ChatOpenAI | ||
| 8 | - | ||
| 9 | -try: | ||
| 10 | - from .llm_image import Video2Frame | ||
| 11 | -except: | ||
| 12 | - from llm_image import Video2Frame | ||
| 13 | - | ||
| 14 | -req_prompt = """ | ||
| 15 | -# Role | ||
| 16 | -你是一名拥有20年经验的足球视频技术分析师，擅长结合**视觉画面**与**解说音频（ASR）**进行跨镜头的事件匹配。你能通过解说员的描述锁定时间背景，并通过视觉特征确认物理细节。 | ||
| 17 | - | ||
| 18 | -# Task | ||
| 19 | -输入包含： | ||
| 20 | -1. 【回放片段】（Replay）：包含视频帧 + **对应的解说文本**。 | ||
| 21 | -2. 【直播进球片段列表】（Live Candidates）：包含视频帧 + **对应的解说文本**。 | ||
| 22 | - | ||
| 23 | -目标：在【直播片段】中找到与【回放片段】属于**同一个进球事件**的片段。如果没有任何片段匹配，返回 null。 | ||
| 24 | - | ||
| 25 | -# Analysis Workflow (多模态分析流程) | ||
| 26 | - | ||
| 27 | -请按照以下步骤进行推理： | ||
| 28 | - | ||
| 29 | -### 第一步：解说词元数据提取 (听觉线索) | ||
| 30 | -分析【回放片段】的解说文本，寻找以下关键信息： | ||
| 31 | -- **时间指代**：解说员是否提到了具体时间？（如“上半场”、“第10分钟”、“开场不久”）。 | ||
| 32 | -- **事件描述**：解说员如何描述这个进球？（如“世界波”、“点球”、“补射”、“乌龙球”）。 | ||
| 33 | -- **球员提及**：解说员念出了谁的名字？ | ||
| 34 | - | ||
| 35 | -### 第二步：视觉物理指纹提取 (视觉线索) | ||
| 36 | -忽略镜头语言（慢放、特写），提取核心物理特征： | ||
| 37 | -- **进攻与射门**：进攻方式（传中/直塞）、射门部位（头/脚）、射门位置。 | ||
| 38 | -- **球路与防守**：球的轨迹（高/低/折射）、门将扑救动作（侧扑/倒地/未动）。 | ||
| 39 | -- **庆祝**：进球后的庆祝动作（仅作为辅助验证）。 | ||
| 40 | - | ||
| 41 | -### 第三步：跨模态匹配与验证 | ||
| 42 | -遍历【直播片段】，结合视觉和听觉进行判断： | ||
| 43 | -- **视觉一致性**：直播片段的动作、轨迹、门将反应是否与回放完全吻合？ | ||
| 44 | -- **听觉一致性**： | ||
| 45 | - - 如果回放解说提到“这是第5分钟的进球”，而直播片段的时间戳是90分钟，**不要直接排除**，而是检查直播片段的解说是否也提到了“回顾第5分钟”或者直播片段的视觉内容确实是第5分钟的动作。 | ||
| 46 | - - **核心原则**：视觉物理特征 > 解说词描述 > 时间戳元数据。 | ||
| 47 | - - **特殊情况**：如果视觉特征高度相似（如同一个球员、同一个角度），但解说词明确说“这是另一个进球（例如：这是他的第二个进球，而回放说是第一个）”，则判定为不匹配。 | ||
| 48 | - | ||
| 49 | -# Logic Constraints (逻辑约束 - 必须严格遵守) | ||
| 50 | - | ||
| 51 | -- **时间单向性原则（铁律）**： | ||
| 52 | - - **回看中的比赛时间一定在比赛片段画面时间之后**。 | ||
| 53 | - - 逻辑：回放是对过去发生事件的回顾。如果【回放片段】中解说提到的比赛时间（或画面显示的比赛时钟）**早于**【直播片段】中解说提到的比赛时间（或画面时钟），则**绝对不可能**是同一个事件。 | ||
| 54 | - - *示例*：回放解说在描述“第10分钟的进球”，而直播片段明确发生在“第5分钟”，则该直播片段**一定不是**目标。 | ||
| 55 | -- **时间陷阱**：回放可能是赛后集锦。如果解说员说“让我们看看**刚才**那个球”或者“**上半场**那个球”，即使当前比赛时间是90分钟，也要去匹配对应时间段的直播片段（或视觉特征）。 | ||
| 56 | -- **同名陷阱**：如果解说提到“又是**凯恩**进球了”，不能只看凯恩，必须看**怎么进的**（头球还是点球）。 | ||
| 57 | -- **无匹配处理**：如果所有候选片段在视觉动作（如射门方式、进球位置）或关键事件逻辑上与回放明显不符，必须判定为无匹配，将 `video_id` 设为 `null`。 | ||
| 58 | - | ||
| 59 | -### 输出要求 | ||
| 60 | - | ||
| 61 | -请仅输出一个JSON格式的结果，不要输出任何分析过程。不要包含 markdown 标记（如 ```json ... ```），不要包含任何解释或额外文本。 | ||
| 62 | -格式如下： | ||
| 63 | -{ | ||
| 64 | - "replay_summary": { | ||
| 65 | - "audio_cues": "解说提到的关键信息（如：'第15分钟', '远射', '德布劳内'）", | ||
| 66 | - "visual_cues": "视觉关键特征（如：'禁区外右脚', '球挂死角', '门将飞身扑救'）" | ||
| 67 | - }, | ||
| 68 | - "reasoning": "综合分析：回放解说提到是'上半场的远射'，视觉显示'17号球员禁区外起脚'。Candidate_1 视觉是'近距离推射'，排除。Candidate_2 视觉是'禁区外远射'，且门将动作一致，虽然直播时间显示是下半场（可能是集锦回顾），但解说也提到了'回顾上半场'，确认为同一事件。", | ||
| 69 | - "video_id": "Candidate_2" | ||
| 70 | -}""" | ||
| 71 | - | ||
| 72 | - | ||
| 73 | -class FootballReplayMatchLive: | ||
| 74 | - def __init__(self, base_url: str, model: str, temperature: float = 0.0, api_key: str = 'no_key', cache_dir:str=None, save_frames_enable:bool=False): | ||
| 75 | - self.base_url = base_url | ||
| 76 | - self.model = model | ||
| 77 | - self.temperature = temperature | ||
| 78 | - self.api_key = api_key | ||
| 79 | - # self.model = ChatOllama(base_url="http://192.168.1.59:11434", model="qwen3.6:35b-a3b-q8_0", temperature=0.7, | ||
| 80 | - # keep_alive=-1, reasoning=False) | ||
| 81 | - # self.model = ChatOpenAI(base_url="http://192.168.1.59:11434/v1", model="qwen3.6:35b-a3b-q8_0", temperature=0.7, | ||
| 82 | - # api_key='no_key') | ||
| 83 | - self.model = ChatOpenAI(base_url=base_url, model=model, temperature=temperature, api_key=api_key, | ||
| 84 | - extra_body={"chat_template_kwargs": {"enable_thinking": False}}) | ||
| 85 | - self.cache_dir = cache_dir | ||
| 86 | - self.video2frame = Video2Frame(cache_dir=cache_dir, save_frames_enable=save_frames_enable) | ||
| 87 | - | ||
| 88 | - def _match_once(self, replay_video_contents: list, live_videos: list[dict], record: list = None): | ||
| 89 | - | ||
| 90 | - if len(live_videos) == 0: | ||
| 91 | - return None | ||
| 92 | - elif len(live_videos) == 1: | ||
| 93 | - live_video_path = live_videos[0].get("url", None) | ||
| 94 | - live_video_id = live_videos[0].get("video_id", os.path.basename(live_video_path)) | ||
| 95 | - asr_text = live_videos[0].get("asr_text", '') | ||
| 96 | - live = { | ||
| 97 | - "video_id": live_video_id, | ||
| 98 | - "video_path": live_video_path, | ||
| 99 | - "asr_text": asr_text, | ||
| 100 | - } | ||
| 101 | - if record is not None: | ||
| 102 | - record.append({"live": live, "llm_result": None, 'live_list': [live]}) | ||
| 103 | - return live | ||
| 104 | - user_contents = [] | ||
| 105 | - # replay_video_contents = video_contents(replay_video["video_path"], "\n【回放片段信息】\n", | ||
| 106 | - # prompt_end=f"\n回放解说内容：{replay_video['asr_text']}\n", | ||
| 107 | - # video_name=os.path.basename(replay_video["video_path"]), | ||
| 108 | - # fps=2, max_frames=999, sampling_mode="head", max_short_edge=480) | ||
| 109 | - live_videos_contents = [] | ||
| 110 | - live_map = {} | ||
| 111 | - live_records = {} | ||
| 112 | - for live_video in live_videos: | ||
| 113 | - live_video_path = live_video.get("url", None) | ||
| 114 | - live_video_id = live_video.get("video_id", os.path.basename(live_video_path)) | ||
| 115 | - event_utc = live_video.get("event_utc", None) | ||
| 116 | - asr_text = live_video.get("asr_text", '') | ||
| 117 | - live_video_contents = live_video.get("llm_contents", None) | ||
| 118 | - if live_video_contents is None: | ||
| 119 | - # event_start = event_utc - timedelta(seconds=10) if event_utc is not None else None | ||
| 120 | - event_start = event_utc - 10000 if event_utc is not None else None | ||
| 121 | - # event_end = event_utc + timedelta(seconds=5) if event_utc is not None else None | ||
| 122 | - event_end = event_utc + 5000 if event_utc is not None else None | ||
| 123 | - live_video_contents = self.video2frame.to_llm_contents( live_video_path, | ||
| 124 | - cache=self.cache_dir, | ||
| 125 | - fps=2, | ||
| 126 | - start=event_start, | ||
| 127 | - end=event_end, | ||
| 128 | - roi=None, | ||
| 129 | - max_px_area=400_000, | ||
| 130 | - prompt_start=f"### 候选片段 video_id: {live_video_id} ###", | ||
| 131 | - prompt_end=f"\n该片段解说内容: {asr_text}\n" | ||
| 132 | - ) | ||
| 133 | - live_map[live_video_id] = { | ||
| 134 | - "video_id": live_video_id, | ||
| 135 | - "video_path": live_video_path, | ||
| 136 | - "event_utc": event_utc, | ||
| 137 | - "asr_text": asr_text, | ||
| 138 | - "contents": live_video_contents, | ||
| 139 | - | ||
| 140 | - } | ||
| 141 | - live_records[live_video_id] = { | ||
| 142 | - "video_id": live_video_id, | ||
| 143 | - "video_path": live_video_path, | ||
| 144 | - "event_utc": event_utc, | ||
| 145 | - "asr_text": asr_text, | ||
| 146 | - } | ||
| 147 | - live_videos_contents.extend(live_video_contents) | ||
| 148 | - | ||
| 149 | - user_contents.extend(replay_video_contents) | ||
| 150 | - user_contents.append({'type': "text", "text": "\n【候选直播片段列表】\n"}) | ||
| 151 | - user_contents.extend(live_videos_contents) | ||
| 152 | - user_contents.append({'type': "text", "text": "\n请根据上述片段进行匹配并按要求输出结果\n"}) | ||
| 153 | - system_message = SystemMessage(content=req_prompt) | ||
| 154 | - user_message = HumanMessage(content=user_contents) | ||
| 155 | - result = self.model.invoke([system_message, user_message]).content | ||
| 156 | - try: | ||
| 157 | - result_json = json.loads(result) | ||
| 158 | - except json.JSONDecodeError: | ||
| 159 | - try: | ||
| 160 | - result_json = json.loads(result.replace("```json", "").replace("```", "")) | ||
| 161 | - except Exception as e: | ||
| 162 | - print("JSON解析失败:", result) | ||
| 163 | - raise e | ||
| 164 | - | ||
| 165 | - video_id = result_json.get("video_id", None) | ||
| 166 | - result_live = live_map.get(video_id, None) | ||
| 167 | - if record is not None: | ||
| 168 | - record.append( | ||
| 169 | - {"live": live_records.get(video_id,None), "llm_result": result_json, 'live_list': list(live_records.values())}) | ||
| 170 | - return result_live | ||
| 171 | - | ||
| 172 | - def _match_batch(self, replay_video_contents: list, live_videos: list[dict], max_parallel: int = 3, cache_path=None, | ||
| 173 | - record: list = None): | ||
| 174 | - """ | ||
| 175 | - Match a replay video with live videos to find the most likely match. | ||
| 176 | - :param max_parallel: | ||
| 177 | - :param replay_video: Path to the replay video.(video_path,asr_text) | ||
| 178 | - :param live_videos: [(video_id,video_path,asr_text)] | ||
| 179 | - :param cache_path: Path to cache the result. | ||
| 180 | - :return: JSON object containing the match result. | ||
| 181 | - """ | ||
| 182 | - if cache_path is not None and os.path.exists(cache_path): | ||
| 183 | - with open(cache_path, 'r', encoding='utf-8') as f: | ||
| 184 | - return json.loads(f.read()) | ||
| 185 | - | ||
| 186 | - # 按照max_parallel对live_videos进行分组 | ||
| 187 | - live_videos_groups = [live_videos[i::max_parallel] for i in range(max_parallel)] | ||
| 188 | - # 过滤掉空的分组 | ||
| 189 | - live_videos_groups = [g for g in live_videos_groups if g] | ||
| 190 | - # 如果group 大于1 并且最后一个group 只有一个元素，将其唯一元素放入倒数第二个group，并删除最后一个group | ||
| 191 | - if len(live_videos_groups) > 1 and len(live_videos_groups[-1]) == 1: | ||
| 192 | - live_videos_groups[-2].append(live_videos_groups[-1][0]) | ||
| 193 | - live_videos_groups.pop() | ||
| 194 | - | ||
| 195 | - if len(live_videos_groups) > 1: | ||
| 196 | - match_result = [] | ||
| 197 | - for live_videos_group in live_videos_groups: | ||
| 198 | - g_live = self._match_once(replay_video_contents, live_videos_group, record) | ||
| 199 | - if g_live is None: | ||
| 200 | - continue | ||
| 201 | - match_result.append(g_live) | ||
| 202 | - if len(match_result) == 0: | ||
| 203 | - return None | ||
| 204 | - else: | ||
| 205 | - return self._match_batch(replay_video_contents, match_result, max_parallel, cache_path, record) | ||
| 206 | - elif len(live_videos_groups) == 1: | ||
| 207 | - return self._match_once(replay_video_contents, live_videos_groups[0], record) | ||
| 208 | - else: | ||
| 209 | - return None | ||
| 210 | - | ||
| 211 | - def match_batch(self, replay_video: dict, live_videos: list[dict], max_parallel: int = 3, cache_dir=None): | ||
| 212 | - | ||
| 213 | - self.cache_dir = cache_dir | ||
| 214 | - cache_path = os.path.join(cache_dir, 'match_live.json') | ||
| 215 | - if cache_path is not None and os.path.exists(cache_path): | ||
| 216 | - try: | ||
| 217 | - with open(cache_path, 'r', encoding='utf-8') as f: | ||
| 218 | - return json.loads(f.read()).get("result", None) | ||
| 219 | - except: | ||
| 220 | - os.remove(cache_path) | ||
| 221 | - | ||
| 222 | - replay_video_path = replay_video.get("url", None) | ||
| 223 | - event_start = replay_video.get("start_utc", None) | ||
| 224 | - event_end = replay_video.get("end_utc", None) | ||
| 225 | - replay_video_contents = self.video2frame.to_llm_contents(replay_video_path, | ||
| 226 | - cache=os.path.join(os.path.dirname(cache_dir), "replay"), | ||
| 227 | - fps=2, | ||
| 228 | - start=event_start, | ||
| 229 | - end=event_end, | ||
| 230 | - roi=None, | ||
| 231 | - max_px_area=400_000, | ||
| 232 | - prompt_start="\n【回放片段信息】\n", | ||
| 233 | - prompt_end=f"\n回放解说内容：无\n" | ||
| 234 | - ) | ||
| 235 | - live_record = [] | ||
| 236 | - result = self._match_batch(replay_video_contents, live_videos, max_parallel, cache_path, live_record) | ||
| 237 | - if result is not None: | ||
| 238 | - result_no_content = { | ||
| 239 | - "video_id": result.get("video_id", None), | ||
| 240 | - "video_path": result.get("video_path", None), | ||
| 241 | - "event_utc": result.get("event_utc", None), | ||
| 242 | - "asr_text": result.get("asr_text", None), | ||
| 243 | - } | ||
| 244 | - else: | ||
| 245 | - result_no_content = None | ||
| 246 | - record = { | ||
| 247 | - "request": { | ||
| 248 | - "replay_video": replay_video, | ||
| 249 | - "live_videos": live_videos, | ||
| 250 | - "max_parallel": max_parallel, | ||
| 251 | - "cache_path": cache_path | ||
| 252 | - }, | ||
| 253 | - "result": result_no_content, | ||
| 254 | - "live_record": live_record | ||
| 255 | - } | ||
| 256 | - if cache_path is not None: | ||
| 257 | - os.makedirs(Path(cache_path).parent, exist_ok=True) | ||
| 258 | - with open(cache_path, 'w', encoding='utf-8') as f: | ||
| 259 | - f.write(json.dumps(record, ensure_ascii=False, indent=4)) | 1 | +import json |
| 2 | +import os.path | ||
| 3 | +from pathlib import Path | ||
| 4 | +from datetime import timedelta | ||
| 5 | + | ||
| 6 | +from langchain_core.messages import SystemMessage, HumanMessage | ||
| 7 | +from langchain_openai import ChatOpenAI | ||
| 8 | + | ||
| 9 | +try: | ||
| 10 | + from .llm_image import Video2Frame | ||
| 11 | +except: | ||
| 12 | + from llm_image import Video2Frame | ||
| 13 | + | ||
| 14 | +req_prompt = """ | ||
| 15 | +# Role | ||
| 16 | +你是一名拥有20年经验的足球视频技术分析师，擅长结合**视觉画面**与**解说音频（ASR）**进行跨镜头的事件匹配。你能通过解说员的描述锁定时间背景，并通过视觉特征确认物理细节。 | ||
| 17 | + | ||
| 18 | +# Task | ||
| 19 | +输入包含： | ||
| 20 | +1. 【回放片段】（Replay）：包含视频帧 + **对应的解说文本**。 | ||
| 21 | +2. 【直播进球片段列表】（Live Candidates）：包含视频帧 + **对应的解说文本**。 | ||
| 22 | + | ||
| 23 | +目标：在【直播片段】中找到与【回放片段】属于**同一个进球事件**的片段。如果没有任何片段匹配，返回 null。 | ||
| 24 | + | ||
| 25 | +# Analysis Workflow (多模态分析流程) | ||
| 26 | + | ||
| 27 | +请按照以下步骤进行推理： | ||
| 28 | + | ||
| 29 | +### 第一步：解说词元数据提取 (听觉线索) | ||
| 30 | +分析【回放片段】的解说文本，寻找以下关键信息： | ||
| 31 | +- **时间指代**：解说员是否提到了具体时间？（如“上半场”、“第10分钟”、“开场不久”）。 | ||
| 32 | +- **事件描述**：解说员如何描述这个进球？（如“世界波”、“点球”、“补射”、“乌龙球”）。 | ||
| 33 | +- **球员提及**：解说员念出了谁的名字？ | ||
| 34 | + | ||
| 35 | +### 第二步：视觉物理指纹提取 (视觉线索) | ||
| 36 | +忽略镜头语言（慢放、特写），提取核心物理特征： | ||
| 37 | +- **进攻与射门**：进攻方式（传中/直塞）、射门部位（头/脚）、射门位置。 | ||
| 38 | +- **球路与防守**：球的轨迹（高/低/折射）、门将扑救动作（侧扑/倒地/未动）。 | ||
| 39 | +- **庆祝**：进球后的庆祝动作（仅作为辅助验证）。 | ||
| 40 | + | ||
| 41 | +### 第三步：跨模态匹配与验证 | ||
| 42 | +遍历【直播片段】，结合视觉和听觉进行判断： | ||
| 43 | +- **视觉一致性**：直播片段的动作、轨迹、门将反应是否与回放完全吻合？ | ||
| 44 | +- **听觉一致性**： | ||
| 45 | + - 如果回放解说提到“这是第5分钟的进球”，而直播片段的时间戳是90分钟，**不要直接排除**，而是检查直播片段的解说是否也提到了“回顾第5分钟”或者直播片段的视觉内容确实是第5分钟的动作。 | ||
| 46 | + - **核心原则**：视觉物理特征 > 解说词描述 > 时间戳元数据。 | ||
| 47 | + - **特殊情况**：如果视觉特征高度相似（如同一个球员、同一个角度），但解说词明确说“这是另一个进球（例如：这是他的第二个进球，而回放说是第一个）”，则判定为不匹配。 | ||
| 48 | + | ||
| 49 | +# Logic Constraints (逻辑约束 - 必须严格遵守) | ||
| 50 | + | ||
| 51 | +- **时间单向性原则（铁律）**： | ||
| 52 | + - **回看中的比赛时间一定在比赛片段画面时间之后**。 | ||
| 53 | + - 逻辑：回放是对过去发生事件的回顾。如果【回放片段】中解说提到的比赛时间（或画面显示的比赛时钟）**早于**【直播片段】中解说提到的比赛时间（或画面时钟），则**绝对不可能**是同一个事件。 | ||
| 54 | + - *示例*：回放解说在描述“第10分钟的进球”，而直播片段明确发生在“第5分钟”，则该直播片段**一定不是**目标。 | ||
| 55 | +- **时间陷阱**：回放可能是赛后集锦。如果解说员说“让我们看看**刚才**那个球”或者“**上半场**那个球”，即使当前比赛时间是90分钟，也要去匹配对应时间段的直播片段（或视觉特征）。 | ||
| 56 | +- **同名陷阱**：如果解说提到“又是**凯恩**进球了”，不能只看凯恩，必须看**怎么进的**（头球还是点球）。 | ||
| 57 | +- **无匹配处理**：如果所有候选片段在视觉动作（如射门方式、进球位置）或关键事件逻辑上与回放明显不符，必须判定为无匹配，将 `video_id` 设为 `null`。 | ||
| 58 | + | ||
| 59 | +### 输出要求 | ||
| 60 | + | ||
| 61 | +请仅输出一个JSON格式的结果，不要输出任何分析过程。不要包含 markdown 标记（如 ```json ... ```），不要包含任何解释或额外文本。 | ||
| 62 | +格式如下： | ||
| 63 | +{ | ||
| 64 | + "replay_summary": { | ||
| 65 | + "audio_cues": "解说提到的关键信息（如：'第15分钟', '远射', '德布劳内'）", | ||
| 66 | + "visual_cues": "视觉关键特征（如：'禁区外右脚', '球挂死角', '门将飞身扑救'）" | ||
| 67 | + }, | ||
| 68 | + "reasoning": "综合分析：回放解说提到是'上半场的远射'，视觉显示'17号球员禁区外起脚'。Candidate_1 视觉是'近距离推射'，排除。Candidate_2 视觉是'禁区外远射'，且门将动作一致，虽然直播时间显示是下半场（可能是集锦回顾），但解说也提到了'回顾上半场'，确认为同一事件。", | ||
| 69 | + "video_id": "Candidate_2" | ||
| 70 | +}""" | ||
| 71 | + | ||
| 72 | + | ||
| 73 | +class FootballReplayMatchLive: | ||
| 74 | + def __init__(self, base_url: str, model: str, temperature: float = 0.0, api_key: str = 'no_key', cache_dir:str=None, save_frames_enable:bool=False): | ||
| 75 | + self.base_url = base_url | ||
| 76 | + self.model = model | ||
| 77 | + self.temperature = temperature | ||
| 78 | + self.api_key = api_key | ||
| 79 | + # self.model = ChatOllama(base_url="http://192.168.1.59:11434", model="qwen3.6:35b-a3b-q8_0", temperature=0.7, | ||
| 80 | + # keep_alive=-1, reasoning=False) | ||
| 81 | + # self.model = ChatOpenAI(base_url="http://192.168.1.59:11434/v1", model="qwen3.6:35b-a3b-q8_0", temperature=0.7, | ||
| 82 | + # api_key='no_key') | ||
| 83 | + self.model = ChatOpenAI(base_url=base_url, model=model, temperature=temperature, api_key=api_key, | ||
| 84 | + extra_body={"chat_template_kwargs": {"enable_thinking": False}}) | ||
| 85 | + self.cache_dir = cache_dir | ||
| 86 | + self.video2frame = Video2Frame(cache_dir=cache_dir, save_frames_enable=save_frames_enable) | ||
| 87 | + | ||
| 88 | + def _match_once(self, replay_video_contents: list, live_videos: list[dict], record: list = None): | ||
| 89 | + | ||
| 90 | + if len(live_videos) == 0: | ||
| 91 | + return None | ||
| 92 | + elif len(live_videos) == 1: | ||
| 93 | + live_video_path = live_videos[0].get("url", None) | ||
| 94 | + live_video_id = live_videos[0].get("video_id", os.path.basename(live_video_path)) | ||
| 95 | + asr_text = live_videos[0].get("asr_text", '') | ||
| 96 | + live = { | ||
| 97 | + "video_id": live_video_id, | ||
| 98 | + "video_path": live_video_path, | ||
| 99 | + "asr_text": asr_text, | ||
| 100 | + } | ||
| 101 | + if record is not None: | ||
| 102 | + record.append({"live": live, "llm_result": None, 'live_list': [live]}) | ||
| 103 | + return live | ||
| 104 | + user_contents = [] | ||
| 105 | + # replay_video_contents = video_contents(replay_video["video_path"], "\n【回放片段信息】\n", | ||
| 106 | + # prompt_end=f"\n回放解说内容：{replay_video['asr_text']}\n", | ||
| 107 | + # video_name=os.path.basename(replay_video["video_path"]), | ||
| 108 | + # fps=2, max_frames=999, sampling_mode="head", max_short_edge=480) | ||
| 109 | + live_videos_contents = [] | ||
| 110 | + live_map = {} | ||
| 111 | + live_records = {} | ||
| 112 | + for live_video in live_videos: | ||
| 113 | + live_video_path = live_video.get("url", None) | ||
| 114 | + live_video_id = live_video.get("video_id", os.path.basename(live_video_path)) | ||
| 115 | + event_utc = live_video.get("event_utc", None) | ||
| 116 | + asr_text = live_video.get("asr_text", '') | ||
| 117 | + live_video_contents = live_video.get("llm_contents", None) | ||
| 118 | + if live_video_contents is None: | ||
| 119 | + # event_start = event_utc - timedelta(seconds=10) if event_utc is not None else None | ||
| 120 | + event_start = event_utc - 10000 if event_utc is not None else None | ||
| 121 | + # event_end = event_utc + timedelta(seconds=5) if event_utc is not None else None | ||
| 122 | + event_end = event_utc + 5000 if event_utc is not None else None | ||
| 123 | + live_video_contents = self.video2frame.to_llm_contents( live_video_path, | ||
| 124 | + cache=self.cache_dir, | ||
| 125 | + fps=2, | ||
| 126 | + start=event_start, | ||
| 127 | + end=event_end, | ||
| 128 | + roi=None, | ||
| 129 | + max_px_area=400_000, | ||
| 130 | + prompt_start=f"### 候选片段 video_id: {live_video_id} ###", | ||
| 131 | + prompt_end=f"\n该片段解说内容: {asr_text}\n" | ||
| 132 | + ) | ||
| 133 | + live_map[live_video_id] = { | ||
| 134 | + "video_id": live_video_id, | ||
| 135 | + "video_path": live_video_path, | ||
| 136 | + "event_utc": event_utc, | ||
| 137 | + "asr_text": asr_text, | ||
| 138 | + "contents": live_video_contents, | ||
| 139 | + | ||
| 140 | + } | ||
| 141 | + live_records[live_video_id] = { | ||
| 142 | + "video_id": live_video_id, | ||
| 143 | + "video_path": live_video_path, | ||
| 144 | + "event_utc": event_utc, | ||
| 145 | + "asr_text": asr_text, | ||
| 146 | + } | ||
| 147 | + live_videos_contents.extend(live_video_contents) | ||
| 148 | + | ||
| 149 | + user_contents.extend(replay_video_contents) | ||
| 150 | + user_contents.append({'type': "text", "text": "\n【候选直播片段列表】\n"}) | ||
| 151 | + user_contents.extend(live_videos_contents) | ||
| 152 | + user_contents.append({'type': "text", "text": "\n请根据上述片段进行匹配并按要求输出结果\n"}) | ||
| 153 | + system_message = SystemMessage(content=req_prompt) | ||
| 154 | + user_message = HumanMessage(content=user_contents) | ||
| 155 | + result = self.model.invoke([system_message, user_message]).content | ||
| 156 | + try: | ||
| 157 | + result_json = json.loads(result) | ||
| 158 | + except json.JSONDecodeError: | ||
| 159 | + try: | ||
| 160 | + result_json = json.loads(result.replace("```json", "").replace("```", "")) | ||
| 161 | + except Exception as e: | ||
| 162 | + print("JSON解析失败:", result) | ||
| 163 | + raise e | ||
| 164 | + | ||
| 165 | + video_id = result_json.get("video_id", None) | ||
| 166 | + result_live = live_map.get(video_id, None) | ||
| 167 | + if record is not None: | ||
| 168 | + record.append( | ||
| 169 | + {"live": live_records.get(video_id,None), "llm_result": result_json, 'live_list': list(live_records.values())}) | ||
| 170 | + return result_live | ||
| 171 | + | ||
| 172 | + def _match_batch(self, replay_video_contents: list, live_videos: list[dict], max_parallel: int = 3, cache_path=None, | ||
| 173 | + record: list = None): | ||
| 174 | + """ | ||
| 175 | + Match a replay video with live videos to find the most likely match. | ||
| 176 | + :param max_parallel: | ||
| 177 | + :param replay_video: Path to the replay video.(video_path,asr_text) | ||
| 178 | + :param live_videos: [(video_id,video_path,asr_text)] | ||
| 179 | + :param cache_path: Path to cache the result. | ||
| 180 | + :return: JSON object containing the match result. | ||
| 181 | + """ | ||
| 182 | + if cache_path is not None and os.path.exists(cache_path): | ||
| 183 | + with open(cache_path, 'r', encoding='utf-8') as f: | ||
| 184 | + return json.loads(f.read()) | ||
| 185 | + | ||
| 186 | + # 按照max_parallel对live_videos进行分组 | ||
| 187 | + live_videos_groups = [live_videos[i::max_parallel] for i in range(max_parallel)] | ||
| 188 | + # 过滤掉空的分组 | ||
| 189 | + live_videos_groups = [g for g in live_videos_groups if g] | ||
| 190 | + # 如果group 大于1 并且最后一个group 只有一个元素，将其唯一元素放入倒数第二个group，并删除最后一个group | ||
| 191 | + if len(live_videos_groups) > 1 and len(live_videos_groups[-1]) == 1: | ||
| 192 | + live_videos_groups[-2].append(live_videos_groups[-1][0]) | ||
| 193 | + live_videos_groups.pop() | ||
| 194 | + | ||
| 195 | + if len(live_videos_groups) > 1: | ||
| 196 | + match_result = [] | ||
| 197 | + for live_videos_group in live_videos_groups: | ||
| 198 | + g_live = self._match_once(replay_video_contents, live_videos_group, record) | ||
| 199 | + if g_live is None: | ||
| 200 | + continue | ||
| 201 | + match_result.append(g_live) | ||
| 202 | + if len(match_result) == 0: | ||
| 203 | + return None | ||
| 204 | + else: | ||
| 205 | + return self._match_batch(replay_video_contents, match_result, max_parallel, cache_path, record) | ||
| 206 | + elif len(live_videos_groups) == 1: | ||
| 207 | + return self._match_once(replay_video_contents, live_videos_groups[0], record) | ||
| 208 | + else: | ||
| 209 | + return None | ||
| 210 | + | ||
| 211 | + def match_batch(self, replay_video: dict, live_videos: list[dict], max_parallel: int = 3, cache_dir=None): | ||
| 212 | + | ||
| 213 | + self.cache_dir = cache_dir | ||
| 214 | + cache_path = os.path.join(cache_dir, 'match_live.json') | ||
| 215 | + if cache_path is not None and os.path.exists(cache_path): | ||
| 216 | + try: | ||
| 217 | + with open(cache_path, 'r', encoding='utf-8') as f: | ||
| 218 | + return json.loads(f.read()).get("result", None) | ||
| 219 | + except: | ||
| 220 | + os.remove(cache_path) | ||
| 221 | + | ||
| 222 | + replay_video_path = replay_video.get("url", None) | ||
| 223 | + event_start = replay_video.get("start_utc", None) | ||
| 224 | + event_end = replay_video.get("end_utc", None) | ||
| 225 | + replay_video_contents = self.video2frame.to_llm_contents(replay_video_path, | ||
| 226 | + cache=os.path.join(os.path.dirname(cache_dir), "replay"), | ||
| 227 | + fps=2, | ||
| 228 | + start=event_start, | ||
| 229 | + end=event_end, | ||
| 230 | + roi=None, | ||
| 231 | + max_px_area=400_000, | ||
| 232 | + prompt_start="\n【回放片段信息】\n", | ||
| 233 | + prompt_end=f"\n回放解说内容：无\n" | ||
| 234 | + ) | ||
| 235 | + live_record = [] | ||
| 236 | + result = self._match_batch(replay_video_contents, live_videos, max_parallel, cache_path, live_record) | ||
| 237 | + if result is not None: | ||
| 238 | + result_no_content = { | ||
| 239 | + "video_id": result.get("video_id", None), | ||
| 240 | + "video_path": result.get("video_path", None), | ||
| 241 | + "event_utc": result.get("event_utc", None), | ||
| 242 | + "asr_text": result.get("asr_text", None), | ||
| 243 | + } | ||
| 244 | + else: | ||
| 245 | + result_no_content = None | ||
| 246 | + record = { | ||
| 247 | + "request": { | ||
| 248 | + "replay_video": replay_video, | ||
| 249 | + "live_videos": live_videos, | ||
| 250 | + "max_parallel": max_parallel, | ||
| 251 | + "cache_path": cache_path | ||
| 252 | + }, | ||
| 253 | + "result": result_no_content, | ||
| 254 | + "live_record": live_record | ||
| 255 | + } | ||
| 256 | + if cache_path is not None: | ||
| 257 | + os.makedirs(Path(cache_path).parent, exist_ok=True) | ||
| 258 | + with open(cache_path, 'w', encoding='utf-8') as f: | ||
| 259 | + f.write(json.dumps(record, ensure_ascii=False, indent=4)) | ||
| 260 | return result_no_content | 260 | return result_no_content |
| 1 | -import json | ||
| 2 | -import os.path | ||
| 3 | -from pathlib import Path | ||
| 4 | - | ||
| 5 | -from langchain_core.messages import SystemMessage, HumanMessage | ||
| 6 | -from langchain_openai import ChatOpenAI | ||
| 7 | - | ||
| 8 | -try: | ||
| 9 | - from .llm_image import Video2Frame | ||
| 10 | -except: | ||
| 11 | - from llm_image import Video2Frame | ||
| 12 | - | ||
| 13 | - | ||
| 14 | -req_prompt = """ | ||
| 15 | -### 角色设定 | ||
| 16 | -你是一位拥有20年经验的足球赛事视频分析专家。你的任务是结合**视频画面**和**解说音频文本**，精准判断视频片段中是否发生了"有效进球"。 | ||
| 17 | - | ||
| 18 | -### 输入内容 | ||
| 19 | - | ||
| 20 | -1. **视频片段**：可能包含赛场、观众、教练、回放、演播室、宣传片、广告等多种画面。 | ||
| 21 | -2. **解说文本**：该片段对应的实时解说内容（ASR）。 | ||
| 22 | - | ||
| 23 | -### 分析逻辑（思维链） | ||
| 24 | - | ||
| 25 | -请综合以下三个维度的信息进行推理： | ||
| 26 | - | ||
| 27 | -#### 1. 场景维度（非比赛画面过滤 - 最高优先级） | ||
| 28 | - | ||
| 29 | -- **非比赛内容识别**：检查画面是否为**宣传片**、**商业广告**、**纯演播室解说**、**集锦混剪**（无连续比赛画面）或* | ||
| 30 | - *静态图文**。 | ||
| 31 | -- **处理规则**：如果视频内容主要是上述非比赛画面，且没有包含明确的实时进球片段，**直接判定为“无进球”** | ||
| 32 | - ，无需进行后续的进球逻辑分析。 | ||
| 33 | - | ||
| 34 | -#### 2. 视觉维度（寻找关键证据） | ||
| 35 | - | ||
| 36 | -- **核心画面**：足球入网、球在网内静止、裁判指中圈。 | ||
| 37 | -- **行为线索**：进攻方疯狂庆祝、防守方抱头懊恼、全场观众起立欢呼。 | ||
| 38 | -- **画面容忍度**：即使画面主要是观众或教练特写，只要上下文（如切入的进球回放）或行为暗示了进球，也应视为进球。 | ||
| 39 | - | ||
| 40 | -#### 3. 听觉维度（解说语义分析） | ||
| 41 | - | ||
| 42 | -- **进球关键词**：寻找如“球进啦”、“Goal”、“得分”、“世界波”、“破门”、“1比0”等肯定性词汇。 | ||
| 43 | -- **情绪语调**：解说员音量突然升高、语速加快、情绪激动（通常伴随进球发生）。 | ||
| 44 | -- **否定排除**：如果解说提到“越位在先”、“进球无效”、“击中横梁”、“偏出”，则视为未进球。 | ||
| 45 | - | ||
| 46 | -### 判定规则 | ||
| 47 | - | ||
| 48 | -- **判定为“无进球”**： | ||
| 49 | - - **画面为宣传片、广告、纯演播室或其他非比赛实时内容；** | ||
| 50 | - - 视觉显示未进（偏出/被扑/中柱）； | ||
| 51 | - - 解说明确表示未进或进球无效； | ||
| 52 | - - 画面与解说均无进球迹象（如普通传球、界外球）。 | ||
| 53 | -- **判定为“进球”**： | ||
| 54 | - - 画面为比赛内容，且视觉清晰显示进球； | ||
| 55 | - - **或** 画面为比赛内容，视觉模糊但解说员明确喊出“球进了”且情绪激动； | ||
| 56 | - - **或** 画面显示庆祝/回放，配合解说确认进球。 | ||
| 57 | - | ||
| 58 | -### 输出要求 | ||
| 59 | - | ||
| 60 | -请仅输出一个JSON格式的结果，不要输出任何分析过程。不要包含 markdown 标记（如 ```json ... ```），不要包含任何解释或额外文本。 | ||
| 61 | -格式如下： | ||
| 62 | -{ | ||
| 63 | -"event_name": "进球" 或 "无进球", | ||
| 64 | -"description": "判定理由，若是宣传片请直接注明" | ||
| 65 | -}""" | ||
| 66 | - | ||
| 67 | - | ||
| 68 | -class FootballReplayVideoEvent: | ||
| 69 | - def __init__(self, base_url: str, model: str, temperature: float = 0.0, api_key: str = 'no_key', cache_dir:str=None, save_frames_enable:bool=False): | ||
| 70 | - self.base_url = base_url | ||
| 71 | - self.model = model | ||
| 72 | - self.temperature = temperature | ||
| 73 | - self.api_key = api_key | ||
| 74 | - # self.model = ChatOllama(base_url="http://192.168.1.59:11434", model="qwen3.6:35b-a3b-q8_0", temperature=0.7, | ||
| 75 | - # keep_alive=-1, reasoning=False) | ||
| 76 | - # self.model = ChatOpenAI(base_url="http://192.168.1.59:11434/v1", model="qwen3.6:35b-a3b-q8_0", temperature=0.7, | ||
| 77 | - # api_key='no_key') | ||
| 78 | - self.model = ChatOpenAI(base_url=base_url, model=model, temperature=temperature, api_key=api_key, | ||
| 79 | - extra_body={"chat_template_kwargs": {"enable_thinking": False}}) | ||
| 80 | - | ||
| 81 | - | ||
| 82 | - self.video2frame = Video2Frame(cache_dir=cache_dir, save_frames_enable=save_frames_enable) | ||
| 83 | - | ||
| 84 | - def video_event(self, replay_pack: dict, asr_text: str = '无', cache_dir=None): | ||
| 85 | - replay_video_path = replay_pack.get("url", None) | ||
| 86 | - | ||
| 87 | - cache_path = os.path.join(cache_dir, 'video_event.json') | ||
| 88 | - if cache_path is not None and os.path.exists(cache_path): | ||
| 89 | - with open(cache_path, 'r', encoding='utf-8') as f: | ||
| 90 | - return json.loads(f.read()) | ||
| 91 | - | ||
| 92 | - event_start = replay_pack.get("start_utc", None) | ||
| 93 | - event_end = replay_pack.get("end_utc", None) | ||
| 94 | - contents = self.video2frame.to_llm_contents(replay_video_path, | ||
| 95 | - cache=cache_dir, | ||
| 96 | - fps=2, | ||
| 97 | - start=event_start, | ||
| 98 | - end=event_end, | ||
| 99 | - roi=None, | ||
| 100 | - max_px_area=400_000, | ||
| 101 | - prompt_start="\n【回放片段信息】\n", | ||
| 102 | - prompt_end=f"\n回放解说内容：{asr_text}\n" | ||
| 103 | - ) | ||
| 104 | - | ||
| 105 | - system_message = SystemMessage(content=req_prompt) | ||
| 106 | - video_message = HumanMessage(content=contents) | ||
| 107 | - asr_message = HumanMessage(content=f"解说内容：{asr_text}") | ||
| 108 | - result = self.model.invoke([system_message, video_message, asr_message]).content | ||
| 109 | - try: | ||
| 110 | - result_json = json.loads(result) | ||
| 111 | - except json.JSONDecodeError: | ||
| 112 | - result_json = json.loads(result.replace("```json", "").replace("```", "")) | ||
| 113 | - if cache_path is not None: | ||
| 114 | - os.makedirs(Path(cache_path).parent, exist_ok=True) | ||
| 115 | - with open(cache_path, 'w', encoding='utf-8') as f: | ||
| 116 | - f.write(result) | ||
| 117 | - return result_json | ||
| 118 | - | ||
| 119 | - | ||
| 120 | -if __name__ == "__main__": | ||
| 121 | - from aabd.base.patched_logging import init_logging | ||
| 122 | - import os | ||
| 123 | - os.environ["APP_LOG_TYPE"] = "console" | ||
| 124 | - init_logging() | ||
| 125 | - fbrv = FootballReplayVideoEvent(base_url="http://192.168.1.59:11434/v1", | ||
| 126 | - model="Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf", | ||
| 127 | - temperature=0.7, | ||
| 128 | - cache_dir="/root/lzw/tmp_0518_replay_cache", | ||
| 129 | - save_frames_enable=True | ||
| 130 | - ) | ||
| 131 | - replay_pack = {"url": "http://video.mam.miguvideo.com/mnt6/fastclip3/wsc/2026/04/29/df760b8d15394cbc9ed0d0a4a9c4b786_1080PS/29133605/vodtmp/3b3e99cf4ca84c3782503d8817242de2.m3u8", | ||
| 132 | - "start_utc": 0, | ||
| 133 | - "end_utc": 999999999, | ||
| 134 | - "asr_text": '无' | ||
| 135 | - } | ||
| 136 | - | ||
| 137 | - replay_goals = fbrv.video_event(replay_pack, cache_path="/root/lzw/tmp_0518_replay_cache") | 1 | +import json |
| 2 | +import os.path | ||
| 3 | +from pathlib import Path | ||
| 4 | + | ||
| 5 | +from langchain_core.messages import SystemMessage, HumanMessage | ||
| 6 | +from langchain_openai import ChatOpenAI | ||
| 7 | + | ||
| 8 | +try: | ||
| 9 | + from .llm_image import Video2Frame | ||
| 10 | +except: | ||
| 11 | + from llm_image import Video2Frame | ||
| 12 | + | ||
| 13 | + | ||
| 14 | +req_prompt = """ | ||
| 15 | +### 角色设定 | ||
| 16 | +你是一位拥有20年经验的足球赛事视频分析专家。你的任务是结合**视频画面**和**解说音频文本**，精准判断视频片段中是否发生了"有效进球"。 | ||
| 17 | + | ||
| 18 | +### 输入内容 | ||
| 19 | + | ||
| 20 | +1. **视频片段**：可能包含赛场、观众、教练、回放、演播室、宣传片、广告等多种画面。 | ||
| 21 | +2. **解说文本**：该片段对应的实时解说内容（ASR）。 | ||
| 22 | + | ||
| 23 | +### 分析逻辑（思维链） | ||
| 24 | + | ||
| 25 | +请综合以下三个维度的信息进行推理： | ||
| 26 | + | ||
| 27 | +#### 1. 场景维度（非比赛画面过滤 - 最高优先级） | ||
| 28 | + | ||
| 29 | +- **非比赛内容识别**：检查画面是否为**宣传片**、**商业广告**、**纯演播室解说**、**集锦混剪**（无连续比赛画面）或* | ||
| 30 | + *静态图文**。 | ||
| 31 | +- **处理规则**：如果视频内容主要是上述非比赛画面，且没有包含明确的实时进球片段，**直接判定为“无进球”** | ||
| 32 | + ，无需进行后续的进球逻辑分析。 | ||
| 33 | + | ||
| 34 | +#### 2. 视觉维度（寻找关键证据） | ||
| 35 | + | ||
| 36 | +- **核心画面**：足球入网、球在网内静止、裁判指中圈。 | ||
| 37 | +- **行为线索**：进攻方疯狂庆祝、防守方抱头懊恼、全场观众起立欢呼。 | ||
| 38 | +- **画面容忍度**：即使画面主要是观众或教练特写，只要上下文（如切入的进球回放）或行为暗示了进球，也应视为进球。 | ||
| 39 | + | ||
| 40 | +#### 3. 听觉维度（解说语义分析） | ||
| 41 | + | ||
| 42 | +- **进球关键词**：寻找如“球进啦”、“Goal”、“得分”、“世界波”、“破门”、“1比0”等肯定性词汇。 | ||
| 43 | +- **情绪语调**：解说员音量突然升高、语速加快、情绪激动（通常伴随进球发生）。 | ||
| 44 | +- **否定排除**：如果解说提到“越位在先”、“进球无效”、“击中横梁”、“偏出”，则视为未进球。 | ||
| 45 | + | ||
| 46 | +### 判定规则 | ||
| 47 | + | ||
| 48 | +- **判定为“无进球”**： | ||
| 49 | + - **画面为宣传片、广告、纯演播室或其他非比赛实时内容；** | ||
| 50 | + - 视觉显示未进（偏出/被扑/中柱）； | ||
| 51 | + - 解说明确表示未进或进球无效； | ||
| 52 | + - 画面与解说均无进球迹象（如普通传球、界外球）。 | ||
| 53 | +- **判定为“进球”**： | ||
| 54 | + - 画面为比赛内容，且视觉清晰显示进球； | ||
| 55 | + - **或** 画面为比赛内容，视觉模糊但解说员明确喊出“球进了”且情绪激动； | ||
| 56 | + - **或** 画面显示庆祝/回放，配合解说确认进球。 | ||
| 57 | + | ||
| 58 | +### 输出要求 | ||
| 59 | + | ||
| 60 | +请仅输出一个JSON格式的结果，不要输出任何分析过程。不要包含 markdown 标记（如 ```json ... ```），不要包含任何解释或额外文本。 | ||
| 61 | +格式如下： | ||
| 62 | +{ | ||
| 63 | +"event_name": "进球" 或 "无进球", | ||
| 64 | +"description": "判定理由，若是宣传片请直接注明" | ||
| 65 | +}""" | ||
| 66 | + | ||
| 67 | + | ||
| 68 | +class FootballReplayVideoEvent: | ||
| 69 | + def __init__(self, base_url: str, model: str, temperature: float = 0.0, api_key: str = 'no_key', cache_dir:str=None, save_frames_enable:bool=False): | ||
| 70 | + self.base_url = base_url | ||
| 71 | + self.model = model | ||
| 72 | + self.temperature = temperature | ||
| 73 | + self.api_key = api_key | ||
| 74 | + # self.model = ChatOllama(base_url="http://192.168.1.59:11434", model="qwen3.6:35b-a3b-q8_0", temperature=0.7, | ||
| 75 | + # keep_alive=-1, reasoning=False) | ||
| 76 | + # self.model = ChatOpenAI(base_url="http://192.168.1.59:11434/v1", model="qwen3.6:35b-a3b-q8_0", temperature=0.7, | ||
| 77 | + # api_key='no_key') | ||
| 78 | + self.model = ChatOpenAI(base_url=base_url, model=model, temperature=temperature, api_key=api_key, | ||
| 79 | + extra_body={"chat_template_kwargs": {"enable_thinking": False}}) | ||
| 80 | + | ||
| 81 | + | ||
| 82 | + self.video2frame = Video2Frame(cache_dir=cache_dir, save_frames_enable=save_frames_enable) | ||
| 83 | + | ||
| 84 | + def video_event(self, replay_pack: dict, asr_text: str = '无', cache_dir=None): | ||
| 85 | + replay_video_path = replay_pack.get("url", None) | ||
| 86 | + | ||
| 87 | + cache_path = os.path.join(cache_dir, 'video_event.json') | ||
| 88 | + if cache_path is not None and os.path.exists(cache_path): | ||
| 89 | + with open(cache_path, 'r', encoding='utf-8') as f: | ||
| 90 | + return json.loads(f.read()) | ||
| 91 | + | ||
| 92 | + event_start = replay_pack.get("start_utc", None) | ||
| 93 | + event_end = replay_pack.get("end_utc", None) | ||
| 94 | + contents = self.video2frame.to_llm_contents(replay_video_path, | ||
| 95 | + cache=cache_dir, | ||
| 96 | + fps=2, | ||
| 97 | + start=event_start, | ||
| 98 | + end=event_end, | ||
| 99 | + roi=None, | ||
| 100 | + max_px_area=400_000, | ||
| 101 | + prompt_start="\n【回放片段信息】\n", | ||
| 102 | + prompt_end=f"\n回放解说内容：{asr_text}\n" | ||
| 103 | + ) | ||
| 104 | + | ||
| 105 | + system_message = SystemMessage(content=req_prompt) | ||
| 106 | + video_message = HumanMessage(content=contents) | ||
| 107 | + asr_message = HumanMessage(content=f"解说内容：{asr_text}") | ||
| 108 | + result = self.model.invoke([system_message, video_message, asr_message]).content | ||
| 109 | + try: | ||
| 110 | + result_json = json.loads(result) | ||
| 111 | + except json.JSONDecodeError: | ||
| 112 | + result_json = json.loads(result.replace("```json", "").replace("```", "")) | ||
| 113 | + if cache_path is not None: | ||
| 114 | + os.makedirs(Path(cache_path).parent, exist_ok=True) | ||
| 115 | + with open(cache_path, 'w', encoding='utf-8') as f: | ||
| 116 | + f.write(result) | ||
| 117 | + return result_json | ||
| 118 | + | ||
| 119 | + | ||
| 120 | +if __name__ == "__main__": | ||
| 121 | + from aabd.base.patched_logging import init_logging | ||
| 122 | + import os | ||
| 123 | + os.environ["APP_LOG_TYPE"] = "console" | ||
| 124 | + init_logging() | ||
| 125 | + fbrv = FootballReplayVideoEvent(base_url="http://192.168.1.59:11434/v1", | ||
| 126 | + model="Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf", | ||
| 127 | + temperature=0.7, | ||
| 128 | + cache_dir="/root/lzw/tmp_0518_replay_cache", | ||
| 129 | + save_frames_enable=True | ||
| 130 | + ) | ||
| 131 | + replay_pack = {"url": "http://video.mam.miguvideo.com/mnt6/fastclip3/wsc/2026/04/29/df760b8d15394cbc9ed0d0a4a9c4b786_1080PS/29133605/vodtmp/3b3e99cf4ca84c3782503d8817242de2.m3u8", | ||
| 132 | + "start_utc": 0, | ||
| 133 | + "end_utc": 999999999, | ||
| 134 | + "asr_text": '无' | ||
| 135 | + } | ||
| 136 | + | ||
| 137 | + replay_goals = fbrv.video_event(replay_pack, cache_path="/root/lzw/tmp_0518_replay_cache") | ||
| 138 | print(replay_goals) | 138 | print(replay_goals) |
| 1 | -import base64 | ||
| 2 | -import shutil | ||
| 3 | -import subprocess | ||
| 4 | -import sys | ||
| 5 | -from pathlib import Path | ||
| 6 | -from typing import Optional | ||
| 7 | -from urllib.parse import urlparse | ||
| 8 | -import cv2 | ||
| 9 | -from aabd.stream_chain import video2frames | ||
| 10 | -import json | ||
| 11 | -import logging | ||
| 12 | -import hashlib | ||
| 13 | - | ||
| 14 | -logger = logging.getLogger(__name__) | ||
| 15 | - | ||
| 16 | -_IS_ABSOLUTE_UTC_THRESHOLD = 1_000_000_000 | ||
| 17 | - | ||
| 18 | - | ||
| 19 | -def download_to_mp4(url: str, output_path: str, duration: Optional[int] = None) -> None: | ||
| 20 | - cmd = ["ffmpeg", "-y", "-fflags", "+discardcorrupt", "-i", url] | ||
| 21 | - if duration is not None: | ||
| 22 | - cmd.extend(["-t", str(duration)]) | ||
| 23 | - cmd.extend(["-c", "copy", "-bsf:a", "aac_adtstoasc", "-movflags", "+faststart", output_path]) | ||
| 24 | - result = subprocess.run(cmd, capture_output=True, text=True, encoding="utf-8", errors="replace") | ||
| 25 | - if result.returncode != 0: | ||
| 26 | - raise RuntimeError(f"ffmpeg 下载失败: {result.stderr}") | ||
| 27 | - | ||
| 28 | - | ||
| 29 | -def resize_frame(frame, max_px_area): | ||
| 30 | - if max_px_area is None: | ||
| 31 | - return frame | ||
| 32 | - h, w = frame.shape[:2] | ||
| 33 | - area = h * w | ||
| 34 | - if area > max_px_area: | ||
| 35 | - scale = (max_px_area / area) ** 0.5 | ||
| 36 | - new_w = int(w * scale) | ||
| 37 | - new_h = int(h * scale) | ||
| 38 | - frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA) | ||
| 39 | - return frame | ||
| 40 | - | ||
| 41 | - | ||
| 42 | -class Video2Frame: | ||
| 43 | - def __init__(self, cache_dir, save_frames_enable=False): | ||
| 44 | - self.cache_dir = Path(cache_dir) | ||
| 45 | - self.save_frames_enable = save_frames_enable | ||
| 46 | - | ||
| 47 | - def get_root_path(self, url, cache=None): | ||
| 48 | - if cache is None: | ||
| 49 | - parsed_url = urlparse(url) | ||
| 50 | - raw = f"{parsed_url.scheme}/{parsed_url.netloc}{parsed_url.path}" | ||
| 51 | - cache = hashlib.md5(raw.encode('utf-8')).hexdigest() | ||
| 52 | - return self.cache_dir / cache | ||
| 53 | - | ||
| 54 | - def iter_mp4(self, file, start, end, fps): | ||
| 55 | - | ||
| 56 | - if (start is not None and start > _IS_ABSOLUTE_UTC_THRESHOLD) or ( | ||
| 57 | - end is not None and end > _IS_ABSOLUTE_UTC_THRESHOLD): | ||
| 58 | - video_start_time = 0 | ||
| 59 | - with video2frames.AVAnyKeyStreamDecoder(file, sei_enable=True, tqdm_enable=False) as (sd, _): | ||
| 60 | - for frame in sd: | ||
| 61 | - if start is not None: | ||
| 62 | - video_start_time = start - frame.get('src_frame_time') | ||
| 63 | - if end is not None: | ||
| 64 | - video_end_time = end - frame.get('src_frame_time') or 0 | ||
| 65 | - break | ||
| 66 | - else: | ||
| 67 | - video_start_time = start if start is not None else 0 | ||
| 68 | - video_end_time = end if end is not None else sys.maxsize | ||
| 69 | - | ||
| 70 | - with video2frames.AVAnyKeyStreamDecoder(file, start=video_start_time, end=video_end_time, control_type='time', | ||
| 71 | - sei_enable=True, max_fps=fps, frame_type='numpy_bgr') as (sd, _): | ||
| 72 | - yield from sd | ||
| 73 | - | ||
| 74 | - def to_frames(self, url, cache=None, fps=None, start=None, end=None, roi=None, max_px_area=None) -> list: | ||
| 75 | - video_root_path = self.get_root_path(url, cache) | ||
| 76 | - | ||
| 77 | - download_video_path = video_root_path / "video.mp4" | ||
| 78 | - | ||
| 79 | - if not download_video_path.exists(): | ||
| 80 | - logger.info(f"download_video: {url} -> {download_video_path}") | ||
| 81 | - video_root_path.mkdir(parents=True, exist_ok=True) | ||
| 82 | - download_to_mp4(url, str(download_video_path), duration=None) | ||
| 83 | - logger.info(f"video_path: {download_video_path}") | ||
| 84 | - cache_dir = video_root_path / 'caches' / f"{start}-{end}-{'_'.join(roi) if roi else 'None'}-{max_px_area}" | ||
| 85 | - cache_frames = cache_dir / "frames" | ||
| 86 | - cache_name = cache_dir / "data.json" | ||
| 87 | - | ||
| 88 | - if cache_name.exists(): | ||
| 89 | - logger.info(f"use_cache: {cache_name}") | ||
| 90 | - # return json.load(cache_name.open('r')) | ||
| 91 | - with cache_name.open('r', encoding='utf-8') as f: | ||
| 92 | - return json.load(f) | ||
| 93 | - else: | ||
| 94 | - cache_dir.mkdir(parents=True, exist_ok=True) | ||
| 95 | - data = { | ||
| 96 | - 'fps': fps, | ||
| 97 | - 'frames': [] | ||
| 98 | - } | ||
| 99 | - if self.save_frames_enable: | ||
| 100 | - cache_frames.mkdir(parents=True, exist_ok=True) | ||
| 101 | - for f in self.iter_mp4(str(download_video_path), start, end, fps): | ||
| 102 | - frame = f['frame'] | ||
| 103 | - time = f['src_frame_time'] | ||
| 104 | - fps = f['fps'] | ||
| 105 | - data['fps'] = fps | ||
| 106 | - if roi is not None: | ||
| 107 | - x1, y1, x2, y2 = roi | ||
| 108 | - frame = frame[y1:y2, x1:x2] | ||
| 109 | - if max_px_area is not None: | ||
| 110 | - frame = resize_frame(frame, max_px_area) | ||
| 111 | - | ||
| 112 | - _, buffer = cv2.imencode('.jpg', frame) | ||
| 113 | - b64_str = base64.b64encode(buffer).decode('utf-8') | ||
| 114 | - | ||
| 115 | - data['frames'].append({'time': time, 'image': b64_str}) | ||
| 116 | - | ||
| 117 | - if self.save_frames_enable: | ||
| 118 | - frame_path = cache_frames / f"{time:015d}.jpg" | ||
| 119 | - cv2.imwrite(frame_path, frame) | ||
| 120 | - | ||
| 121 | - cache_name.parent.mkdir(parents=True, exist_ok=True) | ||
| 122 | - # json.dump(data, cache_name.open('w'), indent=4) | ||
| 123 | - with cache_name.open('w', encoding='utf-8') as f: | ||
| 124 | - json.dump(data, f, indent=4) | ||
| 125 | - logger.info(f"save_cache: {cache_name}") | ||
| 126 | - return data | ||
| 127 | - | ||
| 128 | - def to_llm_contents(self, url, cache=None, fps=None, start=None, end=None, roi=None, max_px_area=None, | ||
| 129 | - prompt_start=None, | ||
| 130 | - prompt_end=None): | ||
| 131 | - data = self.to_frames(url, cache, fps, start, end, roi, max_px_area) | ||
| 132 | - fps = data['fps'] | ||
| 133 | - frames = data['frames'] | ||
| 134 | - | ||
| 135 | - contents = [] | ||
| 136 | - if prompt_start is not None: | ||
| 137 | - contents.append({"type": "text", "text": prompt_start}) | ||
| 138 | - video_prompt = ( | ||
| 139 | - f"以下是从视频中按时间顺序提取的 {len(frames)} 帧画面，fps={fps}，请将它们视为一个连续的视频进行分析。" | ||
| 140 | - ) | ||
| 141 | - contents.append({"type": "text", "text": video_prompt}) | ||
| 142 | - | ||
| 143 | - for frame in frames: | ||
| 144 | - contents.append({ | ||
| 145 | - "type": "image_url", | ||
| 146 | - "image_url": { | ||
| 147 | - "url": f"data:image/jpeg;base64,{frame['image']}" | ||
| 148 | - } | ||
| 149 | - }) | ||
| 150 | - if prompt_end is not None: | ||
| 151 | - contents.append({"type": "text", "text": prompt_end}) | ||
| 152 | - return contents | ||
| 153 | - | ||
| 154 | - | ||
| 155 | -if __name__ == '__main__': | ||
| 156 | - from aabd.base.patched_logging import init_logging | ||
| 157 | - import os | ||
| 158 | - os.environ["APP_LOG_TYPE"] = "console" | ||
| 159 | - init_logging() | ||
| 160 | - # url = "http://video.mam.miguvideo.com/mnt6/fastclip3/wsc/2026/04/29/df760b8d15394cbc9ed0d0a4a9c4b786_1080PS/29133605/vodtmp/3b3e99cf4ca84c3782503d8817242de2.m3u8" | ||
| 161 | - url = rf"/root/lzw/finished/69dd5845dd0412067b8d5587-auto-1776074760653/live/videos/00-14-09-052.mp4" | ||
| 162 | - caceh_dir = r"/root/lzw/aigc-embedding-service/src/football_replay_match/core" | ||
| 163 | - vf = Video2Frame(caceh_dir, save_frames_enable=True) | ||
| 164 | - # vf = Video2Frame(rf'D:\Code\migu\aigc-embedding-service\src\football_replay_match\core', save_frames_enable=True) | ||
| 165 | - # a = vf.to_llm_contents(url, fps=2, start=1777444502543, end=1777444531063, max_px_area=1920 * 1080 // 6) | ||
| 166 | - a = vf.to_llm_contents(url, fps=2, start=0, end=10_000, max_px_area=1920 * 1080 // 6) | ||
| 167 | - from langchain_openai import ChatOpenAI | ||
| 168 | - from langchain_core.messages import SystemMessage, HumanMessage | ||
| 169 | - | ||
| 170 | - model = ChatOpenAI(base_url="http://192.168.1.59:11434/v1", model="Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf", | ||
| 171 | - temperature=0.7, api_key='no_key', | ||
| 172 | - extra_body={"chat_template_kwargs": {"enable_thinking": False}}) | ||
| 173 | - | ||
| 174 | - result = model.invoke([SystemMessage(content="描述一下视频的内容。"), HumanMessage(content=a)]) | ||
| 175 | - print(result) | 1 | +import base64 |
| 2 | +import shutil | ||
| 3 | +import subprocess | ||
| 4 | +import sys | ||
| 5 | +from pathlib import Path | ||
| 6 | +from typing import Optional | ||
| 7 | +from urllib.parse import urlparse | ||
| 8 | +import cv2 | ||
| 9 | +from aabd.stream_chain import video2frames | ||
| 10 | +import json | ||
| 11 | +import logging | ||
| 12 | +import hashlib | ||
| 13 | + | ||
| 14 | +logger = logging.getLogger(__name__) | ||
| 15 | + | ||
| 16 | +_IS_ABSOLUTE_UTC_THRESHOLD = 1_000_000_000 | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +def download_to_mp4(url: str, output_path: str, duration: Optional[int] = None) -> None: | ||
| 20 | + cmd = ["ffmpeg", "-y", "-fflags", "+discardcorrupt", "-i", url] | ||
| 21 | + if duration is not None: | ||
| 22 | + cmd.extend(["-t", str(duration)]) | ||
| 23 | + cmd.extend(["-c", "copy", "-bsf:a", "aac_adtstoasc", "-movflags", "+faststart", output_path]) | ||
| 24 | + result = subprocess.run(cmd, capture_output=True, text=True, encoding="utf-8", errors="replace") | ||
| 25 | + if result.returncode != 0: | ||
| 26 | + raise RuntimeError(f"ffmpeg 下载失败: {result.stderr}") | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +def resize_frame(frame, max_px_area): | ||
| 30 | + if max_px_area is None: | ||
| 31 | + return frame | ||
| 32 | + h, w = frame.shape[:2] | ||
| 33 | + area = h * w | ||
| 34 | + if area > max_px_area: | ||
| 35 | + scale = (max_px_area / area) ** 0.5 | ||
| 36 | + new_w = int(w * scale) | ||
| 37 | + new_h = int(h * scale) | ||
| 38 | + frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA) | ||
| 39 | + return frame | ||
| 40 | + | ||
| 41 | + | ||
| 42 | +class Video2Frame: | ||
| 43 | + def __init__(self, cache_dir, save_frames_enable=False): | ||
| 44 | + self.cache_dir = Path(cache_dir) | ||
| 45 | + self.save_frames_enable = save_frames_enable | ||
| 46 | + | ||
| 47 | + def get_root_path(self, url, cache=None): | ||
| 48 | + if cache is None: | ||
| 49 | + parsed_url = urlparse(url) | ||
| 50 | + raw = f"{parsed_url.scheme}/{parsed_url.netloc}{parsed_url.path}" | ||
| 51 | + cache = hashlib.md5(raw.encode('utf-8')).hexdigest() | ||
| 52 | + return self.cache_dir / cache | ||
| 53 | + | ||
| 54 | + def iter_mp4(self, file, start, end, fps): | ||
| 55 | + | ||
| 56 | + if (start is not None and start > _IS_ABSOLUTE_UTC_THRESHOLD) or ( | ||
| 57 | + end is not None and end > _IS_ABSOLUTE_UTC_THRESHOLD): | ||
| 58 | + video_start_time = 0 | ||
| 59 | + with video2frames.AVAnyKeyStreamDecoder(file, sei_enable=True, tqdm_enable=False) as (sd, _): | ||
| 60 | + for frame in sd: | ||
| 61 | + if start is not None: | ||
| 62 | + video_start_time = start - frame.get('src_frame_time') | ||
| 63 | + if end is not None: | ||
| 64 | + video_end_time = end - frame.get('src_frame_time') or 0 | ||
| 65 | + break | ||
| 66 | + else: | ||
| 67 | + video_start_time = start if start is not None else 0 | ||
| 68 | + video_end_time = end if end is not None else sys.maxsize | ||
| 69 | + | ||
| 70 | + with video2frames.AVAnyKeyStreamDecoder(file, start=video_start_time, end=video_end_time, control_type='time', | ||
| 71 | + sei_enable=True, max_fps=fps, frame_type='numpy_bgr') as (sd, _): | ||
| 72 | + yield from sd | ||
| 73 | + | ||
| 74 | + def to_frames(self, url, cache=None, fps=None, start=None, end=None, roi=None, max_px_area=None) -> list: | ||
| 75 | + video_root_path = self.get_root_path(url, cache) | ||
| 76 | + | ||
| 77 | + download_video_path = video_root_path / "video.mp4" | ||
| 78 | + | ||
| 79 | + if not download_video_path.exists(): | ||
| 80 | + logger.info(f"download_video: {url} -> {download_video_path}") | ||
| 81 | + video_root_path.mkdir(parents=True, exist_ok=True) | ||
| 82 | + download_to_mp4(url, str(download_video_path), duration=None) | ||
| 83 | + logger.info(f"video_path: {download_video_path}") | ||
| 84 | + cache_dir = video_root_path / 'caches' / f"{start}-{end}-{'_'.join(roi) if roi else 'None'}-{max_px_area}" | ||
| 85 | + cache_frames = cache_dir / "frames" | ||
| 86 | + cache_name = cache_dir / "data.json" | ||
| 87 | + | ||
| 88 | + if cache_name.exists(): | ||
| 89 | + logger.info(f"use_cache: {cache_name}") | ||
| 90 | + # return json.load(cache_name.open('r')) | ||
| 91 | + with cache_name.open('r', encoding='utf-8') as f: | ||
| 92 | + return json.load(f) | ||
| 93 | + else: | ||
| 94 | + cache_dir.mkdir(parents=True, exist_ok=True) | ||
| 95 | + data = { | ||
| 96 | + 'fps': fps, | ||
| 97 | + 'frames': [] | ||
| 98 | + } | ||
| 99 | + if self.save_frames_enable: | ||
| 100 | + cache_frames.mkdir(parents=True, exist_ok=True) | ||
| 101 | + for f in self.iter_mp4(str(download_video_path), start, end, fps): | ||
| 102 | + frame = f['frame'] | ||
| 103 | + time = f['src_frame_time'] | ||
| 104 | + fps = f['fps'] | ||
| 105 | + data['fps'] = fps | ||
| 106 | + if roi is not None: | ||
| 107 | + x1, y1, x2, y2 = roi | ||
| 108 | + frame = frame[y1:y2, x1:x2] | ||
| 109 | + if max_px_area is not None: | ||
| 110 | + frame = resize_frame(frame, max_px_area) | ||
| 111 | + | ||
| 112 | + _, buffer = cv2.imencode('.jpg', frame) | ||
| 113 | + b64_str = base64.b64encode(buffer).decode('utf-8') | ||
| 114 | + | ||
| 115 | + data['frames'].append({'time': time, 'image': b64_str}) | ||
| 116 | + | ||
| 117 | + if self.save_frames_enable: | ||
| 118 | + frame_path = cache_frames / f"{time:015d}.jpg" | ||
| 119 | + cv2.imwrite(frame_path, frame) | ||
| 120 | + | ||
| 121 | + cache_name.parent.mkdir(parents=True, exist_ok=True) | ||
| 122 | + # json.dump(data, cache_name.open('w'), indent=4) | ||
| 123 | + with cache_name.open('w', encoding='utf-8') as f: | ||
| 124 | + json.dump(data, f, indent=4) | ||
| 125 | + logger.info(f"save_cache: {cache_name}") | ||
| 126 | + return data | ||
| 127 | + | ||
| 128 | + def to_llm_contents(self, url, cache=None, fps=None, start=None, end=None, roi=None, max_px_area=None, | ||
| 129 | + prompt_start=None, | ||
| 130 | + prompt_end=None): | ||
| 131 | + data = self.to_frames(url, cache, fps, start, end, roi, max_px_area) | ||
| 132 | + fps = data['fps'] | ||
| 133 | + frames = data['frames'] | ||
| 134 | + | ||
| 135 | + contents = [] | ||
| 136 | + if prompt_start is not None: | ||
| 137 | + contents.append({"type": "text", "text": prompt_start}) | ||
| 138 | + video_prompt = ( | ||
| 139 | + f"以下是从视频中按时间顺序提取的 {len(frames)} 帧画面，fps={fps}，请将它们视为一个连续的视频进行分析。" | ||
| 140 | + ) | ||
| 141 | + contents.append({"type": "text", "text": video_prompt}) | ||
| 142 | + | ||
| 143 | + for frame in frames: | ||
| 144 | + contents.append({ | ||
| 145 | + "type": "image_url", | ||
| 146 | + "image_url": { | ||
| 147 | + "url": f"data:image/jpeg;base64,{frame['image']}" | ||
| 148 | + } | ||
| 149 | + }) | ||
| 150 | + if prompt_end is not None: | ||
| 151 | + contents.append({"type": "text", "text": prompt_end}) | ||
| 152 | + return contents | ||
| 153 | + | ||
| 154 | + | ||
| 155 | +if __name__ == '__main__': | ||
| 156 | + from aabd.base.patched_logging import init_logging | ||
| 157 | + import os | ||
| 158 | + os.environ["APP_LOG_TYPE"] = "console" | ||
| 159 | + init_logging() | ||
| 160 | + # url = "http://video.mam.miguvideo.com/mnt6/fastclip3/wsc/2026/04/29/df760b8d15394cbc9ed0d0a4a9c4b786_1080PS/29133605/vodtmp/3b3e99cf4ca84c3782503d8817242de2.m3u8" | ||
| 161 | + url = rf"/root/lzw/finished/69dd5845dd0412067b8d5587-auto-1776074760653/live/videos/00-14-09-052.mp4" | ||
| 162 | + caceh_dir = r"/root/lzw/aigc-embedding-service/src/football_replay_match/core" | ||
| 163 | + vf = Video2Frame(caceh_dir, save_frames_enable=True) | ||
| 164 | + # vf = Video2Frame(rf'D:\Code\migu\aigc-embedding-service\src\football_replay_match\core', save_frames_enable=True) | ||
| 165 | + # a = vf.to_llm_contents(url, fps=2, start=1777444502543, end=1777444531063, max_px_area=1920 * 1080 // 6) | ||
| 166 | + a = vf.to_llm_contents(url, fps=2, start=0, end=10_000, max_px_area=1920 * 1080 // 6) | ||
| 167 | + from langchain_openai import ChatOpenAI | ||
| 168 | + from langchain_core.messages import SystemMessage, HumanMessage | ||
| 169 | + | ||
| 170 | + model = ChatOpenAI(base_url="http://192.168.1.59:11434/v1", model="Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf", | ||
| 171 | + temperature=0.7, api_key='no_key', | ||
| 172 | + extra_body={"chat_template_kwargs": {"enable_thinking": False}}) | ||
| 173 | + | ||
| 174 | + result = model.invoke([SystemMessage(content="描述一下视频的内容。"), HumanMessage(content=a)]) | ||
| 175 | + print(result) |
| 1 | -from football_replay_match_live import FootballReplayMatchLive | ||
| 2 | -# from qwen_asr_util import QwenAsr | ||
| 3 | -from football_replay_video_event_by_llm import FootballReplayVideoEvent | ||
| 4 | -import os | ||
| 5 | -import json | ||
| 6 | - | ||
| 7 | - | ||
| 8 | -def batch_match(): | ||
| 9 | - cache_dir = "/root/lzw/tmp_0518_replay_cache03" | ||
| 10 | - replay_match_live = FootballReplayMatchLive(base_url="http://192.168.1.59:11434/v1", | ||
| 11 | - model="Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf", | ||
| 12 | - temperature=0.7, | ||
| 13 | - cache_dir=cache_dir, | ||
| 14 | - save_frames_enable=True | ||
| 15 | - ) | ||
| 16 | - # qwen_asr = QwenAsr(base_url="http://192.168.1.59:8101/v1", model="Qwen/Qwen3-ASR-1.7B") | ||
| 17 | - fbrv = FootballReplayVideoEvent(base_url="http://192.168.1.59:11434/v1", | ||
| 18 | - model="Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf", | ||
| 19 | - temperature=0.7, | ||
| 20 | - cache_dir=cache_dir, | ||
| 21 | - save_frames_enable=True | ||
| 22 | - ) | ||
| 23 | - | ||
| 24 | - videos_dir = rf"/root/lzw/finished" | ||
| 25 | - for video_name in sorted(os.listdir(videos_dir)): | ||
| 26 | - | ||
| 27 | - live_dir = os.path.join(videos_dir, video_name, 'live') | ||
| 28 | - print(f"live_dir: {live_dir}") | ||
| 29 | - | ||
| 30 | - live_video_dir = os.path.join(live_dir, 'videos') | ||
| 31 | - live_asr_dir = os.path.join(live_dir, 'asr') | ||
| 32 | - | ||
| 33 | - live_packs = [] | ||
| 34 | - for live_video_name in sorted(os.listdir(live_video_dir)): | ||
| 35 | - live_video_path = os.path.join(live_video_dir, live_video_name) | ||
| 36 | - live_asr_path = os.path.join(live_asr_dir, f"{live_video_name}.txt") | ||
| 37 | - # live_asr_text = qwen_asr.asr(live_video_path, cache_path=live_asr_path) | ||
| 38 | - live_asr_text = '' | ||
| 39 | - live_packs.append({ | ||
| 40 | - "video_id": os.path.basename(live_video_path), | ||
| 41 | - "url": live_video_path, | ||
| 42 | - "asr_text": live_asr_text | ||
| 43 | - }) | ||
| 44 | - | ||
| 45 | - replays_dir = os.path.join(videos_dir, video_name, 'replays') | ||
| 46 | - replay_videos_dir = os.path.join(replays_dir, 'videos') | ||
| 47 | - replay_goals_dir = os.path.join(replays_dir, 'goals') | ||
| 48 | - replay_asr_dir = os.path.join(replays_dir, 'asr') | ||
| 49 | - replay_matches_dir = os.path.join(replays_dir, 'matches') | ||
| 50 | - | ||
| 51 | - for replay_video_name in sorted(os.listdir(replay_videos_dir)): | ||
| 52 | - replay_video_path = os.path.join(replay_videos_dir, replay_video_name) | ||
| 53 | - # replay_goals_path = os.path.join(replay_goals_dir, f"{replay_video_name}.json") | ||
| 54 | - | ||
| 55 | - # replay_asr_path = os.path.join(replay_asr_dir, f"{replay_video_name}.txt") | ||
| 56 | - # replay_match_path = os.path.join(replay_matches_dir, f"{replay_video_name}.json") | ||
| 57 | - | ||
| 58 | - # print(100*"*") | ||
| 59 | - # print(replay_match_path) | ||
| 60 | - replay_match_path = None | ||
| 61 | - | ||
| 62 | - # replay_asr_text = qwen_asr.asr(replay_video_path, cache_path=replay_asr_path) | ||
| 63 | - replay_asr_text = '' | ||
| 64 | - replay_pack = {"url": "http://video.mam.miguvideo.com/mnt6/fastclip3/wsc/2026/04/29/df760b8d15394cbc9ed0d0a4a9c4b786_1080PS/29133605/vodtmp/3b3e99cf4ca84c3782503d8817242de2.m3u8", | ||
| 65 | - "start_utc": 1777444502543, | ||
| 66 | - "end_utc": 1777444531063, | ||
| 67 | - "asr_text": '无' | ||
| 68 | - } | ||
| 69 | - | ||
| 70 | - replay_goals = fbrv.video_event(replay_pack, cache_dir=cache_dir) | ||
| 71 | - if replay_goals['event_name'] == '无进球': | ||
| 72 | - continue | ||
| 73 | - # replay_pack = {"video_path": replay_video_path, "asr_text": replay_asr_text} | ||
| 74 | - | ||
| 75 | - try: | ||
| 76 | - import time | ||
| 77 | - start_time = time.time() | ||
| 78 | - print("Start matching...") | ||
| 79 | - print(f"len(live_packs): {len(live_packs)}") | ||
| 80 | - result = replay_match_live.match_batch(replay_pack, live_packs, max_parallel=2, cache_dir=cache_dir) | ||
| 81 | - end_time = time.time() | ||
| 82 | - print(f"Time taken for matching: {end_time - start_time:.2f} seconds") | ||
| 83 | - print(replay_video_path) | ||
| 84 | - print(result) | ||
| 85 | - break | ||
| 86 | - | ||
| 87 | - except Exception as e: | ||
| 88 | - print(f"Error processing {replay_video_path}: {e}") | ||
| 89 | - | ||
| 90 | - print('-' * 20) | ||
| 91 | - | ||
| 92 | - break | ||
| 93 | - | ||
| 94 | - | ||
| 95 | -if __name__ == '__main__': | ||
| 96 | - batch_match() | 1 | +from football_replay_match_live import FootballReplayMatchLive |
| 2 | +# from qwen_asr_util import QwenAsr | ||
| 3 | +from football_replay_video_event_by_llm import FootballReplayVideoEvent | ||
| 4 | +import os | ||
| 5 | +import json | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +def batch_match(): | ||
| 9 | + cache_dir = "/root/lzw/tmp_0518_replay_cache03" | ||
| 10 | + replay_match_live = FootballReplayMatchLive(base_url="http://192.168.1.59:11434/v1", | ||
| 11 | + model="Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf", | ||
| 12 | + temperature=0.7, | ||
| 13 | + cache_dir=cache_dir, | ||
| 14 | + save_frames_enable=True | ||
| 15 | + ) | ||
| 16 | + # qwen_asr = QwenAsr(base_url="http://192.168.1.59:8101/v1", model="Qwen/Qwen3-ASR-1.7B") | ||
| 17 | + fbrv = FootballReplayVideoEvent(base_url="http://192.168.1.59:11434/v1", | ||
| 18 | + model="Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf", | ||
| 19 | + temperature=0.7, | ||
| 20 | + cache_dir=cache_dir, | ||
| 21 | + save_frames_enable=True | ||
| 22 | + ) | ||
| 23 | + | ||
| 24 | + videos_dir = rf"/root/lzw/finished" | ||
| 25 | + for video_name in sorted(os.listdir(videos_dir)): | ||
| 26 | + | ||
| 27 | + live_dir = os.path.join(videos_dir, video_name, 'live') | ||
| 28 | + print(f"live_dir: {live_dir}") | ||
| 29 | + | ||
| 30 | + live_video_dir = os.path.join(live_dir, 'videos') | ||
| 31 | + live_asr_dir = os.path.join(live_dir, 'asr') | ||
| 32 | + | ||
| 33 | + live_packs = [] | ||
| 34 | + for live_video_name in sorted(os.listdir(live_video_dir)): | ||
| 35 | + live_video_path = os.path.join(live_video_dir, live_video_name) | ||
| 36 | + live_asr_path = os.path.join(live_asr_dir, f"{live_video_name}.txt") | ||
| 37 | + # live_asr_text = qwen_asr.asr(live_video_path, cache_path=live_asr_path) | ||
| 38 | + live_asr_text = '' | ||
| 39 | + live_packs.append({ | ||
| 40 | + "video_id": os.path.basename(live_video_path), | ||
| 41 | + "url": live_video_path, | ||
| 42 | + "asr_text": live_asr_text | ||
| 43 | + }) | ||
| 44 | + | ||
| 45 | + replays_dir = os.path.join(videos_dir, video_name, 'replays') | ||
| 46 | + replay_videos_dir = os.path.join(replays_dir, 'videos') | ||
| 47 | + replay_goals_dir = os.path.join(replays_dir, 'goals') | ||
| 48 | + replay_asr_dir = os.path.join(replays_dir, 'asr') | ||
| 49 | + replay_matches_dir = os.path.join(replays_dir, 'matches') | ||
| 50 | + | ||
| 51 | + for replay_video_name in sorted(os.listdir(replay_videos_dir)): | ||
| 52 | + replay_video_path = os.path.join(replay_videos_dir, replay_video_name) | ||
| 53 | + # replay_goals_path = os.path.join(replay_goals_dir, f"{replay_video_name}.json") | ||
| 54 | + | ||
| 55 | + # replay_asr_path = os.path.join(replay_asr_dir, f"{replay_video_name}.txt") | ||
| 56 | + # replay_match_path = os.path.join(replay_matches_dir, f"{replay_video_name}.json") | ||
| 57 | + | ||
| 58 | + # print(100*"*") | ||
| 59 | + # print(replay_match_path) | ||
| 60 | + replay_match_path = None | ||
| 61 | + | ||
| 62 | + # replay_asr_text = qwen_asr.asr(replay_video_path, cache_path=replay_asr_path) | ||
| 63 | + replay_asr_text = '' | ||
| 64 | + replay_pack = {"url": "http://video.mam.miguvideo.com/mnt6/fastclip3/wsc/2026/04/29/df760b8d15394cbc9ed0d0a4a9c4b786_1080PS/29133605/vodtmp/3b3e99cf4ca84c3782503d8817242de2.m3u8", | ||
| 65 | + "start_utc": 1777444502543, | ||
| 66 | + "end_utc": 1777444531063, | ||
| 67 | + "asr_text": '无' | ||
| 68 | + } | ||
| 69 | + | ||
| 70 | + replay_goals = fbrv.video_event(replay_pack, cache_dir=cache_dir) | ||
| 71 | + if replay_goals['event_name'] == '无进球': | ||
| 72 | + continue | ||
| 73 | + # replay_pack = {"video_path": replay_video_path, "asr_text": replay_asr_text} | ||
| 74 | + | ||
| 75 | + try: | ||
| 76 | + import time | ||
| 77 | + start_time = time.time() | ||
| 78 | + print("Start matching...") | ||
| 79 | + print(f"len(live_packs): {len(live_packs)}") | ||
| 80 | + result = replay_match_live.match_batch(replay_pack, live_packs, max_parallel=2, cache_dir=cache_dir) | ||
| 81 | + end_time = time.time() | ||
| 82 | + print(f"Time taken for matching: {end_time - start_time:.2f} seconds") | ||
| 83 | + print(replay_video_path) | ||
| 84 | + print(result) | ||
| 85 | + break | ||
| 86 | + | ||
| 87 | + except Exception as e: | ||
| 88 | + print(f"Error processing {replay_video_path}: {e}") | ||
| 89 | + | ||
| 90 | + print('-' * 20) | ||
| 91 | + | ||
| 92 | + break | ||
| 93 | + | ||
| 94 | + | ||
| 95 | +if __name__ == '__main__': | ||
| 96 | + batch_match() |
| 1 | -from aabd.base.patched_logging import init_logging, get_logger | ||
| 2 | - | ||
| 3 | -init_logging() | ||
| 4 | -logger = get_logger(__name__) | ||
| 5 | -import json | ||
| 6 | -import time | ||
| 7 | - | ||
| 8 | -try: | ||
| 9 | - from config import settings | ||
| 10 | -except: | ||
| 11 | - from .config import settings | ||
| 12 | - | ||
| 13 | - pass | ||
| 14 | -try: | ||
| 15 | - from .core.api import FootballReplayMatch | ||
| 16 | -except: | ||
| 17 | - from core.api import FootballReplayMatch | ||
| 18 | - pass | ||
| 19 | - | ||
| 20 | -from threading import Thread | ||
| 21 | - | ||
| 22 | -from aabd.mq.kafka_client import KafkaMessageIterator, KafkaProducer | ||
| 23 | - | ||
| 24 | - | ||
| 25 | -def get_message_iterator(kafka_config): | ||
| 26 | - return KafkaMessageIterator(bootstrap_servers=kafka_config.get('servers', None), | ||
| 27 | - group_id=kafka_config.get('group_id', None), | ||
| 28 | - topic=kafka_config.get('topic', None), | ||
| 29 | - sasl_plain_username=kafka_config.get('username', None), | ||
| 30 | - sasl_plain_password=kafka_config.get('password', None), | ||
| 31 | - value_deserializer="str") | ||
| 32 | - | ||
| 33 | - | ||
| 34 | -def get_message_producer(kafka_config): | ||
| 35 | - return KafkaProducer(bootstrap_servers=kafka_config.get('servers', None), | ||
| 36 | - sasl_plain_username=kafka_config.get('username', None), | ||
| 37 | - sasl_plain_password=kafka_config.get('password', None)) | ||
| 38 | - | ||
| 39 | - | ||
| 40 | -def get_callback_func(producer, topic): | ||
| 41 | - def callback_func(task_id, call_data): | ||
| 42 | - try: | ||
| 43 | - st = time.time() | ||
| 44 | - producer.send_message_async(topic, call_data, key=task_id, timeout=5) | ||
| 45 | - logger.info(f"sent_message[{task_id}][{time.time() - st:.2f}s]: {call_data}") | ||
| 46 | - except: | ||
| 47 | - logger.exception(f"Error sending message: {call_data}") | ||
| 48 | - | ||
| 49 | - return callback_func | ||
| 50 | - | ||
| 51 | - | ||
| 52 | -def kafka_message_iterator_thread(message_iterator, config, callback_func): | ||
| 53 | - frm = FootballReplayMatch(config) | ||
| 54 | - with message_iterator: | ||
| 55 | - for message in message_iterator: | ||
| 56 | - try: | ||
| 57 | - json_message = json.loads(message) | ||
| 58 | - result = frm.replay_match_event(json_message) | ||
| 59 | - task_id = json_message.get("id") | ||
| 60 | - callback_func(task_id, json.dumps(result, ensure_ascii=False)) | ||
| 61 | - logger.info(f"Processed task_id={task_id}, result={result}") | ||
| 62 | - except Exception: | ||
| 63 | - logger.exception(f"Error processing message: {message}") | ||
| 64 | - | ||
| 65 | - | ||
| 66 | -def main(): | ||
| 67 | - message_iter = get_message_iterator(settings.input_kafka) | ||
| 68 | - kafka_producer = get_message_producer(settings.output_kafka) | ||
| 69 | - | ||
| 70 | - try: | ||
| 71 | - | ||
| 72 | - Thread(target=kafka_message_iterator_thread, | ||
| 73 | - args=(message_iter, settings, | ||
| 74 | - get_callback_func(kafka_producer, settings.output_kafka.get('topic', None)))).start() | ||
| 75 | - | ||
| 76 | - except KeyboardInterrupt: | ||
| 77 | - message_iter.running = False | ||
| 78 | - kafka_producer.close() | ||
| 79 | - | ||
| 80 | - | ||
| 81 | -if __name__ == '__main__': | ||
| 82 | - main() | 1 | +from aabd.base.patched_logging import init_logging, get_logger |
| 2 | + | ||
| 3 | +init_logging() | ||
| 4 | +logger = get_logger(__name__) | ||
| 5 | +import json | ||
| 6 | +import time | ||
| 7 | + | ||
| 8 | +try: | ||
| 9 | + from config import settings | ||
| 10 | +except: | ||
| 11 | + from .config import settings | ||
| 12 | + | ||
| 13 | + pass | ||
| 14 | +try: | ||
| 15 | + from .core.api import FootballReplayMatch | ||
| 16 | +except: | ||
| 17 | + from core.api import FootballReplayMatch | ||
| 18 | + pass | ||
| 19 | + | ||
| 20 | +from threading import Thread | ||
| 21 | + | ||
| 22 | +from aabd.mq.kafka_client import KafkaMessageIterator, KafkaProducer | ||
| 23 | + | ||
| 24 | + | ||
| 25 | +def get_message_iterator(kafka_config): | ||
| 26 | + return KafkaMessageIterator(bootstrap_servers=kafka_config.get('servers', None), | ||
| 27 | + group_id=kafka_config.get('group_id', None), | ||
| 28 | + topic=kafka_config.get('topic', None), | ||
| 29 | + sasl_plain_username=kafka_config.get('username', None), | ||
| 30 | + sasl_plain_password=kafka_config.get('password', None), | ||
| 31 | + value_deserializer="str") | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +def get_message_producer(kafka_config): | ||
| 35 | + return KafkaProducer(bootstrap_servers=kafka_config.get('servers', None), | ||
| 36 | + sasl_plain_username=kafka_config.get('username', None), | ||
| 37 | + sasl_plain_password=kafka_config.get('password', None)) | ||
| 38 | + | ||
| 39 | + | ||
| 40 | +def get_callback_func(producer, topic): | ||
| 41 | + def callback_func(task_id, call_data): | ||
| 42 | + try: | ||
| 43 | + st = time.time() | ||
| 44 | + producer.send_message_async(topic, call_data, key=task_id, timeout=5) | ||
| 45 | + logger.info(f"sent_message[{task_id}][{time.time() - st:.2f}s]: {call_data}") | ||
| 46 | + except: | ||
| 47 | + logger.exception(f"Error sending message: {call_data}") | ||
| 48 | + | ||
| 49 | + return callback_func | ||
| 50 | + | ||
| 51 | + | ||
| 52 | +def kafka_message_iterator_thread(message_iterator, config, callback_func): | ||
| 53 | + frm = FootballReplayMatch(config) | ||
| 54 | + with message_iterator: | ||
| 55 | + for message in message_iterator: | ||
| 56 | + try: | ||
| 57 | + json_message = json.loads(message) | ||
| 58 | + result = frm.replay_match_event(json_message) | ||
| 59 | + task_id = json_message.get("id") | ||
| 60 | + callback_func(task_id, json.dumps(result, ensure_ascii=False)) | ||
| 61 | + logger.info(f"Processed task_id={task_id}, result={result}") | ||
| 62 | + except Exception: | ||
| 63 | + logger.exception(f"Error processing message: {message}") | ||
| 64 | + | ||
| 65 | + | ||
| 66 | +def main(): | ||
| 67 | + message_iter = get_message_iterator(settings.input_kafka) | ||
| 68 | + kafka_producer = get_message_producer(settings.output_kafka) | ||
| 69 | + | ||
| 70 | + try: | ||
| 71 | + | ||
| 72 | + Thread(target=kafka_message_iterator_thread, | ||
| 73 | + args=(message_iter, settings, | ||
| 74 | + get_callback_func(kafka_producer, settings.output_kafka.get('topic', None)))).start() | ||
| 75 | + | ||
| 76 | + except KeyboardInterrupt: | ||
| 77 | + message_iter.running = False | ||
| 78 | + kafka_producer.close() | ||
| 79 | + | ||
| 80 | + | ||
| 81 | +if __name__ == '__main__': | ||
| 82 | + main() |
This diff could not be displayed because it is too large.
-
Please register or login to post a comment