Build lightweight AI agent admin
This commit is contained in:
Codex
@@ -0,0 +1,70 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
media/
|
||||
wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
|
||||
# Virtual Environment
|
||||
venv/
|
||||
ENV/*
|
||||
env/*
|
||||
.venv/
|
||||
|
||||
# IDE
|
||||
.idea/
|
||||
.vscode/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# Environment files - 保护敏感信息
|
||||
env/*
|
||||
!env/example.env
|
||||
|
||||
# Database
|
||||
*.db
|
||||
*.sqlite3
|
||||
|
||||
# Logs
|
||||
*.log
|
||||
logs/
|
||||
|
||||
# Media files
|
||||
media/
|
||||
!media/.gitkeep
|
||||
|
||||
# Chunk uploads temp
|
||||
media/chunk_uploads/
|
||||
|
||||
# Cache
|
||||
.cache/
|
||||
*.pyc
|
||||
|
||||
# Test
|
||||
.pytest_cache/
|
||||
.coverage
|
||||
htmlcov/
|
||||
|
||||
# Alembic
|
||||
alembic/versions/*.pyc
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
@@ -0,0 +1,777 @@
|
||||
# ZQ Platform - FastAPI Backend
|
||||
|
||||
基于 FastAPI 的现代化异步后端服务，使用 SQLAlchemy 异步 ORM + Alembic 数据库迁移 + PostgreSQL。
|
||||
|
||||
## 技术栈
|
||||
|
||||
- **框架**: FastAPI 0.115+
|
||||
- **数据库**: PostgreSQL 16+
|
||||
- **ORM**: SQLAlchemy 2.0+ (异步)
|
||||
- **迁移**: Alembic
|
||||
- **认证**: JWT
|
||||
- **缓存**: Redis
|
||||
- **Python**: 3.12+
|
||||
|
||||
## 项目结构
|
||||
|
||||
```
|
||||
backend-fastapi/
|
||||
├── app/ # 核心应用模块
|
||||
│ ├── config.py # 配置管理
|
||||
│ ├── database.py # 数据库连接
|
||||
│ ├── base_model.py # BaseModel 基类
|
||||
│ ├── base_schema.py # 通用 Schema
|
||||
│ ├── base_service.py # BaseService 基类
|
||||
│ ├── redis.py # Redis 缓存
|
||||
│ └── excel.py # Excel 工具
|
||||
├── core/ # 核心业务模块
|
||||
│ ├── user/ # 用户管理
|
||||
│ ├── role/ # 角色管理
|
||||
│ ├── menu/ # 菜单管理
|
||||
│ ├── dept/ # 部门管理
|
||||
│ ├── permission/ # 权限管理
|
||||
│ └── ...
|
||||
├── scheduler/ # 定时任务模块
|
||||
│ ├── model.py
|
||||
│ ├── service.py
|
||||
│ └── tasks.py
|
||||
├── zq_demo/ # 示例模块
|
||||
│ ├── demo/
|
||||
│ └── demo_cache/
|
||||
├── scripts/ # 工具脚本
|
||||
│ ├── dumpdata.py # 数据导出
|
||||
│ └── loaddata.py # 数据导入
|
||||
├── alembic/ # 数据库迁移
|
||||
│ ├── versions/
|
||||
│ └── env.py
|
||||
├── env/ # 环境配置
|
||||
│ ├── dev.env
|
||||
│ ├── uat.env
|
||||
│ └── prod.env
|
||||
├── main.py # 应用入口
|
||||
├── requirements.txt # 依赖列表
|
||||
└── README.md
|
||||
```
|
||||
|
||||
## 快速开始
|
||||
|
||||
### 1. 环境准备
|
||||
|
||||
```bash
|
||||
# 创建虚拟环境
|
||||
conda create -n zq-fastapi python=3.12
|
||||
conda activate zq-fastapi
|
||||
|
||||
# 或使用 venv
|
||||
python -m venv venv
|
||||
source venv/bin/activate # Linux/Mac
|
||||
# venv\Scripts\activate # Windows
|
||||
```
|
||||
|
||||
### 2. 安装依赖
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### 3. 配置环境变量
|
||||
|
||||
复制环境配置文件：
|
||||
|
||||
```bash
|
||||
cp env/example.env env/dev.env
|
||||
```
|
||||
|
||||
编辑 `env/dev.env`，配置数据库连接：
|
||||
|
||||
```env
|
||||
# 数据库配置
|
||||
DATABASE_URL=postgresql+asyncpg://user:password@localhost:5432/dbname
|
||||
|
||||
# Redis 配置
|
||||
REDIS_HOST=localhost
|
||||
REDIS_PORT=6379
|
||||
REDIS_DB=0
|
||||
|
||||
# JWT 配置
|
||||
SECRET_KEY=your-secret-key-here
|
||||
ALGORITHM=HS256
|
||||
ACCESS_TOKEN_EXPIRE_MINUTES=30
|
||||
```
|
||||
|
||||
### 4. 数据库迁移
|
||||
|
||||
```bash
|
||||
# 首次使用：生成初始迁移
|
||||
alembic revision --autogenerate -m "init tables"
|
||||
|
||||
# 执行迁移
|
||||
alembic upgrade head
|
||||
|
||||
# 导入数据
|
||||
python scripts/loaddata.py db_init.json
|
||||
```
|
||||
|
||||
### 5. 启动服务
|
||||
|
||||
```bash
|
||||
# 开发模式（自动重载）
|
||||
python main.py
|
||||
|
||||
# 或使用 uvicorn
|
||||
uvicorn main:app --reload --host 0.0.0.0 --port 8000
|
||||
```
|
||||
|
||||
### 6. 访问 API 文档
|
||||
|
||||
- **Swagger UI**: http://localhost:8000/docs
|
||||
- **ReDoc**: http://localhost:8000/redoc
|
||||
|
||||
## 数据库操作
|
||||
|
||||
### 迁移命令
|
||||
|
||||
```bash
|
||||
# 查看当前版本
|
||||
alembic current
|
||||
|
||||
# 查看迁移历史
|
||||
alembic history
|
||||
|
||||
# 生成新的迁移文件
|
||||
alembic revision --autogenerate -m "描述信息"
|
||||
|
||||
# 升级到最新版本
|
||||
alembic upgrade head
|
||||
|
||||
# 回滚一个版本
|
||||
alembic downgrade -1
|
||||
|
||||
# 回滚到指定版本
|
||||
alembic downgrade <revision_id>
|
||||
```
|
||||
|
||||
### 数据导入导出
|
||||
|
||||
#### 导出数据（dumpdata.py）
|
||||
|
||||
```bash
|
||||
# 导出所有数据到文件
|
||||
python scripts/dumpdata.py -o db_init.json -f
|
||||
|
||||
# 导出指定模块（如 core）
|
||||
python scripts/dumpdata.py core -o core_data.json -f
|
||||
|
||||
# 导出到标准输出（不指定 -o 参数）
|
||||
python scripts/dumpdata.py > data.json
|
||||
|
||||
# 导出指定模块到标准输出
|
||||
python scripts/dumpdata.py core > core_data.json
|
||||
```
|
||||
|
||||
**参数说明：**
|
||||
|
||||
- `app_name`（位置参数，可选）：指定要导出的应用/模块名称
|
||||
- 例如：`core`、`scheduler`、`zq_demo`
|
||||
- 不指定则导出所有数据
|
||||
|
||||
- `-o, --output`：指定输出文件路径
|
||||
- 例如：`-o db_init.json`
|
||||
- 不指定则输出到标准输出（stdout）
|
||||
|
||||
- `-f, --force`：强制覆盖已存在的文件
|
||||
- 如果输出文件已存在且未使用此参数，脚本会报错并退出
|
||||
- 使用此参数可以强制覆盖现有文件
|
||||
|
||||
**示例：**
|
||||
|
||||
```bash
|
||||
# 导出所有数据，如果文件存在则覆盖
|
||||
python scripts/dumpdata.py -o db_init.json -f
|
||||
|
||||
# 导出 core 模块数据，不覆盖已存在文件（文件存在会报错）
|
||||
python scripts/dumpdata.py core -o core_data.json
|
||||
|
||||
# 导出 scheduler 模块数据到标准输出，然后重定向到文件
|
||||
python scripts/dumpdata.py scheduler > scheduler_data.json
|
||||
```
|
||||
|
||||
#### 导入数据（loaddata.py）
|
||||
|
||||
```bash
|
||||
# 导入数据
|
||||
python scripts/loaddata.py db_init.json
|
||||
|
||||
# 导入多个文件
|
||||
python scripts/loaddata.py core_data.json scheduler_data.json
|
||||
```
|
||||
|
||||
**参数说明：**
|
||||
|
||||
- `files`（位置参数，必需）：要导入的 JSON 文件路径，可以指定多个文件
|
||||
|
||||
## 开发指南
|
||||
|
||||
### 新建模块
|
||||
|
||||
按照以下步骤创建新的业务模块（以 `example` 为例）：
|
||||
|
||||
#### 1. 创建模块目录
|
||||
|
||||
```bash
|
||||
mkdir -p core/example
|
||||
touch core/example/__init__.py
|
||||
touch core/example/model.py
|
||||
touch core/example/schema.py
|
||||
touch core/example/service.py
|
||||
touch core/example/api.py
|
||||
```
|
||||
|
||||
#### 2. 定义模型 (model.py)
|
||||
|
||||
```python
|
||||
from sqlalchemy import Column, String, Boolean
|
||||
from app.base_model import BaseModel
|
||||
|
||||
class Example(BaseModel):
|
||||
__tablename__ = "core_example"
|
||||
|
||||
name = Column(String(100), nullable=False, comment="名称")
|
||||
description = Column(String(500), comment="描述")
|
||||
is_active = Column(Boolean, default=True, comment="是否激活")
|
||||
```
|
||||
|
||||
#### 3. 定义 Schema (schema.py)
|
||||
|
||||
```python
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from typing import Optional
|
||||
from datetime import datetime
|
||||
|
||||
class ExampleBase(BaseModel):
|
||||
name: str
|
||||
description: Optional[str] = None
|
||||
is_active: bool = True
|
||||
|
||||
class ExampleCreate(ExampleBase):
|
||||
pass
|
||||
|
||||
class ExampleUpdate(BaseModel):
|
||||
name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
is_active: Optional[bool] = None
|
||||
|
||||
class ExampleResponse(ExampleBase):
|
||||
id: str
|
||||
sort: int = 0
|
||||
is_deleted: bool = False
|
||||
sys_create_datetime: Optional[datetime] = None
|
||||
sys_update_datetime: Optional[datetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
```
|
||||
|
||||
#### 4. 定义服务 (service.py)
|
||||
|
||||
```python
|
||||
from app.base_service import BaseService
|
||||
from core.example.model import Example
|
||||
from core.example.schema import ExampleCreate, ExampleUpdate
|
||||
|
||||
class ExampleService(BaseService[Example, ExampleCreate, ExampleUpdate]):
|
||||
model = Example
|
||||
```
|
||||
|
||||
#### 5. 定义 API (api.py)
|
||||
|
||||
```python
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from app.database import get_db
|
||||
from app.base_schema import PaginatedResponse, ResponseModel
|
||||
from core.example.schema import ExampleCreate, ExampleUpdate, ExampleResponse
|
||||
from core.example.service import ExampleService
|
||||
|
||||
router = APIRouter(prefix="/example", tags=["示例管理"])
|
||||
|
||||
@router.post("", response_model=ExampleResponse, summary="创建")
|
||||
async def create(data: ExampleCreate, db: AsyncSession = Depends(get_db)):
|
||||
return await ExampleService.create(db=db, data=data)
|
||||
|
||||
@router.get("", response_model=PaginatedResponse[ExampleResponse], summary="获取列表")
|
||||
async def get_list(
|
||||
page: int = Query(default=1, ge=1),
|
||||
page_size: int = Query(default=20, ge=1, le=100, alias="pageSize"),
|
||||
db: AsyncSession = Depends(get_db)
|
||||
):
|
||||
items, total = await ExampleService.get_list(db, page=page, page_size=page_size)
|
||||
return PaginatedResponse(items=items, total=total)
|
||||
|
||||
@router.get("/{record_id}", response_model=ExampleResponse, summary="获取详情")
|
||||
async def get_by_id(record_id: str, db: AsyncSession = Depends(get_db)):
|
||||
result = await ExampleService.get_by_id(db, record_id=record_id)
|
||||
if not result:
|
||||
raise HTTPException(status_code=404, detail="记录不存在")
|
||||
return result
|
||||
|
||||
@router.put("/{record_id}", response_model=ExampleResponse, summary="更新")
|
||||
async def update(record_id: str, data: ExampleUpdate, db: AsyncSession = Depends(get_db)):
|
||||
result = await ExampleService.update(db, record_id=record_id, data=data)
|
||||
if not result:
|
||||
raise HTTPException(status_code=404, detail="记录不存在")
|
||||
return result
|
||||
|
||||
@router.delete("/{record_id}", response_model=ResponseModel, summary="删除")
|
||||
async def delete(record_id: str, db: AsyncSession = Depends(get_db)):
|
||||
success = await ExampleService.delete(db, record_id=record_id)
|
||||
if not success:
|
||||
raise HTTPException(status_code=404, detail="记录不存在")
|
||||
return ResponseModel(message="删除成功")
|
||||
```
|
||||
|
||||
#### 6. 注册路由
|
||||
|
||||
在 `core/router.py` 中添加：
|
||||
|
||||
```python
|
||||
from core.example.api import router as example_router
|
||||
|
||||
router.include_router(example_router)
|
||||
```
|
||||
|
||||
#### 7. 生成数据库迁移
|
||||
|
||||
```bash
|
||||
alembic revision --autogenerate -m "add example table"
|
||||
alembic upgrade head
|
||||
```
|
||||
|
||||
## 核心功能
|
||||
|
||||
### BaseModel
|
||||
|
||||
所有模型继承自 `BaseModel`，自动包含以下字段：
|
||||
|
||||
- `id`: UUID 主键
|
||||
- `sort`: 排序字段
|
||||
- `is_deleted`: 软删除标记
|
||||
- `sys_create_datetime`: 创建时间
|
||||
- `sys_update_datetime`: 更新时间
|
||||
- `sys_creator_id`: 创建人ID
|
||||
- `sys_modifier_id`: 修改人ID
|
||||
|
||||
### BaseService
|
||||
|
||||
提供通用 CRUD 操作：
|
||||
|
||||
- `create()`: 创建记录
|
||||
- `get_by_id()`: 根据ID获取
|
||||
- `get_list()`: 分页查询
|
||||
- `update()`: 更新记录
|
||||
- `delete()`: 删除记录（软删除/硬删除）
|
||||
- `check_unique()`: 唯一性检查
|
||||
- `export_to_excel()`: 导出Excel
|
||||
- `import_from_excel()`: 导入Excel
|
||||
|
||||
### 缓存支持
|
||||
|
||||
使用 Redis 缓存，继承 `CacheService` 获得缓存功能：
|
||||
|
||||
```python
|
||||
from app.cache_service import CacheService
|
||||
|
||||
class ExampleService(CacheService[Example, ExampleCreate, ExampleUpdate]):
|
||||
model = Example
|
||||
cache_prefix = "example"
|
||||
cache_ttl = 3600 # 1小时
|
||||
```
|
||||
|
||||
## 环境配置
|
||||
|
||||
项目支持多环境配置：
|
||||
|
||||
- `env/dev.env`: 开发环境
|
||||
- `env/uat.env`: UAT环境
|
||||
- `env/prod.env`: 生产环境
|
||||
|
||||
通过环境变量 `ENV` 切换：
|
||||
|
||||
```bash
|
||||
export ENV=prod # 使用生产环境配置
|
||||
python main.py
|
||||
```
|
||||
|
||||
## API 规范
|
||||
|
||||
### 响应格式
|
||||
|
||||
成功响应：
|
||||
|
||||
```json
|
||||
{
|
||||
"code": 200,
|
||||
"message": "success",
|
||||
"data": {...}
|
||||
}
|
||||
```
|
||||
|
||||
分页响应：
|
||||
|
||||
```json
|
||||
{
|
||||
"items": [...],
|
||||
"total": 100
|
||||
}
|
||||
```
|
||||
|
||||
错误响应：
|
||||
|
||||
```json
|
||||
{
|
||||
"detail": "错误信息"
|
||||
}
|
||||
```
|
||||
|
||||
### 路由命名规范
|
||||
|
||||
- 使用小写短横线：`/api/core/user-profile`
|
||||
- 静态路由在前：`/api/core/menu/check/name`
|
||||
- 动态路由在后：`/api/core/menu/{menu_id}`
|
||||
|
||||
## 常见问题
|
||||
|
||||
### 1. 迁移文件为空
|
||||
|
||||
确保 `alembic/env.py` 中的 `auto_import_models()` 函数正确扫描了所有模型文件。
|
||||
|
||||
### 2. 路由重定向 307
|
||||
|
||||
检查路由定义，使用 `@router.post("")` 而不是 `@router.post("/")`。
|
||||
|
||||
### 3. 数据库连接失败
|
||||
|
||||
检查 `env/dev.env` 中的 `DATABASE_URL` 配置是否正确。
|
||||
|
||||
# WeasyPrint 安装与配置指南
|
||||
|
||||
本文档介绍如何在不同操作系统上安装和配置 WeasyPrint 及其依赖。
|
||||
|
||||
## 目录
|
||||
|
||||
- [macOS](#macos)
|
||||
- [Linux (Ubuntu/Debian)](#linux-ubuntudebian)
|
||||
- [Linux (CentOS/RHEL)](#linux-centosrhel)
|
||||
- [Windows](#windows)
|
||||
- [验证安装](#验证安装)
|
||||
- [常见问题](#常见问题)
|
||||
|
||||
---
|
||||
|
||||
## macOS
|
||||
|
||||
### 1. 安装系统依赖
|
||||
|
||||
使用 Homebrew 安装所需的系统库：
|
||||
|
||||
```bash
|
||||
brew install pango glib gobject-introspection harfbuzz cairo fontconfig freetype
|
||||
```
|
||||
|
||||
### 2. 安装 Python 依赖
|
||||
|
||||
```bash
|
||||
pip install weasyprint==62.3
|
||||
```
|
||||
|
||||
### 3. 配置环境变量
|
||||
|
||||
WeasyPrint 需要能够找到系统库。将以下内容添加到 `~/.zshrc`（如果使用 bash，则添加到 `~/.bash_profile`）：
|
||||
|
||||
```bash
|
||||
export DYLD_LIBRARY_PATH="/opt/homebrew/opt/glib/lib:/opt/homebrew/opt/pango/lib:/opt/homebrew/opt/harfbuzz/lib:/opt/homebrew/opt/cairo/lib:/opt/homebrew/opt/fontconfig/lib:/opt/homebrew/opt/freetype/lib:$DYLD_LIBRARY_PATH"
|
||||
```
|
||||
|
||||
**自动添加方法：**
|
||||
|
||||
```bash
|
||||
echo 'export DYLD_LIBRARY_PATH="/opt/homebrew/opt/glib/lib:/opt/homebrew/opt/pango/lib:/opt/homebrew/opt/harfbuzz/lib:/opt/homebrew/opt/cairo/lib:/opt/homebrew/opt/fontconfig/lib:/opt/homebrew/opt/freetype/lib:$DYLD_LIBRARY_PATH"' >> ~/.zshrc
|
||||
source ~/.zshrc
|
||||
```
|
||||
|
||||
### 4. 重启终端或 IDE
|
||||
|
||||
关闭并重新打开终端窗口，或者重启 IDE，使环境变量生效。
|
||||
|
||||
---
|
||||
|
||||
## Linux (Ubuntu/Debian)
|
||||
|
||||
### 1. 安装系统依赖
|
||||
|
||||
```bash
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y \
|
||||
libpango-1.0-0 \
|
||||
libpangocairo-1.0-0 \
|
||||
libgdk-pixbuf2.0-0 \
|
||||
libffi-dev \
|
||||
libcairo2 \
|
||||
libglib2.0-0 \
|
||||
libharfbuzz0b \
|
||||
libfontconfig1 \
|
||||
libfreetype6
|
||||
```
|
||||
|
||||
### 2. 安装 Python 依赖
|
||||
|
||||
```bash
|
||||
pip install weasyprint==62.3
|
||||
```
|
||||
|
||||
### 3. 配置环境变量（通常不需要）
|
||||
|
||||
在 Linux 上，系统库通常已经在标准路径中，不需要额外配置环境变量。
|
||||
|
||||
如果遇到库加载问题，可以添加到 `~/.bashrc`：
|
||||
|
||||
```bash
|
||||
export LD_LIBRARY_PATH="/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH"
|
||||
source ~/.bashrc
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Linux (CentOS/RHEL)
|
||||
|
||||
### 1. 安装系统依赖
|
||||
|
||||
```bash
|
||||
sudo yum install -y \
|
||||
pango \
|
||||
pango-devel \
|
||||
cairo \
|
||||
cairo-devel \
|
||||
glib2 \
|
||||
glib2-devel \
|
||||
harfbuzz \
|
||||
harfbuzz-devel \
|
||||
fontconfig \
|
||||
fontconfig-devel \
|
||||
freetype \
|
||||
freetype-devel \
|
||||
libffi-devel
|
||||
```
|
||||
|
||||
或者使用 dnf（CentOS 8+）：
|
||||
|
||||
```bash
|
||||
sudo dnf install -y \
|
||||
pango \
|
||||
pango-devel \
|
||||
cairo \
|
||||
cairo-devel \
|
||||
glib2 \
|
||||
glib2-devel \
|
||||
harfbuzz \
|
||||
harfbuzz-devel \
|
||||
fontconfig \
|
||||
fontconfig-devel \
|
||||
freetype \
|
||||
freetype-devel \
|
||||
libffi-devel
|
||||
```
|
||||
|
||||
### 2. 安装 Python 依赖
|
||||
|
||||
```bash
|
||||
pip install weasyprint==62.3
|
||||
```
|
||||
|
||||
### 3. 配置环境变量（如果需要）
|
||||
|
||||
```bash
|
||||
export LD_LIBRARY_PATH="/usr/lib64:$LD_LIBRARY_PATH"
|
||||
echo 'export LD_LIBRARY_PATH="/usr/lib64:$LD_LIBRARY_PATH"' >> ~/.bashrc
|
||||
source ~/.bashrc
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Windows
|
||||
|
||||
### 方法 1：使用 GTK3 Runtime（推荐）
|
||||
|
||||
1. **下载并安装 GTK3 Runtime**
|
||||
|
||||
访问 [GTK for Windows Runtime](https://github.com/tschoonj/GTK-for-Windows-Runtime-Environment-Installer/releases)，下载最新版本的安装程序（例如 `gtk3-runtime-3.24.31-2022-01-04-ts-win64.exe`）。
|
||||
|
||||
2. **运行安装程序**
|
||||
|
||||
双击安装程序，按照提示完成安装。默认安装路径为 `C:\Program Files\GTK3-Runtime Win64`。
|
||||
|
||||
3. **添加到系统 PATH**
|
||||
|
||||
- 右键点击"此电脑" → "属性" → "高级系统设置" → "环境变量"
|
||||
- 在"系统变量"中找到 `Path`，点击"编辑"
|
||||
- 添加以下路径：
|
||||
```
|
||||
C:\Program Files\GTK3-Runtime Win64\bin
|
||||
```
|
||||
- 点击"确定"保存
|
||||
|
||||
4. **安装 Python 依赖**
|
||||
|
||||
```cmd
|
||||
pip install weasyprint==62.3
|
||||
```
|
||||
|
||||
5. **重启命令提示符或 PowerShell**
|
||||
|
||||
### 方法 2：使用 MSYS2（开发者推荐）
|
||||
|
||||
1. **安装 MSYS2**
|
||||
|
||||
访问 [MSYS2 官网](https://www.msys2.org/)，下载并安装 MSYS2。
|
||||
|
||||
2. **安装依赖包**
|
||||
|
||||
打开 MSYS2 终端，运行：
|
||||
|
||||
```bash
|
||||
pacman -S mingw-w64-x86_64-pango mingw-w64-x86_64-cairo mingw-w64-x86_64-glib2
|
||||
```
|
||||
|
||||
3. **添加到系统 PATH**
|
||||
|
||||
将 MSYS2 的 bin 目录添加到系统 PATH：
|
||||
```
|
||||
C:\msys64\mingw64\bin
|
||||
```
|
||||
|
||||
4. **安装 Python 依赖**
|
||||
|
||||
```cmd
|
||||
pip install weasyprint==62.3
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 验证安装
|
||||
|
||||
运行以下 Python 代码验证 WeasyPrint 是否正确安装：
|
||||
|
||||
```python
|
||||
from weasyprint import HTML
|
||||
|
||||
html_content = """
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<style>
|
||||
body { font-family: "PingFang SC", "Microsoft YaHei", sans-serif; }
|
||||
h1 { color: #333; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>测试中文字体</h1>
|
||||
<p>这是一个测试文档，用于验证 WeasyPrint 是否正确安装。</p>
|
||||
<p>Test English text and 中文文本。</p>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
try:
|
||||
pdf_bytes = HTML(string=html_content).write_pdf()
|
||||
with open('test.pdf', 'wb') as f:
|
||||
f.write(pdf_bytes)
|
||||
print("✓ WeasyPrint 安装成功！已生成 test.pdf")
|
||||
except Exception as e:
|
||||
print(f"✗ WeasyPrint 安装失败：{e}")
|
||||
```
|
||||
|
||||
如果成功，会在当前目录生成 `test.pdf` 文件。
|
||||
|
||||
---
|
||||
|
||||
## 常见问题
|
||||
|
||||
### 1. macOS: `OSError: cannot load library 'libgobject-2.0-0'`
|
||||
|
||||
**原因**：环境变量未正确设置。
|
||||
|
||||
**解决方案**：
|
||||
- 确保已添加环境变量到 `~/.zshrc`
|
||||
- 重启终端或运行 `source ~/.zshrc`
|
||||
- 如果使用 IDE，需要重启 IDE
|
||||
|
||||
### 2. Linux: `ImportError: cannot import name 'HTML' from 'weasyprint'`
|
||||
|
||||
**原因**：系统依赖未安装。
|
||||
|
||||
**解决方案**：
|
||||
```bash
|
||||
sudo apt-get install -y libpango-1.0-0 libpangocairo-1.0-0 libgdk-pixbuf2.0-0
|
||||
```
|
||||
|
||||
### 3. Windows: `OSError: no library called "cairo" was found`
|
||||
|
||||
**原因**：GTK3 Runtime 未安装或未添加到 PATH。
|
||||
|
||||
**解决方案**：
|
||||
- 确保已安装 GTK3 Runtime
|
||||
- 检查 `C:\Program Files\GTK3-Runtime Win64\bin` 是否在系统 PATH 中
|
||||
- 重启命令提示符
|
||||
|
||||
### 4. 中文字体显示为方块
|
||||
|
||||
**原因**：系统缺少中文字体。
|
||||
|
||||
**解决方案**：
|
||||
|
||||
**macOS**：
|
||||
```bash
|
||||
# 系统自带中文字体，通常不需要额外安装
|
||||
```
|
||||
|
||||
**Linux**：
|
||||
```bash
|
||||
sudo apt-get install fonts-noto-cjk fonts-wqy-zenhei
|
||||
```
|
||||
|
||||
**Windows**：
|
||||
- 确保系统已安装中文字体（如微软雅黑、宋体等）
|
||||
- Windows 10/11 默认已包含中文字体
|
||||
|
||||
### 5. Conda 环境中的问题
|
||||
|
||||
如果在 Conda 环境中遇到库加载问题，尝试：
|
||||
|
||||
```bash
|
||||
# 安装 conda-forge 版本
|
||||
conda install -c conda-forge weasyprint
|
||||
```
|
||||
|
||||
或者确保环境变量在激活 Conda 环境后仍然有效。
|
||||
|
||||
---
|
||||
|
||||
## 项目启动
|
||||
|
||||
配置完成后，启动后端服务：
|
||||
|
||||
```bash
|
||||
cd /path/to/backend-fastapi
|
||||
python -m uvicorn main:app --reload
|
||||
```
|
||||
|
||||
如果一切正常，服务应该能够成功启动，并且 PDF 预览功能可以正常使用。
|
||||
|
||||
---
|
||||
|
||||
## 参考链接
|
||||
|
||||
- [WeasyPrint 官方文档](https://doc.courtbouillon.org/weasyprint/stable/)
|
||||
- [WeasyPrint 安装指南](https://doc.courtbouillon.org/weasyprint/stable/first_steps.html#installation)
|
||||
- [WeasyPrint 故障排除](https://doc.courtbouillon.org/weasyprint/stable/first_steps.html#troubleshooting)
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
"""
|
||||
AI 平台模块
|
||||
|
||||
提供类似 Dify 的 LLM 应用开发能力：
|
||||
- 模型配置管理
|
||||
- 对话功能
|
||||
- 工作流编排
|
||||
- 智能体
|
||||
- 工具系统
|
||||
"""
|
||||
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
AI 平台 API
|
||||
"""
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,295 @@
|
||||
"""
|
||||
AI 应用 API
|
||||
"""
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from sqlalchemy import select, func
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.database import get_db
|
||||
from app.base_schema import PaginatedResponse, ResponseModel
|
||||
from ai_platform.models import AIApp, LLMModel
|
||||
from ai_platform.schemas.app_schema import (
|
||||
AppCreate,
|
||||
AppUpdate,
|
||||
AppResponse,
|
||||
AppListResponse,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/apps", tags=["AI-应用"])
|
||||
|
||||
|
||||
@router.get("", response_model=PaginatedResponse[AppListResponse], summary="应用列表")
|
||||
async def list_apps(
|
||||
name: Optional[str] = Query(None, description="名称"),
|
||||
app_type: Optional[str] = Query(None, description="类型"),
|
||||
status: Optional[str] = Query(None, description="状态"),
|
||||
page: int = Query(1, ge=1, description="页码"),
|
||||
page_size: int = Query(20, ge=1, le=100, alias="pageSize", description="每页数量"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""获取应用列表"""
|
||||
query = select(AIApp).where(AIApp.is_deleted == False)
|
||||
|
||||
if name:
|
||||
query = query.where(AIApp.name.ilike(f"%{name}%"))
|
||||
if app_type:
|
||||
query = query.where(AIApp.app_type == app_type)
|
||||
if status:
|
||||
query = query.where(AIApp.status == status)
|
||||
|
||||
# 获取总数
|
||||
count_result = await db.execute(select(func.count()).select_from(query.subquery()))
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
# 分页
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(AIApp.sort.desc(), AIApp.sys_create_datetime.desc())
|
||||
query = query.offset(offset).limit(page_size)
|
||||
|
||||
result = await db.execute(query)
|
||||
apps = result.scalars().all()
|
||||
|
||||
items = []
|
||||
for app in apps:
|
||||
model_name = ""
|
||||
if app.model_id:
|
||||
model_result = await db.execute(
|
||||
select(LLMModel).where(LLMModel.id == app.model_id)
|
||||
)
|
||||
model = model_result.scalar_one_or_none()
|
||||
model_name = model.display_name if model else ""
|
||||
|
||||
items.append({
|
||||
"id": app.id,
|
||||
"name": app.name,
|
||||
"code": app.code,
|
||||
"description": app.description or "",
|
||||
"icon": app.icon or "",
|
||||
"app_type": app.app_type or "chat",
|
||||
"status": app.status or "draft",
|
||||
"model_name": model_name,
|
||||
"is_public": app.is_public or False,
|
||||
"conversation_count": app.conversation_count or 0,
|
||||
"message_count": app.message_count or 0,
|
||||
"sys_create_datetime": app.sys_create_datetime,
|
||||
})
|
||||
|
||||
return PaginatedResponse(items=items, total=total)
|
||||
|
||||
|
||||
@router.get("/published", response_model=List[AppListResponse], summary="获取已发布应用")
|
||||
async def list_published_apps(db: AsyncSession = Depends(get_db)):
|
||||
"""获取已发布的应用列表（用于用户选择）"""
|
||||
query = select(AIApp).where(
|
||||
AIApp.is_deleted == False,
|
||||
AIApp.status == "published"
|
||||
).order_by(AIApp.sort.desc())
|
||||
|
||||
result = await db.execute(query)
|
||||
apps = result.scalars().all()
|
||||
|
||||
items = []
|
||||
for app in apps:
|
||||
model_name = ""
|
||||
if app.model_id:
|
||||
model_result = await db.execute(
|
||||
select(LLMModel).where(LLMModel.id == app.model_id)
|
||||
)
|
||||
model = model_result.scalar_one_or_none()
|
||||
model_name = model.display_name if model else ""
|
||||
|
||||
items.append({
|
||||
"id": app.id,
|
||||
"name": app.name,
|
||||
"code": app.code,
|
||||
"description": app.description or "",
|
||||
"icon": app.icon or "",
|
||||
"app_type": app.app_type or "chat",
|
||||
"status": app.status or "draft",
|
||||
"model_name": model_name,
|
||||
"is_public": app.is_public or False,
|
||||
"conversation_count": app.conversation_count or 0,
|
||||
"message_count": app.message_count or 0,
|
||||
"sys_create_datetime": app.sys_create_datetime,
|
||||
})
|
||||
|
||||
return items
|
||||
|
||||
|
||||
@router.get("/code/{code}", response_model=AppResponse, summary="根据编码获取应用")
|
||||
async def get_app_by_code(code: str, db: AsyncSession = Depends(get_db)):
|
||||
"""根据编码获取应用"""
|
||||
result = await db.execute(
|
||||
select(AIApp).where(AIApp.code == code, AIApp.is_deleted == False)
|
||||
)
|
||||
app = result.scalar_one_or_none()
|
||||
if not app:
|
||||
raise HTTPException(status_code=404, detail="应用不存在")
|
||||
|
||||
return await _build_app_response(app, db)
|
||||
|
||||
|
||||
@router.get("/{app_id}", response_model=AppResponse, summary="应用详情")
|
||||
async def get_app(app_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""获取应用详情"""
|
||||
result = await db.execute(
|
||||
select(AIApp).where(AIApp.id == app_id, AIApp.is_deleted == False)
|
||||
)
|
||||
app = result.scalar_one_or_none()
|
||||
if not app:
|
||||
raise HTTPException(status_code=404, detail="应用不存在")
|
||||
|
||||
return await _build_app_response(app, db)
|
||||
|
||||
|
||||
@router.post("", response_model=AppResponse, summary="创建应用")
|
||||
async def create_app(data: AppCreate, db: AsyncSession = Depends(get_db)):
|
||||
"""创建应用"""
|
||||
# 检查编码是否重复
|
||||
exists_result = await db.execute(
|
||||
select(AIApp).where(AIApp.code == data.code, AIApp.is_deleted == False)
|
||||
)
|
||||
if exists_result.scalar_one_or_none():
|
||||
raise HTTPException(status_code=400, detail=f"应用编码 {data.code} 已存在")
|
||||
|
||||
# 验证模型
|
||||
if data.model_id:
|
||||
model_result = await db.execute(
|
||||
select(LLMModel).where(LLMModel.id == data.model_id, LLMModel.is_deleted == False)
|
||||
)
|
||||
if not model_result.scalar_one_or_none():
|
||||
raise HTTPException(status_code=400, detail="模型不存在")
|
||||
|
||||
app = AIApp(**data.model_dump())
|
||||
db.add(app)
|
||||
await db.commit()
|
||||
await db.refresh(app)
|
||||
|
||||
return await _build_app_response(app, db)
|
||||
|
||||
|
||||
@router.put("/{app_id}", response_model=AppResponse, summary="更新应用")
|
||||
async def update_app(
|
||||
app_id: str,
|
||||
data: AppUpdate,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""更新应用"""
|
||||
result = await db.execute(
|
||||
select(AIApp).where(AIApp.id == app_id, AIApp.is_deleted == False)
|
||||
)
|
||||
app = result.scalar_one_or_none()
|
||||
if not app:
|
||||
raise HTTPException(status_code=404, detail="应用不存在")
|
||||
|
||||
update_data = data.model_dump(exclude_unset=True)
|
||||
|
||||
# 验证模型
|
||||
if "model_id" in update_data and update_data["model_id"]:
|
||||
model_result = await db.execute(
|
||||
select(LLMModel).where(
|
||||
LLMModel.id == update_data["model_id"],
|
||||
LLMModel.is_deleted == False
|
||||
)
|
||||
)
|
||||
if not model_result.scalar_one_or_none():
|
||||
raise HTTPException(status_code=400, detail="模型不存在")
|
||||
|
||||
for key, value in update_data.items():
|
||||
setattr(app, key, value)
|
||||
|
||||
await db.commit()
|
||||
await db.refresh(app)
|
||||
|
||||
return await _build_app_response(app, db)
|
||||
|
||||
|
||||
@router.delete("/{app_id}", response_model=ResponseModel, summary="删除应用")
|
||||
async def delete_app(app_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""删除应用"""
|
||||
result = await db.execute(
|
||||
select(AIApp).where(AIApp.id == app_id, AIApp.is_deleted == False)
|
||||
)
|
||||
app = result.scalar_one_or_none()
|
||||
if not app:
|
||||
raise HTTPException(status_code=404, detail="应用不存在")
|
||||
|
||||
app.is_deleted = True
|
||||
await db.commit()
|
||||
|
||||
return ResponseModel(message="删除成功")
|
||||
|
||||
|
||||
@router.post("/{app_id}/publish", response_model=AppResponse, summary="发布应用")
|
||||
async def publish_app(app_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""发布应用"""
|
||||
result = await db.execute(
|
||||
select(AIApp).where(AIApp.id == app_id, AIApp.is_deleted == False)
|
||||
)
|
||||
app = result.scalar_one_or_none()
|
||||
if not app:
|
||||
raise HTTPException(status_code=404, detail="应用不存在")
|
||||
|
||||
app.status = "published"
|
||||
await db.commit()
|
||||
await db.refresh(app)
|
||||
|
||||
return await _build_app_response(app, db)
|
||||
|
||||
|
||||
@router.post("/{app_id}/disable", response_model=AppResponse, summary="停用应用")
|
||||
async def disable_app(app_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""停用应用"""
|
||||
result = await db.execute(
|
||||
select(AIApp).where(AIApp.id == app_id, AIApp.is_deleted == False)
|
||||
)
|
||||
app = result.scalar_one_or_none()
|
||||
if not app:
|
||||
raise HTTPException(status_code=404, detail="应用不存在")
|
||||
|
||||
app.status = "disabled"
|
||||
await db.commit()
|
||||
await db.refresh(app)
|
||||
|
||||
return await _build_app_response(app, db)
|
||||
|
||||
|
||||
async def _build_app_response(app: AIApp, db: AsyncSession) -> dict:
|
||||
"""构建应用输出"""
|
||||
model_name = ""
|
||||
if app.model_id:
|
||||
model_result = await db.execute(
|
||||
select(LLMModel).where(LLMModel.id == app.model_id)
|
||||
)
|
||||
model = model_result.scalar_one_or_none()
|
||||
model_name = model.display_name if model else ""
|
||||
|
||||
return {
|
||||
"id": app.id,
|
||||
"name": app.name,
|
||||
"code": app.code,
|
||||
"description": app.description or "",
|
||||
"icon": app.icon or "",
|
||||
"app_type": app.app_type or "chat",
|
||||
"status": app.status or "draft",
|
||||
"model_id": app.model_id,
|
||||
"model_name": model_name,
|
||||
"system_prompt": app.system_prompt or "",
|
||||
"temperature": app.temperature or 0.7,
|
||||
"top_p": app.top_p or 1.0,
|
||||
"max_tokens": app.max_tokens or 2048,
|
||||
"opening_statement": app.opening_statement or "",
|
||||
"suggested_questions": app.suggested_questions or [],
|
||||
"workflow_definition": app.workflow_definition or {},
|
||||
"is_public": app.is_public or False,
|
||||
"conversation_count": app.conversation_count or 0,
|
||||
"message_count": app.message_count or 0,
|
||||
"sort": app.sort or 0,
|
||||
"sys_create_datetime": app.sys_create_datetime,
|
||||
"sys_update_datetime": app.sys_update_datetime,
|
||||
}
|
||||
@@ -0,0 +1,337 @@
|
||||
"""
|
||||
对话 API
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from fastapi.responses import StreamingResponse
|
||||
from sqlalchemy import select, func
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.database import get_db
|
||||
from app.base_schema import PaginatedResponse, ResponseModel
|
||||
from utils.context import get_current_user_id_from_context
|
||||
from ai_platform.models import Conversation, Message, AIApp
|
||||
from ai_platform.schemas.chat_schema import (
|
||||
ConversationCreate,
|
||||
ConversationUpdate,
|
||||
ConversationResponse,
|
||||
ConversationListResponse,
|
||||
MessageResponse,
|
||||
SendMessageInput,
|
||||
MessageFeedbackInput,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/chat", tags=["AI-对话"])
|
||||
|
||||
|
||||
@router.get("/conversations", response_model=PaginatedResponse[ConversationListResponse], summary="对话列表")
|
||||
async def list_conversations(
|
||||
app_id: str = Query(..., description="应用 ID"),
|
||||
page: int = Query(1, ge=1, description="页码"),
|
||||
page_size: int = Query(20, ge=1, le=100, alias="pageSize", description="每页数量"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""获取对话列表"""
|
||||
query = select(Conversation).where(
|
||||
Conversation.app_id == app_id,
|
||||
Conversation.is_deleted == False
|
||||
)
|
||||
|
||||
# 获取总数
|
||||
count_result = await db.execute(select(func.count()).select_from(query.subquery()))
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
# 分页
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(
|
||||
Conversation.is_pinned.desc(),
|
||||
Conversation.sys_update_datetime.desc()
|
||||
)
|
||||
query = query.offset(offset).limit(page_size)
|
||||
|
||||
result = await db.execute(query)
|
||||
conversations = result.scalars().all()
|
||||
|
||||
items = [
|
||||
{
|
||||
"id": c.id,
|
||||
"title": c.title or "",
|
||||
"message_count": c.message_count or 0,
|
||||
"is_pinned": c.is_pinned or False,
|
||||
"sys_update_datetime": c.sys_update_datetime,
|
||||
}
|
||||
for c in conversations
|
||||
]
|
||||
|
||||
return PaginatedResponse(items=items, total=total)
|
||||
|
||||
|
||||
@router.post("/conversations", response_model=ConversationResponse, summary="创建对话")
|
||||
async def create_conversation(data: ConversationCreate, db: AsyncSession = Depends(get_db)):
|
||||
"""创建对话"""
|
||||
# 验证应用
|
||||
app_result = await db.execute(
|
||||
select(AIApp).where(AIApp.id == data.app_id, AIApp.is_deleted == False)
|
||||
)
|
||||
app = app_result.scalar_one_or_none()
|
||||
if not app:
|
||||
raise HTTPException(status_code=400, detail="应用不存在")
|
||||
|
||||
user_id = get_current_user_id_from_context()
|
||||
if not user_id:
|
||||
raise HTTPException(status_code=401, detail="未提供认证凭据")
|
||||
|
||||
conversation = Conversation(
|
||||
app_id=data.app_id,
|
||||
user_id=user_id,
|
||||
title=data.title or "新对话",
|
||||
)
|
||||
db.add(conversation)
|
||||
await db.commit()
|
||||
await db.refresh(conversation)
|
||||
|
||||
return await _build_conversation_response(conversation, db)
|
||||
|
||||
|
||||
@router.get("/conversations/{conversation_id}", response_model=ConversationResponse, summary="对话详情")
|
||||
async def get_conversation(conversation_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""获取对话详情"""
|
||||
result = await db.execute(
|
||||
select(Conversation).where(
|
||||
Conversation.id == conversation_id,
|
||||
Conversation.is_deleted == False
|
||||
)
|
||||
)
|
||||
conversation = result.scalar_one_or_none()
|
||||
if not conversation:
|
||||
raise HTTPException(status_code=404, detail="对话不存在")
|
||||
|
||||
return await _build_conversation_response(conversation, db)
|
||||
|
||||
|
||||
@router.put("/conversations/{conversation_id}", response_model=ConversationResponse, summary="更新对话")
|
||||
async def update_conversation(
|
||||
conversation_id: str,
|
||||
data: ConversationUpdate,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""更新对话"""
|
||||
result = await db.execute(
|
||||
select(Conversation).where(
|
||||
Conversation.id == conversation_id,
|
||||
Conversation.is_deleted == False
|
||||
)
|
||||
)
|
||||
conversation = result.scalar_one_or_none()
|
||||
if not conversation:
|
||||
raise HTTPException(status_code=404, detail="对话不存在")
|
||||
|
||||
if data.title is not None:
|
||||
conversation.title = data.title
|
||||
if data.is_pinned is not None:
|
||||
conversation.is_pinned = data.is_pinned
|
||||
|
||||
await db.commit()
|
||||
await db.refresh(conversation)
|
||||
|
||||
return await _build_conversation_response(conversation, db)
|
||||
|
||||
|
||||
@router.delete("/conversations/{conversation_id}", response_model=ResponseModel, summary="删除对话")
|
||||
async def delete_conversation(conversation_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""删除对话"""
|
||||
result = await db.execute(
|
||||
select(Conversation).where(
|
||||
Conversation.id == conversation_id,
|
||||
Conversation.is_deleted == False
|
||||
)
|
||||
)
|
||||
conversation = result.scalar_one_or_none()
|
||||
if not conversation:
|
||||
raise HTTPException(status_code=404, detail="对话不存在")
|
||||
|
||||
conversation.is_deleted = True
|
||||
await db.commit()
|
||||
|
||||
return ResponseModel(message="删除成功")
|
||||
|
||||
|
||||
@router.get("/conversations/{conversation_id}/messages", response_model=List[MessageResponse], summary="获取消息列表")
|
||||
async def get_messages(
|
||||
conversation_id: str,
|
||||
limit: int = Query(50, ge=1, le=200, description="限制数量"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""获取对话消息"""
|
||||
# 验证对话存在
|
||||
conv_result = await db.execute(
|
||||
select(Conversation).where(
|
||||
Conversation.id == conversation_id,
|
||||
Conversation.is_deleted == False
|
||||
)
|
||||
)
|
||||
if not conv_result.scalar_one_or_none():
|
||||
raise HTTPException(status_code=404, detail="对话不存在")
|
||||
|
||||
query = select(Message).where(
|
||||
Message.conversation_id == conversation_id,
|
||||
Message.is_deleted == False
|
||||
).order_by(Message.sys_create_datetime).limit(limit)
|
||||
|
||||
result = await db.execute(query)
|
||||
messages = result.scalars().all()
|
||||
|
||||
return [
|
||||
{
|
||||
"id": m.id,
|
||||
"role": m.role,
|
||||
"content": m.content or "",
|
||||
"status": m.status or "completed",
|
||||
"prompt_tokens": m.prompt_tokens or 0,
|
||||
"completion_tokens": m.completion_tokens or 0,
|
||||
"total_tokens": m.total_tokens or 0,
|
||||
"model_name": m.model_name or "",
|
||||
"latency": m.latency or 0,
|
||||
"error_message": m.error_message or "",
|
||||
"feedback": m.feedback or "",
|
||||
"sys_create_datetime": m.sys_create_datetime,
|
||||
}
|
||||
for m in messages
|
||||
]
|
||||
|
||||
|
||||
@router.post("/conversations/{conversation_id}/messages", response_model=MessageResponse, summary="发送消息")
|
||||
async def send_message(
|
||||
conversation_id: str,
|
||||
data: SendMessageInput,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""发送消息并获取 AI 回复"""
|
||||
from ai_platform.services.chat_service import ChatService
|
||||
|
||||
# 验证对话存在
|
||||
conv_result = await db.execute(
|
||||
select(Conversation).where(
|
||||
Conversation.id == conversation_id,
|
||||
Conversation.is_deleted == False
|
||||
)
|
||||
)
|
||||
conversation = conv_result.scalar_one_or_none()
|
||||
if not conversation:
|
||||
raise HTTPException(status_code=404, detail="对话不存在")
|
||||
|
||||
# 使用ChatService发送消息
|
||||
chat_service = ChatService(db)
|
||||
_, assistant_message = await chat_service.send_message(
|
||||
conversation_id=conversation_id,
|
||||
content=data.content,
|
||||
)
|
||||
|
||||
return {
|
||||
"id": assistant_message.id,
|
||||
"role": assistant_message.role,
|
||||
"content": assistant_message.content or "",
|
||||
"status": assistant_message.status or "completed",
|
||||
"prompt_tokens": assistant_message.prompt_tokens or 0,
|
||||
"completion_tokens": assistant_message.completion_tokens or 0,
|
||||
"total_tokens": assistant_message.total_tokens or 0,
|
||||
"model_name": assistant_message.model_name or "",
|
||||
"latency": assistant_message.latency or 0,
|
||||
"error_message": assistant_message.error_message or "",
|
||||
"feedback": assistant_message.feedback or "",
|
||||
"sys_create_datetime": assistant_message.sys_create_datetime,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/conversations/{conversation_id}/messages/stream", summary="流式发送消息")
|
||||
async def send_message_stream(
|
||||
conversation_id: str,
|
||||
data: SendMessageInput,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""发送消息并获取 AI 流式回复（SSE）"""
|
||||
from ai_platform.services.chat_service import ChatService
|
||||
|
||||
# 验证对话存在
|
||||
conv_result = await db.execute(
|
||||
select(Conversation).where(
|
||||
Conversation.id == conversation_id,
|
||||
Conversation.is_deleted == False
|
||||
)
|
||||
)
|
||||
conversation = conv_result.scalar_one_or_none()
|
||||
if not conversation:
|
||||
raise HTTPException(status_code=404, detail="对话不存在")
|
||||
|
||||
async def generate():
|
||||
chat_service = ChatService(db)
|
||||
async for event in chat_service.send_message_stream(
|
||||
conversation_id=conversation_id,
|
||||
content=data.content,
|
||||
):
|
||||
yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
return StreamingResponse(
|
||||
generate(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@router.post("/messages/{message_id}/feedback", response_model=ResponseModel, summary="消息反馈")
|
||||
async def message_feedback(
|
||||
message_id: str,
|
||||
data: MessageFeedbackInput,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""对消息进行反馈"""
|
||||
result = await db.execute(
|
||||
select(Message).where(
|
||||
Message.id == message_id,
|
||||
Message.is_deleted == False
|
||||
)
|
||||
)
|
||||
message = result.scalar_one_or_none()
|
||||
if not message:
|
||||
raise HTTPException(status_code=404, detail="消息不存在")
|
||||
|
||||
if data.feedback not in ("like", "dislike", ""):
|
||||
raise HTTPException(status_code=400, detail="无效的反馈类型")
|
||||
|
||||
message.feedback = data.feedback
|
||||
await db.commit()
|
||||
|
||||
return ResponseModel(message="反馈成功")
|
||||
|
||||
|
||||
async def _build_conversation_response(conversation: Conversation, db: AsyncSession) -> dict:
|
||||
"""构建对话输出"""
|
||||
app_name = ""
|
||||
if conversation.app_id:
|
||||
app_result = await db.execute(
|
||||
select(AIApp).where(AIApp.id == conversation.app_id)
|
||||
)
|
||||
app = app_result.scalar_one_or_none()
|
||||
app_name = app.name if app else ""
|
||||
|
||||
return {
|
||||
"id": conversation.id,
|
||||
"app_id": conversation.app_id,
|
||||
"app_name": app_name,
|
||||
"title": conversation.title or "",
|
||||
"message_count": conversation.message_count or 0,
|
||||
"total_tokens": conversation.total_tokens or 0,
|
||||
"is_pinned": conversation.is_pinned or False,
|
||||
"sort": conversation.sort or 0,
|
||||
"sys_create_datetime": conversation.sys_create_datetime,
|
||||
"sys_update_datetime": conversation.sys_update_datetime,
|
||||
}
|
||||
@@ -0,0 +1,295 @@
|
||||
"""
|
||||
LLM 模型 API
|
||||
"""
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Body
|
||||
from sqlalchemy import select, func
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.database import get_db
|
||||
from app.base_schema import PaginatedResponse, ResponseModel
|
||||
from ai_platform.models import LLMModel, LLMProvider
|
||||
from ai_platform.schemas.model_schema import (
|
||||
ModelCreate,
|
||||
ModelUpdate,
|
||||
ModelResponse,
|
||||
ModelListResponse,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/model", tags=["AI-模型"])
|
||||
|
||||
|
||||
@router.get("/list", response_model=PaginatedResponse[ModelListResponse], summary="模型列表")
|
||||
async def list_models(
|
||||
provider_id: Optional[str] = Query(None, description="提供商 ID"),
|
||||
model_type: Optional[str] = Query(None, description="模型类型"),
|
||||
is_active: Optional[bool] = Query(None, description="是否启用"),
|
||||
page: int = Query(1, ge=1, description="页码"),
|
||||
page_size: int = Query(20, ge=1, le=100, alias="pageSize", description="每页数量"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""获取模型列表"""
|
||||
query = select(LLMModel).where(LLMModel.is_deleted == False)
|
||||
|
||||
if provider_id:
|
||||
query = query.where(LLMModel.provider_id == provider_id)
|
||||
if model_type:
|
||||
query = query.where(LLMModel.model_type == model_type)
|
||||
if is_active is not None:
|
||||
query = query.where(LLMModel.is_active == is_active)
|
||||
|
||||
# 获取总数
|
||||
count_result = await db.execute(select(func.count()).select_from(query.subquery()))
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
# 分页
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(LLMModel.sort.desc(), LLMModel.sys_create_datetime.desc())
|
||||
query = query.offset(offset).limit(page_size)
|
||||
|
||||
result = await db.execute(query)
|
||||
models = result.scalars().all()
|
||||
|
||||
# 获取提供商名称
|
||||
items = []
|
||||
for model in models:
|
||||
provider_result = await db.execute(
|
||||
select(LLMProvider).where(LLMProvider.id == model.provider_id)
|
||||
)
|
||||
provider = provider_result.scalar_one_or_none()
|
||||
items.append({
|
||||
"id": model.id,
|
||||
"provider_id": model.provider_id,
|
||||
"provider_name": provider.name if provider else "",
|
||||
"model_name": model.model_name,
|
||||
"display_name": model.display_name,
|
||||
"model_type": model.model_type or "chat",
|
||||
"is_active": model.is_active,
|
||||
})
|
||||
|
||||
return PaginatedResponse(items=items, total=total)
|
||||
|
||||
|
||||
@router.get("/active", response_model=List[ModelListResponse], summary="获取可用模型列表")
|
||||
async def list_active_models(db: AsyncSession = Depends(get_db)):
|
||||
"""获取所有可用的模型（用于选择器）"""
|
||||
query = select(LLMModel).where(
|
||||
LLMModel.is_deleted == False,
|
||||
LLMModel.is_active == True
|
||||
).order_by(LLMModel.display_name)
|
||||
|
||||
result = await db.execute(query)
|
||||
models = result.scalars().all()
|
||||
|
||||
items = []
|
||||
for model in models:
|
||||
provider_result = await db.execute(
|
||||
select(LLMProvider).where(
|
||||
LLMProvider.id == model.provider_id,
|
||||
LLMProvider.is_deleted == False,
|
||||
LLMProvider.is_active == True
|
||||
)
|
||||
)
|
||||
provider = provider_result.scalar_one_or_none()
|
||||
if provider:
|
||||
items.append({
|
||||
"id": model.id,
|
||||
"provider_id": model.provider_id,
|
||||
"provider_name": provider.name,
|
||||
"model_name": model.model_name,
|
||||
"display_name": model.display_name,
|
||||
"model_type": model.model_type or "chat",
|
||||
"is_active": model.is_active,
|
||||
})
|
||||
|
||||
return items
|
||||
|
||||
|
||||
@router.get("/{model_id}", response_model=ModelResponse, summary="模型详情")
|
||||
async def get_model(model_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""获取模型详情"""
|
||||
result = await db.execute(
|
||||
select(LLMModel).where(
|
||||
LLMModel.id == model_id,
|
||||
LLMModel.is_deleted == False
|
||||
)
|
||||
)
|
||||
model = result.scalar_one_or_none()
|
||||
if not model:
|
||||
raise HTTPException(status_code=404, detail="模型不存在")
|
||||
|
||||
# 获取提供商名称
|
||||
provider_result = await db.execute(
|
||||
select(LLMProvider).where(LLMProvider.id == model.provider_id)
|
||||
)
|
||||
provider = provider_result.scalar_one_or_none()
|
||||
|
||||
return _build_model_response(model, provider)
|
||||
|
||||
|
||||
@router.post("", response_model=ModelResponse, summary="创建模型")
|
||||
async def create_model(data: ModelCreate, db: AsyncSession = Depends(get_db)):
|
||||
"""创建模型"""
|
||||
# 验证提供商
|
||||
provider_result = await db.execute(
|
||||
select(LLMProvider).where(
|
||||
LLMProvider.id == data.provider_id,
|
||||
LLMProvider.is_deleted == False
|
||||
)
|
||||
)
|
||||
provider = provider_result.scalar_one_or_none()
|
||||
if not provider:
|
||||
raise HTTPException(status_code=400, detail="提供商不存在")
|
||||
|
||||
# 检查模型名称是否重复
|
||||
exists_result = await db.execute(
|
||||
select(LLMModel).where(
|
||||
LLMModel.provider_id == data.provider_id,
|
||||
LLMModel.model_name == data.model_name,
|
||||
LLMModel.is_deleted == False
|
||||
)
|
||||
)
|
||||
if exists_result.scalar_one_or_none():
|
||||
raise HTTPException(status_code=400, detail=f"模型 {data.model_name} 已存在")
|
||||
|
||||
model = LLMModel(**data.model_dump())
|
||||
db.add(model)
|
||||
await db.commit()
|
||||
await db.refresh(model)
|
||||
|
||||
return _build_model_response(model, provider)
|
||||
|
||||
|
||||
@router.post("/batch", response_model=ResponseModel, summary="批量创建模型")
|
||||
async def batch_create_models(
|
||||
provider_id: str = Query(..., description="提供商 ID"),
|
||||
models: List[dict] = Body(..., description="模型列表"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""批量创建模型（从默认模型列表）"""
|
||||
provider_result = await db.execute(
|
||||
select(LLMProvider).where(
|
||||
LLMProvider.id == provider_id,
|
||||
LLMProvider.is_deleted == False
|
||||
)
|
||||
)
|
||||
provider = provider_result.scalar_one_or_none()
|
||||
if not provider:
|
||||
raise HTTPException(status_code=400, detail="提供商不存在")
|
||||
|
||||
created_count = 0
|
||||
for model_data in models:
|
||||
model_name = model_data.get("model_name", "")
|
||||
if not model_name:
|
||||
continue
|
||||
|
||||
# 跳过已存在的
|
||||
exists_result = await db.execute(
|
||||
select(LLMModel).where(
|
||||
LLMModel.provider_id == provider_id,
|
||||
LLMModel.model_name == model_name,
|
||||
LLMModel.is_deleted == False
|
||||
)
|
||||
)
|
||||
if exists_result.scalar_one_or_none():
|
||||
continue
|
||||
|
||||
model = LLMModel(
|
||||
provider_id=provider_id,
|
||||
model_name=model_name,
|
||||
display_name=model_data.get("display_name", model_name),
|
||||
model_type=model_data.get("model_type", "chat"),
|
||||
max_tokens=model_data.get("max_tokens", 4096),
|
||||
context_window=model_data.get("context_window", 4096),
|
||||
supports_vision=model_data.get("supports_vision", False),
|
||||
supports_function_call=model_data.get("supports_function_call", False),
|
||||
input_price=model_data.get("input_price", 0),
|
||||
output_price=model_data.get("output_price", 0),
|
||||
is_active=True,
|
||||
)
|
||||
db.add(model)
|
||||
created_count += 1
|
||||
|
||||
await db.commit()
|
||||
return ResponseModel(message=f"成功创建 {created_count} 个模型")
|
||||
|
||||
|
||||
@router.put("/{model_id}", response_model=ModelResponse, summary="更新模型")
|
||||
async def update_model(
|
||||
model_id: str,
|
||||
data: ModelUpdate,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""更新模型"""
|
||||
result = await db.execute(
|
||||
select(LLMModel).where(
|
||||
LLMModel.id == model_id,
|
||||
LLMModel.is_deleted == False
|
||||
)
|
||||
)
|
||||
model = result.scalar_one_or_none()
|
||||
if not model:
|
||||
raise HTTPException(status_code=404, detail="模型不存在")
|
||||
|
||||
update_data = data.model_dump(exclude_unset=True)
|
||||
for key, value in update_data.items():
|
||||
setattr(model, key, value)
|
||||
|
||||
await db.commit()
|
||||
await db.refresh(model)
|
||||
|
||||
# 获取提供商名称
|
||||
provider_result = await db.execute(
|
||||
select(LLMProvider).where(LLMProvider.id == model.provider_id)
|
||||
)
|
||||
provider = provider_result.scalar_one_or_none()
|
||||
|
||||
return _build_model_response(model, provider)
|
||||
|
||||
|
||||
@router.delete("/{model_id}", response_model=ResponseModel, summary="删除模型")
|
||||
async def delete_model(model_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""删除模型"""
|
||||
result = await db.execute(
|
||||
select(LLMModel).where(
|
||||
LLMModel.id == model_id,
|
||||
LLMModel.is_deleted == False
|
||||
)
|
||||
)
|
||||
model = result.scalar_one_or_none()
|
||||
if not model:
|
||||
raise HTTPException(status_code=404, detail="模型不存在")
|
||||
|
||||
model.is_deleted = True
|
||||
await db.commit()
|
||||
|
||||
return ResponseModel(message="删除成功")
|
||||
|
||||
|
||||
def _build_model_response(model: LLMModel, provider: Optional[LLMProvider] = None) -> dict:
|
||||
"""构建模型输出"""
|
||||
return {
|
||||
"id": model.id,
|
||||
"provider_id": model.provider_id,
|
||||
"provider_name": provider.name if provider else "",
|
||||
"model_name": model.model_name,
|
||||
"display_name": model.display_name,
|
||||
"model_type": model.model_type or "chat",
|
||||
"max_tokens": model.max_tokens or 4096,
|
||||
"context_window": model.context_window or 4096,
|
||||
"default_temperature": model.default_temperature or 0.7,
|
||||
"default_top_p": model.default_top_p or 1.0,
|
||||
"input_price": model.input_price or 0,
|
||||
"output_price": model.output_price or 0,
|
||||
"supports_vision": model.supports_vision or False,
|
||||
"supports_function_call": model.supports_function_call or False,
|
||||
"supports_streaming": model.supports_streaming if model.supports_streaming is not None else True,
|
||||
"is_active": model.is_active,
|
||||
"sort": model.sort or 0,
|
||||
"sys_create_datetime": model.sys_create_datetime,
|
||||
"sys_update_datetime": model.sys_update_datetime,
|
||||
}
|
||||
@@ -0,0 +1,270 @@
|
||||
"""
|
||||
LLM 提供商 API
|
||||
"""
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from sqlalchemy import select, func
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.database import get_db
|
||||
from app.base_schema import PaginatedResponse, ResponseModel
|
||||
from ai_platform.models import LLMProvider
|
||||
from ai_platform.schemas.provider_schema import (
|
||||
ProviderCreate,
|
||||
ProviderUpdate,
|
||||
ProviderResponse,
|
||||
ProviderListResponse,
|
||||
ProviderTypeResponse,
|
||||
)
|
||||
from ai_platform.providers import ProviderRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/provider", tags=["AI-提供商"])
|
||||
|
||||
|
||||
@router.get("/types", response_model=List[ProviderTypeResponse], summary="获取提供商类型列表")
|
||||
async def get_provider_types():
|
||||
"""获取所有支持的提供商类型"""
|
||||
return ProviderRegistry.get_all_types()
|
||||
|
||||
|
||||
@router.get("/list", response_model=PaginatedResponse[ProviderListResponse], summary="提供商列表")
|
||||
async def list_providers(
|
||||
name: Optional[str] = Query(None, description="名称"),
|
||||
provider_type: Optional[str] = Query(None, description="类型"),
|
||||
is_active: Optional[bool] = Query(None, description="是否启用"),
|
||||
page: int = Query(1, ge=1, description="页码"),
|
||||
page_size: int = Query(20, ge=1, le=100, alias="pageSize", description="每页数量"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""获取提供商列表"""
|
||||
query = select(LLMProvider).where(LLMProvider.is_deleted == False)
|
||||
|
||||
if name:
|
||||
query = query.where(LLMProvider.name.ilike(f"%{name}%"))
|
||||
if provider_type:
|
||||
query = query.where(LLMProvider.provider_type == provider_type)
|
||||
if is_active is not None:
|
||||
query = query.where(LLMProvider.is_active == is_active)
|
||||
|
||||
# 获取总数
|
||||
count_result = await db.execute(select(func.count()).select_from(query.subquery()))
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
# 分页
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(LLMProvider.sort.desc(), LLMProvider.sys_create_datetime.desc())
|
||||
query = query.offset(offset).limit(page_size)
|
||||
|
||||
result = await db.execute(query)
|
||||
items = result.scalars().all()
|
||||
|
||||
return PaginatedResponse(items=items, total=total)
|
||||
|
||||
|
||||
@router.get("/{provider_id}", response_model=ProviderResponse, summary="提供商详情")
|
||||
async def get_provider(provider_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""获取提供商详情"""
|
||||
result = await db.execute(
|
||||
select(LLMProvider).where(
|
||||
LLMProvider.id == provider_id,
|
||||
LLMProvider.is_deleted == False
|
||||
)
|
||||
)
|
||||
provider = result.scalar_one_or_none()
|
||||
if not provider:
|
||||
raise HTTPException(status_code=404, detail="提供商不存在")
|
||||
|
||||
return _build_provider_response(provider)
|
||||
|
||||
|
||||
@router.post("", response_model=ProviderResponse, summary="创建提供商")
|
||||
async def create_provider(data: ProviderCreate, db: AsyncSession = Depends(get_db)):
|
||||
"""创建提供商"""
|
||||
provider = LLMProvider(**data.model_dump())
|
||||
db.add(provider)
|
||||
await db.commit()
|
||||
await db.refresh(provider)
|
||||
|
||||
return _build_provider_response(provider)
|
||||
|
||||
|
||||
@router.put("/{provider_id}", response_model=ProviderResponse, summary="更新提供商")
|
||||
async def update_provider(
|
||||
provider_id: str,
|
||||
data: ProviderUpdate,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""更新提供商"""
|
||||
result = await db.execute(
|
||||
select(LLMProvider).where(
|
||||
LLMProvider.id == provider_id,
|
||||
LLMProvider.is_deleted == False
|
||||
)
|
||||
)
|
||||
provider = result.scalar_one_or_none()
|
||||
if not provider:
|
||||
raise HTTPException(status_code=404, detail="提供商不存在")
|
||||
|
||||
update_data = data.model_dump(exclude_unset=True)
|
||||
for key, value in update_data.items():
|
||||
setattr(provider, key, value)
|
||||
|
||||
await db.commit()
|
||||
await db.refresh(provider)
|
||||
|
||||
return _build_provider_response(provider)
|
||||
|
||||
|
||||
@router.delete("/{provider_id}", response_model=ResponseModel, summary="删除提供商")
|
||||
async def delete_provider(provider_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""删除提供商"""
|
||||
result = await db.execute(
|
||||
select(LLMProvider).where(
|
||||
LLMProvider.id == provider_id,
|
||||
LLMProvider.is_deleted == False
|
||||
)
|
||||
)
|
||||
provider = result.scalar_one_or_none()
|
||||
if not provider:
|
||||
raise HTTPException(status_code=404, detail="提供商不存在")
|
||||
|
||||
provider.is_deleted = True
|
||||
await db.commit()
|
||||
|
||||
return ResponseModel(message="删除成功")
|
||||
|
||||
|
||||
@router.post("/{provider_id}/test", summary="测试提供商连接")
|
||||
async def test_provider(provider_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""测试提供商连接"""
|
||||
result = await db.execute(
|
||||
select(LLMProvider).where(
|
||||
LLMProvider.id == provider_id,
|
||||
LLMProvider.is_deleted == False
|
||||
)
|
||||
)
|
||||
provider = result.scalar_one_or_none()
|
||||
if not provider:
|
||||
raise HTTPException(status_code=404, detail="提供商不存在")
|
||||
|
||||
# 创建提供商实例
|
||||
provider_instance = ProviderRegistry.create_instance(
|
||||
provider_type=provider.provider_type,
|
||||
api_key=provider.api_key,
|
||||
api_base=provider.api_base,
|
||||
ollama_host=provider.ollama_host,
|
||||
)
|
||||
|
||||
if not provider_instance:
|
||||
raise HTTPException(status_code=400, detail=f"不支持的提供商类型: {provider.provider_type}")
|
||||
|
||||
# 验证配置
|
||||
if not provider_instance.validate_config():
|
||||
raise HTTPException(status_code=400, detail="配置无效，请检查 API Key")
|
||||
|
||||
# 尝试获取模型列表（如果支持）
|
||||
try:
|
||||
models = provider_instance.get_available_models()
|
||||
return {
|
||||
"success": True,
|
||||
"message": "连接成功",
|
||||
"models": models,
|
||||
}
|
||||
except Exception as e:
|
||||
logger.warning(f"获取模型列表失败: {e}")
|
||||
return {
|
||||
"success": True,
|
||||
"message": "连接成功（无法获取模型列表）",
|
||||
"models": [],
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{provider_id}/default-models", summary="获取默认模型列表")
|
||||
async def get_default_models(provider_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""获取提供商的默认模型列表"""
|
||||
result = await db.execute(
|
||||
select(LLMProvider).where(
|
||||
LLMProvider.id == provider_id,
|
||||
LLMProvider.is_deleted == False
|
||||
)
|
||||
)
|
||||
provider = result.scalar_one_or_none()
|
||||
if not provider:
|
||||
raise HTTPException(status_code=404, detail="提供商不存在")
|
||||
|
||||
return ProviderRegistry.get_default_models(provider.provider_type)
|
||||
|
||||
|
||||
@router.get("/{provider_id}/fetch-models", summary="在线拉取提供商模型列表")
|
||||
async def fetch_provider_models(provider_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""
|
||||
通过提供商 API 在线拉取最新模型列表。
|
||||
如果在线拉取失败，自动 fallback 到硬编码的默认模型列表。
|
||||
"""
|
||||
result = await db.execute(
|
||||
select(LLMProvider).where(
|
||||
LLMProvider.id == provider_id,
|
||||
LLMProvider.is_deleted == False
|
||||
)
|
||||
)
|
||||
provider = result.scalar_one_or_none()
|
||||
if not provider:
|
||||
raise HTTPException(status_code=404, detail="提供商不存在")
|
||||
|
||||
# 构建 kwargs（Ollama 需要 ollama_host）
|
||||
kwargs = {}
|
||||
api_base = provider.api_base or ''
|
||||
if provider.provider_type == 'ollama' and provider.ollama_host:
|
||||
api_base = provider.ollama_host
|
||||
|
||||
# 尝试在线拉取
|
||||
try:
|
||||
models = await ProviderRegistry.fetch_models_from_api(
|
||||
provider_type=provider.provider_type,
|
||||
api_key=provider.api_key or '',
|
||||
api_base=api_base,
|
||||
**kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f'在线拉取模型列表异常 [provider={provider.name}, type={provider.provider_type}]: {e}'
|
||||
)
|
||||
models = []
|
||||
|
||||
source = 'api'
|
||||
if not models:
|
||||
logger.warning(
|
||||
f'在线拉取模型列表为空，fallback 到默认列表 '
|
||||
f'[provider={provider.name}, type={provider.provider_type}, api_base={api_base or "(empty)"}]'
|
||||
)
|
||||
models = ProviderRegistry.get_default_models(provider.provider_type)
|
||||
source = 'default'
|
||||
|
||||
return {
|
||||
"source": source,
|
||||
"models": models,
|
||||
}
|
||||
|
||||
|
||||
def _build_provider_response(provider: LLMProvider) -> dict:
|
||||
"""构建提供商输出"""
|
||||
return {
|
||||
"id": provider.id,
|
||||
"name": provider.name,
|
||||
"provider_type": provider.provider_type,
|
||||
"api_key_masked": provider.get_api_key_masked(),
|
||||
"api_base": provider.api_base or "",
|
||||
"api_version": provider.api_version or "",
|
||||
"ollama_host": provider.ollama_host or "",
|
||||
"description": provider.description or "",
|
||||
"is_active": provider.is_active,
|
||||
"quota_limit": provider.quota_limit or 0,
|
||||
"quota_used": provider.quota_used or 0,
|
||||
"sort": provider.sort or 0,
|
||||
"sys_create_datetime": provider.sys_create_datetime,
|
||||
"sys_update_datetime": provider.sys_update_datetime,
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
"""
|
||||
语音识别 API
|
||||
"""
|
||||
from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile
|
||||
from fastapi.responses import Response
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.database import get_db
|
||||
|
||||
router = APIRouter(prefix="/speech", tags=["AI-语音"])
|
||||
|
||||
|
||||
@router.post("/transcribe", summary="语音转文字")
|
||||
async def transcribe(
|
||||
audio: UploadFile = File(..., description="音频文件"),
|
||||
language: str = Query("zh", description="语言"),
|
||||
provider: str = Query("dashscope", description="提供商"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""
|
||||
语音转文字（ASR）
|
||||
|
||||
将音频文件转换为文字，支持：
|
||||
- dashscope: 阿里云百炼（默认，推荐）
|
||||
- openai: OpenAI Whisper
|
||||
|
||||
支持的音频格式：wav, mp3, webm, pcm, opus
|
||||
"""
|
||||
from ai_platform.services.speech_service import SpeechService
|
||||
|
||||
service = SpeechService(db=db)
|
||||
await service._resolve_dashscope_api_key()
|
||||
audio_content = await audio.read()
|
||||
result = service.transcribe(
|
||||
audio_file=audio_content,
|
||||
language=language,
|
||||
provider=provider,
|
||||
)
|
||||
|
||||
if result["success"]:
|
||||
return {
|
||||
"text": result["text"],
|
||||
"duration": result.get("duration", 0),
|
||||
}
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail=result["error"])
|
||||
|
||||
|
||||
@router.post("/tts", summary="文字转语音")
|
||||
async def text_to_speech(
|
||||
text: str = Query(..., description="要转换的文字"),
|
||||
voice: str = Query("sambert-zhichu-v1", description="声音"),
|
||||
provider: str = Query("dashscope", description="提供商"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""
|
||||
文字转语音（TTS）
|
||||
|
||||
将文字转换为语音，返回音频文件
|
||||
|
||||
DashScope 可用声音：
|
||||
- sambert-zhichu-v1: 知厨（男声）
|
||||
- sambert-zhimiao-emo-v1: 知妙（女声，带情感）
|
||||
- sambert-zhiying-v1: 知莺（女声）
|
||||
|
||||
OpenAI 可用声音：
|
||||
- alloy, echo, fable, onyx, nova, shimmer
|
||||
"""
|
||||
from ai_platform.services.speech_service import SpeechService
|
||||
|
||||
service = SpeechService(db=db)
|
||||
await service._resolve_dashscope_api_key()
|
||||
result = service.text_to_speech(
|
||||
text=text,
|
||||
voice=voice,
|
||||
provider=provider,
|
||||
)
|
||||
|
||||
if result["success"]:
|
||||
return Response(
|
||||
content=result["audio_data"],
|
||||
media_type=result["content_type"],
|
||||
headers={"Content-Disposition": 'attachment; filename="speech.wav"'},
|
||||
)
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail=result["error"])
|
||||
@@ -0,0 +1,720 @@
|
||||
"""
|
||||
AI 工作流 API
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, List, Optional
|
||||
from datetime import datetime
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from fastapi.responses import StreamingResponse
|
||||
from sqlalchemy import select, func, or_, and_
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.database import get_db
|
||||
from app.base_schema import PaginatedResponse, ResponseModel
|
||||
from ai_platform.models import AIWorkflow, AIWorkflowVersion, AIWorkflowRun
|
||||
from core.application.model import Application
|
||||
from ai_platform.schemas.workflow_schema import (
|
||||
WorkflowCreate,
|
||||
WorkflowUpdate,
|
||||
WorkflowResponse,
|
||||
WorkflowListResponse,
|
||||
WorkflowRunInput,
|
||||
WorkflowRunResponse,
|
||||
WorkflowRunListResponse,
|
||||
WorkflowImportCheckIn,
|
||||
WorkflowImportCheckOut,
|
||||
WorkflowImportIn,
|
||||
NodeSchemaResponse,
|
||||
)
|
||||
from ai_platform.services.workflow_import_export import (
|
||||
WorkflowImportExportException,
|
||||
export_config as export_workflow_config,
|
||||
check_import as check_workflow_import,
|
||||
import_config as import_workflow_config,
|
||||
)
|
||||
from ai_platform.nodes.registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/workflow", tags=["AI-工作流"])
|
||||
|
||||
|
||||
# ============ 节点 Schema API ============
|
||||
|
||||
@router.get("/nodes/schemas", response_model=List[NodeSchemaResponse], summary="获取节点 Schema 列表")
|
||||
async def get_node_schemas():
|
||||
"""获取所有已注册节点的 Schema"""
|
||||
return NodeRegistry.get_all_schemas()
|
||||
|
||||
|
||||
@router.get("/nodes/schemas/by-category", summary="按分类获取节点 Schema")
|
||||
async def get_node_schemas_by_category():
|
||||
"""按分类获取所有已注册节点的 Schema"""
|
||||
return NodeRegistry.get_schemas_by_category()
|
||||
|
||||
|
||||
# ============ 工作流 API ============
|
||||
|
||||
@router.get("/list", response_model=PaginatedResponse[WorkflowListResponse], summary="工作流列表")
|
||||
async def list_workflows(
|
||||
name: Optional[str] = Query(None, description="名称"),
|
||||
status: Optional[str] = Query(None, description="状态"),
|
||||
workflow_type: Optional[str] = Query(None, description="工作流类型"),
|
||||
application_id: Optional[str] = Query(None, alias="applicationId", description="所属应用ID"),
|
||||
page: int = Query(1, ge=1, description="页码"),
|
||||
page_size: int = Query(20, ge=1, le=1000, alias="pageSize", description="每页数量"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""获取工作流列表（自动应用数据权限）"""
|
||||
from ai_platform.services.workflow_service import AIWorkflowService
|
||||
from app.data_scope_utils import get_data_scope_filter, apply_data_scope_to_conditions
|
||||
|
||||
conditions = [AIWorkflow.is_deleted == False]
|
||||
|
||||
if application_id:
|
||||
conditions.append(or_(
|
||||
AIWorkflow.application_id == application_id,
|
||||
and_(AIWorkflow.application_id.is_(None), AIWorkflow.is_global == True)
|
||||
))
|
||||
if name:
|
||||
conditions.append(AIWorkflow.name.ilike(f"%{name}%"))
|
||||
if status:
|
||||
conditions.append(AIWorkflow.status == status)
|
||||
if workflow_type:
|
||||
conditions.append(AIWorkflow.workflow_type == workflow_type)
|
||||
|
||||
# 获取数据权限过滤条件并应用
|
||||
from ai_platform.services.workflow_service import RESOURCE_TYPE
|
||||
data_scope_filter = await get_data_scope_filter(db, RESOURCE_TYPE)
|
||||
scope_conditions = apply_data_scope_to_conditions(AIWorkflow, data_scope_filter)
|
||||
conditions.extend(scope_conditions)
|
||||
|
||||
# 获取总数
|
||||
query = select(AIWorkflow).where(and_(*conditions))
|
||||
count_result = await db.execute(select(func.count()).select_from(query.subquery()))
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
# 分页
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(AIWorkflow.sort.desc(), AIWorkflow.sys_create_datetime.desc())
|
||||
query = query.offset(offset).limit(page_size)
|
||||
|
||||
result = await db.execute(query)
|
||||
workflows = result.scalars().all()
|
||||
|
||||
# 批量查询应用名称
|
||||
app_ids = list({w.application_id for w in workflows if w.application_id})
|
||||
app_name_map = {}
|
||||
if app_ids:
|
||||
app_result = await db.execute(
|
||||
select(Application.id, Application.name).where(Application.id.in_(app_ids))
|
||||
)
|
||||
app_name_map = {row.id: row.name for row in app_result}
|
||||
|
||||
items = [_build_workflow_list_response(w, app_name_map.get(w.application_id, "")) for w in workflows]
|
||||
|
||||
return PaginatedResponse(items=items, total=total)
|
||||
|
||||
|
||||
# 全局运行记录路由须注册在 /{workflow_id} 之前，避免 "runs" 被当作 workflow_id
|
||||
@router.get("/runs", response_model=PaginatedResponse[WorkflowRunListResponse], summary="全局工作流运行记录")
|
||||
async def list_all_workflow_runs(
|
||||
workflow_id: Optional[str] = Query(None, alias="workflowId", description="工作流ID"),
|
||||
status: Optional[str] = Query(None, description="运行状态"),
|
||||
trigger_type: Optional[str] = Query(None, alias="triggerType", description="触发来源"),
|
||||
page: int = Query(1, ge=1, description="页码"),
|
||||
page_size: int = Query(20, ge=1, le=100, alias="pageSize", description="每页数量"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""获取全局工作流运行记录（可按工作流、状态、触发来源筛选）"""
|
||||
conditions = [AIWorkflowRun.is_deleted == False]
|
||||
if workflow_id:
|
||||
conditions.append(AIWorkflowRun.workflow_id == workflow_id)
|
||||
if status:
|
||||
conditions.append(AIWorkflowRun.status == status)
|
||||
if trigger_type:
|
||||
conditions.append(AIWorkflowRun.trigger_type == trigger_type)
|
||||
|
||||
query = select(AIWorkflowRun).where(*conditions).order_by(
|
||||
AIWorkflowRun.sys_create_datetime.desc()
|
||||
)
|
||||
|
||||
count_result = await db.execute(select(func.count()).select_from(query.subquery()))
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
offset = (page - 1) * page_size
|
||||
result = await db.execute(query.offset(offset).limit(page_size))
|
||||
runs = result.scalars().all()
|
||||
|
||||
workflow_ids = {r.workflow_id for r in runs if r.workflow_id}
|
||||
workflow_name_map: dict[str, str] = {}
|
||||
if workflow_ids:
|
||||
wf_result = await db.execute(
|
||||
select(AIWorkflow.id, AIWorkflow.name).where(AIWorkflow.id.in_(workflow_ids))
|
||||
)
|
||||
workflow_name_map = {row[0]: row[1] for row in wf_result.all()}
|
||||
|
||||
items = [
|
||||
_build_run_list_item(r, workflow_name_map.get(r.workflow_id, ""))
|
||||
for r in runs
|
||||
]
|
||||
return PaginatedResponse(items=items, total=total)
|
||||
|
||||
|
||||
@router.get("/runs/{run_id}", response_model=WorkflowRunResponse, summary="运行记录详情")
|
||||
async def get_workflow_run(run_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""获取运行记录详情"""
|
||||
result = await db.execute(
|
||||
select(AIWorkflowRun).where(AIWorkflowRun.id == run_id, AIWorkflowRun.is_deleted == False)
|
||||
)
|
||||
run = result.scalar_one_or_none()
|
||||
if not run:
|
||||
raise HTTPException(status_code=404, detail="运行记录不存在")
|
||||
|
||||
workflow_result = await db.execute(
|
||||
select(AIWorkflow).where(AIWorkflow.id == run.workflow_id)
|
||||
)
|
||||
workflow = workflow_result.scalar_one_or_none()
|
||||
workflow_name = workflow.name if workflow else ""
|
||||
|
||||
return _build_run_detail(run, workflow_name)
|
||||
|
||||
|
||||
@router.post("/runs/{run_id}/stop", response_model=ResponseModel, summary="停止运行")
|
||||
async def stop_workflow_run(run_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""停止工作流运行"""
|
||||
result = await db.execute(
|
||||
select(AIWorkflowRun).where(AIWorkflowRun.id == run_id, AIWorkflowRun.is_deleted == False)
|
||||
)
|
||||
run = result.scalar_one_or_none()
|
||||
if not run:
|
||||
raise HTTPException(status_code=404, detail="运行记录不存在")
|
||||
|
||||
if run.status not in ("pending", "running"):
|
||||
raise HTTPException(status_code=400, detail="工作流已完成，无法停止")
|
||||
|
||||
run.status = "stopped"
|
||||
await db.commit()
|
||||
|
||||
return ResponseModel(message="已停止")
|
||||
|
||||
|
||||
class ResumeWorkflowInput(BaseModel):
|
||||
"""恢复工作流输入"""
|
||||
user_input: Any = Field(..., description="用户输入（可以是字符串、布尔值、对象等）")
|
||||
|
||||
|
||||
@router.post("/runs/{run_id}/resume", summary="恢复工作流运行")
|
||||
async def resume_workflow_run(
|
||||
run_id: str,
|
||||
data: ResumeWorkflowInput,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""恢复等待中的工作流运行（SSE）"""
|
||||
import json
|
||||
from ai_platform.services.workflow_service import AIWorkflowService
|
||||
|
||||
async def generate():
|
||||
workflow_service = AIWorkflowService(db)
|
||||
async for event in workflow_service.resume_workflow_stream_async(
|
||||
run_id=run_id,
|
||||
user_input=data.user_input,
|
||||
):
|
||||
yield f"data: {json.dumps(event, ensure_ascii=False, default=str)}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
return StreamingResponse(
|
||||
generate(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@router.post("/runs/{run_id}/resume/stream", summary="流式恢复工作流运行")
|
||||
async def resume_workflow_run_stream(
|
||||
run_id: str,
|
||||
data: ResumeWorkflowInput,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""流式恢复等待中的工作流运行（SSE）"""
|
||||
import json
|
||||
from ai_platform.services.workflow_service import AIWorkflowService
|
||||
|
||||
async def generate():
|
||||
workflow_service = AIWorkflowService(db)
|
||||
async for event in workflow_service.resume_workflow_stream_async(
|
||||
run_id=run_id,
|
||||
user_input=data.user_input,
|
||||
):
|
||||
yield f"data: {json.dumps(event, ensure_ascii=False, default=str)}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
return StreamingResponse(
|
||||
generate(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@router.get("/code/{code}", response_model=WorkflowResponse, summary="根据编码获取工作流")
|
||||
async def get_workflow_by_code(code: str, db: AsyncSession = Depends(get_db)):
|
||||
"""根据编码获取工作流"""
|
||||
result = await db.execute(
|
||||
select(AIWorkflow).where(AIWorkflow.code == code, AIWorkflow.is_deleted == False)
|
||||
)
|
||||
workflow = result.scalar_one_or_none()
|
||||
if not workflow:
|
||||
raise HTTPException(status_code=404, detail="工作流不存在")
|
||||
return _build_workflow_response(workflow)
|
||||
|
||||
|
||||
@router.get("/{workflow_id}", response_model=WorkflowResponse, summary="工作流详情")
|
||||
async def get_workflow(workflow_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""获取工作流详情"""
|
||||
result = await db.execute(
|
||||
select(AIWorkflow).where(AIWorkflow.id == workflow_id, AIWorkflow.is_deleted == False)
|
||||
)
|
||||
workflow = result.scalar_one_or_none()
|
||||
if not workflow:
|
||||
raise HTTPException(status_code=404, detail="工作流不存在")
|
||||
return _build_workflow_response(workflow)
|
||||
|
||||
|
||||
@router.post("", response_model=WorkflowResponse, summary="创建工作流")
|
||||
async def create_workflow(data: WorkflowCreate, db: AsyncSession = Depends(get_db)):
|
||||
"""创建工作流"""
|
||||
# 检查编码是否重复
|
||||
exists_result = await db.execute(
|
||||
select(AIWorkflow).where(AIWorkflow.code == data.code, AIWorkflow.is_deleted == False)
|
||||
)
|
||||
if exists_result.scalar_one_or_none():
|
||||
raise HTTPException(status_code=400, detail=f"工作流编码 {data.code} 已存在")
|
||||
|
||||
workflow = AIWorkflow(**data.model_dump())
|
||||
db.add(workflow)
|
||||
await db.commit()
|
||||
await db.refresh(workflow)
|
||||
|
||||
return _build_workflow_response(workflow)
|
||||
|
||||
|
||||
@router.put("/{workflow_id}", response_model=WorkflowResponse, summary="更新工作流")
|
||||
async def update_workflow(
|
||||
workflow_id: str,
|
||||
data: WorkflowUpdate,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""更新工作流"""
|
||||
result = await db.execute(
|
||||
select(AIWorkflow).where(AIWorkflow.id == workflow_id, AIWorkflow.is_deleted == False)
|
||||
)
|
||||
workflow = result.scalar_one_or_none()
|
||||
if not workflow:
|
||||
raise HTTPException(status_code=404, detail="工作流不存在")
|
||||
|
||||
# 如果更新了 code，检查编码是否重复（排除自身）
|
||||
if data.code and data.code != workflow.code:
|
||||
exists_result = await db.execute(
|
||||
select(AIWorkflow).where(
|
||||
AIWorkflow.code == data.code,
|
||||
AIWorkflow.id != workflow_id,
|
||||
AIWorkflow.is_deleted == False
|
||||
)
|
||||
)
|
||||
if exists_result.scalar_one_or_none():
|
||||
raise HTTPException(status_code=400, detail=f"工作流编码 {data.code} 已存在")
|
||||
|
||||
update_data = data.model_dump(exclude_unset=True)
|
||||
for key, value in update_data.items():
|
||||
setattr(workflow, key, value)
|
||||
|
||||
# 更新草稿版本号
|
||||
workflow.version = (workflow.version or 0) + 1
|
||||
|
||||
await db.commit()
|
||||
await db.refresh(workflow)
|
||||
|
||||
return _build_workflow_response(workflow)
|
||||
|
||||
|
||||
@router.delete("/{workflow_id}", response_model=ResponseModel, summary="删除工作流")
|
||||
async def delete_workflow(workflow_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""删除工作流"""
|
||||
result = await db.execute(
|
||||
select(AIWorkflow).where(AIWorkflow.id == workflow_id, AIWorkflow.is_deleted == False)
|
||||
)
|
||||
workflow = result.scalar_one_or_none()
|
||||
if not workflow:
|
||||
raise HTTPException(status_code=404, detail="工作流不存在")
|
||||
|
||||
workflow.is_deleted = True
|
||||
await db.commit()
|
||||
|
||||
return ResponseModel(message="删除成功")
|
||||
|
||||
|
||||
@router.post("/{workflow_id}/copy", response_model=WorkflowResponse, summary="复制工作流")
|
||||
async def copy_workflow(workflow_id: str, db: AsyncSession = Depends(get_db)):
|
||||
"""复制工作流"""
|
||||
# 获取原工作流
|
||||
result = await db.execute(
|
||||
select(AIWorkflow).where(AIWorkflow.id == workflow_id, AIWorkflow.is_deleted == False)
|
||||
)
|
||||
workflow = result.scalar_one_or_none()
|
||||
if not workflow:
|
||||
raise HTTPException(status_code=404, detail="工作流不存在")
|
||||
|
||||
# 生成新的编码（原编码 + _copy + 时间戳）
|
||||
import time
|
||||
timestamp = int(time.time() * 1000)
|
||||
new_code = f"{workflow.code}_copy_{timestamp}"
|
||||
|
||||
# 检查编码是否重复（理论上不会，但保险起见）
|
||||
exists_result = await db.execute(
|
||||
select(AIWorkflow).where(AIWorkflow.code == new_code, AIWorkflow.is_deleted == False)
|
||||
)
|
||||
if exists_result.scalar_one_or_none():
|
||||
new_code = f"{workflow.code}_copy_{timestamp}_{int(time.time())}"
|
||||
|
||||
# 创建新工作流
|
||||
new_workflow = AIWorkflow(
|
||||
name=f"{workflow.name} (副本)",
|
||||
code=new_code,
|
||||
description=workflow.description,
|
||||
workflow_type=workflow.workflow_type,
|
||||
definition=workflow.definition, # 复制工作流定义
|
||||
input_variables=workflow.input_variables,
|
||||
output_variables=workflow.output_variables,
|
||||
status="draft", # 新工作流默认为草稿状态
|
||||
version=1,
|
||||
published_version=None,
|
||||
published_at=None,
|
||||
published_definition=None,
|
||||
)
|
||||
db.add(new_workflow)
|
||||
await db.commit()
|
||||
await db.refresh(new_workflow)
|
||||
|
||||
return _build_workflow_response(new_workflow)
|
||||
|
||||
|
||||
@router.get("/{workflow_id}/export", summary="导出工作流配置")
|
||||
async def export_workflow(
|
||||
workflow_id: str,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""导出工作流配置为 JSON（草稿 definition）"""
|
||||
try:
|
||||
config = await export_workflow_config(db, workflow_id)
|
||||
content = json.dumps(config, ensure_ascii=False, indent=2)
|
||||
return StreamingResponse(
|
||||
iter([content]),
|
||||
media_type="application/json",
|
||||
headers={
|
||||
"Content-Disposition": f'attachment; filename="{config["code"]}.json"'
|
||||
},
|
||||
)
|
||||
except WorkflowImportExportException as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/import/check", response_model=WorkflowImportCheckOut, summary="导入预检查")
|
||||
async def check_import_workflow(
|
||||
data: WorkflowImportCheckIn,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""导入预检查：检查工作流编码是否冲突"""
|
||||
try:
|
||||
return await check_workflow_import(db, data.code)
|
||||
except WorkflowImportExportException as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/import", response_model=WorkflowResponse, summary="导入工作流配置")
|
||||
async def import_workflow(
|
||||
data: WorkflowImportIn,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""导入工作流配置"""
|
||||
try:
|
||||
workflow = await import_workflow_config(db, data.model_dump())
|
||||
await db.commit()
|
||||
await db.refresh(workflow)
|
||||
return _build_workflow_response(workflow)
|
||||
except WorkflowImportExportException as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/{workflow_id}/publish", response_model=WorkflowResponse, summary="发布工作流")
|
||||
async def publish_workflow(
|
||||
workflow_id: str,
|
||||
description: str = Query("", description="版本说明"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""发布工作流"""
|
||||
result = await db.execute(
|
||||
select(AIWorkflow).where(AIWorkflow.id == workflow_id, AIWorkflow.is_deleted == False)
|
||||
)
|
||||
workflow = result.scalar_one_or_none()
|
||||
if not workflow:
|
||||
raise HTTPException(status_code=404, detail="工作流不存在")
|
||||
|
||||
# 创建版本记录
|
||||
new_version = (workflow.published_version or 0) + 1
|
||||
version = AIWorkflowVersion(
|
||||
workflow_id=workflow_id,
|
||||
version=new_version,
|
||||
definition=workflow.definition or {},
|
||||
description=description,
|
||||
published_at=datetime.now(),
|
||||
)
|
||||
db.add(version)
|
||||
|
||||
# 更新工作流
|
||||
workflow.status = "published"
|
||||
workflow.published_version = new_version
|
||||
workflow.published_at = datetime.now()
|
||||
workflow.published_definition = workflow.definition
|
||||
|
||||
await db.commit()
|
||||
await db.refresh(workflow)
|
||||
|
||||
return _build_workflow_response(workflow)
|
||||
|
||||
|
||||
@router.get("/{workflow_id}/versions", summary="获取版本历史")
|
||||
async def list_workflow_versions(
|
||||
workflow_id: str,
|
||||
page: int = Query(1, ge=1, description="页码"),
|
||||
page_size: int = Query(20, ge=1, le=100, alias="pageSize", description="每页数量"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""获取工作流版本历史"""
|
||||
query = select(AIWorkflowVersion).where(
|
||||
AIWorkflowVersion.workflow_id == workflow_id,
|
||||
AIWorkflowVersion.is_deleted == False
|
||||
).order_by(AIWorkflowVersion.version.desc())
|
||||
|
||||
# 获取总数
|
||||
count_result = await db.execute(select(func.count()).select_from(query.subquery()))
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
# 分页
|
||||
offset = (page - 1) * page_size
|
||||
query = query.offset(offset).limit(page_size)
|
||||
|
||||
result = await db.execute(query)
|
||||
versions = result.scalars().all()
|
||||
|
||||
items = [
|
||||
{
|
||||
"id": v.id,
|
||||
"version": v.version,
|
||||
"description": v.description or "",
|
||||
"published_at": v.published_at,
|
||||
"run_count": v.run_count or 0,
|
||||
"success_count": v.success_count or 0,
|
||||
}
|
||||
for v in versions
|
||||
]
|
||||
|
||||
return PaginatedResponse(items=items, total=total)
|
||||
|
||||
|
||||
@router.get("/{workflow_id}/runs", response_model=PaginatedResponse[WorkflowRunListResponse], summary="工作流运行记录")
|
||||
async def list_workflow_runs(
|
||||
workflow_id: str,
|
||||
page: int = Query(1, ge=1, description="页码"),
|
||||
page_size: int = Query(20, ge=1, le=100, alias="pageSize", description="每页数量"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""获取工作流运行记录"""
|
||||
query = select(AIWorkflowRun).where(
|
||||
AIWorkflowRun.workflow_id == workflow_id,
|
||||
AIWorkflowRun.is_deleted == False
|
||||
).order_by(AIWorkflowRun.sys_create_datetime.desc())
|
||||
|
||||
# 获取总数
|
||||
count_result = await db.execute(select(func.count()).select_from(query.subquery()))
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
# 分页
|
||||
offset = (page - 1) * page_size
|
||||
query = query.offset(offset).limit(page_size)
|
||||
|
||||
result = await db.execute(query)
|
||||
runs = result.scalars().all()
|
||||
|
||||
# 获取工作流名称
|
||||
workflow_result = await db.execute(
|
||||
select(AIWorkflow).where(AIWorkflow.id == workflow_id)
|
||||
)
|
||||
workflow = workflow_result.scalar_one_or_none()
|
||||
workflow_name = workflow.name if workflow else ""
|
||||
|
||||
items = [
|
||||
_build_run_list_item(r, workflow_name)
|
||||
for r in runs
|
||||
]
|
||||
|
||||
return PaginatedResponse(items=items, total=total)
|
||||
|
||||
|
||||
def _build_run_list_item(run: AIWorkflowRun, workflow_name: str = "") -> dict:
|
||||
return {
|
||||
"id": run.id,
|
||||
"workflow_id": run.workflow_id,
|
||||
"workflow_name": workflow_name,
|
||||
"status": run.status or "pending",
|
||||
"trigger_type": run.trigger_type or "api",
|
||||
"total_steps": run.total_steps or 0,
|
||||
"total_tokens": run.total_tokens or 0,
|
||||
"elapsed_time": run.elapsed_time or 0,
|
||||
"error_message": (run.error_message or "")[:200],
|
||||
"started_at": run.started_at,
|
||||
"completed_at": run.completed_at,
|
||||
}
|
||||
|
||||
|
||||
def _build_run_detail(run: AIWorkflowRun, workflow_name: str = "") -> dict:
|
||||
return {
|
||||
"id": run.id,
|
||||
"workflow_id": run.workflow_id,
|
||||
"workflow_name": workflow_name,
|
||||
"status": run.status or "pending",
|
||||
"trigger_type": run.trigger_type or "api",
|
||||
"use_draft": bool(run.use_draft),
|
||||
"workflow_version": run.workflow_version,
|
||||
"definition_snapshot": run.definition_snapshot or {},
|
||||
"inputs": run.inputs or {},
|
||||
"outputs": run.outputs or {},
|
||||
"execution_log": run.execution_log or [],
|
||||
"current_node_id": run.current_node_id or "",
|
||||
"waiting_config": run.waiting_config or {},
|
||||
"error_message": run.error_message or "",
|
||||
"total_tokens": run.total_tokens or 0,
|
||||
"total_steps": run.total_steps or 0,
|
||||
"elapsed_time": run.elapsed_time or 0,
|
||||
"started_at": run.started_at,
|
||||
"completed_at": run.completed_at,
|
||||
"sys_create_datetime": run.sys_create_datetime,
|
||||
}
|
||||
|
||||
|
||||
class WorkflowRunInput(BaseModel):
|
||||
"""工作流运行输入"""
|
||||
inputs: dict = Field(default_factory=dict, description="输入变量")
|
||||
use_draft: bool = Field(default=False, description="是否使用草稿版本")
|
||||
|
||||
|
||||
@router.post("/{workflow_id}/run", summary="运行工作流")
|
||||
async def run_workflow(
|
||||
workflow_id: str,
|
||||
data: WorkflowRunInput,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""运行工作流（非流式）"""
|
||||
from ai_platform.services.workflow_service import AIWorkflowService
|
||||
|
||||
workflow_service = AIWorkflowService(db)
|
||||
run = await workflow_service.run_workflow(
|
||||
workflow_id=workflow_id,
|
||||
inputs=data.inputs,
|
||||
use_draft=data.use_draft,
|
||||
trigger_type='api',
|
||||
)
|
||||
|
||||
return {
|
||||
"id": run.id,
|
||||
"workflow_id": run.workflow_id,
|
||||
"status": run.status,
|
||||
"outputs": run.outputs or {},
|
||||
"execution_log": run.execution_log or [],
|
||||
"total_tokens": run.total_tokens or 0,
|
||||
"total_steps": run.total_steps or 0,
|
||||
"elapsed_time": run.elapsed_time or 0,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/{workflow_id}/run/stream", summary="流式运行工作流")
|
||||
async def run_workflow_stream(
|
||||
workflow_id: str,
|
||||
data: WorkflowRunInput,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""流式运行工作流（SSE）"""
|
||||
import json
|
||||
from ai_platform.services.workflow_service import AIWorkflowService
|
||||
|
||||
async def generate():
|
||||
workflow_service = AIWorkflowService(db)
|
||||
trigger_type = 'editor_draft' if data.use_draft else 'editor_published'
|
||||
async for event in workflow_service.run_workflow_stream_async(
|
||||
workflow_id=workflow_id,
|
||||
inputs=data.inputs,
|
||||
use_draft=data.use_draft,
|
||||
trigger_type=trigger_type,
|
||||
):
|
||||
yield f"data: {json.dumps(event, ensure_ascii=False, default=str)}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
return StreamingResponse(
|
||||
generate(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _build_workflow_response(workflow: AIWorkflow) -> dict:
|
||||
"""构建工作流输出"""
|
||||
return {
|
||||
"id": workflow.id,
|
||||
"name": workflow.name,
|
||||
"code": workflow.code,
|
||||
"workflow_type": workflow.workflow_type or "general",
|
||||
"description": workflow.description or "",
|
||||
"status": workflow.status or "draft",
|
||||
"version": workflow.version or 1,
|
||||
"published_version": workflow.published_version,
|
||||
"published_at": workflow.published_at,
|
||||
"published_definition": workflow.published_definition,
|
||||
"definition": workflow.definition or {},
|
||||
"input_variables": workflow.input_variables or [],
|
||||
"output_variables": workflow.output_variables or [],
|
||||
"run_count": workflow.run_count or 0,
|
||||
"success_count": workflow.success_count or 0,
|
||||
"sort": workflow.sort or 0,
|
||||
"sys_create_datetime": workflow.sys_create_datetime,
|
||||
"sys_update_datetime": workflow.sys_update_datetime,
|
||||
}
|
||||
|
||||
|
||||
def _build_workflow_list_response(workflow: AIWorkflow, application_name: str = "") -> dict:
|
||||
"""构建工作流列表输出"""
|
||||
return {
|
||||
"id": workflow.id,
|
||||
"application_id": workflow.application_id,
|
||||
"application_name": application_name,
|
||||
"is_global": workflow.is_global or False,
|
||||
"name": workflow.name,
|
||||
"code": workflow.code,
|
||||
"workflow_type": workflow.workflow_type or "general",
|
||||
"description": workflow.description or "",
|
||||
"status": workflow.status or "draft",
|
||||
"version": workflow.version or 1,
|
||||
"run_count": workflow.run_count or 0,
|
||||
"success_count": workflow.success_count or 0,
|
||||
"sys_create_datetime": workflow.sys_create_datetime,
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
AI 知识库模块
|
||||
"""
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,45 @@
|
||||
"""
|
||||
文档分块策略
|
||||
"""
|
||||
from .base import BaseChunker, ChunkResult
|
||||
from .recursive import RecursiveChunker
|
||||
from .markdown import MarkdownChunker
|
||||
from .fixed import FixedChunker
|
||||
from .qa_chunker import QAChunker
|
||||
from .sentence import SentenceChunker
|
||||
from .auto import AutoChunker
|
||||
|
||||
__all__ = [
|
||||
'BaseChunker',
|
||||
'ChunkResult',
|
||||
'RecursiveChunker',
|
||||
'MarkdownChunker',
|
||||
'FixedChunker',
|
||||
'QAChunker',
|
||||
'SentenceChunker',
|
||||
'AutoChunker',
|
||||
'get_chunker',
|
||||
]
|
||||
|
||||
|
||||
def get_chunker(strategy: str, chunk_size: int = 500, chunk_overlap: int = 50, separator: str = None, **kwargs) -> BaseChunker:
|
||||
"""
|
||||
根据策略名称获取分块器实例
|
||||
|
||||
Args:
|
||||
strategy: 分块策略名称（recursive/markdown/fixed/qa/sentence/auto）
|
||||
chunk_size: 分块大小
|
||||
chunk_overlap: 分块重叠
|
||||
separator: 自定义分隔符
|
||||
**kwargs: 额外参数（如 QAChunker 的 llm_caller）
|
||||
"""
|
||||
chunkers = {
|
||||
'recursive': RecursiveChunker,
|
||||
'markdown': MarkdownChunker,
|
||||
'fixed': FixedChunker,
|
||||
'qa': QAChunker,
|
||||
'sentence': SentenceChunker,
|
||||
'auto': AutoChunker,
|
||||
}
|
||||
chunker_cls = chunkers.get(strategy, RecursiveChunker)
|
||||
return chunker_cls(chunk_size=chunk_size, chunk_overlap=chunk_overlap, separator=separator, **kwargs)
|
||||
@@ -0,0 +1,90 @@
|
||||
"""
|
||||
自动分块策略
|
||||
|
||||
根据文件类型自动选择最佳分块器。
|
||||
参考 Dify 的 auto 分块模式。
|
||||
"""
|
||||
import logging
|
||||
from typing import Dict, Any, List, Optional
|
||||
|
||||
from .base import BaseChunker, ChunkResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 文件类型 → 推荐分块策略
|
||||
FILE_TYPE_STRATEGY_MAP = {
|
||||
# Markdown 文件使用 Markdown 分块器
|
||||
'md': 'markdown',
|
||||
'markdown': 'markdown',
|
||||
# 代码文件使用按句子分块（按行/语句边界）
|
||||
'py': 'sentence',
|
||||
'js': 'sentence',
|
||||
'ts': 'sentence',
|
||||
'java': 'sentence',
|
||||
'go': 'sentence',
|
||||
'rs': 'sentence',
|
||||
'c': 'sentence',
|
||||
'cpp': 'sentence',
|
||||
'h': 'sentence',
|
||||
# 纯文本使用按句子分块
|
||||
'txt': 'sentence',
|
||||
# CSV/Excel 使用固定大小（表格数据按行分割更合理）
|
||||
'csv': 'fixed',
|
||||
'xlsx': 'fixed',
|
||||
'xls': 'fixed',
|
||||
# HTML 使用 Markdown 分块器（HTML 结构类似）
|
||||
'html': 'markdown',
|
||||
'htm': 'markdown',
|
||||
# 其他文档类型使用递归分块
|
||||
'pdf': 'recursive',
|
||||
'docx': 'recursive',
|
||||
'doc': 'recursive',
|
||||
'pptx': 'recursive',
|
||||
'ppt': 'recursive',
|
||||
}
|
||||
|
||||
|
||||
class AutoChunker(BaseChunker):
|
||||
"""
|
||||
自动分块器
|
||||
|
||||
根据文档的文件类型自动选择最佳分块策略。
|
||||
metadata 中需要包含 'file_type' 字段。
|
||||
"""
|
||||
|
||||
def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
|
||||
"""自动选择分块策略并执行"""
|
||||
if not text or not text.strip():
|
||||
return []
|
||||
|
||||
metadata = metadata or {}
|
||||
file_type = metadata.get('file_type', '').lower().lstrip('.')
|
||||
|
||||
# 根据文件类型选择策略
|
||||
strategy = FILE_TYPE_STRATEGY_MAP.get(file_type, 'recursive')
|
||||
|
||||
logger.info(f"AutoChunker: file_type={file_type} → strategy={strategy}")
|
||||
|
||||
# 动态创建对应的分块器
|
||||
chunker = self._get_chunker(strategy)
|
||||
return chunker.chunk(text, metadata)
|
||||
|
||||
def _get_chunker(self, strategy: str) -> BaseChunker:
|
||||
"""获取对应策略的分块器实例"""
|
||||
from .recursive import RecursiveChunker
|
||||
from .markdown import MarkdownChunker
|
||||
from .fixed import FixedChunker
|
||||
from .sentence import SentenceChunker
|
||||
|
||||
chunkers = {
|
||||
'recursive': RecursiveChunker,
|
||||
'markdown': MarkdownChunker,
|
||||
'fixed': FixedChunker,
|
||||
'sentence': SentenceChunker,
|
||||
}
|
||||
cls = chunkers.get(strategy, RecursiveChunker)
|
||||
return cls(
|
||||
chunk_size=self.chunk_size,
|
||||
chunk_overlap=self.chunk_overlap,
|
||||
separator=self.separator,
|
||||
)
|
||||
@@ -0,0 +1,92 @@
|
||||
"""
|
||||
分块策略基类
|
||||
"""
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, Any, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChunkResult:
|
||||
"""分块结果"""
|
||||
content: str
|
||||
position: int = 0
|
||||
char_count: int = 0
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.char_count:
|
||||
self.char_count = len(self.content)
|
||||
|
||||
|
||||
class BaseChunker(ABC):
|
||||
"""
|
||||
分块策略基类
|
||||
|
||||
所有分块策略必须继承此类并实现 chunk 方法
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunk_size: int = 500,
|
||||
chunk_overlap: int = 50,
|
||||
separator: Optional[str] = None,
|
||||
):
|
||||
self.chunk_size = chunk_size
|
||||
self.chunk_overlap = chunk_overlap
|
||||
self.separator = separator
|
||||
|
||||
@abstractmethod
|
||||
def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
|
||||
"""
|
||||
将文本分块
|
||||
|
||||
Args:
|
||||
text: 原始文本
|
||||
metadata: 文档元数据
|
||||
|
||||
Returns:
|
||||
分块结果列表
|
||||
"""
|
||||
pass
|
||||
|
||||
def _clean_text(self, text: str) -> str:
|
||||
"""清理文本：去除多余空白"""
|
||||
import re
|
||||
# 合并连续空行为单个空行
|
||||
text = re.sub(r'\n{3,}', '\n\n', text)
|
||||
# 去除行尾空白
|
||||
text = '\n'.join(line.rstrip() for line in text.split('\n'))
|
||||
return text.strip()
|
||||
|
||||
def _merge_small_chunks(self, chunks: List[str], min_size: int = 50) -> List[str]:
|
||||
"""合并过小的分块"""
|
||||
if not chunks:
|
||||
return []
|
||||
|
||||
merged = []
|
||||
buffer = ""
|
||||
|
||||
for chunk in chunks:
|
||||
if not chunk.strip():
|
||||
continue
|
||||
if buffer and len(buffer) + len(chunk) <= self.chunk_size:
|
||||
buffer = buffer + "\n" + chunk
|
||||
elif buffer and len(buffer) < min_size:
|
||||
buffer = buffer + "\n" + chunk
|
||||
else:
|
||||
if buffer:
|
||||
merged.append(buffer)
|
||||
buffer = chunk
|
||||
|
||||
if buffer:
|
||||
# 最后一个 buffer 如果太小，合并到前一个
|
||||
if merged and len(buffer) < min_size:
|
||||
merged[-1] = merged[-1] + "\n" + buffer
|
||||
else:
|
||||
merged.append(buffer)
|
||||
|
||||
return merged
|
||||
@@ -0,0 +1,57 @@
|
||||
"""
|
||||
固定大小分块策略
|
||||
|
||||
按固定字符数分割文本，最简单的分块方式
|
||||
"""
|
||||
import logging
|
||||
from typing import Dict, Any, List, Optional
|
||||
|
||||
from .base import BaseChunker, ChunkResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FixedChunker(BaseChunker):
|
||||
"""
|
||||
固定大小分块器
|
||||
|
||||
按固定字符数分割文本，相邻分块之间有 overlap 重叠
|
||||
"""
|
||||
|
||||
def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
|
||||
"""固定大小分块"""
|
||||
if not text or not text.strip():
|
||||
return []
|
||||
|
||||
text = self._clean_text(text)
|
||||
metadata = metadata or {}
|
||||
|
||||
if len(text) <= self.chunk_size:
|
||||
return [ChunkResult(
|
||||
content=text,
|
||||
position=0,
|
||||
metadata={**metadata},
|
||||
)]
|
||||
|
||||
chunks = []
|
||||
start = 0
|
||||
position = 0
|
||||
step = self.chunk_size - self.chunk_overlap
|
||||
|
||||
while start < len(text):
|
||||
end = min(start + self.chunk_size, len(text))
|
||||
chunk_text = text[start:end].strip()
|
||||
|
||||
if chunk_text:
|
||||
chunks.append(ChunkResult(
|
||||
content=chunk_text,
|
||||
position=position,
|
||||
metadata={**metadata},
|
||||
))
|
||||
position += 1
|
||||
|
||||
start += step
|
||||
if step <= 0:
|
||||
break
|
||||
|
||||
return chunks
|
||||
@@ -0,0 +1,172 @@
|
||||
"""
|
||||
Markdown 结构化分块策略
|
||||
|
||||
按 Markdown 标题层级分割文档，保留文档结构信息
|
||||
"""
|
||||
import re
|
||||
import logging
|
||||
from typing import Dict, Any, List, Optional
|
||||
|
||||
from .base import BaseChunker, ChunkResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MarkdownChunker(BaseChunker):
|
||||
"""
|
||||
Markdown 分块器
|
||||
|
||||
按标题层级分割 Markdown 文档，每个标题下的内容作为一个分块
|
||||
如果单个标题下的内容超过 chunk_size，则使用递归分割
|
||||
"""
|
||||
|
||||
def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
|
||||
"""按 Markdown 标题分块"""
|
||||
if not text or not text.strip():
|
||||
return []
|
||||
|
||||
text = self._clean_text(text)
|
||||
metadata = metadata or {}
|
||||
|
||||
# 按标题分割
|
||||
sections = self._split_by_headers(text)
|
||||
|
||||
# 处理每个 section
|
||||
raw_chunks = []
|
||||
for section in sections:
|
||||
header = section.get('header', '')
|
||||
content = section.get('content', '')
|
||||
level = section.get('level', 0)
|
||||
|
||||
if not content.strip():
|
||||
continue
|
||||
|
||||
# 组合标题和内容
|
||||
full_text = f"{header}\n{content}" if header else content
|
||||
|
||||
if len(full_text) <= self.chunk_size:
|
||||
raw_chunks.append({
|
||||
'content': full_text.strip(),
|
||||
'metadata': {
|
||||
**metadata,
|
||||
'header': header,
|
||||
'header_level': level,
|
||||
}
|
||||
})
|
||||
else:
|
||||
# 内容超长，递归分割
|
||||
sub_chunks = self._split_long_section(content, header)
|
||||
for i, sub in enumerate(sub_chunks):
|
||||
raw_chunks.append({
|
||||
'content': sub.strip(),
|
||||
'metadata': {
|
||||
**metadata,
|
||||
'header': header,
|
||||
'header_level': level,
|
||||
'sub_chunk': i,
|
||||
}
|
||||
})
|
||||
|
||||
# 合并过小的分块
|
||||
merged = self._merge_small_section_chunks(raw_chunks)
|
||||
|
||||
# 构建结果
|
||||
results = []
|
||||
for i, item in enumerate(merged):
|
||||
if item['content'].strip():
|
||||
results.append(ChunkResult(
|
||||
content=item['content'],
|
||||
position=i,
|
||||
metadata=item.get('metadata', {}),
|
||||
))
|
||||
|
||||
return results
|
||||
|
||||
def _split_by_headers(self, text: str) -> List[Dict[str, Any]]:
|
||||
"""按 Markdown 标题分割"""
|
||||
# 匹配 Markdown 标题: # Title, ## Title, ### Title 等
|
||||
header_pattern = re.compile(r'^(#{1,6})\s+(.+)$', re.MULTILINE)
|
||||
|
||||
sections = []
|
||||
last_end = 0
|
||||
last_header = ''
|
||||
last_level = 0
|
||||
|
||||
for match in header_pattern.finditer(text):
|
||||
# 保存前一个 section 的内容
|
||||
if last_end > 0 or match.start() > 0:
|
||||
content = text[last_end:match.start()]
|
||||
if content.strip() or last_header:
|
||||
sections.append({
|
||||
'header': last_header,
|
||||
'content': content.strip(),
|
||||
'level': last_level,
|
||||
})
|
||||
|
||||
last_header = match.group(0)
|
||||
last_level = len(match.group(1))
|
||||
last_end = match.end()
|
||||
|
||||
# 最后一个 section
|
||||
remaining = text[last_end:]
|
||||
if remaining.strip() or last_header:
|
||||
sections.append({
|
||||
'header': last_header,
|
||||
'content': remaining.strip(),
|
||||
'level': last_level,
|
||||
})
|
||||
|
||||
# 如果没有找到任何标题，整个文本作为一个 section
|
||||
if not sections:
|
||||
sections.append({
|
||||
'header': '',
|
||||
'content': text.strip(),
|
||||
'level': 0,
|
||||
})
|
||||
|
||||
return sections
|
||||
|
||||
def _split_long_section(self, content: str, header: str = '') -> List[str]:
|
||||
"""分割超长的 section 内容"""
|
||||
from .recursive import RecursiveChunker
|
||||
|
||||
chunker = RecursiveChunker(
|
||||
chunk_size=self.chunk_size,
|
||||
chunk_overlap=self.chunk_overlap,
|
||||
)
|
||||
results = chunker.chunk(content)
|
||||
|
||||
chunks = []
|
||||
for i, result in enumerate(results):
|
||||
# 第一个分块带上标题
|
||||
if i == 0 and header:
|
||||
chunks.append(f"{header}\n{result.content}")
|
||||
else:
|
||||
chunks.append(result.content)
|
||||
|
||||
return chunks if chunks else [content]
|
||||
|
||||
def _merge_small_section_chunks(self, chunks: List[Dict], min_size: int = 80) -> List[Dict]:
|
||||
"""合并过小的 section 分块"""
|
||||
if not chunks:
|
||||
return []
|
||||
|
||||
merged = []
|
||||
buffer = None
|
||||
|
||||
for chunk in chunks:
|
||||
if buffer is None:
|
||||
buffer = chunk
|
||||
elif len(buffer['content']) < min_size and len(buffer['content']) + len(chunk['content']) <= self.chunk_size:
|
||||
buffer['content'] = buffer['content'] + "\n\n" + chunk['content']
|
||||
else:
|
||||
merged.append(buffer)
|
||||
buffer = chunk
|
||||
|
||||
if buffer:
|
||||
if merged and len(buffer['content']) < min_size:
|
||||
merged[-1]['content'] = merged[-1]['content'] + "\n\n" + buffer['content']
|
||||
else:
|
||||
merged.append(buffer)
|
||||
|
||||
return merged
|
||||
@@ -0,0 +1,197 @@
|
||||
"""
|
||||
Q&A 自动拆分分块策略
|
||||
|
||||
使用 LLM 将文档内容自动拆分为问答对。
|
||||
每个分段的 content 存储 question，metadata 中存储 answer。
|
||||
检索时用 question 做向量匹配，返回 answer 作为上下文。
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
from typing import Dict, Any, List, Optional
|
||||
|
||||
from .base import BaseChunker, ChunkResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Q&A 拆分的系统提示词
|
||||
QA_SYSTEM_PROMPT = """你是一个专业的知识库问答对生成助手。请根据给定的文本内容，生成高质量的问答对（Q&A pairs）。
|
||||
|
||||
要求：
|
||||
1. 问题应该是用户可能会问的自然语言问题
|
||||
2. 答案应该准确、完整，直接来源于原文
|
||||
3. 每个问答对应该覆盖文本中的一个独立知识点
|
||||
4. 问题要具体明确，避免过于宽泛
|
||||
5. 答案要简洁但完整，包含必要的上下文
|
||||
|
||||
请以 JSON 数组格式输出，每个元素包含 question 和 answer 字段：
|
||||
```json
|
||||
[
|
||||
{"question": "问题1", "answer": "答案1"},
|
||||
{"question": "问题2", "answer": "答案2"}
|
||||
]
|
||||
```
|
||||
|
||||
只输出 JSON 数组，不要输出其他内容。"""
|
||||
|
||||
|
||||
class QAChunker(BaseChunker):
|
||||
"""
|
||||
Q&A 自动拆分分块器
|
||||
|
||||
使用 LLM 将文本拆分为问答对。
|
||||
需要在初始化时传入 LLM 调用函数。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunk_size: int = 500,
|
||||
chunk_overlap: int = 50,
|
||||
separator: Optional[str] = None,
|
||||
llm_caller: Optional[Any] = None,
|
||||
):
|
||||
super().__init__(chunk_size, chunk_overlap, separator)
|
||||
self._llm_caller = llm_caller
|
||||
|
||||
def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
|
||||
"""
|
||||
同步分块（Q&A 模式不支持同步调用，返回空列表）
|
||||
请使用 chunk_async 方法。
|
||||
"""
|
||||
logger.warning("QAChunker.chunk() 不支持同步调用，请使用 chunk_async()")
|
||||
return []
|
||||
|
||||
async def chunk_async(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
|
||||
"""
|
||||
异步分块：使用 LLM 将文本拆分为 Q&A 对
|
||||
|
||||
Args:
|
||||
text: 原始文本
|
||||
metadata: 文档元数据
|
||||
|
||||
Returns:
|
||||
分块结果列表，每个 ChunkResult 的 content 为 question，
|
||||
metadata 中包含 answer 和 chunk_mode='qa'
|
||||
"""
|
||||
if not text or not text.strip():
|
||||
return []
|
||||
|
||||
metadata = metadata or {}
|
||||
text = self._clean_text(text)
|
||||
|
||||
# 如果文本太长，先按段落粗分再逐段生成 Q&A
|
||||
max_input_size = self.chunk_size * 8 # LLM 输入上限
|
||||
if len(text) > max_input_size:
|
||||
segments = self._split_for_qa(text, max_input_size)
|
||||
else:
|
||||
segments = [text]
|
||||
|
||||
all_results = []
|
||||
position = 0
|
||||
|
||||
for segment in segments:
|
||||
qa_pairs = await self._generate_qa_pairs(segment)
|
||||
for qa in qa_pairs:
|
||||
question = qa.get('question', '').strip()
|
||||
answer = qa.get('answer', '').strip()
|
||||
if not question or not answer:
|
||||
continue
|
||||
|
||||
all_results.append(ChunkResult(
|
||||
content=question,
|
||||
position=position,
|
||||
metadata={
|
||||
**metadata,
|
||||
'answer': answer,
|
||||
'chunk_mode': 'qa',
|
||||
},
|
||||
))
|
||||
position += 1
|
||||
|
||||
logger.info(f"Q&A 拆分完成: {len(all_results)} 个问答对")
|
||||
return all_results
|
||||
|
||||
async def _generate_qa_pairs(self, text: str) -> List[Dict[str, str]]:
|
||||
"""调用 LLM 生成 Q&A 对"""
|
||||
if not self._llm_caller:
|
||||
logger.error("QAChunker: 未配置 LLM 调用函数")
|
||||
return []
|
||||
|
||||
try:
|
||||
user_prompt = f"请根据以下文本生成问答对：\n\n{text}"
|
||||
|
||||
response_text = await self._llm_caller(
|
||||
system_prompt=QA_SYSTEM_PROMPT,
|
||||
user_prompt=user_prompt,
|
||||
)
|
||||
|
||||
if not response_text:
|
||||
return []
|
||||
|
||||
# 解析 JSON 响应
|
||||
return self._parse_qa_response(response_text)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Q&A 生成失败: {e}")
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _parse_qa_response(response_text: str) -> List[Dict[str, str]]:
|
||||
"""解析 LLM 返回的 Q&A JSON"""
|
||||
try:
|
||||
# 尝试直接解析
|
||||
result = json.loads(response_text)
|
||||
if isinstance(result, list):
|
||||
return [
|
||||
item for item in result
|
||||
if isinstance(item, dict) and 'question' in item and 'answer' in item
|
||||
]
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# 尝试提取 JSON 代码块
|
||||
import re
|
||||
json_match = re.search(r'```(?:json)?\s*\n?(.*?)\n?```', response_text, re.DOTALL)
|
||||
if json_match:
|
||||
try:
|
||||
result = json.loads(json_match.group(1))
|
||||
if isinstance(result, list):
|
||||
return [
|
||||
item for item in result
|
||||
if isinstance(item, dict) and 'question' in item and 'answer' in item
|
||||
]
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# 尝试找到 [ ... ] 部分
|
||||
bracket_match = re.search(r'\[.*\]', response_text, re.DOTALL)
|
||||
if bracket_match:
|
||||
try:
|
||||
result = json.loads(bracket_match.group(0))
|
||||
if isinstance(result, list):
|
||||
return [
|
||||
item for item in result
|
||||
if isinstance(item, dict) and 'question' in item and 'answer' in item
|
||||
]
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
logger.warning(f"无法解析 Q&A 响应: {response_text[:200]}")
|
||||
return []
|
||||
|
||||
def _split_for_qa(self, text: str, max_size: int) -> List[str]:
|
||||
"""将长文本按段落分割为适合 LLM 处理的片段"""
|
||||
paragraphs = text.split('\n\n')
|
||||
segments = []
|
||||
current = ""
|
||||
|
||||
for para in paragraphs:
|
||||
if current and len(current) + len(para) + 2 > max_size:
|
||||
segments.append(current.strip())
|
||||
current = para
|
||||
else:
|
||||
current = current + "\n\n" + para if current else para
|
||||
|
||||
if current.strip():
|
||||
segments.append(current.strip())
|
||||
|
||||
return segments
|
||||
@@ -0,0 +1,160 @@
|
||||
"""
|
||||
递归字符分块策略
|
||||
|
||||
最常用的分块策略，按照分隔符层级递归分割文本
|
||||
优先按段落 → 句子 → 字符的顺序分割
|
||||
"""
|
||||
import logging
|
||||
from typing import Dict, Any, List, Optional
|
||||
|
||||
from .base import BaseChunker, ChunkResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 默认分隔符层级（从大到小）
|
||||
DEFAULT_SEPARATORS = [
|
||||
"\n\n", # 段落
|
||||
"\n", # 换行
|
||||
"。", # 中文句号
|
||||
"！", # 中文感叹号
|
||||
"？", # 中文问号
|
||||
"；", # 中文分号
|
||||
". ", # 英文句号
|
||||
"! ", # 英文感叹号
|
||||
"? ", # 英文问号
|
||||
"; ", # 英文分号
|
||||
"，", # 中文逗号
|
||||
", ", # 英文逗号
|
||||
" ", # 空格
|
||||
"", # 逐字符
|
||||
]
|
||||
|
||||
|
||||
class RecursiveChunker(BaseChunker):
|
||||
"""
|
||||
递归字符分块器
|
||||
|
||||
按分隔符层级递归分割文本，确保每个分块不超过 chunk_size，
|
||||
相邻分块之间有 chunk_overlap 的重叠
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunk_size: int = 500,
|
||||
chunk_overlap: int = 50,
|
||||
separator: Optional[str] = None,
|
||||
separators: Optional[List[str]] = None,
|
||||
):
|
||||
super().__init__(chunk_size, chunk_overlap, separator)
|
||||
if separator:
|
||||
self.separators = [separator] + DEFAULT_SEPARATORS
|
||||
elif separators:
|
||||
self.separators = separators
|
||||
else:
|
||||
self.separators = DEFAULT_SEPARATORS
|
||||
|
||||
def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
|
||||
"""递归分块"""
|
||||
if not text or not text.strip():
|
||||
return []
|
||||
|
||||
text = self._clean_text(text)
|
||||
metadata = metadata or {}
|
||||
|
||||
# 递归分割
|
||||
raw_chunks = self._recursive_split(text, self.separators)
|
||||
|
||||
# 合并过小的分块
|
||||
raw_chunks = self._merge_small_chunks(raw_chunks)
|
||||
|
||||
# 添加重叠
|
||||
chunks_with_overlap = self._add_overlap(raw_chunks)
|
||||
|
||||
# 构建结果
|
||||
results = []
|
||||
for i, content in enumerate(chunks_with_overlap):
|
||||
if content.strip():
|
||||
results.append(ChunkResult(
|
||||
content=content.strip(),
|
||||
position=i,
|
||||
metadata={**metadata},
|
||||
))
|
||||
|
||||
return results
|
||||
|
||||
def _recursive_split(self, text: str, separators: List[str]) -> List[str]:
|
||||
"""递归分割文本"""
|
||||
if len(text) <= self.chunk_size:
|
||||
return [text] if text.strip() else []
|
||||
|
||||
# 找到合适的分隔符
|
||||
separator = ""
|
||||
for sep in separators:
|
||||
if sep == "":
|
||||
separator = sep
|
||||
break
|
||||
if sep in text:
|
||||
separator = sep
|
||||
break
|
||||
|
||||
# 按分隔符分割
|
||||
if separator:
|
||||
splits = text.split(separator)
|
||||
else:
|
||||
# 逐字符分割
|
||||
splits = list(text)
|
||||
|
||||
# 合并分割结果，确保不超过 chunk_size
|
||||
chunks = []
|
||||
current = ""
|
||||
|
||||
for split in splits:
|
||||
piece = split if not separator else split
|
||||
test_piece = current + separator + piece if current else piece
|
||||
|
||||
if len(test_piece) <= self.chunk_size:
|
||||
current = test_piece
|
||||
else:
|
||||
if current:
|
||||
chunks.append(current)
|
||||
# 如果单个片段超过 chunk_size，递归处理
|
||||
if len(piece) > self.chunk_size:
|
||||
remaining_separators = separators[separators.index(separator) + 1:] if separator in separators else separators[1:]
|
||||
if remaining_separators:
|
||||
sub_chunks = self._recursive_split(piece, remaining_separators)
|
||||
chunks.extend(sub_chunks)
|
||||
current = ""
|
||||
else:
|
||||
# 没有更小的分隔符了，强制截断
|
||||
for j in range(0, len(piece), self.chunk_size):
|
||||
chunks.append(piece[j:j + self.chunk_size])
|
||||
current = ""
|
||||
else:
|
||||
current = piece
|
||||
|
||||
if current:
|
||||
chunks.append(current)
|
||||
|
||||
return chunks
|
||||
|
||||
def _add_overlap(self, chunks: List[str]) -> List[str]:
|
||||
"""为相邻分块添加重叠"""
|
||||
if self.chunk_overlap <= 0 or len(chunks) <= 1:
|
||||
return chunks
|
||||
|
||||
result = []
|
||||
for i, chunk in enumerate(chunks):
|
||||
if i == 0:
|
||||
result.append(chunk)
|
||||
else:
|
||||
# 从前一个分块的末尾取 overlap 字符作为前缀
|
||||
prev = chunks[i - 1]
|
||||
overlap_text = prev[-self.chunk_overlap:] if len(prev) > self.chunk_overlap else prev
|
||||
# 确保合并后不超过 chunk_size 太多
|
||||
combined = overlap_text + "\n" + chunk
|
||||
if len(combined) <= self.chunk_size * 1.2:
|
||||
result.append(combined)
|
||||
else:
|
||||
result.append(chunk)
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,105 @@
|
||||
"""
|
||||
按句子分块策略
|
||||
|
||||
按句号/问号/感叹号等句子边界分割文本，
|
||||
然后将小句子合并到不超过 chunk_size 的分块中。
|
||||
参考 Dify 的 sentence 分块模式。
|
||||
"""
|
||||
import logging
|
||||
import re
|
||||
from typing import Dict, Any, List, Optional
|
||||
|
||||
from .base import BaseChunker, ChunkResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 句子分隔符正则（中英文句号/问号/感叹号/分号）
|
||||
SENTENCE_PATTERN = re.compile(
|
||||
r'(?<=[。！？；.!?;])\s*'
|
||||
)
|
||||
|
||||
|
||||
class SentenceChunker(BaseChunker):
|
||||
"""
|
||||
按句子分块器
|
||||
|
||||
先按句子边界分割文本，再将相邻句子合并为不超过 chunk_size 的分块。
|
||||
保证每个分块都是完整句子的组合，不会在句子中间截断。
|
||||
"""
|
||||
|
||||
def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
|
||||
"""按句子分块"""
|
||||
if not text or not text.strip():
|
||||
return []
|
||||
|
||||
text = self._clean_text(text)
|
||||
metadata = metadata or {}
|
||||
|
||||
# 按句子边界分割
|
||||
sentences = SENTENCE_PATTERN.split(text)
|
||||
sentences = [s.strip() for s in sentences if s.strip()]
|
||||
|
||||
if not sentences:
|
||||
return [ChunkResult(content=text, position=0, metadata={**metadata})]
|
||||
|
||||
# 合并句子为分块（不超过 chunk_size）
|
||||
chunks = []
|
||||
current = ""
|
||||
position = 0
|
||||
|
||||
for sentence in sentences:
|
||||
# 如果单个句子就超过 chunk_size，强制作为独立分块
|
||||
if len(sentence) > self.chunk_size:
|
||||
if current:
|
||||
chunks.append(current)
|
||||
current = ""
|
||||
chunks.append(sentence)
|
||||
continue
|
||||
|
||||
test = current + sentence if not current else current + " " + sentence
|
||||
if len(test) <= self.chunk_size:
|
||||
current = test
|
||||
else:
|
||||
if current:
|
||||
chunks.append(current)
|
||||
current = sentence
|
||||
|
||||
if current:
|
||||
chunks.append(current)
|
||||
|
||||
# 合并过小的分块
|
||||
chunks = self._merge_small_chunks(chunks)
|
||||
|
||||
# 添加重叠
|
||||
if self.chunk_overlap > 0 and len(chunks) > 1:
|
||||
chunks = self._add_sentence_overlap(chunks)
|
||||
|
||||
# 构建结果
|
||||
results = []
|
||||
for i, content in enumerate(chunks):
|
||||
if content.strip():
|
||||
results.append(ChunkResult(
|
||||
content=content.strip(),
|
||||
position=i,
|
||||
metadata={**metadata},
|
||||
))
|
||||
|
||||
return results
|
||||
|
||||
def _add_sentence_overlap(self, chunks: List[str]) -> List[str]:
|
||||
"""为相邻分块添加句子级重叠"""
|
||||
result = [chunks[0]]
|
||||
for i in range(1, len(chunks)):
|
||||
prev = chunks[i - 1]
|
||||
# 从前一个分块取最后一个句子作为重叠
|
||||
prev_sentences = SENTENCE_PATTERN.split(prev)
|
||||
prev_sentences = [s.strip() for s in prev_sentences if s.strip()]
|
||||
if prev_sentences:
|
||||
overlap = prev_sentences[-1]
|
||||
if len(overlap) <= self.chunk_overlap:
|
||||
combined = overlap + " " + chunks[i]
|
||||
if len(combined) <= self.chunk_size * 1.2:
|
||||
result.append(combined)
|
||||
continue
|
||||
result.append(chunks[i])
|
||||
return result
|
||||
@@ -0,0 +1,16 @@
|
||||
"""
|
||||
知识库数据模型
|
||||
"""
|
||||
from .knowledge_base_model import KnowledgeBase
|
||||
from .document_model import KnowledgeDocument
|
||||
from .segment_model import KnowledgeSegment
|
||||
from .annotation_model import KnowledgeAnnotation
|
||||
from .retrieval_log_model import KnowledgeRetrievalLog
|
||||
|
||||
__all__ = [
|
||||
'KnowledgeBase',
|
||||
'KnowledgeDocument',
|
||||
'KnowledgeSegment',
|
||||
'KnowledgeAnnotation',
|
||||
'KnowledgeRetrievalLog',
|
||||
]
|
||||
@@ -0,0 +1,34 @@
|
||||
"""
|
||||
知识库标注模型（Q&A 对）
|
||||
|
||||
手动添加的高优先级问答对，检索时优先匹配。
|
||||
"""
|
||||
from sqlalchemy import Column, String, Text, Integer, Boolean, Index
|
||||
|
||||
from app.base_model import BaseModel
|
||||
|
||||
|
||||
class KnowledgeAnnotation(BaseModel):
|
||||
"""
|
||||
知识库标注（Q&A 对）
|
||||
|
||||
用户手动添加的问答对，检索时优先匹配 question，返回 answer。
|
||||
"""
|
||||
__tablename__ = "ai_knowledge_annotation"
|
||||
|
||||
knowledge_base_id = Column(String(21), nullable=False, index=True, comment="所属知识库ID（逻辑外键关联ai_knowledge_base）")
|
||||
|
||||
# Q&A 内容
|
||||
question = Column(Text, nullable=False, comment="问题")
|
||||
answer = Column(Text, nullable=False, comment="答案")
|
||||
|
||||
# 向量化状态
|
||||
embedding_status = Column(String(20), default="pending", comment="向量化状态: pending/completed/failed")
|
||||
|
||||
# 状态
|
||||
enabled = Column(Boolean, default=True, comment="是否启用")
|
||||
hit_count = Column(Integer, default=0, comment="命中次数")
|
||||
|
||||
__table_args__ = (
|
||||
Index('idx_annotation_kb_enabled', 'knowledge_base_id', 'enabled'),
|
||||
)
|
||||
@@ -0,0 +1,39 @@
|
||||
"""
|
||||
知识库文档模型
|
||||
"""
|
||||
from sqlalchemy import Column, String, Text, Integer, Boolean, DateTime, BigInteger
|
||||
|
||||
from app.base_model import BaseModel
|
||||
|
||||
|
||||
class KnowledgeDocument(BaseModel):
|
||||
"""
|
||||
知识库文档
|
||||
|
||||
记录上传到知识库的文档信息及处理状态
|
||||
"""
|
||||
__tablename__ = "ai_knowledge_document"
|
||||
|
||||
knowledge_base_id = Column(String(21), nullable=False, index=True, comment="所属知识库ID（逻辑外键关联ai_knowledge_base）")
|
||||
file_id = Column(String(21), nullable=True, comment="关联文件ID（逻辑外键关联core_file_manager）")
|
||||
name = Column(String(255), nullable=False, comment="文档名称")
|
||||
file_type = Column(String(20), nullable=True, comment="文件类型: pdf/docx/txt/md/xlsx/csv/html/pptx")
|
||||
file_size = Column(BigInteger, default=0, comment="文件大小（字节）")
|
||||
content_hash = Column(String(64), nullable=True, index=True, comment="内容MD5（用于去重）")
|
||||
|
||||
# 处理结果
|
||||
segment_count = Column(Integer, default=0, comment="分段数量")
|
||||
token_count = Column(Integer, default=0, comment="Token 总数")
|
||||
char_count = Column(Integer, default=0, comment="字符总数")
|
||||
|
||||
# 处理状态
|
||||
status = Column(String(20), default="pending", index=True, comment="状态: pending/indexing/completed/failed/disabled")
|
||||
error_message = Column(Text, nullable=True, comment="错误信息")
|
||||
indexing_started_at = Column(DateTime, nullable=True, comment="索引开始时间")
|
||||
indexing_completed_at = Column(DateTime, nullable=True, comment="索引完成时间")
|
||||
|
||||
# 去重
|
||||
duplicate_warning = Column(Text, nullable=True, comment="内容重复警告（跨知识库检测）")
|
||||
|
||||
# 是否启用
|
||||
enabled = Column(Boolean, default=True, comment="是否启用（禁用后不参与检索）")
|
||||
@@ -0,0 +1,55 @@
|
||||
"""
|
||||
知识库模型
|
||||
"""
|
||||
from sqlalchemy import Column, String, Text, Integer, Float, Boolean, JSON
|
||||
|
||||
from app.base_model import BaseModel
|
||||
|
||||
|
||||
class KnowledgeBase(BaseModel):
|
||||
"""
|
||||
知识库
|
||||
|
||||
管理文档集合，配置分块策略和检索参数
|
||||
"""
|
||||
__tablename__ = "ai_knowledge_base"
|
||||
|
||||
application_id = Column(String(21), nullable=True, index=True, comment="所属应用ID（逻辑外键关联core_application）")
|
||||
is_global = Column(Boolean, default=False, comment="是否在子应用中可见")
|
||||
name = Column(String(100), nullable=False, comment="知识库名称")
|
||||
code = Column(String(100), nullable=False, unique=True, comment="知识库编码")
|
||||
description = Column(Text, nullable=True, comment="描述")
|
||||
icon = Column(String(50), default="", comment="图标")
|
||||
|
||||
# Embedding 配置
|
||||
embedding_model_id = Column(String(21), nullable=True, comment="Embedding 模型ID（逻辑外键关联ai_llm_model）")
|
||||
embedding_dimensions = Column(Integer, default=1536, comment="向量维度")
|
||||
|
||||
# 分块策略
|
||||
chunk_strategy = Column(String(20), default="recursive", comment="分块策略: recursive/semantic/markdown/fixed")
|
||||
chunk_size = Column(Integer, default=500, comment="分块大小（字符数）")
|
||||
chunk_overlap = Column(Integer, default=50, comment="分块重叠大小（字符数）")
|
||||
separator = Column(String(50), nullable=True, comment="自定义分隔符")
|
||||
|
||||
# 检索配置
|
||||
retrieval_mode = Column(String(20), default="hybrid", comment="检索模式: vector/fulltext/hybrid")
|
||||
top_k = Column(Integer, default=5, comment="检索返回数量")
|
||||
score_threshold = Column(Float, default=0.5, comment="相似度阈值（0-1）")
|
||||
rerank_enabled = Column(Boolean, default=False, comment="是否启用重排序")
|
||||
rerank_model_id = Column(String(21), nullable=True, comment="重排序模型ID")
|
||||
retrieval_weight = Column(Float, default=1.0, comment="检索权重（多知识库检索时的加权系数，0.1-10.0）")
|
||||
|
||||
# 预处理规则
|
||||
process_rules = Column(JSON, nullable=True, comment="预处理规则（清洗配置）")
|
||||
|
||||
# 索引模式
|
||||
indexing_technique = Column(String(20), default="high_quality", comment="索引模式: high_quality/economy")
|
||||
|
||||
# 统计
|
||||
document_count = Column(Integer, default=0, comment="文档数量")
|
||||
segment_count = Column(Integer, default=0, comment="分段数量")
|
||||
total_token_count = Column(Integer, default=0, comment="总 Token 数")
|
||||
total_char_count = Column(Integer, default=0, comment="总字符数")
|
||||
|
||||
# 状态
|
||||
status = Column(String(20), default="active", comment="状态: active/disabled")
|
||||
@@ -0,0 +1,41 @@
|
||||
"""
|
||||
知识库检索日志模型
|
||||
|
||||
记录每次检索的查询、结果、耗时等信息，用于分析检索质量。
|
||||
"""
|
||||
from sqlalchemy import Column, String, Text, Integer, Float, JSON, Index
|
||||
|
||||
from app.base_model import BaseModel
|
||||
|
||||
|
||||
class KnowledgeRetrievalLog(BaseModel):
|
||||
"""
|
||||
知识库检索日志
|
||||
|
||||
记录每次检索请求的完整信息，用于检索质量分析和优化。
|
||||
"""
|
||||
__tablename__ = "ai_knowledge_retrieval_log"
|
||||
|
||||
# 检索请求
|
||||
query = Column(Text, nullable=False, comment="查询文本")
|
||||
knowledge_base_ids = Column(JSON, nullable=False, comment="检索的知识库ID列表")
|
||||
retrieval_mode = Column(String(20), default="hybrid", comment="检索模式: vector/fulltext/hybrid")
|
||||
top_k = Column(Integer, default=5, comment="请求的返回数量")
|
||||
score_threshold = Column(Float, default=0.5, comment="相似度阈值")
|
||||
|
||||
# 检索结果
|
||||
result_count = Column(Integer, default=0, comment="实际返回结果数")
|
||||
results = Column(JSON, nullable=True, comment="检索结果摘要（segment_id/score/kb_id）")
|
||||
rerank_applied = Column(String(5), default="false", comment="是否应用了重排序")
|
||||
|
||||
# 性能
|
||||
elapsed_time = Column(Integer, default=0, comment="耗时（毫秒）")
|
||||
|
||||
# 来源
|
||||
source = Column(String(50), nullable=True, comment="调用来源: api/workflow/chat")
|
||||
user_id = Column(String(21), nullable=True, comment="操作用户ID")
|
||||
|
||||
__table_args__ = (
|
||||
Index('idx_retrieval_log_query_time', 'sys_create_datetime'),
|
||||
Index('idx_retrieval_log_user', 'user_id'),
|
||||
)
|
||||
@@ -0,0 +1,47 @@
|
||||
"""
|
||||
知识库分段模型
|
||||
"""
|
||||
from sqlalchemy import Column, String, Text, Integer, Boolean, JSON, Index
|
||||
|
||||
from app.base_model import BaseModel
|
||||
|
||||
|
||||
class KnowledgeSegment(BaseModel):
|
||||
"""
|
||||
知识库分段（Chunk）
|
||||
|
||||
文档经过分块后的最小检索单元
|
||||
向量数据存储在 Qdrant 向量数据库中，此表只存业务数据
|
||||
"""
|
||||
__tablename__ = "ai_knowledge_segment"
|
||||
|
||||
knowledge_base_id = Column(String(21), nullable=False, index=True, comment="所属知识库ID（逻辑外键关联ai_knowledge_base）")
|
||||
document_id = Column(String(21), nullable=False, index=True, comment="所属文档ID（逻辑外键关联ai_knowledge_document）")
|
||||
|
||||
# 内容
|
||||
position = Column(Integer, default=0, comment="在文档中的位置序号")
|
||||
content = Column(Text, nullable=False, comment="分段文本内容")
|
||||
answer = Column(Text, nullable=True, comment="Q&A 模式的答案内容")
|
||||
token_count = Column(Integer, default=0, comment="Token 数")
|
||||
char_count = Column(Integer, default=0, comment="字符数")
|
||||
word_count = Column(Integer, default=0, comment="词数")
|
||||
|
||||
# 元数据
|
||||
page_number = Column(Integer, nullable=True, comment="来源页码（PDF/PPT）")
|
||||
keywords = Column(JSON, nullable=True, comment="关键词列表（用于全文检索增强）")
|
||||
extra_metadata = Column(JSON, nullable=True, comment="元数据（标题/来源等）")
|
||||
|
||||
# 向量化状态（向量数据存在 Qdrant 中，这里只记录状态）
|
||||
embedding_status = Column(String(20), default="pending", comment="向量化状态: pending/completed/failed")
|
||||
|
||||
# 父子分段（Small-to-Big）
|
||||
parent_segment_id = Column(String(21), nullable=True, index=True, comment="父分段ID（逻辑外键，用于 Small-to-Big 检索）")
|
||||
|
||||
# 状态
|
||||
enabled = Column(Boolean, default=True, comment="是否启用（禁用后不参与检索）")
|
||||
hit_count = Column(Integer, default=0, comment="命中次数")
|
||||
|
||||
__table_args__ = (
|
||||
Index('idx_segment_kb_doc', 'knowledge_base_id', 'document_id'),
|
||||
Index('idx_segment_kb_enabled', 'knowledge_base_id', 'enabled'),
|
||||
)
|
||||
@@ -0,0 +1,6 @@
|
||||
"""
|
||||
知识库 Schema
|
||||
"""
|
||||
from .knowledge_base_schema import *
|
||||
from .document_schema import *
|
||||
from .segment_schema import *
|
||||
@@ -0,0 +1,37 @@
|
||||
"""
|
||||
知识库标注 Schema
|
||||
"""
|
||||
from typing import Optional, List
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, Field, ConfigDict
|
||||
|
||||
from app.base_schema import CSTDatetime
|
||||
|
||||
|
||||
class AnnotationCreateInput(BaseModel):
|
||||
"""创建标注"""
|
||||
question: str = Field(..., min_length=1, description="问题")
|
||||
answer: str = Field(..., min_length=1, description="答案")
|
||||
|
||||
|
||||
class AnnotationUpdateInput(BaseModel):
|
||||
"""更新标注"""
|
||||
question: Optional[str] = Field(None, min_length=1, description="问题")
|
||||
answer: Optional[str] = Field(None, min_length=1, description="答案")
|
||||
enabled: Optional[bool] = Field(None, description="是否启用")
|
||||
|
||||
|
||||
class AnnotationResponse(BaseModel):
|
||||
"""标注输出"""
|
||||
id: str
|
||||
knowledge_base_id: str
|
||||
question: str
|
||||
answer: str
|
||||
embedding_status: str = "pending"
|
||||
enabled: bool = True
|
||||
hit_count: int = 0
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
sys_update_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
@@ -0,0 +1,62 @@
|
||||
"""
|
||||
知识库文档 Schema
|
||||
"""
|
||||
from typing import Optional, List
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, Field, ConfigDict
|
||||
|
||||
from app.base_schema import CSTDatetime
|
||||
|
||||
|
||||
class DocumentUploadInput(BaseModel):
|
||||
"""文档上传输入（通过文件管理系统上传后传入 file_id）"""
|
||||
file_id: str = Field(..., description="文件ID（来自文件管理系统）")
|
||||
name: Optional[str] = Field(None, description="文档名称（不传则使用文件名）")
|
||||
|
||||
|
||||
class DocumentBatchUploadInput(BaseModel):
|
||||
"""批量文档上传"""
|
||||
file_ids: List[str] = Field(..., min_length=1, description="文件ID列表")
|
||||
|
||||
|
||||
class DocumentResponse(BaseModel):
|
||||
"""文档输出"""
|
||||
id: str
|
||||
knowledge_base_id: str
|
||||
file_id: Optional[str] = None
|
||||
name: str
|
||||
file_type: str = ""
|
||||
file_size: int = 0
|
||||
content_hash: str = ""
|
||||
segment_count: int = 0
|
||||
token_count: int = 0
|
||||
char_count: int = 0
|
||||
status: str = "pending"
|
||||
error_message: str = ""
|
||||
duplicate_warning: Optional[str] = None
|
||||
enabled: bool = True
|
||||
indexing_started_at: Optional[CSTDatetime] = None
|
||||
indexing_completed_at: Optional[CSTDatetime] = None
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
sys_update_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class DocumentListResponse(BaseModel):
|
||||
"""文档列表输出"""
|
||||
id: str
|
||||
knowledge_base_id: str
|
||||
file_id: Optional[str] = None
|
||||
name: str
|
||||
file_type: str = ""
|
||||
file_size: int = 0
|
||||
segment_count: int = 0
|
||||
token_count: int = 0
|
||||
status: str = "pending"
|
||||
duplicate_warning: Optional[str] = None
|
||||
enabled: bool = True
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
@@ -0,0 +1,115 @@
|
||||
"""
|
||||
知识库 Schema
|
||||
"""
|
||||
from typing import Optional, List, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, Field, ConfigDict
|
||||
|
||||
from app.base_schema import CSTDatetime
|
||||
|
||||
|
||||
class KnowledgeBaseCreate(BaseModel):
|
||||
"""创建知识库"""
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
application_id: Optional[str] = Field(None, description="所属应用ID")
|
||||
is_global: bool = Field(default=False, description="是否在子应用中可见")
|
||||
name: str = Field(..., max_length=100, description="知识库名称")
|
||||
code: str = Field(..., max_length=100, description="知识库编码")
|
||||
description: Optional[str] = Field(None, description="描述")
|
||||
icon: str = Field(default="", description="图标")
|
||||
embedding_model_id: Optional[str] = Field(None, description="Embedding 模型ID")
|
||||
embedding_dimensions: int = Field(default=1536, description="向量维度")
|
||||
chunk_strategy: str = Field(default="recursive", description="分块策略")
|
||||
chunk_size: int = Field(default=500, ge=100, le=4000, description="分块大小")
|
||||
chunk_overlap: int = Field(default=50, ge=0, le=500, description="分块重叠")
|
||||
separator: Optional[str] = Field(None, description="自定义分隔符")
|
||||
retrieval_mode: str = Field(default="hybrid", description="检索模式")
|
||||
top_k: int = Field(default=5, ge=1, le=20, description="检索数量")
|
||||
score_threshold: float = Field(default=0.5, ge=0, le=1, description="相似度阈值")
|
||||
rerank_enabled: bool = Field(default=False, description="是否启用重排序")
|
||||
rerank_model_id: Optional[str] = Field(None, description="重排序模型ID")
|
||||
retrieval_weight: float = Field(default=1.0, ge=0.1, le=10.0, description="检索权重")
|
||||
process_rules: Optional[Dict[str, Any]] = Field(None, description="预处理规则")
|
||||
indexing_technique: str = Field(default="high_quality", description="索引模式: high_quality/economy")
|
||||
|
||||
|
||||
class KnowledgeBaseUpdate(BaseModel):
|
||||
"""更新知识库"""
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
icon: Optional[str] = None
|
||||
embedding_model_id: Optional[str] = None
|
||||
embedding_dimensions: Optional[int] = None
|
||||
chunk_strategy: Optional[str] = None
|
||||
chunk_size: Optional[int] = Field(None, ge=100, le=4000)
|
||||
chunk_overlap: Optional[int] = Field(None, ge=0, le=500)
|
||||
separator: Optional[str] = None
|
||||
retrieval_mode: Optional[str] = None
|
||||
top_k: Optional[int] = Field(None, ge=1, le=20)
|
||||
score_threshold: Optional[float] = Field(None, ge=0, le=1)
|
||||
rerank_enabled: Optional[bool] = None
|
||||
rerank_model_id: Optional[str] = None
|
||||
retrieval_weight: Optional[float] = Field(None, ge=0.1, le=10.0)
|
||||
process_rules: Optional[Dict[str, Any]] = None
|
||||
indexing_technique: Optional[str] = None
|
||||
status: Optional[str] = None
|
||||
is_global: Optional[bool] = None
|
||||
|
||||
|
||||
class KnowledgeBaseResponse(BaseModel):
|
||||
"""知识库详情输出"""
|
||||
id: str
|
||||
application_id: Optional[str] = None
|
||||
is_global: bool = False
|
||||
name: str
|
||||
code: str
|
||||
description: str = ""
|
||||
icon: str = ""
|
||||
embedding_model_id: Optional[str] = None
|
||||
embedding_model_name: str = ""
|
||||
embedding_dimensions: int = 1536
|
||||
chunk_strategy: str = "recursive"
|
||||
chunk_size: int = 500
|
||||
chunk_overlap: int = 50
|
||||
separator: Optional[str] = None
|
||||
retrieval_mode: str = "hybrid"
|
||||
top_k: int = 5
|
||||
score_threshold: float = 0.5
|
||||
rerank_enabled: bool = False
|
||||
rerank_model_id: Optional[str] = None
|
||||
retrieval_weight: float = 1.0
|
||||
process_rules: Optional[Dict[str, Any]] = None
|
||||
indexing_technique: str = "high_quality"
|
||||
document_count: int = 0
|
||||
segment_count: int = 0
|
||||
total_token_count: int = 0
|
||||
total_char_count: int = 0
|
||||
status: str = "active"
|
||||
sort: int = 0
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
sys_update_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True, protected_namespaces=())
|
||||
|
||||
|
||||
class KnowledgeBaseListResponse(BaseModel):
|
||||
"""知识库列表输出"""
|
||||
id: str
|
||||
application_id: Optional[str] = None
|
||||
application_name: str = ""
|
||||
is_global: bool = False
|
||||
name: str
|
||||
code: str
|
||||
description: str = ""
|
||||
icon: str = ""
|
||||
embedding_model_name: str = ""
|
||||
document_count: int = 0
|
||||
segment_count: int = 0
|
||||
status: str = "active"
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True, protected_namespaces=())
|
||||
@@ -0,0 +1,28 @@
|
||||
"""
|
||||
检索日志 Schema
|
||||
"""
|
||||
from typing import Optional, List, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, Field, ConfigDict
|
||||
|
||||
from app.base_schema import CSTDatetime
|
||||
|
||||
|
||||
class RetrievalLogResponse(BaseModel):
|
||||
"""检索日志输出"""
|
||||
id: str
|
||||
query: str
|
||||
knowledge_base_ids: List[str] = []
|
||||
retrieval_mode: str = "hybrid"
|
||||
top_k: int = 5
|
||||
score_threshold: float = 0.5
|
||||
result_count: int = 0
|
||||
results: Optional[List[Dict[str, Any]]] = None
|
||||
rerank_applied: str = "false"
|
||||
elapsed_time: int = 0
|
||||
source: Optional[str] = None
|
||||
user_id: Optional[str] = None
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
@@ -0,0 +1,140 @@
|
||||
"""
|
||||
知识库分段 Schema
|
||||
"""
|
||||
from typing import Optional, List, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, Field, ConfigDict
|
||||
|
||||
from app.base_schema import CSTDatetime
|
||||
|
||||
|
||||
class SegmentResponse(BaseModel):
|
||||
"""分段输出"""
|
||||
id: str
|
||||
knowledge_base_id: str
|
||||
document_id: str
|
||||
document_name: str = ""
|
||||
position: int = 0
|
||||
content: str
|
||||
answer: Optional[str] = None
|
||||
token_count: int = 0
|
||||
char_count: int = 0
|
||||
word_count: int = 0
|
||||
page_number: Optional[int] = None
|
||||
keywords: Optional[List[str]] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
embedding_status: str = "pending"
|
||||
enabled: bool = True
|
||||
hit_count: int = 0
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class SegmentListResponse(BaseModel):
|
||||
"""分段列表输出"""
|
||||
id: str
|
||||
document_id: str
|
||||
document_name: str = ""
|
||||
position: int = 0
|
||||
content: str
|
||||
answer: Optional[str] = None
|
||||
token_count: int = 0
|
||||
char_count: int = 0
|
||||
word_count: int = 0
|
||||
page_number: Optional[int] = None
|
||||
keywords: Optional[List[str]] = None
|
||||
extra_metadata: Optional[Dict[str, Any]] = None
|
||||
enabled: bool = True
|
||||
hit_count: int = 0
|
||||
embedding_status: str = "pending"
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class SegmentUpdateInput(BaseModel):
|
||||
"""更新分段"""
|
||||
content: Optional[str] = Field(None, description="分段内容")
|
||||
keywords: Optional[List[str]] = Field(None, description="关键词")
|
||||
enabled: Optional[bool] = Field(None, description="是否启用")
|
||||
extra_metadata: Optional[Dict[str, Any]] = Field(None, description="元数据")
|
||||
|
||||
|
||||
class SegmentCreateInput(BaseModel):
|
||||
"""手动创建分段"""
|
||||
content: str = Field(..., min_length=1, description="分段内容")
|
||||
answer: Optional[str] = Field(None, description="Q&A 模式的答案")
|
||||
keywords: Optional[List[str]] = Field(None, description="关键词")
|
||||
|
||||
|
||||
class ChunkPreviewInput(BaseModel):
|
||||
"""分块预览输入"""
|
||||
file_id: str = Field(..., description="文件ID")
|
||||
chunk_strategy: str = Field(default="recursive", description="分块策略")
|
||||
chunk_size: int = Field(default=500, ge=100, le=4000, description="分块大小")
|
||||
chunk_overlap: int = Field(default=50, ge=0, le=500, description="分块重叠")
|
||||
separator: Optional[str] = Field(None, description="自定义分隔符")
|
||||
process_rules: Optional[Dict[str, Any]] = Field(None, description="预处理规则")
|
||||
|
||||
|
||||
class ChunkPreviewItem(BaseModel):
|
||||
"""分块预览结果项"""
|
||||
position: int = 0
|
||||
content: str = ""
|
||||
char_count: int = 0
|
||||
token_count: int = 0
|
||||
word_count: int = 0
|
||||
answer: Optional[str] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class ChunkPreviewResponse(BaseModel):
|
||||
"""分块预览响应"""
|
||||
chunks: List[ChunkPreviewItem] = Field(default_factory=list)
|
||||
total: int = 0
|
||||
strategy: str = ""
|
||||
chunk_size: int = 0
|
||||
chunk_overlap: int = 0
|
||||
|
||||
|
||||
class RetrievalInput(BaseModel):
|
||||
"""检索输入"""
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
query: str = Field(..., min_length=1, description="查询文本")
|
||||
knowledge_base_ids: List[str] = Field(..., min_length=1, description="知识库ID列表")
|
||||
top_k: int = Field(default=5, ge=1, le=20, description="返回数量")
|
||||
score_threshold: float = Field(default=0.5, ge=0, le=1, description="相似度阈值")
|
||||
retrieval_mode: Optional[str] = Field(None, description="检索模式（不传则使用知识库配置）")
|
||||
rerank_enabled: Optional[bool] = Field(None, description="是否启用重排序（不传则使用知识库配置）")
|
||||
rerank_model_id: Optional[str] = Field(None, description="重排序模型ID（不传则使用知识库配置）")
|
||||
metadata_filter: Optional[Dict[str, Any]] = Field(None, description="元数据过滤条件")
|
||||
|
||||
|
||||
class RetrievalResult(BaseModel):
|
||||
"""检索结果"""
|
||||
segment_id: str
|
||||
document_id: str
|
||||
document_name: str = ""
|
||||
knowledge_base_id: str
|
||||
knowledge_base_name: str = ""
|
||||
content: str
|
||||
score: float = 0.0
|
||||
token_count: int = 0
|
||||
page_number: Optional[int] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
keywords: Optional[List[str]] = None
|
||||
match_source: Optional[str] = Field(None, description="命中来源: vector/fulltext/annotation")
|
||||
parent_content: Optional[str] = Field(None, description="父分段内容（Small-to-Big 模式）")
|
||||
|
||||
|
||||
class RetrievalResponse(BaseModel):
|
||||
"""检索响应"""
|
||||
results: List[RetrievalResult] = Field(default_factory=list)
|
||||
total: int = 0
|
||||
query: str = ""
|
||||
elapsed_time: int = 0
|
||||
retrieval_mode: str = ""
|
||||
rerank_applied: bool = False
|
||||
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
知识库服务
|
||||
"""
|
||||
@@ -0,0 +1,159 @@
|
||||
"""
|
||||
文档预处理/清洗服务
|
||||
|
||||
参考 Dify 的 DatasetProcessRule，支持可配置的文本清洗规则。
|
||||
在文本提取之后、分块之前执行。
|
||||
"""
|
||||
import logging
|
||||
import re
|
||||
from typing import Dict, Any, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 默认预处理规则
|
||||
DEFAULT_PROCESS_RULES: Dict[str, Any] = {
|
||||
"pre_processing_rules": [
|
||||
{"id": "remove_extra_spaces", "enabled": True},
|
||||
{"id": "remove_urls_emails", "enabled": False},
|
||||
{"id": "remove_html_tags", "enabled": False},
|
||||
{"id": "remove_consecutive_newlines", "enabled": True},
|
||||
{"id": "remove_trailing_whitespace", "enabled": True},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
class CleaningService:
|
||||
"""
|
||||
文本清洗服务
|
||||
|
||||
支持的清洗规则：
|
||||
- remove_extra_spaces: 合并连续空格为单个空格
|
||||
- remove_urls_emails: 移除 URL 和邮箱地址
|
||||
- remove_html_tags: 移除 HTML 标签
|
||||
- remove_consecutive_newlines: 合并连续空行（3+）为双空行
|
||||
- remove_trailing_whitespace: 去除行尾空白
|
||||
"""
|
||||
|
||||
# 规则处理器映射
|
||||
RULE_PROCESSORS = {
|
||||
"remove_extra_spaces": "_remove_extra_spaces",
|
||||
"remove_urls_emails": "_remove_urls_emails",
|
||||
"remove_html_tags": "_remove_html_tags",
|
||||
"remove_consecutive_newlines": "_remove_consecutive_newlines",
|
||||
"remove_trailing_whitespace": "_remove_trailing_whitespace",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def clean(cls, text: str, process_rules: Optional[Dict[str, Any]] = None) -> str:
|
||||
"""
|
||||
根据预处理规则清洗文本
|
||||
|
||||
Args:
|
||||
text: 原始文本
|
||||
process_rules: 预处理规则配置，为 None 则使用默认规则
|
||||
|
||||
Returns:
|
||||
清洗后的文本
|
||||
"""
|
||||
if not text:
|
||||
return text
|
||||
|
||||
rules = process_rules or DEFAULT_PROCESS_RULES
|
||||
pre_rules = rules.get("pre_processing_rules", [])
|
||||
|
||||
original_length = len(text)
|
||||
|
||||
for rule in pre_rules:
|
||||
rule_id = rule.get("id", "")
|
||||
enabled = rule.get("enabled", False)
|
||||
|
||||
if not enabled:
|
||||
continue
|
||||
|
||||
processor_name = cls.RULE_PROCESSORS.get(rule_id)
|
||||
if not processor_name:
|
||||
logger.warning(f"未知的预处理规则: {rule_id}")
|
||||
continue
|
||||
|
||||
processor = getattr(cls, processor_name, None)
|
||||
if processor:
|
||||
text = processor(text)
|
||||
|
||||
cleaned_length = len(text)
|
||||
if original_length != cleaned_length:
|
||||
logger.info(
|
||||
f"文本清洗完成: {original_length} -> {cleaned_length} 字符 "
|
||||
f"(减少 {original_length - cleaned_length})"
|
||||
)
|
||||
|
||||
return text.strip()
|
||||
|
||||
@staticmethod
|
||||
def _remove_extra_spaces(text: str) -> str:
|
||||
"""合并连续空格为单个空格（保留换行符）"""
|
||||
# 只处理同一行内的连续空格，不影响换行
|
||||
lines = text.split('\n')
|
||||
cleaned_lines = []
|
||||
for line in lines:
|
||||
cleaned_lines.append(re.sub(r'[ \t]+', ' ', line))
|
||||
return '\n'.join(cleaned_lines)
|
||||
|
||||
@staticmethod
|
||||
def _remove_urls_emails(text: str) -> str:
|
||||
"""移除 URL 和邮箱地址"""
|
||||
# 移除 URL
|
||||
text = re.sub(
|
||||
r'https?://[^\s<>"{}|\\^`\[\]]+',
|
||||
'',
|
||||
text,
|
||||
)
|
||||
# 移除邮箱
|
||||
text = re.sub(
|
||||
r'[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}',
|
||||
'',
|
||||
text,
|
||||
)
|
||||
return text
|
||||
|
||||
@staticmethod
|
||||
def _remove_html_tags(text: str) -> str:
|
||||
"""移除 HTML 标签，保留文本内容"""
|
||||
# 移除 script 和 style 标签及其内容
|
||||
text = re.sub(r'<script[^>]*>.*?</script>', '', text, flags=re.DOTALL | re.IGNORECASE)
|
||||
text = re.sub(r'<style[^>]*>.*?</style>', '', text, flags=re.DOTALL | re.IGNORECASE)
|
||||
# 移除所有 HTML 标签
|
||||
text = re.sub(r'<[^>]+>', '', text)
|
||||
# 解码常见 HTML 实体
|
||||
text = text.replace(' ', ' ')
|
||||
text = text.replace('<', '<')
|
||||
text = text.replace('>', '>')
|
||||
text = text.replace('&', '&')
|
||||
text = text.replace('"', '"')
|
||||
text = text.replace(''', "'")
|
||||
return text
|
||||
|
||||
@staticmethod
|
||||
def _remove_consecutive_newlines(text: str) -> str:
|
||||
"""合并连续空行（3个以上换行）为双换行"""
|
||||
return re.sub(r'\n{3,}', '\n\n', text)
|
||||
|
||||
@staticmethod
|
||||
def _remove_trailing_whitespace(text: str) -> str:
|
||||
"""去除每行行尾空白"""
|
||||
return '\n'.join(line.rstrip() for line in text.split('\n'))
|
||||
|
||||
@classmethod
|
||||
def get_default_rules(cls) -> Dict[str, Any]:
|
||||
"""获取默认预处理规则"""
|
||||
return DEFAULT_PROCESS_RULES.copy()
|
||||
|
||||
@classmethod
|
||||
def get_available_rules(cls) -> List[Dict[str, str]]:
|
||||
"""获取所有可用的预处理规则"""
|
||||
return [
|
||||
{"id": "remove_extra_spaces", "label": "合并连续空格"},
|
||||
{"id": "remove_urls_emails", "label": "移除 URL 和邮箱"},
|
||||
{"id": "remove_html_tags", "label": "移除 HTML 标签"},
|
||||
{"id": "remove_consecutive_newlines", "label": "合并连续空行"},
|
||||
{"id": "remove_trailing_whitespace", "label": "去除行尾空白"},
|
||||
]
|
||||
@@ -0,0 +1,425 @@
|
||||
"""
|
||||
文档服务
|
||||
|
||||
文档上传、管理、状态控制
|
||||
"""
|
||||
import logging
|
||||
from typing import Optional, List, Tuple
|
||||
|
||||
from sqlalchemy import select, func, update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from ai_platform.knowledge.models import KnowledgeBase, KnowledgeDocument, KnowledgeSegment
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DocumentService:
|
||||
"""文档服务"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self._db = db
|
||||
|
||||
async def get_list(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
page: int = 1,
|
||||
page_size: int = 20,
|
||||
name: Optional[str] = None,
|
||||
status: Optional[str] = None,
|
||||
) -> Tuple[List[KnowledgeDocument], int]:
|
||||
"""获取文档列表"""
|
||||
query = select(KnowledgeDocument).where(
|
||||
KnowledgeDocument.knowledge_base_id == knowledge_base_id,
|
||||
KnowledgeDocument.is_deleted == False,
|
||||
)
|
||||
|
||||
if name:
|
||||
query = query.where(KnowledgeDocument.name.ilike(f"%{name}%"))
|
||||
if status:
|
||||
query = query.where(KnowledgeDocument.status == status)
|
||||
|
||||
count_result = await self._db.execute(select(func.count()).select_from(query.subquery()))
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(KnowledgeDocument.sys_create_datetime.desc())
|
||||
query = query.offset(offset).limit(page_size)
|
||||
|
||||
result = await self._db.execute(query)
|
||||
items = result.scalars().all()
|
||||
|
||||
return items, total
|
||||
|
||||
async def get_by_id(self, doc_id: str) -> Optional[KnowledgeDocument]:
|
||||
"""获取文档详情"""
|
||||
result = await self._db.execute(
|
||||
select(KnowledgeDocument).where(
|
||||
KnowledgeDocument.id == doc_id,
|
||||
KnowledgeDocument.is_deleted == False
|
||||
)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def add_document(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
file_id: str,
|
||||
name: Optional[str] = None,
|
||||
) -> KnowledgeDocument:
|
||||
"""
|
||||
添加文档到知识库
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
file_id: 文件管理系统中的文件 ID
|
||||
name: 文档名称（不传则从文件信息获取）
|
||||
"""
|
||||
# 验证知识库存在
|
||||
kb_result = await self._db.execute(
|
||||
select(KnowledgeBase).where(
|
||||
KnowledgeBase.id == knowledge_base_id,
|
||||
KnowledgeBase.is_deleted == False
|
||||
)
|
||||
)
|
||||
kb = kb_result.scalar_one_or_none()
|
||||
if not kb:
|
||||
raise ValueError('知识库不存在')
|
||||
|
||||
# 获取文件信息
|
||||
from core.file_manager.model import FileManager
|
||||
file_result = await self._db.execute(
|
||||
select(FileManager).where(
|
||||
FileManager.id == file_id,
|
||||
FileManager.is_deleted == False
|
||||
)
|
||||
)
|
||||
file_info = file_result.scalar_one_or_none()
|
||||
if not file_info:
|
||||
raise ValueError('文件不存在')
|
||||
|
||||
# 检查是否已添加（通过 file_id 去重）
|
||||
existing = await self._db.execute(
|
||||
select(KnowledgeDocument).where(
|
||||
KnowledgeDocument.knowledge_base_id == knowledge_base_id,
|
||||
KnowledgeDocument.file_id == file_id,
|
||||
KnowledgeDocument.is_deleted == False,
|
||||
)
|
||||
)
|
||||
if existing.scalars().first():
|
||||
raise ValueError('该文件已添加到知识库')
|
||||
|
||||
# 通过文件 MD5 检测内容重复（跨知识库）
|
||||
duplicate_warning = None
|
||||
if file_info.md5:
|
||||
dup_result = await self._db.execute(
|
||||
select(KnowledgeDocument).where(
|
||||
KnowledgeDocument.content_hash == file_info.md5,
|
||||
KnowledgeDocument.is_deleted == False,
|
||||
KnowledgeDocument.knowledge_base_id != knowledge_base_id,
|
||||
).limit(1)
|
||||
)
|
||||
dup_doc = dup_result.scalar_one_or_none()
|
||||
if dup_doc:
|
||||
duplicate_warning = f'该文件内容与其他知识库中的文档 "{dup_doc.name}" 重复'
|
||||
logger.info(f'文档内容重复检测: file_id={file_id}, 重复文档={dup_doc.id}')
|
||||
|
||||
# 同知识库内容去重（严格阻止）
|
||||
same_kb_dup = await self._db.execute(
|
||||
select(KnowledgeDocument).where(
|
||||
KnowledgeDocument.content_hash == file_info.md5,
|
||||
KnowledgeDocument.knowledge_base_id == knowledge_base_id,
|
||||
KnowledgeDocument.is_deleted == False,
|
||||
).limit(1)
|
||||
)
|
||||
if same_kb_dup.scalar_one_or_none():
|
||||
raise ValueError('该知识库中已存在相同内容的文档')
|
||||
|
||||
doc = KnowledgeDocument(
|
||||
knowledge_base_id=knowledge_base_id,
|
||||
file_id=file_id,
|
||||
name=name or file_info.name,
|
||||
file_type=file_info.file_ext or '',
|
||||
file_size=file_info.size or 0,
|
||||
content_hash=file_info.md5 or '',
|
||||
status='pending',
|
||||
duplicate_warning=duplicate_warning,
|
||||
)
|
||||
self._db.add(doc)
|
||||
await self._db.commit()
|
||||
await self._db.refresh(doc)
|
||||
|
||||
return doc
|
||||
|
||||
async def batch_add_documents(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
file_ids: List[str],
|
||||
) -> List[KnowledgeDocument]:
|
||||
"""批量添加文档"""
|
||||
docs = []
|
||||
for file_id in file_ids:
|
||||
try:
|
||||
doc = await self.add_document(knowledge_base_id, file_id)
|
||||
docs.append(doc)
|
||||
except ValueError as e:
|
||||
logger.warning(f'添加文档失败 (file_id={file_id}): {e}')
|
||||
continue
|
||||
return docs
|
||||
|
||||
async def delete_document(self, doc_id: str) -> bool:
|
||||
"""删除文档（软删除，同时删除分段 + 清理 Qdrant 向量）"""
|
||||
doc = await self.get_by_id(doc_id)
|
||||
if not doc:
|
||||
return False
|
||||
|
||||
doc.is_deleted = True
|
||||
|
||||
# 软删除关联分段
|
||||
await self._db.execute(
|
||||
update(KnowledgeSegment).where(
|
||||
KnowledgeSegment.document_id == doc_id
|
||||
).values(is_deleted=True)
|
||||
)
|
||||
|
||||
# 更新知识库统计
|
||||
from ai_platform.knowledge.services.indexing_service import IndexingService
|
||||
indexing_service = IndexingService(self._db)
|
||||
await indexing_service._update_kb_stats(doc.knowledge_base_id)
|
||||
|
||||
await self._db.commit()
|
||||
|
||||
# 从 Qdrant 删除该文档的所有向量
|
||||
try:
|
||||
from ai_platform.knowledge.vector_store import get_vector_store
|
||||
vector_store = get_vector_store()
|
||||
await vector_store.delete_by_filter(
|
||||
doc.knowledge_base_id,
|
||||
filter_conditions={'document_id': doc_id},
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f'从 Qdrant 删除文档向量失败: {e}')
|
||||
|
||||
return True
|
||||
|
||||
async def toggle_document(self, doc_id: str, enabled: bool) -> Optional[KnowledgeDocument]:
|
||||
"""启用/禁用文档"""
|
||||
doc = await self.get_by_id(doc_id)
|
||||
if not doc:
|
||||
return None
|
||||
|
||||
doc.enabled = enabled
|
||||
|
||||
# 同时启用/禁用关联分段
|
||||
await self._db.execute(
|
||||
update(KnowledgeSegment).where(
|
||||
KnowledgeSegment.document_id == doc_id,
|
||||
KnowledgeSegment.is_deleted == False,
|
||||
).values(enabled=enabled)
|
||||
)
|
||||
|
||||
await self._db.commit()
|
||||
await self._db.refresh(doc)
|
||||
return doc
|
||||
|
||||
async def get_segments(
|
||||
self,
|
||||
document_id: str,
|
||||
page: int = 1,
|
||||
page_size: int = 20,
|
||||
keyword: Optional[str] = None,
|
||||
) -> Tuple[List[KnowledgeSegment], int]:
|
||||
"""获取文档的分段列表"""
|
||||
query = select(KnowledgeSegment).where(
|
||||
KnowledgeSegment.document_id == document_id,
|
||||
KnowledgeSegment.is_deleted == False,
|
||||
)
|
||||
|
||||
if keyword:
|
||||
query = query.where(KnowledgeSegment.content.ilike(f"%{keyword}%"))
|
||||
|
||||
count_result = await self._db.execute(select(func.count()).select_from(query.subquery()))
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(KnowledgeSegment.position.asc())
|
||||
query = query.offset(offset).limit(page_size)
|
||||
|
||||
result = await self._db.execute(query)
|
||||
items = result.scalars().all()
|
||||
|
||||
return items, total
|
||||
|
||||
async def get_kb_segments(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
page: int = 1,
|
||||
page_size: int = 20,
|
||||
keyword: Optional[str] = None,
|
||||
enabled: Optional[bool] = None,
|
||||
embedding_status: Optional[str] = None,
|
||||
metadata_key: Optional[str] = None,
|
||||
metadata_value: Optional[str] = None,
|
||||
) -> Tuple[List[KnowledgeSegment], int]:
|
||||
"""获取知识库的所有分段"""
|
||||
from app.db_compat import json_extract, json_has_key
|
||||
|
||||
query = select(KnowledgeSegment).where(
|
||||
KnowledgeSegment.knowledge_base_id == knowledge_base_id,
|
||||
KnowledgeSegment.is_deleted == False,
|
||||
)
|
||||
|
||||
if keyword:
|
||||
query = query.where(KnowledgeSegment.content.ilike(f"%{keyword}%"))
|
||||
if enabled is not None:
|
||||
query = query.where(KnowledgeSegment.enabled == enabled)
|
||||
if embedding_status:
|
||||
query = query.where(KnowledgeSegment.embedding_status == embedding_status)
|
||||
if metadata_key:
|
||||
if metadata_value:
|
||||
query = query.where(
|
||||
json_extract(KnowledgeSegment.extra_metadata, metadata_key) == metadata_value
|
||||
)
|
||||
else:
|
||||
query = query.where(
|
||||
json_has_key(KnowledgeSegment.extra_metadata, metadata_key)
|
||||
)
|
||||
|
||||
count_result = await self._db.execute(select(func.count()).select_from(query.subquery()))
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(KnowledgeSegment.document_id, KnowledgeSegment.position.asc())
|
||||
query = query.offset(offset).limit(page_size)
|
||||
|
||||
result = await self._db.execute(query)
|
||||
items = result.scalars().all()
|
||||
|
||||
return items, total
|
||||
|
||||
async def update_segment(
|
||||
self,
|
||||
segment_id: str,
|
||||
content: Optional[str] = None,
|
||||
keywords: Optional[List[str]] = None,
|
||||
enabled: Optional[bool] = None,
|
||||
extra_metadata: Optional[dict] = None,
|
||||
) -> Optional[KnowledgeSegment]:
|
||||
"""更新分段"""
|
||||
result = await self._db.execute(
|
||||
select(KnowledgeSegment).where(
|
||||
KnowledgeSegment.id == segment_id,
|
||||
KnowledgeSegment.is_deleted == False
|
||||
)
|
||||
)
|
||||
segment = result.scalar_one_or_none()
|
||||
if not segment:
|
||||
return None
|
||||
|
||||
need_reindex = False
|
||||
if content is not None and content != segment.content:
|
||||
segment.content = content
|
||||
segment.char_count = len(content)
|
||||
segment.embedding_status = 'pending'
|
||||
need_reindex = True
|
||||
if keywords is not None:
|
||||
segment.keywords = keywords
|
||||
if enabled is not None:
|
||||
segment.enabled = enabled
|
||||
if extra_metadata is not None:
|
||||
segment.extra_metadata = extra_metadata
|
||||
|
||||
await self._db.commit()
|
||||
|
||||
# 如果内容变更，重新向量化
|
||||
if need_reindex:
|
||||
from ai_platform.knowledge.services.indexing_service import IndexingService
|
||||
indexing_service = IndexingService(self._db)
|
||||
await indexing_service.index_segment(segment.knowledge_base_id, segment_id)
|
||||
|
||||
await self._db.refresh(segment)
|
||||
return segment
|
||||
|
||||
async def add_segment(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
document_id: str,
|
||||
content: str,
|
||||
keywords: Optional[List[str]] = None,
|
||||
) -> KnowledgeSegment:
|
||||
"""手动添加分段"""
|
||||
# 获取当前最大 position
|
||||
max_pos_result = await self._db.execute(
|
||||
select(func.max(KnowledgeSegment.position)).where(
|
||||
KnowledgeSegment.document_id == document_id,
|
||||
KnowledgeSegment.is_deleted == False,
|
||||
)
|
||||
)
|
||||
max_pos = max_pos_result.scalar() or 0
|
||||
|
||||
segment = KnowledgeSegment(
|
||||
knowledge_base_id=knowledge_base_id,
|
||||
document_id=document_id,
|
||||
position=max_pos + 1,
|
||||
content=content,
|
||||
char_count=len(content),
|
||||
keywords=keywords,
|
||||
embedding_status='pending',
|
||||
enabled=True,
|
||||
)
|
||||
self._db.add(segment)
|
||||
await self._db.commit()
|
||||
await self._db.refresh(segment)
|
||||
|
||||
# 向量化
|
||||
from ai_platform.knowledge.services.indexing_service import IndexingService
|
||||
indexing_service = IndexingService(self._db)
|
||||
await indexing_service.index_segment(knowledge_base_id, segment.id)
|
||||
|
||||
# 更新统计
|
||||
await indexing_service._update_kb_stats(knowledge_base_id)
|
||||
await self._db.commit()
|
||||
|
||||
await self._db.refresh(segment)
|
||||
return segment
|
||||
|
||||
async def delete_segment(self, segment_id: str) -> bool:
|
||||
"""删除分段（软删除 + 清理 Qdrant 向量）"""
|
||||
from sqlalchemy import update
|
||||
|
||||
# 先查询获取 kb_id
|
||||
result = await self._db.execute(
|
||||
select(KnowledgeSegment.knowledge_base_id).where(
|
||||
KnowledgeSegment.id == segment_id,
|
||||
KnowledgeSegment.is_deleted == False
|
||||
)
|
||||
)
|
||||
row = result.first()
|
||||
if not row:
|
||||
return False
|
||||
|
||||
kb_id = str(row[0])
|
||||
|
||||
# 直接 SQL UPDATE 避免并发场景下的 StaleDataError
|
||||
await self._db.execute(
|
||||
update(KnowledgeSegment)
|
||||
.where(KnowledgeSegment.id == segment_id)
|
||||
.values(is_deleted=True)
|
||||
)
|
||||
await self._db.commit()
|
||||
|
||||
# 从 Qdrant 删除该分段的向量
|
||||
try:
|
||||
from ai_platform.knowledge.vector_store import get_vector_store
|
||||
vector_store = get_vector_store()
|
||||
await vector_store.delete(kb_id, [str(segment_id)])
|
||||
except Exception as e:
|
||||
logger.warning(f'从 Qdrant 删除分段向量失败: {e}')
|
||||
|
||||
# 更新统计
|
||||
from ai_platform.knowledge.services.indexing_service import IndexingService
|
||||
indexing_service = IndexingService(self._db)
|
||||
await indexing_service._update_kb_stats(kb_id)
|
||||
await self._db.commit()
|
||||
|
||||
return True
|
||||
@@ -0,0 +1,206 @@
|
||||
"""
|
||||
Embedding 服务
|
||||
|
||||
通过 OpenAI 兼容的 Embeddings API 将文本转换为向量
|
||||
支持所有兼容 OpenAI 接口的提供商（OpenAI、Qwen、Ollama 等）
|
||||
"""
|
||||
import logging
|
||||
import math
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 单次批量请求的默认最大文本数
|
||||
DEFAULT_BATCH_SIZE = 50
|
||||
|
||||
# 不同提供商的批次大小限制
|
||||
PROVIDER_BATCH_SIZE = {
|
||||
'qwen': 10, # 阿里云 DashScope 限制 10
|
||||
'dashscope': 10,
|
||||
'siliconflow': 10, # 硅基流动限制较小
|
||||
'ollama': 1, # Ollama 通常逐条处理
|
||||
}
|
||||
|
||||
|
||||
|
||||
class EmbeddingService:
|
||||
"""
|
||||
Embedding 服务
|
||||
|
||||
通过模型 ID 获取对应的提供商，调用 OpenAI 兼容的 Embeddings API
|
||||
"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self._db = db
|
||||
self._client_cache = {}
|
||||
|
||||
async def _get_client_and_model(self, model_id: str):
|
||||
"""
|
||||
根据模型 ID 获取异步客户端和模型名称
|
||||
|
||||
Returns:
|
||||
(async_client, model_name, max_tokens, provider_type)
|
||||
"""
|
||||
from ai_platform.models import LLMModel, LLMProvider
|
||||
|
||||
result = await self._db.execute(
|
||||
select(LLMModel).where(
|
||||
LLMModel.id == model_id,
|
||||
LLMModel.is_active == True,
|
||||
LLMModel.is_deleted == False
|
||||
)
|
||||
)
|
||||
model = result.scalar_one_or_none()
|
||||
if not model:
|
||||
raise ValueError(f'Embedding 模型不存在或已禁用: {model_id}')
|
||||
if model.model_type != 'embedding':
|
||||
raise ValueError(f'模型 {model.display_name} 不是 Embedding 类型')
|
||||
|
||||
provider_result = await self._db.execute(
|
||||
select(LLMProvider).where(
|
||||
LLMProvider.id == model.provider_id,
|
||||
LLMProvider.is_active == True,
|
||||
LLMProvider.is_deleted == False
|
||||
)
|
||||
)
|
||||
provider = provider_result.scalar_one_or_none()
|
||||
if not provider:
|
||||
raise ValueError('Embedding 模型对应的提供商不存在或已禁用')
|
||||
|
||||
cache_key = str(provider.id)
|
||||
if cache_key not in self._client_cache:
|
||||
import httpx
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
# 根据提供商类型确定 base_url
|
||||
if provider.provider_type == 'ollama':
|
||||
base_url = (provider.ollama_host or 'http://localhost:11434').rstrip('/') + '/v1'
|
||||
else:
|
||||
from ai_platform.providers.registry import ProviderRegistry
|
||||
provider_class = ProviderRegistry.get(provider.provider_type)
|
||||
default_base = getattr(provider_class, 'DEFAULT_API_BASE', 'https://api.openai.com/v1') if provider_class else 'https://api.openai.com/v1'
|
||||
base_url = provider.api_base or default_base
|
||||
|
||||
self._client_cache[cache_key] = AsyncOpenAI(
|
||||
api_key=provider.api_key or 'ollama',
|
||||
base_url=base_url,
|
||||
timeout=httpx.Timeout(120.0, connect=30.0),
|
||||
max_retries=5,
|
||||
)
|
||||
|
||||
return self._client_cache[cache_key], model.model_name, model.context_window or 8191, provider.provider_type
|
||||
|
||||
async def embed_text(
|
||||
self,
|
||||
model_id: str,
|
||||
text: str,
|
||||
dimensions: Optional[int] = None,
|
||||
) -> List[float]:
|
||||
"""
|
||||
将单个文本转换为向量
|
||||
|
||||
Args:
|
||||
model_id: Embedding 模型 ID
|
||||
text: 文本内容
|
||||
dimensions: 向量维度（可选，部分模型支持）
|
||||
|
||||
Returns:
|
||||
向量列表 List[float]
|
||||
"""
|
||||
results = await self.embed_texts(model_id, [text], dimensions)
|
||||
return results[0]
|
||||
|
||||
async def embed_texts(
|
||||
self,
|
||||
model_id: str,
|
||||
texts: List[str],
|
||||
dimensions: Optional[int] = None,
|
||||
) -> List[List[float]]:
|
||||
"""
|
||||
批量将文本转换为向量
|
||||
|
||||
Args:
|
||||
model_id: Embedding 模型 ID
|
||||
texts: 文本列表
|
||||
dimensions: 向量维度（可选）
|
||||
|
||||
Returns:
|
||||
向量列表 List[List[float]]
|
||||
"""
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
client, model_name, max_tokens, provider_type = await self._get_client_and_model(model_id)
|
||||
|
||||
# 根据提供商类型确定批次大小
|
||||
batch_size = PROVIDER_BATCH_SIZE.get(provider_type, DEFAULT_BATCH_SIZE)
|
||||
|
||||
# 预处理：截断超长文本
|
||||
processed_texts = []
|
||||
for text in texts:
|
||||
if not text or not text.strip():
|
||||
processed_texts.append(" ")
|
||||
else:
|
||||
# 粗略估算 token 数（中文约 1 字 = 1.5 token，英文约 4 字符 = 1 token）
|
||||
# 保守截断到 max_tokens * 2 个字符
|
||||
max_chars = max_tokens * 2
|
||||
if len(text) > max_chars:
|
||||
processed_texts.append(text[:max_chars])
|
||||
else:
|
||||
processed_texts.append(text)
|
||||
|
||||
# 分批处理
|
||||
all_embeddings = [None] * len(processed_texts)
|
||||
total_batches = math.ceil(len(processed_texts) / batch_size)
|
||||
|
||||
for batch_idx in range(total_batches):
|
||||
start = batch_idx * batch_size
|
||||
end = min(start + batch_size, len(processed_texts))
|
||||
batch_texts = processed_texts[start:end]
|
||||
|
||||
try:
|
||||
kwargs = {
|
||||
'model': model_name,
|
||||
'input': batch_texts,
|
||||
}
|
||||
# 仅对明确支持 dimensions 参数的模型传递该参数
|
||||
if dimensions:
|
||||
model_lower = model_name.lower()
|
||||
# OpenAI text-embedding-3 系列原生支持任意 dimensions
|
||||
if 'text-embedding-3' in model_lower:
|
||||
kwargs['dimensions'] = dimensions
|
||||
# DashScope text-embedding-v3 只接受 [64,128,256,512,768,1024]
|
||||
elif 'text-embedding-v3' in model_lower and dimensions in (64, 128, 256, 512, 768, 1024):
|
||||
kwargs['dimensions'] = dimensions
|
||||
|
||||
response = await client.embeddings.create(**kwargs)
|
||||
|
||||
for item in response.data:
|
||||
all_embeddings[start + item.index] = item.embedding
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'Embedding 批次 {batch_idx + 1}/{total_batches} 失败: {e}')
|
||||
raise ValueError(f'Embedding 调用失败: {str(e)}')
|
||||
|
||||
# 检查是否所有向量都已生成
|
||||
for i, emb in enumerate(all_embeddings):
|
||||
if emb is None:
|
||||
raise ValueError(f'第 {i} 个文本的向量未生成')
|
||||
|
||||
return all_embeddings
|
||||
|
||||
async def get_embedding_dimensions(self, model_id: str) -> int:
|
||||
"""
|
||||
获取模型的向量维度（通过嵌入一个测试文本来检测）
|
||||
|
||||
Args:
|
||||
model_id: Embedding 模型 ID
|
||||
|
||||
Returns:
|
||||
向量维度
|
||||
"""
|
||||
test_embedding = await self.embed_text(model_id, "test")
|
||||
return len(test_embedding)
|
||||
@@ -0,0 +1,91 @@
|
||||
"""
|
||||
索引进度推送服务
|
||||
|
||||
通过 Redis Pub/Sub 推送索引进度，前端通过 SSE 订阅。
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Redis 频道前缀
|
||||
CHANNEL_PREFIX = "knowledge:indexing:progress:"
|
||||
|
||||
|
||||
class IndexingProgressService:
|
||||
"""索引进度推送服务"""
|
||||
|
||||
@staticmethod
|
||||
def _channel(knowledge_base_id: str) -> str:
|
||||
return f"{CHANNEL_PREFIX}{knowledge_base_id}"
|
||||
|
||||
@classmethod
|
||||
async def publish(
|
||||
cls,
|
||||
knowledge_base_id: str,
|
||||
document_id: str,
|
||||
step: str,
|
||||
progress: float,
|
||||
message: str = "",
|
||||
document_name: str = "",
|
||||
error: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
发布索引进度事件
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
document_id: 文档 ID
|
||||
step: 当前步骤 (extracting/cleaning/chunking/vectorizing/completed/failed)
|
||||
progress: 进度 0.0 ~ 1.0
|
||||
message: 进度描述
|
||||
document_name: 文档名称
|
||||
error: 错误信息（仅 failed 步骤）
|
||||
"""
|
||||
try:
|
||||
from utils.redis import RedisClient
|
||||
client = await RedisClient.get_client()
|
||||
|
||||
event = {
|
||||
"document_id": document_id,
|
||||
"document_name": document_name,
|
||||
"step": step,
|
||||
"progress": round(progress, 2),
|
||||
"message": message,
|
||||
}
|
||||
if error:
|
||||
event["error"] = error
|
||||
|
||||
await client.publish(
|
||||
cls._channel(knowledge_base_id),
|
||||
json.dumps(event, ensure_ascii=False),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"发布索引进度失败: {e}")
|
||||
|
||||
@classmethod
|
||||
async def subscribe(cls, knowledge_base_id: str):
|
||||
"""
|
||||
订阅索引进度事件（异步生成器，用于 SSE）
|
||||
|
||||
Yields:
|
||||
dict: 进度事件
|
||||
"""
|
||||
from utils.redis import RedisClient
|
||||
client = await RedisClient.get_client()
|
||||
pubsub = client.pubsub()
|
||||
channel = cls._channel(knowledge_base_id)
|
||||
|
||||
await pubsub.subscribe(channel)
|
||||
try:
|
||||
async for message in pubsub.listen():
|
||||
if message["type"] == "message":
|
||||
try:
|
||||
data = json.loads(message["data"])
|
||||
yield data
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
continue
|
||||
finally:
|
||||
await pubsub.unsubscribe(channel)
|
||||
await pubsub.close()
|
||||
@@ -0,0 +1,540 @@
|
||||
"""
|
||||
文档索引服务
|
||||
|
||||
负责文档处理管道：文本提取 → 分块 → 向量化 → 入库
|
||||
分段数据存入业务数据库，向量数据存入 Qdrant 向量数据库
|
||||
"""
|
||||
import hashlib
|
||||
import logging
|
||||
import math
|
||||
from datetime import datetime
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
from sqlalchemy import select, func, delete
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from ai_platform.knowledge.models import KnowledgeBase, KnowledgeDocument, KnowledgeSegment
|
||||
from ai_platform.knowledge.chunking import get_chunker
|
||||
from ai_platform.knowledge.services.embedding_service import EmbeddingService
|
||||
from ai_platform.knowledge.vector_store import get_vector_store, VectorPoint
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 向量化批次大小
|
||||
EMBEDDING_BATCH_SIZE = 50
|
||||
|
||||
|
||||
class IndexingService:
|
||||
"""
|
||||
文档索引服务
|
||||
|
||||
处理管道：
|
||||
1. 从文件管理系统提取文本内容
|
||||
2. 按知识库配置的策略分块
|
||||
3. 调用 Embedding 模型向量化
|
||||
4. 将分段写入业务数据库，向量写入 Qdrant
|
||||
"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self._db = db
|
||||
self._embedding_service = EmbeddingService(db)
|
||||
self._vector_store = get_vector_store()
|
||||
|
||||
async def index_document(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
document_id: str,
|
||||
) -> Tuple[int, int]:
|
||||
"""
|
||||
索引单个文档
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
document_id: 文档 ID
|
||||
|
||||
Returns:
|
||||
(segment_count, token_count) 分段数和 Token 数
|
||||
"""
|
||||
# 1. 获取知识库配置
|
||||
kb_result = await self._db.execute(
|
||||
select(KnowledgeBase).where(
|
||||
KnowledgeBase.id == knowledge_base_id,
|
||||
KnowledgeBase.is_deleted == False
|
||||
)
|
||||
)
|
||||
kb = kb_result.scalar_one_or_none()
|
||||
if not kb:
|
||||
raise ValueError(f'知识库不存在: {knowledge_base_id}')
|
||||
|
||||
if not kb.embedding_model_id:
|
||||
raise ValueError('知识库未配置 Embedding 模型')
|
||||
|
||||
# 2. 获取文档
|
||||
doc_result = await self._db.execute(
|
||||
select(KnowledgeDocument).where(
|
||||
KnowledgeDocument.id == document_id,
|
||||
KnowledgeDocument.is_deleted == False
|
||||
)
|
||||
)
|
||||
doc = doc_result.scalar_one_or_none()
|
||||
if not doc:
|
||||
raise ValueError(f'文档不存在: {document_id}')
|
||||
|
||||
# 更新状态为 indexing
|
||||
doc.status = 'indexing'
|
||||
doc.indexing_started_at = datetime.now()
|
||||
doc.error_message = None
|
||||
await self._db.commit()
|
||||
|
||||
try:
|
||||
from ai_platform.knowledge.services.indexing_progress_service import IndexingProgressService
|
||||
|
||||
# 3. 提取文本
|
||||
await IndexingProgressService.publish(
|
||||
knowledge_base_id, document_id, step='extracting', progress=0.1,
|
||||
message='正在提取文本内容...', document_name=doc.name,
|
||||
)
|
||||
text_content = await self._extract_text(doc.file_id)
|
||||
if not text_content or not text_content.strip():
|
||||
raise ValueError('文档内容为空，无法索引')
|
||||
|
||||
# 3.5 预处理/清洗
|
||||
await IndexingProgressService.publish(
|
||||
knowledge_base_id, document_id, step='cleaning', progress=0.2,
|
||||
message='正在预处理/清洗文本...', document_name=doc.name,
|
||||
)
|
||||
from ai_platform.knowledge.services.cleaning_service import CleaningService
|
||||
text_content = CleaningService.clean(text_content, kb.process_rules)
|
||||
|
||||
# 计算内容哈希（用于去重）
|
||||
content_hash = hashlib.md5(text_content.encode('utf-8')).hexdigest()
|
||||
doc.content_hash = content_hash
|
||||
|
||||
# 4. 分块
|
||||
await IndexingProgressService.publish(
|
||||
knowledge_base_id, document_id, step='chunking', progress=0.3,
|
||||
message='正在分块...', document_name=doc.name,
|
||||
)
|
||||
chunk_strategy = kb.chunk_strategy or 'recursive'
|
||||
chunk_kwargs = {}
|
||||
|
||||
# Q&A 模式需要 LLM 调用函数
|
||||
if chunk_strategy == 'qa':
|
||||
chunk_kwargs['llm_caller'] = self._create_llm_caller(kb)
|
||||
|
||||
chunker = get_chunker(
|
||||
strategy=chunk_strategy,
|
||||
chunk_size=kb.chunk_size or 500,
|
||||
chunk_overlap=kb.chunk_overlap or 50,
|
||||
separator=kb.separator,
|
||||
**chunk_kwargs,
|
||||
)
|
||||
|
||||
doc_metadata = {
|
||||
'document_id': document_id,
|
||||
'document_name': doc.name,
|
||||
'file_type': doc.file_type,
|
||||
}
|
||||
|
||||
# Q&A 模式使用异步分块
|
||||
if chunk_strategy == 'qa' and hasattr(chunker, 'chunk_async'):
|
||||
chunks = await chunker.chunk_async(text_content, metadata=doc_metadata)
|
||||
else:
|
||||
chunks = chunker.chunk(text_content, metadata=doc_metadata)
|
||||
|
||||
if not chunks:
|
||||
raise ValueError('文档分块结果为空')
|
||||
|
||||
# 5. 删除旧的分段（重新索引场景）
|
||||
await self._delete_document_segments(document_id, knowledge_base_id)
|
||||
|
||||
# 判断索引模式
|
||||
is_economy = (kb.indexing_technique == 'economy')
|
||||
|
||||
# 6. 确保 Qdrant collection 存在（经济模式跳过）
|
||||
if not is_economy:
|
||||
# 自动检测并修正 embedding 维度
|
||||
try:
|
||||
real_dim = await self._embedding_service.get_embedding_dimensions(kb.embedding_model_id)
|
||||
if real_dim != kb.embedding_dimensions:
|
||||
logger.info(f'修正 embedding 维度: {kb.embedding_dimensions} → {real_dim}')
|
||||
kb.embedding_dimensions = real_dim
|
||||
await self._db.commit()
|
||||
except Exception as e:
|
||||
logger.warning(f'自动检测 embedding 维度失败: {e}')
|
||||
|
||||
vector_size = kb.embedding_dimensions or 1536
|
||||
await self._vector_store.ensure_collection(knowledge_base_id, vector_size)
|
||||
|
||||
# 7. 入库（分批处理）
|
||||
segment_count = 0
|
||||
total_token_count = 0
|
||||
total_char_count = 0
|
||||
failed_embedding_count = 0
|
||||
|
||||
total_batches = math.ceil(len(chunks) / EMBEDDING_BATCH_SIZE)
|
||||
|
||||
for batch_idx in range(total_batches):
|
||||
start = batch_idx * EMBEDDING_BATCH_SIZE
|
||||
end = min(start + EMBEDDING_BATCH_SIZE, len(chunks))
|
||||
batch_chunks = chunks[start:end]
|
||||
|
||||
# 高质量模式：批量向量化；经济模式：跳过
|
||||
if is_economy:
|
||||
embeddings = [None] * len(batch_chunks)
|
||||
else:
|
||||
batch_texts = [c.content for c in batch_chunks]
|
||||
try:
|
||||
embeddings = await self._embedding_service.embed_texts(
|
||||
model_id=kb.embedding_model_id,
|
||||
texts=batch_texts,
|
||||
dimensions=kb.embedding_dimensions if kb.embedding_dimensions else None,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f'向量化批次 {batch_idx + 1}/{total_batches} 失败: {e}')
|
||||
embeddings = [None] * len(batch_texts)
|
||||
failed_embedding_count += len(batch_texts)
|
||||
|
||||
# 创建分段记录（业务数据库）+ 收集向量点（Qdrant）
|
||||
vector_points = []
|
||||
for i, chunk in enumerate(batch_chunks):
|
||||
embedding = embeddings[i] if i < len(embeddings) else None
|
||||
char_count = len(chunk.content)
|
||||
token_count = self._estimate_tokens(chunk.content)
|
||||
word_count = self._count_words(chunk.content)
|
||||
|
||||
# 自动提取关键词（高质量和经济模式均提取，增强全文检索）
|
||||
keywords = chunk.metadata.get('keywords')
|
||||
if not keywords:
|
||||
keywords = self._extract_keywords(chunk.content)
|
||||
|
||||
# 经济模式下 embedding_status 标记为 'skipped'
|
||||
if is_economy:
|
||||
emb_status = 'skipped'
|
||||
else:
|
||||
emb_status = 'completed' if embedding else 'failed'
|
||||
|
||||
segment = KnowledgeSegment(
|
||||
knowledge_base_id=knowledge_base_id,
|
||||
document_id=document_id,
|
||||
position=start + i,
|
||||
content=chunk.content,
|
||||
answer=chunk.metadata.get('answer'),
|
||||
token_count=token_count,
|
||||
char_count=char_count,
|
||||
word_count=word_count,
|
||||
page_number=chunk.metadata.get('page_number'),
|
||||
keywords=keywords,
|
||||
extra_metadata=chunk.metadata,
|
||||
embedding_status=emb_status,
|
||||
enabled=True,
|
||||
)
|
||||
self._db.add(segment)
|
||||
await self._db.flush()
|
||||
|
||||
# 收集向量点，稍后批量写入 Qdrant（经济模式跳过）
|
||||
if embedding and not is_economy:
|
||||
vector_points.append(VectorPoint(
|
||||
id=str(segment.id),
|
||||
vector=embedding,
|
||||
payload={
|
||||
'document_id': document_id,
|
||||
'knowledge_base_id': knowledge_base_id,
|
||||
'position': start + i,
|
||||
},
|
||||
))
|
||||
|
||||
segment_count += 1
|
||||
total_token_count += token_count
|
||||
total_char_count += char_count
|
||||
|
||||
# 提交业务数据库
|
||||
await self._db.commit()
|
||||
|
||||
# 批量写入 Qdrant（经济模式跳过）
|
||||
if vector_points:
|
||||
await self._vector_store.upsert(knowledge_base_id, vector_points)
|
||||
|
||||
step_label = '关键词提取中' if is_economy else '向量化中'
|
||||
batch_progress = 0.3 + 0.6 * (end / len(chunks))
|
||||
await IndexingProgressService.publish(
|
||||
knowledge_base_id, document_id, step='vectorizing',
|
||||
progress=batch_progress,
|
||||
message=f'{step_label} {end}/{len(chunks)}',
|
||||
document_name=doc.name,
|
||||
)
|
||||
logger.info(f'文档 {doc.name} 索引进度: {end}/{len(chunks)}')
|
||||
|
||||
# 7. 更新文档状态
|
||||
if not is_economy and failed_embedding_count > 0:
|
||||
if failed_embedding_count >= segment_count:
|
||||
doc.status = 'failed'
|
||||
doc.error_message = f'所有 {segment_count} 个分段向量化失败'
|
||||
else:
|
||||
doc.status = 'completed'
|
||||
doc.error_message = f'{failed_embedding_count}/{segment_count} 个分段向量化失败'
|
||||
else:
|
||||
doc.status = 'completed'
|
||||
doc.segment_count = segment_count
|
||||
doc.token_count = total_token_count
|
||||
doc.char_count = total_char_count
|
||||
doc.indexing_completed_at = datetime.now()
|
||||
|
||||
# 8. 更新知识库统计
|
||||
await self._update_kb_stats(knowledge_base_id)
|
||||
|
||||
await self._db.commit()
|
||||
|
||||
await IndexingProgressService.publish(
|
||||
knowledge_base_id, document_id, step='completed', progress=1.0,
|
||||
message=f'索引完成: {segment_count} 个分段',
|
||||
document_name=doc.name,
|
||||
)
|
||||
logger.info(f'文档 {doc.name} 索引完成: {segment_count} 个分段, {total_token_count} tokens')
|
||||
return segment_count, total_token_count
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'文档索引失败: {e}')
|
||||
doc.status = 'failed'
|
||||
doc.error_message = str(e)[:500]
|
||||
await self._db.commit()
|
||||
await IndexingProgressService.publish(
|
||||
knowledge_base_id, document_id, step='failed', progress=0.0,
|
||||
message='索引失败', document_name=doc.name,
|
||||
error=str(e)[:200],
|
||||
)
|
||||
raise
|
||||
|
||||
async def reindex_document(self, knowledge_base_id: str, document_id: str) -> Tuple[int, int]:
|
||||
"""重新索引文档（删除旧分段后重新处理）"""
|
||||
return await self.index_document(knowledge_base_id, document_id)
|
||||
|
||||
async def index_segment(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
segment_id: str,
|
||||
) -> bool:
|
||||
"""
|
||||
为单个分段生成向量（用于手动添加或更新分段后）
|
||||
"""
|
||||
kb_result = await self._db.execute(
|
||||
select(KnowledgeBase).where(
|
||||
KnowledgeBase.id == knowledge_base_id,
|
||||
KnowledgeBase.is_deleted == False
|
||||
)
|
||||
)
|
||||
kb = kb_result.scalar_one_or_none()
|
||||
if not kb or not kb.embedding_model_id:
|
||||
return False
|
||||
|
||||
seg_result = await self._db.execute(
|
||||
select(KnowledgeSegment).where(
|
||||
KnowledgeSegment.id == segment_id,
|
||||
KnowledgeSegment.is_deleted == False
|
||||
)
|
||||
)
|
||||
segment = seg_result.scalar_one_or_none()
|
||||
if not segment:
|
||||
return False
|
||||
|
||||
try:
|
||||
embedding = await self._embedding_service.embed_text(
|
||||
model_id=kb.embedding_model_id,
|
||||
text=segment.content,
|
||||
dimensions=kb.embedding_dimensions if kb.embedding_dimensions else None,
|
||||
)
|
||||
|
||||
# 确保 collection 存在
|
||||
vector_size = kb.embedding_dimensions or 1536
|
||||
await self._vector_store.ensure_collection(knowledge_base_id, vector_size)
|
||||
|
||||
# 写入 Qdrant
|
||||
await self._vector_store.upsert(knowledge_base_id, [VectorPoint(
|
||||
id=str(segment_id),
|
||||
vector=embedding,
|
||||
payload={
|
||||
'document_id': str(segment.document_id),
|
||||
'knowledge_base_id': knowledge_base_id,
|
||||
'position': segment.position or 0,
|
||||
},
|
||||
)])
|
||||
|
||||
segment.embedding_status = 'completed'
|
||||
await self._db.commit()
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f'分段向量化失败: {e}')
|
||||
segment.embedding_status = 'failed'
|
||||
await self._db.commit()
|
||||
return False
|
||||
|
||||
async def _extract_text(self, file_id: str) -> str:
|
||||
"""从文件管理系统提取文本内容（启用 OCR 支持图片和扫描版 PDF）"""
|
||||
from core.file_manager.service import FileManagerService
|
||||
|
||||
text_content = await FileManagerService.get_file_text_content(
|
||||
self._db, file_id, enable_ocr=True
|
||||
)
|
||||
if not text_content:
|
||||
raise ValueError('无法提取文件文本内容')
|
||||
return text_content
|
||||
|
||||
async def _delete_document_segments(self, document_id: str, knowledge_base_id: str):
|
||||
"""删除文档的所有分段（业务数据库 + Qdrant）"""
|
||||
# 先从 Qdrant 删除该文档的所有向量
|
||||
try:
|
||||
await self._vector_store.delete_by_filter(
|
||||
knowledge_base_id,
|
||||
filter_conditions={'document_id': document_id},
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f'从 Qdrant 删除文档向量失败: {e}')
|
||||
|
||||
# 再从业务数据库删除分段记录
|
||||
await self._db.execute(
|
||||
delete(KnowledgeSegment).where(
|
||||
KnowledgeSegment.document_id == document_id
|
||||
)
|
||||
)
|
||||
|
||||
async def _update_kb_stats(self, knowledge_base_id: str):
|
||||
"""更新知识库统计信息"""
|
||||
# 文档数
|
||||
doc_count_result = await self._db.execute(
|
||||
select(func.count()).select_from(KnowledgeDocument).where(
|
||||
KnowledgeDocument.knowledge_base_id == knowledge_base_id,
|
||||
KnowledgeDocument.is_deleted == False,
|
||||
)
|
||||
)
|
||||
doc_count = doc_count_result.scalar() or 0
|
||||
|
||||
# 分段数和 Token 数
|
||||
seg_stats = await self._db.execute(
|
||||
select(
|
||||
func.count(),
|
||||
func.coalesce(func.sum(KnowledgeSegment.token_count), 0),
|
||||
func.coalesce(func.sum(KnowledgeSegment.char_count), 0),
|
||||
).where(
|
||||
KnowledgeSegment.knowledge_base_id == knowledge_base_id,
|
||||
KnowledgeSegment.is_deleted == False,
|
||||
)
|
||||
)
|
||||
row = seg_stats.one()
|
||||
seg_count = row[0] or 0
|
||||
total_tokens = row[1] or 0
|
||||
total_chars = row[2] or 0
|
||||
|
||||
# 更新知识库
|
||||
kb_result = await self._db.execute(
|
||||
select(KnowledgeBase).where(KnowledgeBase.id == knowledge_base_id)
|
||||
)
|
||||
kb = kb_result.scalar_one_or_none()
|
||||
if kb:
|
||||
kb.document_count = doc_count
|
||||
kb.segment_count = seg_count
|
||||
kb.total_token_count = total_tokens
|
||||
kb.total_char_count = total_chars
|
||||
|
||||
def _create_llm_caller(self, kb: KnowledgeBase):
|
||||
"""
|
||||
创建 Q&A 分块所需的 LLM 调用函数
|
||||
|
||||
使用知识库所属应用中配置的第一个 chat 类型模型。
|
||||
"""
|
||||
db = self._db
|
||||
|
||||
async def llm_caller(system_prompt: str, user_prompt: str) -> str:
|
||||
from ai_platform.services.llm_service import LLMService
|
||||
|
||||
# 查找可用的 chat 模型
|
||||
from ai_platform.models import LLMModel
|
||||
model_result = await db.execute(
|
||||
select(LLMModel).where(
|
||||
LLMModel.model_type == 'chat',
|
||||
LLMModel.is_active == True,
|
||||
LLMModel.is_deleted == False,
|
||||
).limit(1)
|
||||
)
|
||||
model = model_result.scalar_one_or_none()
|
||||
if not model:
|
||||
raise ValueError('未找到可用的 chat 模型，无法进行 Q&A 拆分')
|
||||
|
||||
llm_service = LLMService(db)
|
||||
messages = []
|
||||
if system_prompt:
|
||||
messages.append({'role': 'system', 'content': system_prompt})
|
||||
messages.append({'role': 'user', 'content': user_prompt})
|
||||
|
||||
response = await llm_service.chat_async(
|
||||
model_id=str(model.id),
|
||||
messages=messages,
|
||||
temperature=0.3,
|
||||
max_tokens=4096,
|
||||
)
|
||||
return response.content or ''
|
||||
|
||||
return llm_caller
|
||||
|
||||
@staticmethod
|
||||
def _extract_keywords(text: str, max_keywords: int = 10) -> List[str]:
|
||||
"""
|
||||
从文本中提取关键词（经济模式使用）
|
||||
|
||||
使用简单的 TF 统计提取高频词，无需外部依赖。
|
||||
"""
|
||||
import re
|
||||
if not text:
|
||||
return []
|
||||
|
||||
# 中文分词（简单按标点和空格分割）
|
||||
# 提取中文词组（2-4字）和英文单词
|
||||
chinese_words = re.findall(r'[\u4e00-\u9fff]{2,4}', text)
|
||||
english_words = [w.lower() for w in re.findall(r'[a-zA-Z]{3,}', text)]
|
||||
|
||||
all_words = chinese_words + english_words
|
||||
|
||||
# 停用词（简单列表）
|
||||
stop_words = {
|
||||
'的', '了', '在', '是', '我', '有', '和', '就', '不', '人', '都', '一',
|
||||
'一个', '上', '也', '很', '到', '说', '要', '去', '你', '会', '着',
|
||||
'没有', '看', '好', '自己', '这', '他', '她', '它', '我们', '他们',
|
||||
'可以', '这个', '那个', '什么', '如果', '因为', '所以', '但是', '而且',
|
||||
'the', 'and', 'for', 'are', 'but', 'not', 'you', 'all', 'can',
|
||||
'had', 'her', 'was', 'one', 'our', 'out', 'has', 'have', 'been',
|
||||
'this', 'that', 'with', 'from', 'they', 'will', 'would', 'there',
|
||||
}
|
||||
|
||||
# 词频统计
|
||||
word_freq = {}
|
||||
for word in all_words:
|
||||
if word in stop_words or len(word) < 2:
|
||||
continue
|
||||
word_freq[word] = word_freq.get(word, 0) + 1
|
||||
|
||||
# 按频率排序取 top
|
||||
sorted_words = sorted(word_freq.items(), key=lambda x: x[1], reverse=True)
|
||||
return [w for w, _ in sorted_words[:max_keywords]]
|
||||
|
||||
@staticmethod
|
||||
def _count_words(text: str) -> int:
|
||||
"""计算词数（中文按字计算，英文按空格分词）"""
|
||||
if not text:
|
||||
return 0
|
||||
import re
|
||||
# 中文字符数
|
||||
chinese_chars = len(re.findall(r'[\u4e00-\u9fff]', text))
|
||||
# 英文单词数
|
||||
english_words = len(re.findall(r'[a-zA-Z]+', text))
|
||||
return chinese_chars + english_words
|
||||
|
||||
@staticmethod
|
||||
def _estimate_tokens(text: str) -> int:
|
||||
"""粗略估算文本的 Token 数"""
|
||||
if not text:
|
||||
return 0
|
||||
# 中文约 1 字 = 1.5 token，英文约 4 字符 = 1 token
|
||||
# 简单混合估算
|
||||
chinese_chars = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
|
||||
other_chars = len(text) - chinese_chars
|
||||
return int(chinese_chars * 1.5 + other_chars / 4)
|
||||
@@ -0,0 +1,267 @@
|
||||
"""
|
||||
知识库服务
|
||||
|
||||
知识库 CRUD 操作
|
||||
|
||||
数据权限：
|
||||
- 使用 get_list_with_data_scope() 自动应用数据权限
|
||||
- 支持本人、本部门、本部门及下级、全部等数据范围
|
||||
"""
|
||||
import logging
|
||||
from typing import Optional, List, Tuple
|
||||
|
||||
from sqlalchemy import select, func, or_, and_
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from ai_platform.knowledge.models import KnowledgeBase, KnowledgeDocument, KnowledgeSegment
|
||||
from ai_platform.knowledge.schemas.knowledge_base_schema import KnowledgeBaseCreate, KnowledgeBaseUpdate
|
||||
from app.data_scope_utils import get_data_scope_filter, apply_data_scope_to_conditions
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 资源类型（用于数据权限配置）
|
||||
RESOURCE_TYPE = "knowledge_base"
|
||||
RESOURCE_DISPLAY_NAME = "知识库管理"
|
||||
|
||||
|
||||
class KnowledgeService:
|
||||
"""知识库服务"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self._db = db
|
||||
|
||||
async def get_list(
|
||||
self,
|
||||
page: int = 1,
|
||||
page_size: int = 20,
|
||||
name: Optional[str] = None,
|
||||
status: Optional[str] = None,
|
||||
application_id: Optional[str] = None,
|
||||
) -> Tuple[List[KnowledgeBase], int]:
|
||||
"""获取知识库列表"""
|
||||
query = select(KnowledgeBase).where(KnowledgeBase.is_deleted == False)
|
||||
|
||||
if application_id:
|
||||
query = query.where(or_(
|
||||
KnowledgeBase.application_id == application_id,
|
||||
and_(KnowledgeBase.application_id.is_(None), KnowledgeBase.is_global == True)
|
||||
))
|
||||
if name:
|
||||
query = query.where(KnowledgeBase.name.ilike(f"%{name}%"))
|
||||
if status:
|
||||
query = query.where(KnowledgeBase.status == status)
|
||||
|
||||
# 总数
|
||||
count_result = await self._db.execute(select(func.count()).select_from(query.subquery()))
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
# 分页
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(KnowledgeBase.sort.desc(), KnowledgeBase.sys_create_datetime.desc())
|
||||
query = query.offset(offset).limit(page_size)
|
||||
|
||||
result = await self._db.execute(query)
|
||||
items = result.scalars().all()
|
||||
|
||||
return items, total
|
||||
|
||||
async def get_list_with_data_scope(
|
||||
self,
|
||||
page: int = 1,
|
||||
page_size: int = 20,
|
||||
name: Optional[str] = None,
|
||||
status: Optional[str] = None,
|
||||
application_id: Optional[str] = None,
|
||||
) -> Tuple[List[KnowledgeBase], int]:
|
||||
"""
|
||||
获取知识库列表（带数据权限过滤）
|
||||
|
||||
自动从上下文获取当前用户信息，应用数据权限过滤
|
||||
"""
|
||||
conditions = [KnowledgeBase.is_deleted == False]
|
||||
|
||||
if application_id:
|
||||
conditions.append(or_(
|
||||
KnowledgeBase.application_id == application_id,
|
||||
and_(KnowledgeBase.application_id.is_(None), KnowledgeBase.is_global == True)
|
||||
))
|
||||
if name:
|
||||
conditions.append(KnowledgeBase.name.ilike(f"%{name}%"))
|
||||
if status:
|
||||
conditions.append(KnowledgeBase.status == status)
|
||||
|
||||
# 获取数据权限过滤条件并应用
|
||||
data_scope_filter = await get_data_scope_filter(self._db, RESOURCE_TYPE)
|
||||
scope_conditions = apply_data_scope_to_conditions(KnowledgeBase, data_scope_filter)
|
||||
conditions.extend(scope_conditions)
|
||||
|
||||
# 总数
|
||||
query = select(KnowledgeBase).where(and_(*conditions))
|
||||
count_result = await self._db.execute(select(func.count()).select_from(query.subquery()))
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
# 分页
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(KnowledgeBase.sort.desc(), KnowledgeBase.sys_create_datetime.desc())
|
||||
query = query.offset(offset).limit(page_size)
|
||||
|
||||
result = await self._db.execute(query)
|
||||
items = result.scalars().all()
|
||||
|
||||
return items, total
|
||||
|
||||
async def get_by_id(self, kb_id: str) -> Optional[KnowledgeBase]:
|
||||
"""获取知识库详情"""
|
||||
result = await self._db.execute(
|
||||
select(KnowledgeBase).where(
|
||||
KnowledgeBase.id == kb_id,
|
||||
KnowledgeBase.is_deleted == False
|
||||
)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def get_by_code(self, code: str) -> Optional[KnowledgeBase]:
|
||||
"""根据编码获取知识库"""
|
||||
result = await self._db.execute(
|
||||
select(KnowledgeBase).where(
|
||||
KnowledgeBase.code == code,
|
||||
KnowledgeBase.is_deleted == False
|
||||
)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def create(self, data: KnowledgeBaseCreate) -> KnowledgeBase:
|
||||
"""创建知识库"""
|
||||
# 检查编码唯一性
|
||||
existing = await self.get_by_code(data.code)
|
||||
if existing:
|
||||
raise ValueError(f'知识库编码 {data.code} 已存在')
|
||||
|
||||
kb_data = data.model_dump()
|
||||
|
||||
# 自动检测 embedding 模型的真实维度
|
||||
if data.embedding_model_id:
|
||||
try:
|
||||
from ai_platform.knowledge.services.embedding_service import EmbeddingService
|
||||
embedding_service = EmbeddingService(self._db)
|
||||
real_dim = await embedding_service.get_embedding_dimensions(data.embedding_model_id)
|
||||
kb_data['embedding_dimensions'] = real_dim
|
||||
logger.info(f'自动检测 embedding 维度: {real_dim}')
|
||||
except Exception as e:
|
||||
logger.warning(f'自动检测 embedding 维度失败，使用默认值: {e}')
|
||||
|
||||
kb = KnowledgeBase(**kb_data)
|
||||
|
||||
# 自动填充创建人和部门
|
||||
from utils.context import get_current_user_info_from_context
|
||||
user_info = get_current_user_info_from_context()
|
||||
if user_info:
|
||||
if not kb.sys_creator_id:
|
||||
kb.sys_creator_id = user_info.get('user_id')
|
||||
if not kb.sys_dept_id and user_info.get('dept_id'):
|
||||
kb.sys_dept_id = user_info.get('dept_id')
|
||||
|
||||
self._db.add(kb)
|
||||
await self._db.commit()
|
||||
await self._db.refresh(kb)
|
||||
return kb
|
||||
|
||||
async def update(self, kb_id: str, data: KnowledgeBaseUpdate) -> Optional[KnowledgeBase]:
|
||||
"""更新知识库"""
|
||||
kb = await self.get_by_id(kb_id)
|
||||
if not kb:
|
||||
return None
|
||||
|
||||
update_data = data.model_dump(exclude_unset=True)
|
||||
|
||||
# 如果更换了 embedding 模型，自动重新检测维度
|
||||
new_model_id = update_data.get('embedding_model_id')
|
||||
if new_model_id and new_model_id != kb.embedding_model_id:
|
||||
try:
|
||||
from ai_platform.knowledge.services.embedding_service import EmbeddingService
|
||||
embedding_service = EmbeddingService(self._db)
|
||||
real_dim = await embedding_service.get_embedding_dimensions(new_model_id)
|
||||
update_data['embedding_dimensions'] = real_dim
|
||||
logger.info(f'更换模型后自动检测 embedding 维度: {real_dim}')
|
||||
except Exception as e:
|
||||
logger.warning(f'自动检测 embedding 维度失败: {e}')
|
||||
|
||||
for key, value in update_data.items():
|
||||
setattr(kb, key, value)
|
||||
|
||||
await self._db.commit()
|
||||
await self._db.refresh(kb)
|
||||
return kb
|
||||
|
||||
async def delete(self, kb_id: str) -> bool:
|
||||
"""删除知识库（软删除 + 清理 Qdrant collection）"""
|
||||
kb = await self.get_by_id(kb_id)
|
||||
if not kb:
|
||||
return False
|
||||
|
||||
kb.is_deleted = True
|
||||
|
||||
# 同时软删除所有文档和分段
|
||||
doc_result = await self._db.execute(
|
||||
select(KnowledgeDocument).where(
|
||||
KnowledgeDocument.knowledge_base_id == kb_id,
|
||||
KnowledgeDocument.is_deleted == False
|
||||
)
|
||||
)
|
||||
docs = doc_result.scalars().all()
|
||||
for doc in docs:
|
||||
doc.is_deleted = True
|
||||
|
||||
# 软删除分段
|
||||
from sqlalchemy import update
|
||||
await self._db.execute(
|
||||
update(KnowledgeSegment).where(
|
||||
KnowledgeSegment.knowledge_base_id == kb_id
|
||||
).values(is_deleted=True)
|
||||
)
|
||||
|
||||
await self._db.commit()
|
||||
|
||||
# 删除 Qdrant 中对应的 collection
|
||||
try:
|
||||
from ai_platform.knowledge.vector_store import get_vector_store
|
||||
vector_store = get_vector_store()
|
||||
await vector_store.delete_collection(kb_id)
|
||||
except Exception as e:
|
||||
logger.warning(f'删除 Qdrant collection 失败: {e}')
|
||||
|
||||
return True
|
||||
|
||||
async def get_simple_list(self, application_id: Optional[str] = None) -> List[dict]:
|
||||
"""获取知识库简单列表（用于下拉选择）"""
|
||||
query = select(
|
||||
KnowledgeBase.id,
|
||||
KnowledgeBase.name,
|
||||
KnowledgeBase.code,
|
||||
KnowledgeBase.document_count,
|
||||
KnowledgeBase.segment_count,
|
||||
).where(
|
||||
KnowledgeBase.is_deleted == False,
|
||||
KnowledgeBase.status == 'active',
|
||||
)
|
||||
|
||||
if application_id:
|
||||
query = query.where(or_(
|
||||
KnowledgeBase.application_id == application_id,
|
||||
and_(KnowledgeBase.application_id.is_(None), KnowledgeBase.is_global == True)
|
||||
))
|
||||
|
||||
query = query.order_by(KnowledgeBase.sort.desc(), KnowledgeBase.sys_create_datetime.desc())
|
||||
result = await self._db.execute(query)
|
||||
rows = result.all()
|
||||
|
||||
return [
|
||||
{
|
||||
'id': row.id,
|
||||
'name': row.name,
|
||||
'code': row.code,
|
||||
'document_count': row.document_count or 0,
|
||||
'segment_count': row.segment_count or 0,
|
||||
}
|
||||
for row in rows
|
||||
]
|
||||
@@ -0,0 +1,182 @@
|
||||
"""
|
||||
Rerank 重排序服务
|
||||
|
||||
通过 Rerank 模型对检索结果进行重新排序，提升检索质量。
|
||||
支持 Jina/Cohere 风格的 Rerank API（大多数提供商兼容此接口）。
|
||||
"""
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RerankResult:
|
||||
"""重排序结果"""
|
||||
index: int
|
||||
relevance_score: float
|
||||
|
||||
|
||||
class RerankService:
|
||||
"""
|
||||
Rerank 重排序服务
|
||||
|
||||
通过模型 ID 获取对应的提供商，调用 Rerank API 对文档进行重排序。
|
||||
支持两种 API 风格：
|
||||
- Jina/Cohere 风格：POST /v1/rerank
|
||||
- OpenAI 兼容风格（部分提供商）
|
||||
"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self._db = db
|
||||
self._client_cache = {}
|
||||
|
||||
async def _get_client_config(self, model_id: str):
|
||||
"""
|
||||
根据模型 ID 获取 API 配置
|
||||
|
||||
Returns:
|
||||
(base_url, api_key, model_name)
|
||||
"""
|
||||
from ai_platform.models import LLMModel, LLMProvider
|
||||
|
||||
result = await self._db.execute(
|
||||
select(LLMModel).where(
|
||||
LLMModel.id == model_id,
|
||||
LLMModel.is_active == True,
|
||||
LLMModel.is_deleted == False
|
||||
)
|
||||
)
|
||||
model = result.scalar_one_or_none()
|
||||
if not model:
|
||||
raise ValueError(f'Rerank 模型不存在或已禁用: {model_id}')
|
||||
if model.model_type != 'rerank':
|
||||
raise ValueError(f'模型 {model.display_name} 不是 Rerank 类型')
|
||||
|
||||
provider_result = await self._db.execute(
|
||||
select(LLMProvider).where(
|
||||
LLMProvider.id == model.provider_id,
|
||||
LLMProvider.is_active == True,
|
||||
LLMProvider.is_deleted == False
|
||||
)
|
||||
)
|
||||
provider = provider_result.scalar_one_or_none()
|
||||
if not provider:
|
||||
raise ValueError('Rerank 模型对应的提供商不存在或已禁用')
|
||||
|
||||
if provider.provider_type == 'ollama':
|
||||
base_url = (provider.ollama_host or 'http://localhost:11434').rstrip('/') + '/v1'
|
||||
else:
|
||||
base_url = provider.api_base or 'https://api.openai.com/v1'
|
||||
|
||||
api_key = provider.api_key or 'ollama'
|
||||
|
||||
return base_url, api_key, model.model_name
|
||||
|
||||
async def rerank(
|
||||
self,
|
||||
model_id: str,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_n: Optional[int] = None,
|
||||
) -> List[RerankResult]:
|
||||
"""
|
||||
对文档列表进行重排序
|
||||
|
||||
Args:
|
||||
model_id: Rerank 模型 ID
|
||||
query: 查询文本
|
||||
documents: 待排序的文档列表
|
||||
top_n: 返回前 N 个结果（默认返回全部）
|
||||
|
||||
Returns:
|
||||
按相关性降序排列的 RerankResult 列表
|
||||
"""
|
||||
if not documents:
|
||||
return []
|
||||
|
||||
if top_n is None:
|
||||
top_n = len(documents)
|
||||
|
||||
base_url, api_key, model_name = await self._get_client_config(model_id)
|
||||
|
||||
try:
|
||||
return await self._call_rerank_api(
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
model_name=model_name,
|
||||
query=query,
|
||||
documents=documents,
|
||||
top_n=top_n,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f'Rerank 调用失败: {e}')
|
||||
raise ValueError(f'Rerank 调用失败: {str(e)}')
|
||||
|
||||
async def _call_rerank_api(
|
||||
self,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
model_name: str,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_n: int,
|
||||
) -> List[RerankResult]:
|
||||
"""
|
||||
调用 Rerank API（Jina/Cohere 兼容风格）
|
||||
|
||||
POST {base_url}/rerank
|
||||
{
|
||||
"model": "...",
|
||||
"query": "...",
|
||||
"documents": ["...", "..."],
|
||||
"top_n": 5
|
||||
}
|
||||
|
||||
Response:
|
||||
{
|
||||
"results": [
|
||||
{"index": 0, "relevance_score": 0.95},
|
||||
{"index": 2, "relevance_score": 0.87},
|
||||
...
|
||||
]
|
||||
}
|
||||
"""
|
||||
import httpx
|
||||
|
||||
url = base_url.rstrip('/') + '/rerank'
|
||||
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': f'Bearer {api_key}',
|
||||
}
|
||||
|
||||
payload = {
|
||||
'model': model_name,
|
||||
'query': query,
|
||||
'documents': documents,
|
||||
'top_n': top_n,
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=60) as client:
|
||||
response = await client.post(url, json=payload, headers=headers)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
# 解析结果（兼容 Jina/Cohere/通义千问 等格式）
|
||||
raw_results = data.get('results', [])
|
||||
results = []
|
||||
for item in raw_results:
|
||||
results.append(RerankResult(
|
||||
index=item.get('index', 0),
|
||||
relevance_score=item.get('relevance_score', 0.0),
|
||||
))
|
||||
|
||||
# 按相关性降序排序
|
||||
results.sort(key=lambda r: r.relevance_score, reverse=True)
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,666 @@
|
||||
"""
|
||||
检索服务
|
||||
|
||||
支持向量检索、全文检索、混合检索（RRF 融合）
|
||||
向量检索通过 Qdrant 向量数据库实现，全文检索通过业务数据库 SQL 实现
|
||||
"""
|
||||
import logging
|
||||
import time
|
||||
from typing import List, Optional, Dict, Any
|
||||
|
||||
from sqlalchemy import select, func
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from ai_platform.knowledge.models import KnowledgeBase, KnowledgeDocument, KnowledgeSegment
|
||||
from ai_platform.knowledge.services.embedding_service import EmbeddingService
|
||||
from ai_platform.knowledge.schemas.segment_schema import RetrievalResult
|
||||
from ai_platform.knowledge.vector_store import get_vector_store
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# RRF 融合常数
|
||||
RRF_K = 60
|
||||
|
||||
|
||||
class RetrievalService:
|
||||
"""
|
||||
检索服务
|
||||
|
||||
支持三种检索模式：
|
||||
- vector: 纯向量检索（通过 Qdrant）
|
||||
- fulltext: 纯全文检索（通过业务数据库 LIKE + 关键词匹配）
|
||||
- hybrid: 混合检索（向量 + 全文，RRF 融合排序）
|
||||
"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self._db = db
|
||||
self._embedding_service = EmbeddingService(db)
|
||||
self._vector_store = get_vector_store()
|
||||
|
||||
async def retrieve(
|
||||
self,
|
||||
query: str,
|
||||
knowledge_base_ids: List[str],
|
||||
top_k: int = 5,
|
||||
score_threshold: float = 0.5,
|
||||
retrieval_mode: Optional[str] = None,
|
||||
rerank_enabled: Optional[bool] = None,
|
||||
rerank_model_id: Optional[str] = None,
|
||||
metadata_filter: Optional[Dict[str, Any]] = None,
|
||||
) -> List[RetrievalResult]:
|
||||
"""
|
||||
检索知识库
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
knowledge_base_ids: 知识库 ID 列表
|
||||
top_k: 返回数量
|
||||
score_threshold: 相似度阈值
|
||||
retrieval_mode: 检索模式（不传则使用第一个知识库的配置）
|
||||
rerank_enabled: 是否启用重排序（不传则使用知识库配置）
|
||||
rerank_model_id: 重排序模型 ID（不传则使用知识库配置）
|
||||
metadata_filter: 元数据过滤条件
|
||||
|
||||
Returns:
|
||||
检索结果列表
|
||||
"""
|
||||
if not query or not knowledge_base_ids:
|
||||
return []
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# 获取知识库配置
|
||||
kb_map = await self._get_knowledge_bases(knowledge_base_ids)
|
||||
if not kb_map:
|
||||
return []
|
||||
|
||||
first_kb = list(kb_map.values())[0]
|
||||
|
||||
# 经济模式强制使用全文检索
|
||||
is_economy = getattr(first_kb, 'indexing_technique', 'high_quality') == 'economy'
|
||||
|
||||
# 确定检索模式
|
||||
if is_economy:
|
||||
retrieval_mode = 'fulltext'
|
||||
elif not retrieval_mode:
|
||||
retrieval_mode = first_kb.retrieval_mode or 'hybrid'
|
||||
|
||||
# 确定 rerank 配置（参数优先，否则使用知识库配置；经济模式禁用 rerank）
|
||||
if is_economy:
|
||||
rerank_enabled = False
|
||||
else:
|
||||
if rerank_enabled is None:
|
||||
rerank_enabled = first_kb.rerank_enabled or False
|
||||
if not rerank_model_id:
|
||||
rerank_model_id = first_kb.rerank_model_id
|
||||
|
||||
# 获取 embedding 模型（使用第一个知识库的配置）
|
||||
embedding_model_id = first_kb.embedding_model_id
|
||||
|
||||
# 如果启用了 rerank，初始检索多取一些候选结果
|
||||
candidate_multiplier = 3 if rerank_enabled and rerank_model_id else 2
|
||||
candidate_top_k = top_k * candidate_multiplier
|
||||
|
||||
results = []
|
||||
|
||||
if retrieval_mode == 'vector':
|
||||
results = await self._vector_search(
|
||||
query, knowledge_base_ids, embedding_model_id,
|
||||
top_k=candidate_top_k, score_threshold=score_threshold,
|
||||
dimensions=first_kb.embedding_dimensions,
|
||||
metadata_filter=metadata_filter,
|
||||
)
|
||||
for r in results:
|
||||
r.match_source = 'vector'
|
||||
elif retrieval_mode == 'fulltext':
|
||||
results = await self._fulltext_search(
|
||||
query, knowledge_base_ids, top_k=candidate_top_k,
|
||||
metadata_filter=metadata_filter,
|
||||
)
|
||||
for r in results:
|
||||
r.match_source = 'fulltext'
|
||||
elif retrieval_mode == 'hybrid':
|
||||
# 混合检索：向量 + 全文，RRF 融合
|
||||
vector_results = await self._vector_search(
|
||||
query, knowledge_base_ids, embedding_model_id,
|
||||
top_k=candidate_top_k, score_threshold=score_threshold,
|
||||
dimensions=first_kb.embedding_dimensions,
|
||||
metadata_filter=metadata_filter,
|
||||
)
|
||||
for r in vector_results:
|
||||
r.match_source = 'vector'
|
||||
fulltext_results = await self._fulltext_search(
|
||||
query, knowledge_base_ids, top_k=candidate_top_k,
|
||||
metadata_filter=metadata_filter,
|
||||
)
|
||||
for r in fulltext_results:
|
||||
r.match_source = 'fulltext'
|
||||
results = self._rrf_merge(vector_results, fulltext_results)
|
||||
|
||||
# 多知识库权重加权
|
||||
if len(kb_map) > 1:
|
||||
for r in results:
|
||||
kb = kb_map.get(r.knowledge_base_id)
|
||||
weight = getattr(kb, 'retrieval_weight', 1.0) or 1.0 if kb else 1.0
|
||||
if weight != 1.0:
|
||||
r.score = round(r.score * weight, 4)
|
||||
results.sort(key=lambda x: x.score, reverse=True)
|
||||
|
||||
# 过滤低分结果（rerank 前先粗筛）
|
||||
if not rerank_enabled:
|
||||
results = [r for r in results if r.score >= score_threshold]
|
||||
|
||||
# Rerank 重排序
|
||||
if rerank_enabled and rerank_model_id and results:
|
||||
results = await self._rerank_results(query, results, rerank_model_id, top_k)
|
||||
# rerank 后再按阈值过滤
|
||||
results = [r for r in results if r.score >= score_threshold]
|
||||
|
||||
# 截断到 top_k
|
||||
results = results[:top_k]
|
||||
|
||||
# 内容级去重（多知识库检索时可能有重复内容）
|
||||
results = self._deduplicate_results(results)
|
||||
|
||||
# 标注优先匹配：将匹配到的标注结果插入到最前面
|
||||
annotation_results = await self._match_annotations(
|
||||
query, knowledge_base_ids, embedding_model_id,
|
||||
score_threshold=score_threshold,
|
||||
dimensions=first_kb.embedding_dimensions,
|
||||
)
|
||||
if annotation_results:
|
||||
for r in annotation_results:
|
||||
r.match_source = 'annotation'
|
||||
# 标注结果置顶，去重后合并
|
||||
existing_ids = {r.segment_id for r in annotation_results}
|
||||
results = annotation_results + [r for r in results if r.segment_id not in existing_ids]
|
||||
results = results[:top_k]
|
||||
|
||||
# 填充知识库名称和文档名称
|
||||
await self._fill_names(results, kb_map)
|
||||
|
||||
# 填充父分段内容（Small-to-Big 模式）
|
||||
await self._fill_parent_content(results)
|
||||
|
||||
# 更新命中次数
|
||||
segment_ids = [r.segment_id for r in results]
|
||||
if segment_ids:
|
||||
await self._update_hit_counts(segment_ids)
|
||||
|
||||
elapsed = int((time.time() - start_time) * 1000)
|
||||
rerank_info = ', rerank=ON' if rerank_enabled else ''
|
||||
logger.info(f'检索完成: {len(results)} 条结果, 耗时 {elapsed}ms, 模式={retrieval_mode}{rerank_info}')
|
||||
|
||||
return results
|
||||
|
||||
async def _rerank_results(
|
||||
self,
|
||||
query: str,
|
||||
results: List[RetrievalResult],
|
||||
rerank_model_id: str,
|
||||
top_n: int,
|
||||
) -> List[RetrievalResult]:
|
||||
"""使用 Rerank 模型对检索结果重排序"""
|
||||
from ai_platform.knowledge.services.rerank_service import RerankService
|
||||
|
||||
try:
|
||||
rerank_service = RerankService(self._db)
|
||||
documents = [r.content for r in results]
|
||||
|
||||
rerank_results = await rerank_service.rerank(
|
||||
model_id=rerank_model_id,
|
||||
query=query,
|
||||
documents=documents,
|
||||
top_n=top_n,
|
||||
)
|
||||
|
||||
# 按 rerank 分数重新排列结果
|
||||
reranked = []
|
||||
for rr in rerank_results:
|
||||
if 0 <= rr.index < len(results):
|
||||
result = results[rr.index]
|
||||
result.score = round(rr.relevance_score, 4)
|
||||
reranked.append(result)
|
||||
|
||||
logger.info(f'Rerank 完成: {len(results)} -> {len(reranked)} 条结果')
|
||||
return reranked
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f'Rerank 失败，使用原始排序: {e}')
|
||||
return results
|
||||
|
||||
async def _vector_search(
|
||||
self,
|
||||
query: str,
|
||||
knowledge_base_ids: List[str],
|
||||
embedding_model_id: str,
|
||||
top_k: int = 10,
|
||||
score_threshold: float = 0.0,
|
||||
dimensions: Optional[int] = None,
|
||||
metadata_filter: Optional[Dict[str, Any]] = None,
|
||||
) -> List[RetrievalResult]:
|
||||
"""向量检索（通过 Qdrant 余弦相似度）"""
|
||||
if not embedding_model_id:
|
||||
logger.warning('未配置 Embedding 模型，跳过向量检索')
|
||||
return []
|
||||
|
||||
try:
|
||||
query_embedding = await self._embedding_service.embed_text(
|
||||
model_id=embedding_model_id,
|
||||
text=query,
|
||||
dimensions=dimensions,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f'查询向量化失败: {e}')
|
||||
return []
|
||||
|
||||
# 对每个知识库分别搜索（每个知识库对应一个 Qdrant collection）
|
||||
all_hits = []
|
||||
for kb_id in knowledge_base_ids:
|
||||
hits = await self._vector_store.search(
|
||||
knowledge_base_id=kb_id,
|
||||
query_vector=query_embedding,
|
||||
top_k=top_k,
|
||||
score_threshold=score_threshold,
|
||||
)
|
||||
all_hits.extend(hits)
|
||||
|
||||
if not all_hits:
|
||||
return []
|
||||
|
||||
# 按分数排序
|
||||
all_hits.sort(key=lambda h: h.score, reverse=True)
|
||||
all_hits = all_hits[:top_k]
|
||||
|
||||
# 从业务数据库获取分段详情
|
||||
segment_ids = [h.id for h in all_hits]
|
||||
score_map = {h.id: h.score for h in all_hits}
|
||||
|
||||
seg_result = await self._db.execute(
|
||||
select(KnowledgeSegment).where(
|
||||
KnowledgeSegment.id.in_(segment_ids),
|
||||
KnowledgeSegment.is_deleted == False,
|
||||
KnowledgeSegment.enabled == True,
|
||||
)
|
||||
)
|
||||
segments = {str(s.id): s for s in seg_result.scalars().all()}
|
||||
|
||||
results = []
|
||||
for hit in all_hits:
|
||||
seg = segments.get(hit.id)
|
||||
if not seg:
|
||||
continue
|
||||
# Q&A 模式：question 用于匹配，返回 answer 作为 content
|
||||
content = seg.answer if seg.answer else seg.content
|
||||
meta = dict(seg.extra_metadata) if seg.extra_metadata else {}
|
||||
if seg.answer:
|
||||
meta['question'] = seg.content
|
||||
meta['chunk_mode'] = 'qa'
|
||||
results.append(RetrievalResult(
|
||||
segment_id=str(seg.id),
|
||||
document_id=str(seg.document_id),
|
||||
knowledge_base_id=str(seg.knowledge_base_id),
|
||||
content=content,
|
||||
score=round(score_map.get(hit.id, 0.0), 4),
|
||||
token_count=seg.token_count or 0,
|
||||
page_number=seg.page_number,
|
||||
metadata=meta,
|
||||
keywords=seg.keywords,
|
||||
))
|
||||
|
||||
return results
|
||||
|
||||
async def _fulltext_search(
|
||||
self,
|
||||
query: str,
|
||||
knowledge_base_ids: List[str],
|
||||
top_k: int = 10,
|
||||
metadata_filter: Optional[Dict[str, Any]] = None,
|
||||
) -> List[RetrievalResult]:
|
||||
"""全文检索（基于 ORM LIKE 和关键词匹配，兼容所有数据库）"""
|
||||
import re
|
||||
keywords = re.split(r'[\s,，。.!！?？;；、]+', query)
|
||||
keywords = [k.strip() for k in keywords if k.strip() and len(k.strip()) >= 2]
|
||||
|
||||
if not keywords:
|
||||
keywords = [query.strip()]
|
||||
|
||||
# 使用 SQLAlchemy ORM 构建查询（兼容 PG / MySQL 等）
|
||||
from sqlalchemy import or_
|
||||
keyword_conditions = [
|
||||
func.lower(KnowledgeSegment.content).contains(kw.lower())
|
||||
for kw in keywords[:5]
|
||||
]
|
||||
|
||||
conditions = [
|
||||
KnowledgeSegment.knowledge_base_id.in_(knowledge_base_ids),
|
||||
KnowledgeSegment.is_deleted == False,
|
||||
KnowledgeSegment.enabled == True,
|
||||
or_(*keyword_conditions),
|
||||
]
|
||||
|
||||
# 元数据过滤
|
||||
if metadata_filter:
|
||||
conditions.extend(self._build_metadata_conditions(metadata_filter))
|
||||
|
||||
stmt = (
|
||||
select(KnowledgeSegment)
|
||||
.where(*conditions)
|
||||
.order_by(KnowledgeSegment.char_count.asc())
|
||||
.limit(top_k)
|
||||
)
|
||||
|
||||
result = await self._db.execute(stmt)
|
||||
segments = result.scalars().all()
|
||||
|
||||
# 计算简单的关键词匹配分数
|
||||
results = []
|
||||
for seg in segments:
|
||||
content_lower = seg.content.lower()
|
||||
match_count = sum(1 for kw in keywords if kw.lower() in content_lower)
|
||||
score = match_count / len(keywords) if keywords else 0
|
||||
# Q&A 模式：返回 answer 作为 content
|
||||
content = seg.answer if seg.answer else seg.content
|
||||
meta = dict(seg.extra_metadata) if seg.extra_metadata else {}
|
||||
if seg.answer:
|
||||
meta['question'] = seg.content
|
||||
meta['chunk_mode'] = 'qa'
|
||||
results.append(RetrievalResult(
|
||||
segment_id=str(seg.id),
|
||||
document_id=str(seg.document_id),
|
||||
knowledge_base_id=str(seg.knowledge_base_id),
|
||||
content=content,
|
||||
score=round(score, 4),
|
||||
token_count=seg.token_count or 0,
|
||||
page_number=seg.page_number,
|
||||
metadata=meta,
|
||||
keywords=seg.keywords,
|
||||
))
|
||||
|
||||
results.sort(key=lambda x: x.score, reverse=True)
|
||||
return results
|
||||
|
||||
def _rrf_merge(
|
||||
self,
|
||||
vector_results: List[RetrievalResult],
|
||||
fulltext_results: List[RetrievalResult],
|
||||
) -> List[RetrievalResult]:
|
||||
"""
|
||||
RRF (Reciprocal Rank Fusion) 融合排序
|
||||
|
||||
RRF_score = sum(1 / (k + rank_i)) for each result list
|
||||
"""
|
||||
scores = {} # segment_id -> (rrf_score, result)
|
||||
|
||||
# 向量检索结果排名
|
||||
for rank, result in enumerate(vector_results):
|
||||
rrf_score = 1.0 / (RRF_K + rank + 1)
|
||||
if result.segment_id in scores:
|
||||
old_score, old_result = scores[result.segment_id]
|
||||
scores[result.segment_id] = (old_score + rrf_score, old_result)
|
||||
else:
|
||||
scores[result.segment_id] = (rrf_score, result)
|
||||
|
||||
# 全文检索结果排名
|
||||
for rank, result in enumerate(fulltext_results):
|
||||
rrf_score = 1.0 / (RRF_K + rank + 1)
|
||||
if result.segment_id in scores:
|
||||
old_score, old_result = scores[result.segment_id]
|
||||
scores[result.segment_id] = (old_score + rrf_score, old_result)
|
||||
else:
|
||||
scores[result.segment_id] = (rrf_score, result)
|
||||
|
||||
# 按 RRF 分数排序
|
||||
sorted_items = sorted(scores.values(), key=lambda x: x[0], reverse=True)
|
||||
|
||||
if not sorted_items:
|
||||
return []
|
||||
|
||||
# 归一化分数到 0-1
|
||||
# RRF 单条结果的理论最大分数为 2/(k+1)（同时出现在两个列表的第一名）
|
||||
# 使用理论最大值归一化，避免单条结果被归一化为 100%
|
||||
theoretical_max = 2.0 / (RRF_K + 1)
|
||||
|
||||
results = []
|
||||
for rrf_score, result in sorted_items:
|
||||
normalized_score = min(rrf_score / theoretical_max, 1.0)
|
||||
result.score = round(normalized_score, 4)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def _deduplicate_results(results: List[RetrievalResult], similarity_threshold: float = 0.95) -> List[RetrievalResult]:
|
||||
"""
|
||||
内容级去重（多知识库检索时可能有重复内容）
|
||||
|
||||
使用内容前 200 字符的相似度判断是否重复，保留分数最高的。
|
||||
"""
|
||||
if len(results) <= 1:
|
||||
return results
|
||||
|
||||
deduplicated = []
|
||||
seen_contents = []
|
||||
|
||||
for r in results:
|
||||
content_key = r.content[:200].strip().lower()
|
||||
is_dup = False
|
||||
for seen in seen_contents:
|
||||
# 简单的字符重叠率判断
|
||||
if content_key == seen:
|
||||
is_dup = True
|
||||
break
|
||||
# 如果前 200 字符有 95% 以上重叠，视为重复
|
||||
shorter = min(len(content_key), len(seen))
|
||||
if shorter > 0:
|
||||
common = sum(1 for a, b in zip(content_key, seen) if a == b)
|
||||
if common / shorter >= similarity_threshold:
|
||||
is_dup = True
|
||||
break
|
||||
if not is_dup:
|
||||
deduplicated.append(r)
|
||||
seen_contents.append(content_key)
|
||||
|
||||
return deduplicated
|
||||
|
||||
async def _fill_parent_content(self, results: List[RetrievalResult]):
|
||||
"""填充父分段内容（Small-to-Big 模式）"""
|
||||
if not results:
|
||||
return
|
||||
|
||||
# 获取所有 segment_id，查询是否有 parent_segment_id
|
||||
segment_ids = [r.segment_id for r in results if r.segment_id]
|
||||
if not segment_ids:
|
||||
return
|
||||
|
||||
seg_result = await self._db.execute(
|
||||
select(KnowledgeSegment.id, KnowledgeSegment.parent_segment_id).where(
|
||||
KnowledgeSegment.id.in_(segment_ids),
|
||||
KnowledgeSegment.is_deleted == False,
|
||||
)
|
||||
)
|
||||
parent_map = {}
|
||||
for row in seg_result:
|
||||
if row.parent_segment_id:
|
||||
parent_map[str(row.id)] = row.parent_segment_id
|
||||
|
||||
if not parent_map:
|
||||
return
|
||||
|
||||
# 批量获取父分段内容
|
||||
parent_ids = list(set(parent_map.values()))
|
||||
parent_result = await self._db.execute(
|
||||
select(KnowledgeSegment.id, KnowledgeSegment.content).where(
|
||||
KnowledgeSegment.id.in_(parent_ids),
|
||||
KnowledgeSegment.is_deleted == False,
|
||||
)
|
||||
)
|
||||
parent_content_map = {str(row.id): row.content for row in parent_result}
|
||||
|
||||
# 填充到结果中
|
||||
for r in results:
|
||||
parent_id = parent_map.get(r.segment_id)
|
||||
if parent_id:
|
||||
r.parent_content = parent_content_map.get(str(parent_id))
|
||||
|
||||
@staticmethod
|
||||
def _build_metadata_conditions(metadata_filter: Dict[str, Any]) -> list:
|
||||
"""构建元数据过滤条件（基于 JSON 字段，跨数据库兼容）"""
|
||||
from app.db_compat import json_extract
|
||||
|
||||
conditions = []
|
||||
for key, value in metadata_filter.items():
|
||||
if value is not None:
|
||||
# 使用跨数据库兼容的 json_extract 函数
|
||||
try:
|
||||
conditions.append(
|
||||
json_extract(KnowledgeSegment.extra_metadata, key) == str(value)
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return conditions
|
||||
|
||||
async def _get_knowledge_bases(self, kb_ids: List[str]) -> Dict[str, KnowledgeBase]:
|
||||
"""批量获取知识库"""
|
||||
result = await self._db.execute(
|
||||
select(KnowledgeBase).where(
|
||||
KnowledgeBase.id.in_(kb_ids),
|
||||
KnowledgeBase.is_deleted == False,
|
||||
)
|
||||
)
|
||||
kbs = result.scalars().all()
|
||||
return {str(kb.id): kb for kb in kbs}
|
||||
|
||||
async def _fill_names(self, results: List[RetrievalResult], kb_map: Dict[str, KnowledgeBase]):
|
||||
"""填充知识库名称和文档名称"""
|
||||
if not results:
|
||||
return
|
||||
|
||||
# 获取文档名称
|
||||
doc_ids = list({r.document_id for r in results})
|
||||
doc_result = await self._db.execute(
|
||||
select(KnowledgeDocument.id, KnowledgeDocument.name).where(
|
||||
KnowledgeDocument.id.in_(doc_ids)
|
||||
)
|
||||
)
|
||||
doc_name_map = {row.id: row.name for row in doc_result}
|
||||
|
||||
for result in results:
|
||||
result.document_name = doc_name_map.get(result.document_id, '')
|
||||
kb = kb_map.get(result.knowledge_base_id)
|
||||
result.knowledge_base_name = kb.name if kb else ''
|
||||
|
||||
async def _update_hit_counts(self, segment_ids: List[str]):
|
||||
"""更新分段命中次数（兼容所有数据库）"""
|
||||
if not segment_ids:
|
||||
return
|
||||
try:
|
||||
from sqlalchemy import update
|
||||
await self._db.execute(
|
||||
update(KnowledgeSegment)
|
||||
.where(KnowledgeSegment.id.in_(segment_ids))
|
||||
.values(hit_count=func.coalesce(KnowledgeSegment.hit_count, 0) + 1)
|
||||
)
|
||||
await self._db.commit()
|
||||
except Exception as e:
|
||||
logger.warning(f'更新命中次数失败: {e}')
|
||||
|
||||
async def _match_annotations(
|
||||
self,
|
||||
query: str,
|
||||
knowledge_base_ids: List[str],
|
||||
embedding_model_id: Optional[str],
|
||||
score_threshold: float = 0.5,
|
||||
dimensions: Optional[int] = None,
|
||||
max_results: int = 3,
|
||||
) -> List[RetrievalResult]:
|
||||
"""
|
||||
匹配标注（Q&A 对）
|
||||
|
||||
通过向量相似度匹配标注的 question，返回对应的 answer。
|
||||
标注结果优先级高于普通分段。
|
||||
"""
|
||||
from ai_platform.knowledge.models import KnowledgeAnnotation
|
||||
|
||||
if not embedding_model_id:
|
||||
return []
|
||||
|
||||
try:
|
||||
# 向量化查询
|
||||
query_embedding = await self._embedding_service.embed_text(
|
||||
model_id=embedding_model_id,
|
||||
text=query,
|
||||
dimensions=dimensions if dimensions else None,
|
||||
)
|
||||
|
||||
# 在 Qdrant 中搜索标注向量（payload.type == 'annotation'）
|
||||
all_hits = []
|
||||
for kb_id in knowledge_base_ids:
|
||||
try:
|
||||
hits = await self._vector_store.search(
|
||||
knowledge_base_id=kb_id,
|
||||
query_vector=query_embedding,
|
||||
top_k=max_results,
|
||||
score_threshold=score_threshold,
|
||||
filter_conditions={'type': 'annotation'},
|
||||
)
|
||||
all_hits.extend(hits)
|
||||
except Exception as e:
|
||||
logger.warning(f'标注向量搜索失败 (kb={kb_id}): {e}')
|
||||
continue
|
||||
|
||||
if not all_hits:
|
||||
return []
|
||||
|
||||
# 按分数排序取 top
|
||||
all_hits.sort(key=lambda h: h.score, reverse=True)
|
||||
all_hits = all_hits[:max_results]
|
||||
|
||||
# 从数据库获取标注详情
|
||||
annotation_ids = [h.id for h in all_hits]
|
||||
score_map = {h.id: h.score for h in all_hits}
|
||||
|
||||
ann_result = await self._db.execute(
|
||||
select(KnowledgeAnnotation).where(
|
||||
KnowledgeAnnotation.id.in_(annotation_ids),
|
||||
KnowledgeAnnotation.is_deleted == False,
|
||||
KnowledgeAnnotation.enabled == True,
|
||||
)
|
||||
)
|
||||
annotations = {str(a.id): a for a in ann_result.scalars().all()}
|
||||
|
||||
results = []
|
||||
for hit in all_hits:
|
||||
ann = annotations.get(hit.id)
|
||||
if not ann:
|
||||
continue
|
||||
# 标注结果：content 返回 answer，segment_id 用 annotation id
|
||||
results.append(RetrievalResult(
|
||||
segment_id=str(ann.id),
|
||||
document_id='',
|
||||
document_name='[Q&A]',
|
||||
knowledge_base_id=str(ann.knowledge_base_id),
|
||||
content=ann.answer,
|
||||
score=round(score_map.get(hit.id, 0.0), 4),
|
||||
token_count=0,
|
||||
metadata={'type': 'annotation', 'question': ann.question},
|
||||
))
|
||||
|
||||
# 更新标注命中次数
|
||||
if annotation_ids:
|
||||
try:
|
||||
from sqlalchemy import update
|
||||
await self._db.execute(
|
||||
update(KnowledgeAnnotation)
|
||||
.where(KnowledgeAnnotation.id.in_(annotation_ids))
|
||||
.values(hit_count=func.coalesce(KnowledgeAnnotation.hit_count, 0) + 1)
|
||||
)
|
||||
await self._db.commit()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f'标注匹配失败: {e}')
|
||||
return []
|
||||
@@ -0,0 +1,45 @@
|
||||
"""
|
||||
向量存储模块
|
||||
|
||||
提供可插拔的向量存储后端，支持 Qdrant 等专业向量数据库。
|
||||
与业务数据库完全解耦，segment 表只存业务数据，向量数据存在向量数据库中。
|
||||
"""
|
||||
from ai_platform.knowledge.vector_store.base import BaseVectorStore, VectorPoint, VectorSearchResult
|
||||
from ai_platform.knowledge.vector_store.qdrant_store import QdrantVectorStore
|
||||
|
||||
__all__ = [
|
||||
'BaseVectorStore',
|
||||
'VectorPoint',
|
||||
'VectorSearchResult',
|
||||
'QdrantVectorStore',
|
||||
'get_vector_store',
|
||||
]
|
||||
|
||||
# 单例缓存
|
||||
_vector_store_instance: BaseVectorStore | None = None
|
||||
|
||||
|
||||
def get_vector_store() -> BaseVectorStore:
|
||||
"""
|
||||
工厂函数：根据配置获取向量存储实例（单例）
|
||||
"""
|
||||
global _vector_store_instance
|
||||
if _vector_store_instance is not None:
|
||||
return _vector_store_instance
|
||||
|
||||
from app.config import settings
|
||||
|
||||
store_type = getattr(settings, 'VECTOR_STORE_TYPE', 'qdrant')
|
||||
|
||||
if store_type == 'qdrant':
|
||||
_vector_store_instance = QdrantVectorStore(
|
||||
host=getattr(settings, 'QDRANT_HOST', 'localhost'),
|
||||
port=getattr(settings, 'QDRANT_PORT', 6333),
|
||||
api_key=getattr(settings, 'QDRANT_API_KEY', None),
|
||||
grpc_port=getattr(settings, 'QDRANT_GRPC_PORT', 6334),
|
||||
prefer_grpc=getattr(settings, 'QDRANT_PREFER_GRPC', False),
|
||||
)
|
||||
else:
|
||||
raise ValueError(f'不支持的向量存储类型: {store_type}')
|
||||
|
||||
return _vector_store_instance
|
||||
@@ -0,0 +1,138 @@
|
||||
"""
|
||||
向量存储抽象基类
|
||||
|
||||
定义向量存储的统一接口，所有向量存储后端必须实现这些方法。
|
||||
"""
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VectorPoint:
|
||||
"""向量数据点"""
|
||||
id: str
|
||||
vector: List[float]
|
||||
payload: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VectorSearchResult:
|
||||
"""向量搜索结果"""
|
||||
id: str
|
||||
score: float
|
||||
payload: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class BaseVectorStore(ABC):
|
||||
"""
|
||||
向量存储抽象基类
|
||||
|
||||
每个知识库对应一个 collection，collection 名称格式: kb_{knowledge_base_id}
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def collection_name(knowledge_base_id: str) -> str:
|
||||
"""生成 collection 名称"""
|
||||
return f"kb_{knowledge_base_id.replace('-', '_')}"
|
||||
|
||||
@abstractmethod
|
||||
async def ensure_collection(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
vector_size: int,
|
||||
) -> None:
|
||||
"""
|
||||
确保 collection 存在，不存在则创建
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
vector_size: 向量维度
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def delete_collection(self, knowledge_base_id: str) -> None:
|
||||
"""
|
||||
删除 collection（删除知识库时调用）
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def upsert(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
points: List[VectorPoint],
|
||||
) -> None:
|
||||
"""
|
||||
批量写入/更新向量
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
points: 向量数据点列表
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def delete(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
point_ids: List[str],
|
||||
) -> None:
|
||||
"""
|
||||
批量删除向量
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
point_ids: 要删除的向量 ID 列表
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def search(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
query_vector: List[float],
|
||||
top_k: int = 10,
|
||||
score_threshold: float = 0.0,
|
||||
filter_conditions: Optional[Dict[str, Any]] = None,
|
||||
) -> List[VectorSearchResult]:
|
||||
"""
|
||||
向量相似度搜索
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
query_vector: 查询向量
|
||||
top_k: 返回数量
|
||||
score_threshold: 最低相似度阈值
|
||||
filter_conditions: 过滤条件（如 {"document_id": "xxx"}）
|
||||
|
||||
Returns:
|
||||
搜索结果列表，按相似度降序排列
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def delete_by_filter(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
filter_conditions: Dict[str, Any],
|
||||
) -> None:
|
||||
"""
|
||||
按条件删除向量（如删除某个文档的所有向量）
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
filter_conditions: 过滤条件（如 {"document_id": "xxx"}）
|
||||
"""
|
||||
...
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""健康检查"""
|
||||
return True
|
||||
@@ -0,0 +1,260 @@
|
||||
"""
|
||||
Qdrant 向量存储实现
|
||||
|
||||
使用 Qdrant 作为向量数据库后端，通过 qdrant-client 进行交互。
|
||||
每个知识库对应一个 Qdrant collection。
|
||||
"""
|
||||
import logging
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
from qdrant_client.models import (
|
||||
Distance,
|
||||
FieldCondition,
|
||||
Filter,
|
||||
FilterSelector,
|
||||
MatchValue,
|
||||
PointIdsList,
|
||||
PointStruct,
|
||||
VectorParams,
|
||||
)
|
||||
|
||||
from ai_platform.knowledge.vector_store.base import (
|
||||
BaseVectorStore,
|
||||
VectorPoint,
|
||||
VectorSearchResult,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class QdrantVectorStore(BaseVectorStore):
|
||||
"""
|
||||
Qdrant 向量存储
|
||||
|
||||
特性：
|
||||
- 高性能向量检索（HNSW 索引）
|
||||
- 支持 payload 过滤
|
||||
- 支持 REST 和 gRPC 协议
|
||||
- 与业务数据库完全解耦
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "localhost",
|
||||
port: int = 6333,
|
||||
api_key: Optional[str] = None,
|
||||
grpc_port: int = 6334,
|
||||
prefer_grpc: bool = False,
|
||||
):
|
||||
self._host = host
|
||||
self._port = port
|
||||
self._api_key = api_key
|
||||
self._grpc_port = grpc_port
|
||||
self._prefer_grpc = prefer_grpc
|
||||
self._client: Optional[AsyncQdrantClient] = None
|
||||
|
||||
async def _get_client(self) -> AsyncQdrantClient:
|
||||
"""获取或创建 Qdrant 客户端（懒初始化）"""
|
||||
if self._client is None:
|
||||
# 如果 host 已包含协议前缀，直接作为 url 使用；
|
||||
# 否则拼接 http:// 避免 qdrant-client 对非 localhost 域名自动走 HTTPS
|
||||
if self._host.startswith("http://") or self._host.startswith("https://"):
|
||||
url = f"{self._host}:{self._port}"
|
||||
else:
|
||||
url = f"http://{self._host}:{self._port}"
|
||||
self._client = AsyncQdrantClient(
|
||||
url=url,
|
||||
api_key=self._api_key,
|
||||
grpc_port=self._grpc_port,
|
||||
prefer_grpc=self._prefer_grpc,
|
||||
timeout=30,
|
||||
)
|
||||
return self._client
|
||||
|
||||
async def ensure_collection(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
vector_size: int,
|
||||
) -> None:
|
||||
"""确保 collection 存在且维度匹配"""
|
||||
client = await self._get_client()
|
||||
name = self.collection_name(knowledge_base_id)
|
||||
|
||||
collections = await client.get_collections()
|
||||
existing_names = {c.name for c in collections.collections}
|
||||
|
||||
if name in existing_names:
|
||||
# 检查已有 collection 的维度是否匹配
|
||||
info = await client.get_collection(collection_name=name)
|
||||
existing_size = info.config.params.vectors.size
|
||||
if existing_size != vector_size:
|
||||
logger.warning(
|
||||
f"Qdrant collection {name} 维度不匹配: "
|
||||
f"已有={existing_size}, 期望={vector_size}，删除重建"
|
||||
)
|
||||
await client.delete_collection(collection_name=name)
|
||||
else:
|
||||
return
|
||||
|
||||
await client.create_collection(
|
||||
collection_name=name,
|
||||
vectors_config=VectorParams(
|
||||
size=vector_size,
|
||||
distance=Distance.COSINE,
|
||||
),
|
||||
)
|
||||
# 创建 payload 索引，加速过滤查询
|
||||
await client.create_payload_index(
|
||||
collection_name=name,
|
||||
field_name="document_id",
|
||||
field_schema="keyword",
|
||||
)
|
||||
logger.info(f"Qdrant collection 已创建: {name} (dim={vector_size})")
|
||||
|
||||
async def delete_collection(self, knowledge_base_id: str) -> None:
|
||||
"""删除 collection"""
|
||||
client = await self._get_client()
|
||||
name = self.collection_name(knowledge_base_id)
|
||||
try:
|
||||
await client.delete_collection(collection_name=name)
|
||||
logger.info(f"Qdrant collection 已删除: {name}")
|
||||
except Exception as e:
|
||||
logger.warning(f"删除 Qdrant collection 失败: {name}, {e}")
|
||||
|
||||
async def upsert(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
points: List[VectorPoint],
|
||||
) -> None:
|
||||
"""批量写入/更新向量"""
|
||||
if not points:
|
||||
return
|
||||
|
||||
client = await self._get_client()
|
||||
name = self.collection_name(knowledge_base_id)
|
||||
|
||||
qdrant_points = [
|
||||
PointStruct(
|
||||
id=self._to_uuid(p.id),
|
||||
vector=p.vector,
|
||||
payload={**p.payload, 'segment_id': p.id},
|
||||
)
|
||||
for p in points
|
||||
]
|
||||
|
||||
# Qdrant 单次 upsert 建议不超过 100 个点
|
||||
batch_size = 100
|
||||
for i in range(0, len(qdrant_points), batch_size):
|
||||
batch = qdrant_points[i:i + batch_size]
|
||||
await client.upsert(
|
||||
collection_name=name,
|
||||
points=batch,
|
||||
)
|
||||
|
||||
async def delete(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
point_ids: List[str],
|
||||
) -> None:
|
||||
"""批量删除向量"""
|
||||
if not point_ids:
|
||||
return
|
||||
|
||||
client = await self._get_client()
|
||||
name = self.collection_name(knowledge_base_id)
|
||||
|
||||
uuid_ids = [self._to_uuid(pid) for pid in point_ids]
|
||||
await client.delete(
|
||||
collection_name=name,
|
||||
points_selector=PointIdsList(points=uuid_ids),
|
||||
)
|
||||
|
||||
async def search(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
query_vector: List[float],
|
||||
top_k: int = 10,
|
||||
score_threshold: float = 0.0,
|
||||
filter_conditions: Optional[Dict[str, Any]] = None,
|
||||
) -> List[VectorSearchResult]:
|
||||
"""向量相似度搜索"""
|
||||
client = await self._get_client()
|
||||
name = self.collection_name(knowledge_base_id)
|
||||
|
||||
# 构建过滤条件
|
||||
query_filter = self._build_filter(filter_conditions) if filter_conditions else None
|
||||
|
||||
try:
|
||||
results = await client.search(
|
||||
collection_name=name,
|
||||
query_vector=query_vector,
|
||||
limit=top_k,
|
||||
score_threshold=score_threshold,
|
||||
query_filter=query_filter,
|
||||
with_payload=True,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Qdrant 搜索失败: {e}")
|
||||
return []
|
||||
|
||||
return [
|
||||
VectorSearchResult(
|
||||
id=(hit.payload or {}).get('segment_id', str(hit.id)),
|
||||
score=hit.score,
|
||||
payload=hit.payload or {},
|
||||
)
|
||||
for hit in results
|
||||
]
|
||||
|
||||
async def delete_by_filter(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
filter_conditions: Dict[str, Any],
|
||||
) -> None:
|
||||
"""按条件删除向量"""
|
||||
client = await self._get_client()
|
||||
name = self.collection_name(knowledge_base_id)
|
||||
|
||||
query_filter = self._build_filter(filter_conditions)
|
||||
if query_filter:
|
||||
await client.delete(
|
||||
collection_name=name,
|
||||
points_selector=FilterSelector(filter=query_filter),
|
||||
)
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""健康检查"""
|
||||
try:
|
||||
client = await self._get_client()
|
||||
await client.get_collections()
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Qdrant 健康检查失败: {e}")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _to_uuid(string_id: str) -> str:
|
||||
"""将任意字符串 ID 确定性转换为 UUID5（Qdrant 要求 point ID 为 UUID 或整数）"""
|
||||
return str(uuid.uuid5(uuid.NAMESPACE_DNS, string_id))
|
||||
|
||||
@staticmethod
|
||||
def _build_filter(conditions: Dict[str, Any]) -> Optional[Filter]:
|
||||
"""构建 Qdrant 过滤条件"""
|
||||
if not conditions:
|
||||
return None
|
||||
|
||||
must = []
|
||||
for key, value in conditions.items():
|
||||
if isinstance(value, list):
|
||||
# 列表值：任一匹配（OR 语义），用 should 包裹后作为一个 must 条件
|
||||
should_conditions = [
|
||||
FieldCondition(key=key, match=MatchValue(value=v))
|
||||
for v in value
|
||||
]
|
||||
must.append(Filter(should=should_conditions))
|
||||
else:
|
||||
must.append(FieldCondition(key=key, match=MatchValue(value=value)))
|
||||
|
||||
return Filter(must=must) if must else None
|
||||
@@ -0,0 +1,29 @@
|
||||
"""
|
||||
AI 平台数据模型
|
||||
"""
|
||||
from .provider import LLMProvider
|
||||
from .model import LLMModel
|
||||
from .app import AIApp
|
||||
from .conversation import Conversation, Message
|
||||
from .workflow import AIWorkflow, AIWorkflowVersion, AIWorkflowRun
|
||||
from .prompt_template import PromptTemplate
|
||||
from .agent import Agent, AgentConversation, AgentMessage
|
||||
from ai_platform.knowledge.models import KnowledgeBase, KnowledgeDocument, KnowledgeSegment
|
||||
|
||||
__all__ = [
|
||||
'LLMProvider',
|
||||
'LLMModel',
|
||||
'AIApp',
|
||||
'Conversation',
|
||||
'Message',
|
||||
'AIWorkflow',
|
||||
'AIWorkflowVersion',
|
||||
'AIWorkflowRun',
|
||||
'PromptTemplate',
|
||||
'Agent',
|
||||
'AgentConversation',
|
||||
'AgentMessage',
|
||||
'KnowledgeBase',
|
||||
'KnowledgeDocument',
|
||||
'KnowledgeSegment',
|
||||
]
|
||||
@@ -0,0 +1,95 @@
|
||||
"""
|
||||
智能体模型
|
||||
"""
|
||||
from sqlalchemy import Column, String, Text, Float, Integer, Boolean, JSON, Index
|
||||
|
||||
from app.base_model import BaseModel
|
||||
|
||||
|
||||
class Agent(BaseModel):
|
||||
"""
|
||||
智能体定义
|
||||
|
||||
智能体是一个能够自主决策、调用工具、多轮推理的 AI 实体
|
||||
支持两种模式：
|
||||
- autonomous: 自主规划模式 - Agent 自动拆解任务并执行
|
||||
- dialog_flow: 对话流模式 - 按预定义流程与用户交互
|
||||
"""
|
||||
__tablename__ = "ai_agent"
|
||||
|
||||
application_id = Column(String(21), nullable=True, index=True, comment="所属应用ID（逻辑外键关联core_application）")
|
||||
is_global = Column(Boolean, default=False, comment="是否在子应用中可见")
|
||||
name = Column(String(100), nullable=False, comment="智能体名称")
|
||||
code = Column(String(100), unique=True, nullable=False, comment="智能体编码")
|
||||
description = Column(Text, default="", comment="智能体描述")
|
||||
avatar = Column(String(500), default="", comment="头像 URL")
|
||||
mode = Column(String(20), default="autonomous", comment="运行模式: autonomous/dialog_flow")
|
||||
status = Column(String(20), default="draft", comment="状态: draft/published/disabled")
|
||||
persona = Column(JSON, default=dict, comment="人设配置")
|
||||
system_prompt = Column(Text, default="", comment="系统提示词")
|
||||
model_id = Column(String(21), nullable=True, index=True, comment="默认模型ID（逻辑外键关联ai_llm_model）")
|
||||
temperature = Column(Float, default=0.7, comment="温度参数（0-2）")
|
||||
top_p = Column(Float, default=1.0, comment="top_p 参数")
|
||||
max_tokens = Column(Integer, default=4096, comment="最大输出 Token")
|
||||
max_iterations = Column(Integer, default=10, comment="最大推理轮数")
|
||||
welcome_message = Column(Text, default="", comment="开场白")
|
||||
suggested_questions = Column(JSON, default=list, comment="推荐问题列表")
|
||||
workflow_id = Column(String(21), nullable=True, index=True, comment="关联工作流ID（逻辑外键关联ai_workflow）")
|
||||
enable_memory = Column(Boolean, default=False, comment="是否启用对话记忆")
|
||||
memory_window = Column(Integer, default=10, comment="记忆窗口大小（最近 N 轮对话）")
|
||||
enable_streaming = Column(Boolean, default=True, comment="是否启用流式输出（自主规划模式）")
|
||||
knowledge_base_ids = Column(JSON, default=list, comment="关联的知识库ID列表")
|
||||
knowledge_config = Column(JSON, default=dict, comment="知识库检索配置(top_k/score_threshold/retrieval_mode等)")
|
||||
is_public = Column(Boolean, default=False, comment="是否公开")
|
||||
conversation_count = Column(Integer, default=0, comment="对话数量")
|
||||
message_count = Column(Integer, default=0, comment="消息数量")
|
||||
total_tokens = Column(Integer, default=0, comment="总 Token 消耗")
|
||||
|
||||
|
||||
class AgentConversation(BaseModel):
|
||||
"""
|
||||
智能体对话
|
||||
"""
|
||||
__tablename__ = "ai_agent_conversation"
|
||||
|
||||
agent_id = Column(String(21), nullable=False, index=True, comment="智能体ID（逻辑外键关联ai_agent）")
|
||||
user_id = Column(String(21), nullable=True, index=True, comment="用户ID（逻辑外键关联core_user）")
|
||||
title = Column(String(200), default="", comment="对话标题")
|
||||
summary = Column(Text, default="", comment="对话摘要")
|
||||
workflow_run_id = Column(String(21), nullable=True, comment="工作流运行实例ID（逻辑外键关联ai_workflow_run）")
|
||||
waiting_node_id = Column(String(100), default="", comment="等待输入的节点 ID")
|
||||
extra_data = Column(JSON, default=dict, comment="元数据")
|
||||
message_count = Column(Integer, default=0, comment="消息数量")
|
||||
total_tokens = Column(Integer, default=0, comment="总 Token 消耗")
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_ai_agent_conversation_agent_user", "agent_id", "user_id"),
|
||||
)
|
||||
|
||||
|
||||
class AgentMessage(BaseModel):
|
||||
"""
|
||||
智能体消息
|
||||
|
||||
记录对话中的每条消息，包括用户消息、助手回复、工具调用等
|
||||
"""
|
||||
__tablename__ = "ai_agent_message"
|
||||
|
||||
conversation_id = Column(String(21), nullable=False, index=True, comment="对话ID（逻辑外键关联ai_agent_conversation）")
|
||||
role = Column(String(20), nullable=False, comment="角色: user/assistant/tool/system")
|
||||
content = Column(Text, default="", comment="消息内容")
|
||||
attachments = Column(JSON, default=list, comment="附件列表 [{id, type, name, url, mime_type, size}]")
|
||||
status = Column(String(20), default="completed", comment="状态: pending/completed/failed")
|
||||
reasoning_steps = Column(JSON, default=list, comment="推理步骤")
|
||||
tool_calls = Column(JSON, default=list, comment="工具调用记录")
|
||||
prompt_tokens = Column(Integer, default=0, comment="提示 Token")
|
||||
completion_tokens = Column(Integer, default=0, comment="生成 Token")
|
||||
total_tokens = Column(Integer, default=0, comment="总 Token")
|
||||
elapsed_time = Column(Integer, default=0, comment="耗时（毫秒）")
|
||||
error_message = Column(Text, default="", comment="错误信息")
|
||||
feedback = Column(String(20), default="", comment="用户反馈（like/dislike）")
|
||||
feedback_content = Column(Text, default="", comment="反馈内容")
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_ai_agent_message_conversation_role", "conversation_id", "role"),
|
||||
)
|
||||
@@ -0,0 +1,37 @@
|
||||
"""
|
||||
AI 应用模型
|
||||
"""
|
||||
from sqlalchemy import Column, String, Text, Float, Integer, Boolean, JSON
|
||||
|
||||
from app.base_model import BaseModel
|
||||
|
||||
|
||||
class AIApp(BaseModel):
|
||||
"""
|
||||
AI 应用
|
||||
|
||||
支持的应用类型：
|
||||
- chat: 聊天助手
|
||||
- completion: 文本生成
|
||||
- workflow: 工作流应用
|
||||
- agent: Agent 应用（预留）
|
||||
"""
|
||||
__tablename__ = "ai_app"
|
||||
|
||||
name = Column(String(100), nullable=False, comment="应用名称")
|
||||
code = Column(String(100), unique=True, nullable=False, comment="应用编码")
|
||||
description = Column(Text, default="", comment="应用描述")
|
||||
icon = Column(String(100), default="", comment="应用图标")
|
||||
app_type = Column(String(20), default="chat", comment="应用类型: chat/completion/workflow/agent")
|
||||
status = Column(String(20), default="draft", comment="状态: draft/published/disabled")
|
||||
model_id = Column(String(21), nullable=True, index=True, comment="默认模型ID（逻辑外键关联ai_llm_model）")
|
||||
system_prompt = Column(Text, default="", comment="系统提示词")
|
||||
temperature = Column(Float, default=0.7, comment="温度参数")
|
||||
top_p = Column(Float, default=1.0, comment="top_p 参数")
|
||||
max_tokens = Column(Integer, default=2048, comment="最大输出 Token")
|
||||
workflow_definition = Column(JSON, default=dict, comment="工作流定义")
|
||||
opening_statement = Column(Text, default="", comment="开场白")
|
||||
suggested_questions = Column(JSON, default=list, comment="建议问题列表")
|
||||
is_public = Column(Boolean, default=False, comment="是否公开")
|
||||
conversation_count = Column(Integer, default=0, comment="对话数量")
|
||||
message_count = Column(Integer, default=0, comment="消息数量")
|
||||
@@ -0,0 +1,50 @@
|
||||
"""
|
||||
对话模型
|
||||
"""
|
||||
from sqlalchemy import Column, String, Text, Float, Integer, Boolean, Index
|
||||
|
||||
from app.base_model import BaseModel
|
||||
|
||||
|
||||
class Conversation(BaseModel):
|
||||
"""
|
||||
对话
|
||||
"""
|
||||
__tablename__ = "ai_conversation"
|
||||
|
||||
app_id = Column(String(21), nullable=False, index=True, comment="所属应用ID（逻辑外键关联ai_app）")
|
||||
user_id = Column(String(21), nullable=False, index=True, comment="用户ID（逻辑外键关联core_user）")
|
||||
title = Column(String(200), default="", comment="对话标题")
|
||||
model_override_id = Column(String(21), nullable=True, comment="覆盖模型ID（逻辑外键关联ai_llm_model）")
|
||||
temperature_override = Column(Float, nullable=True, comment="覆盖温度参数")
|
||||
message_count = Column(Integer, default=0, comment="消息数量")
|
||||
total_tokens = Column(Integer, default=0, comment="总 Token 数")
|
||||
is_pinned = Column(Boolean, default=False, comment="是否置顶")
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_ai_conversation_user_app", "user_id", "app_id"),
|
||||
)
|
||||
|
||||
|
||||
class Message(BaseModel):
|
||||
"""
|
||||
消息
|
||||
"""
|
||||
__tablename__ = "ai_message"
|
||||
|
||||
conversation_id = Column(String(21), nullable=False, index=True, comment="所属对话ID（逻辑外键关联ai_conversation）")
|
||||
role = Column(String(20), nullable=False, comment="角色: system/user/assistant")
|
||||
content = Column(Text, nullable=False, comment="消息内容")
|
||||
status = Column(String(20), default="completed", comment="状态: pending/completed/failed/stopped")
|
||||
prompt_tokens = Column(Integer, default=0, comment="提示 Token 数")
|
||||
completion_tokens = Column(Integer, default=0, comment="补全 Token 数")
|
||||
total_tokens = Column(Integer, default=0, comment="总 Token 数")
|
||||
model_name = Column(String(100), default="", comment="使用的模型名称")
|
||||
latency = Column(Integer, default=0, comment="响应耗时（毫秒）")
|
||||
error_message = Column(Text, default="", comment="错误信息")
|
||||
parent_message_id = Column(String(21), nullable=True, comment="父消息ID（逻辑外键关联自身）")
|
||||
feedback = Column(String(20), default="", comment="用户反馈: like/dislike")
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_ai_message_conversation_created", "conversation_id", "sys_create_datetime"),
|
||||
)
|
||||
@@ -0,0 +1,30 @@
|
||||
"""
|
||||
LLM 模型配置
|
||||
"""
|
||||
from sqlalchemy import Column, String, Integer, Float, Boolean, Numeric
|
||||
|
||||
from app.base_model import BaseModel
|
||||
|
||||
|
||||
class LLMModel(BaseModel):
|
||||
"""
|
||||
LLM 模型配置
|
||||
|
||||
每个提供商可以配置多个模型
|
||||
"""
|
||||
__tablename__ = "ai_llm_model"
|
||||
|
||||
provider_id = Column(String(21), nullable=False, index=True, comment="所属提供商ID（逻辑外键关联ai_llm_provider）")
|
||||
model_name = Column(String(100), nullable=False, comment="模型名称（API 调用时使用）")
|
||||
display_name = Column(String(100), nullable=False, comment="显示名称")
|
||||
model_type = Column(String(20), default="chat", comment="模型类型: chat/completion/embedding/rerank")
|
||||
max_tokens = Column(Integer, default=4096, comment="最大 Token 数")
|
||||
context_window = Column(Integer, default=4096, comment="上下文窗口大小")
|
||||
default_temperature = Column(Float, default=0.7, comment="默认温度参数")
|
||||
default_top_p = Column(Float, default=1.0, comment="默认 top_p 参数")
|
||||
input_price = Column(Numeric(10, 6), default=0, comment="输入价格（每 1K tokens）")
|
||||
output_price = Column(Numeric(10, 6), default=0, comment="输出价格（每 1K tokens）")
|
||||
is_active = Column(Boolean, default=True, comment="是否启用")
|
||||
supports_vision = Column(Boolean, default=False, comment="是否支持视觉（图片输入）")
|
||||
supports_function_call = Column(Boolean, default=False, comment="是否支持函数调用")
|
||||
supports_streaming = Column(Boolean, default=True, comment="是否支持流式输出")
|
||||
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
Prompt 模板模型
|
||||
"""
|
||||
from sqlalchemy import Column, String, Text, Integer, Boolean, JSON
|
||||
|
||||
from app.base_model import BaseModel
|
||||
|
||||
|
||||
class PromptTemplate(BaseModel):
|
||||
"""
|
||||
Prompt 模板
|
||||
|
||||
用于管理和复用 Prompt
|
||||
"""
|
||||
__tablename__ = "ai_prompt_template"
|
||||
|
||||
name = Column(String(100), nullable=False, comment="模板名称")
|
||||
code = Column(String(100), unique=True, nullable=False, comment="模板编码")
|
||||
category = Column(String(20), default="system", comment="分类: system/user/assistant/few_shot")
|
||||
description = Column(Text, default="", comment="描述")
|
||||
content = Column(Text, nullable=False, comment="模板内容")
|
||||
variables = Column(JSON, default=list, comment="变量定义列表")
|
||||
tags = Column(JSON, default=list, comment="标签列表")
|
||||
is_public = Column(Boolean, default=False, comment="是否公开")
|
||||
usage_count = Column(Integer, default=0, comment="使用次数")
|
||||
|
||||
def render(self, variables: dict) -> str:
|
||||
"""
|
||||
渲染模板
|
||||
|
||||
Args:
|
||||
variables: 变量字典
|
||||
|
||||
Returns:
|
||||
渲染后的内容
|
||||
"""
|
||||
content = self.content or ""
|
||||
for key, value in variables.items():
|
||||
content = content.replace(f'{{{{{key}}}}}', str(value))
|
||||
return content
|
||||
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
LLM 提供商模型
|
||||
"""
|
||||
from sqlalchemy import Column, String, Text, Boolean, Integer
|
||||
|
||||
from app.base_model import BaseModel
|
||||
|
||||
|
||||
class LLMProvider(BaseModel):
|
||||
"""
|
||||
LLM 提供商配置
|
||||
|
||||
支持的提供商类型：
|
||||
- openai: OpenAI (GPT-3.5, GPT-4, etc.)
|
||||
- claude: Anthropic Claude
|
||||
- qwen: 阿里通义千问
|
||||
- ollama: 本地 Ollama
|
||||
- azure_openai: Azure OpenAI
|
||||
- zhipu: 智谱 AI
|
||||
- moonshot: Moonshot (Kimi)
|
||||
- deepseek: DeepSeek
|
||||
"""
|
||||
__tablename__ = "ai_llm_provider"
|
||||
|
||||
name = Column(String(100), nullable=False, comment="提供商名称")
|
||||
provider_type = Column(String(50), nullable=False, comment="提供商类型")
|
||||
api_key = Column(Text, default="", comment="API Key（加密存储）")
|
||||
api_base = Column(String(500), default="", comment="API 地址（可选，用于自定义端点）")
|
||||
api_version = Column(String(50), default="", comment="API 版本（Azure OpenAI 专用）")
|
||||
ollama_host = Column(String(200), default="http://localhost:11434", comment="Ollama 服务地址")
|
||||
is_active = Column(Boolean, default=True, comment="是否启用")
|
||||
description = Column(Text, default="", comment="描述")
|
||||
quota_limit = Column(Integer, default=0, comment="配额限制（0 表示无限制）")
|
||||
quota_used = Column(Integer, default=0, comment="已使用配额")
|
||||
|
||||
def get_api_key_masked(self) -> str:
|
||||
"""获取脱敏的 API Key"""
|
||||
if not self.api_key:
|
||||
return ''
|
||||
if len(self.api_key) <= 8:
|
||||
return '*' * len(self.api_key)
|
||||
return self.api_key[:4] + '*' * (len(self.api_key) - 8) + self.api_key[-4:]
|
||||
@@ -0,0 +1,83 @@
|
||||
"""
|
||||
AI 工作流模型
|
||||
"""
|
||||
from sqlalchemy import Column, String, Text, Integer, Boolean, DateTime, JSON, Index
|
||||
|
||||
from app.base_model import BaseModel
|
||||
|
||||
|
||||
class AIWorkflow(BaseModel):
|
||||
"""
|
||||
AI 工作流定义
|
||||
|
||||
独立于 AIApp 的工作流定义，可以被多个应用引用
|
||||
"""
|
||||
__tablename__ = "ai_workflow"
|
||||
|
||||
application_id = Column(String(21), nullable=True, index=True, comment="所属应用ID（逻辑外键关联core_application）")
|
||||
is_global = Column(Boolean, default=False, comment="是否在子应用中可见")
|
||||
name = Column(String(100), nullable=False, comment="工作流名称")
|
||||
code = Column(String(100), unique=True, nullable=False, comment="工作流编码")
|
||||
workflow_type = Column(String(30), default="general", comment="工作流类型: general/application/form/report/data_process/automation")
|
||||
description = Column(Text, default="", comment="描述")
|
||||
status = Column(String(20), default="draft", comment="状态: draft/published/disabled")
|
||||
version = Column(Integer, default=1, comment="当前草稿版本号")
|
||||
published_version = Column(Integer, nullable=True, comment="已发布的版本号")
|
||||
published_at = Column(DateTime, nullable=True, comment="最后发布时间")
|
||||
published_definition = Column(JSON, default=dict, comment="已发布版本的工作流定义")
|
||||
definition = Column(JSON, default=dict, comment="工作流定义（草稿）")
|
||||
input_variables = Column(JSON, default=list, comment="输入变量定义")
|
||||
output_variables = Column(JSON, default=list, comment="输出变量定义")
|
||||
run_count = Column(Integer, default=0, comment="运行次数")
|
||||
success_count = Column(Integer, default=0, comment="成功次数")
|
||||
|
||||
|
||||
class AIWorkflowVersion(BaseModel):
|
||||
"""
|
||||
AI 工作流版本历史
|
||||
|
||||
每次发布时创建一条版本记录
|
||||
"""
|
||||
__tablename__ = "ai_workflow_version"
|
||||
|
||||
workflow_id = Column(String(21), nullable=False, index=True, comment="工作流ID（逻辑外键关联ai_workflow）")
|
||||
version = Column(Integer, nullable=False, comment="版本号")
|
||||
definition = Column(JSON, default=dict, comment="该版本的工作流定义")
|
||||
description = Column(Text, default="", comment="版本说明")
|
||||
published_by_id = Column(String(21), nullable=True, comment="发布人ID（逻辑外键关联core_user）")
|
||||
published_at = Column(DateTime, nullable=True, comment="发布时间")
|
||||
run_count = Column(Integer, default=0, comment="运行次数")
|
||||
success_count = Column(Integer, default=0, comment="成功次数")
|
||||
|
||||
|
||||
class AIWorkflowRun(BaseModel):
|
||||
"""
|
||||
AI 工作流运行记录
|
||||
"""
|
||||
__tablename__ = "ai_workflow_run"
|
||||
|
||||
workflow_id = Column(String(21), nullable=False, index=True, comment="工作流ID（逻辑外键关联ai_workflow）")
|
||||
app_id = Column(String(21), nullable=True, index=True, comment="关联应用ID（逻辑外键关联ai_app）")
|
||||
conversation_id = Column(String(21), nullable=True, comment="关联对话ID（逻辑外键关联ai_conversation）")
|
||||
user_id = Column(String(21), nullable=True, index=True, comment="执行用户ID（逻辑外键关联core_user）")
|
||||
status = Column(String(20), default="pending", comment="状态: pending/running/waiting/completed/failed/stopped")
|
||||
trigger_type = Column(String(30), default="api", comment="触发来源: editor_draft/editor_published/agent/api/form_button")
|
||||
use_draft = Column(Boolean, default=False, comment="是否使用草稿定义执行")
|
||||
workflow_version = Column(Integer, nullable=True, comment="执行时发布版本号，草稿运行为空")
|
||||
definition_snapshot = Column(JSON, default=dict, comment="运行开始时的工作流定义快照")
|
||||
inputs = Column(JSON, default=dict, comment="输入数据")
|
||||
outputs = Column(JSON, default=dict, comment="输出数据")
|
||||
execution_log = Column(JSON, default=list, comment="执行日志")
|
||||
current_node_id = Column(String(100), default="", comment="当前节点 ID")
|
||||
waiting_config = Column(JSON, default=dict, comment="等待用户输入的配置")
|
||||
error_message = Column(Text, default="", comment="错误信息")
|
||||
total_tokens = Column(Integer, default=0, comment="总 Token 数")
|
||||
total_steps = Column(Integer, default=0, comment="总步骤数")
|
||||
elapsed_time = Column(Integer, default=0, comment="总耗时（毫秒）")
|
||||
started_at = Column(DateTime, nullable=True, comment="开始时间")
|
||||
completed_at = Column(DateTime, nullable=True, comment="完成时间")
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_ai_workflow_run_workflow_status", "workflow_id", "status"),
|
||||
Index("ix_ai_workflow_run_user_status", "user_id", "status"),
|
||||
)
|
||||
@@ -0,0 +1,18 @@
|
||||
"""
|
||||
AI 工作流节点系统
|
||||
|
||||
提供可扩展的节点架构，支持：
|
||||
- 内置节点（LLM、条件、代码等）
|
||||
- 自定义节点扩展
|
||||
- 知识库节点（预留）
|
||||
- 工具节点（预留）
|
||||
"""
|
||||
from .base import BaseNode, NodeContext, NodeResult
|
||||
from .registry import NodeRegistry
|
||||
|
||||
__all__ = [
|
||||
'BaseNode',
|
||||
'NodeContext',
|
||||
'NodeResult',
|
||||
'NodeRegistry',
|
||||
]
|
||||
@@ -0,0 +1,394 @@
|
||||
"""
|
||||
节点基类
|
||||
"""
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class NodeContext:
|
||||
"""
|
||||
节点执行上下文
|
||||
|
||||
包含节点执行所需的所有信息
|
||||
"""
|
||||
# 工作流运行实例
|
||||
workflow_run_id: str = ''
|
||||
|
||||
# 变量存储（所有节点共享）
|
||||
variables: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# 用户输入
|
||||
user_input: str = ''
|
||||
|
||||
# 当前用户
|
||||
user_id: str = ''
|
||||
|
||||
# 对话历史（用于 LLM 节点）
|
||||
conversation_history: List[Dict[str, str]] = field(default_factory=list)
|
||||
|
||||
# 节点配置
|
||||
node_config: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# 上一个节点的输出
|
||||
previous_output: Any = None
|
||||
|
||||
# 元数据
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# 数据库会话（用于需要数据库访问的节点）
|
||||
db_session: Any = None
|
||||
|
||||
def get_variable(self, name: str, default: Any = None) -> Any:
|
||||
"""获取变量"""
|
||||
return self.variables.get(name, default)
|
||||
|
||||
def set_variable(self, name: str, value: Any) -> None:
|
||||
"""设置变量"""
|
||||
self.variables[name] = value
|
||||
|
||||
def resolve_template(self, template: str) -> str:
|
||||
"""
|
||||
解析模板中的变量引用
|
||||
|
||||
支持格式：
|
||||
- {{variable_name}} - 直接变量引用
|
||||
- {{variable_name[0]}} - 数组索引访问
|
||||
- {{variable_name[-1]}} - 负数索引（最后一个）
|
||||
- {{NodeID.key}} - 节点输出引用（如果存在）
|
||||
- {{NodeID.key}}.property - 访问解析结果的嵌套属性
|
||||
- {{NodeID.key[0].property}} - 数组索引 + 属性访问
|
||||
"""
|
||||
import re
|
||||
import json
|
||||
result = template
|
||||
|
||||
# 匹配 {{...}} 格式的变量引用，以及可选的后续属性访问 .property1.property2...
|
||||
pattern = r'\{\{([^}]+)\}\}((?:\.[a-zA-Z_][a-zA-Z0-9_]*)*)'
|
||||
|
||||
def get_nested_value(obj, path: str):
|
||||
"""从对象中获取嵌套属性值，支持数组索引"""
|
||||
if not path:
|
||||
return obj
|
||||
|
||||
# 移除开头的点
|
||||
if path.startswith('.'):
|
||||
path = path[1:]
|
||||
|
||||
current = obj
|
||||
# 使用正则分割路径，支持 .property 和 [index] 格式
|
||||
# 例如: "items[0].name" -> ["items", "[0]", "name"]
|
||||
parts = re.split(r'(?=\[)|\.', path)
|
||||
|
||||
for part in parts:
|
||||
if not part:
|
||||
continue
|
||||
|
||||
# 如果是字符串，尝试解析为 JSON
|
||||
if isinstance(current, str):
|
||||
try:
|
||||
current = json.loads(current)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return None
|
||||
|
||||
# 检查是否是数组索引 [n]
|
||||
index_match = re.match(r'\[(-?\d+)\]', part)
|
||||
if index_match:
|
||||
index = int(index_match.group(1))
|
||||
if isinstance(current, (list, tuple)):
|
||||
try:
|
||||
current = current[index]
|
||||
except IndexError:
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
# 从字典中获取属性
|
||||
elif isinstance(current, dict):
|
||||
if part in current:
|
||||
current = current[part]
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
|
||||
return current
|
||||
|
||||
def replace_var(match):
|
||||
var_ref = match.group(1).strip()
|
||||
extra_path = match.group(2) or '' # 额外的属性路径，如 .customer_name
|
||||
|
||||
value = None
|
||||
index_path = '' # 变量名后的索引/属性路径
|
||||
|
||||
# 检查是否有数组索引 [n]，如果有则分割
|
||||
# 例如: "loop_results[0]" -> var_name="loop_results", index_path="[0]"
|
||||
# 例如: "llm-123.llm_response" -> var_name="llm-123.llm_response", index_path=""
|
||||
bracket_pos = var_ref.find('[')
|
||||
if bracket_pos > 0:
|
||||
var_name = var_ref[:bracket_pos]
|
||||
index_path = var_ref[bracket_pos:]
|
||||
else:
|
||||
var_name = var_ref
|
||||
|
||||
# 尝试解析 NodeID.key.subkey... 格式（支持多层属性访问）
|
||||
if '.' in var_name:
|
||||
parts = var_name.split('.')
|
||||
node_id = parts[0]
|
||||
remaining_path = '.'.join(parts[1:]) # 剩余路径，如 "item.module_name"
|
||||
|
||||
# 先尝试从节点输出命名空间获取
|
||||
node_outputs = self.variables.get(f'_node_{node_id}')
|
||||
if isinstance(node_outputs, dict):
|
||||
# 尝试获取第一层 key
|
||||
first_key = parts[1] if len(parts) > 1 else None
|
||||
if first_key and first_key in node_outputs:
|
||||
value = node_outputs[first_key]
|
||||
# 如果还有更多层级，继续递归获取
|
||||
if len(parts) > 2:
|
||||
nested_path = '.'.join(parts[2:])
|
||||
nested_value = get_nested_value(value, nested_path)
|
||||
if nested_value is not None:
|
||||
value = nested_value
|
||||
|
||||
# 回退：尝试直接从变量获取（node_id 作为变量名）
|
||||
if value is None and node_id in self.variables:
|
||||
node_data = self.variables[node_id]
|
||||
if isinstance(node_data, dict):
|
||||
first_key = parts[1] if len(parts) > 1 else None
|
||||
if first_key and first_key in node_data:
|
||||
value = node_data[first_key]
|
||||
if len(parts) > 2:
|
||||
nested_path = '.'.join(parts[2:])
|
||||
nested_value = get_nested_value(value, nested_path)
|
||||
if nested_value is not None:
|
||||
value = nested_value
|
||||
|
||||
# 再回退：直接从顶层变量获取完整路径
|
||||
if value is None and remaining_path in self.variables:
|
||||
value = self.variables[remaining_path]
|
||||
else:
|
||||
# 直接变量引用
|
||||
if var_name in self.variables:
|
||||
value = self.variables[var_name]
|
||||
|
||||
# 如果找到了值，处理索引路径和额外的属性路径
|
||||
if value is not None:
|
||||
# 合并索引路径和额外路径
|
||||
full_path = index_path + extra_path
|
||||
if full_path:
|
||||
nested_value = get_nested_value(value, full_path)
|
||||
if nested_value is not None:
|
||||
return str(nested_value)
|
||||
# 嵌套属性未找到，返回原始值
|
||||
return str(value)
|
||||
return str(value)
|
||||
|
||||
# 未找到变量，保持原样
|
||||
return match.group(0)
|
||||
|
||||
result = re.sub(pattern, replace_var, result)
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class NodeResult:
|
||||
"""
|
||||
节点执行结果
|
||||
"""
|
||||
# 是否成功
|
||||
success: bool = True
|
||||
|
||||
# 输出数据
|
||||
output: Any = None
|
||||
|
||||
# 输出变量（会合并到上下文的 variables 中）
|
||||
output_variables: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# 错误信息
|
||||
error: str = ''
|
||||
|
||||
# 下一个节点 ID（用于条件分支）
|
||||
next_node_id: str = ''
|
||||
|
||||
# Token 使用（LLM 节点）
|
||||
tokens_used: int = 0
|
||||
|
||||
# 耗时（毫秒）
|
||||
elapsed_time: int = 0
|
||||
|
||||
# 元数据
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# ========== 对话流相关 ==========
|
||||
|
||||
# 是否等待用户输入（对话流模式）
|
||||
waiting_for_input: bool = False
|
||||
|
||||
# 等待配置（描述需要什么类型的输入）
|
||||
waiting_config: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# 事件列表（如发送消息）
|
||||
events: List[Dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
# ========== 设计预览相关 ==========
|
||||
|
||||
# 设计预览数据（用于工作流中显示设计结果并允许编辑）
|
||||
preview: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class BaseNode(ABC):
|
||||
"""
|
||||
节点基类
|
||||
|
||||
所有节点必须继承此类并实现 execute 方法
|
||||
"""
|
||||
|
||||
# 节点类型标识（必须唯一）
|
||||
node_type: str = ''
|
||||
|
||||
# 节点显示名称
|
||||
node_name: str = ''
|
||||
|
||||
# 节点分类
|
||||
node_category: str = 'basic' # basic, llm, logic, data, tool, knowledge
|
||||
|
||||
# 节点图标
|
||||
node_icon: str = ''
|
||||
|
||||
# 节点描述
|
||||
node_description: str = ''
|
||||
|
||||
# 输入参数定义
|
||||
inputs: List[Dict[str, Any]] = []
|
||||
|
||||
# 输出参数定义
|
||||
outputs: List[Dict[str, Any]] = []
|
||||
|
||||
# 是否支持多个输出分支
|
||||
supports_branches: bool = False
|
||||
|
||||
def __init__(self, config: Dict[str, Any] = None):
|
||||
"""
|
||||
初始化节点
|
||||
|
||||
Args:
|
||||
config: 节点配置
|
||||
"""
|
||||
self.config = config or {}
|
||||
|
||||
@abstractmethod
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""
|
||||
执行节点
|
||||
|
||||
Args:
|
||||
context: 执行上下文
|
||||
|
||||
Returns:
|
||||
NodeResult
|
||||
"""
|
||||
pass
|
||||
|
||||
async def execute_async(self, context: NodeContext) -> NodeResult:
|
||||
"""
|
||||
异步执行节点（默认调用同步方法）
|
||||
|
||||
Args:
|
||||
context: 执行上下文
|
||||
|
||||
Returns:
|
||||
NodeResult
|
||||
"""
|
||||
return self.execute(context)
|
||||
|
||||
def validate_config(self) -> tuple:
|
||||
"""
|
||||
验证节点配置
|
||||
|
||||
Returns:
|
||||
(is_valid, error_message)
|
||||
"""
|
||||
return True, ''
|
||||
|
||||
@classmethod
|
||||
def get_schema(cls) -> Dict[str, Any]:
|
||||
"""
|
||||
获取节点 Schema（供前端渲染）
|
||||
|
||||
Returns:
|
||||
节点 Schema
|
||||
"""
|
||||
return {
|
||||
'type': cls.node_type,
|
||||
'name': cls.node_name,
|
||||
'category': cls.node_category,
|
||||
'icon': cls.node_icon,
|
||||
'description': cls.node_description,
|
||||
'inputs': cls.inputs,
|
||||
'outputs': cls.outputs,
|
||||
'supports_branches': cls.supports_branches,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""
|
||||
获取节点配置 Schema（供前端表单渲染）
|
||||
|
||||
Returns:
|
||||
配置 Schema
|
||||
"""
|
||||
return {}
|
||||
|
||||
def resolve_require_confirmation(self, context: NodeContext, default: bool = True) -> bool:
|
||||
"""
|
||||
解析 require_confirmation 配置
|
||||
|
||||
支持三种模式：
|
||||
1. 布尔值：直接使用 True/False
|
||||
2. 字符串 'always'/'never'：始终确认/从不确认
|
||||
3. 变量引用：{{variable_name}} 格式，解析变量值作为布尔值
|
||||
|
||||
Args:
|
||||
context: 节点执行上下文
|
||||
default: 默认值
|
||||
|
||||
Returns:
|
||||
是否需要确认
|
||||
"""
|
||||
value = self.config.get('require_confirmation', default)
|
||||
|
||||
# 布尔值直接返回
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
|
||||
# 字符串处理
|
||||
if isinstance(value, str):
|
||||
value_lower = value.lower().strip()
|
||||
|
||||
# 固定模式
|
||||
if value_lower in ('always', 'true', '1', 'yes'):
|
||||
return True
|
||||
if value_lower in ('never', 'false', '0', 'no'):
|
||||
return False
|
||||
|
||||
# 变量引用模式：{{variable_name}}
|
||||
if '{{' in value and '}}' in value:
|
||||
resolved = context.resolve_template(value)
|
||||
# 解析后的值转换为布尔值
|
||||
if isinstance(resolved, bool):
|
||||
return resolved
|
||||
if isinstance(resolved, str):
|
||||
resolved_lower = resolved.lower().strip()
|
||||
if resolved_lower in ('true', '1', 'yes'):
|
||||
return True
|
||||
if resolved_lower in ('false', '0', 'no'):
|
||||
return False
|
||||
# 非空字符串视为 True
|
||||
return bool(resolved and resolved != value)
|
||||
|
||||
# 其他情况返回默认值
|
||||
return default
|
||||
@@ -0,0 +1,69 @@
|
||||
"""
|
||||
内置节点
|
||||
"""
|
||||
|
||||
# Node modules are imported by ai_platform.nodes.registry so optional nodes can
|
||||
# degrade independently when their runtime dependencies are unavailable.
|
||||
|
||||
__all__ = [
|
||||
'StartNode',
|
||||
'EndNode',
|
||||
'LLMNode',
|
||||
'ConditionNode',
|
||||
'CodeNode',
|
||||
'HttpNode',
|
||||
'TemplateNode',
|
||||
'VariableNode',
|
||||
'ParallelNode',
|
||||
'MergeNode',
|
||||
'BaseDatabaseNode',
|
||||
'DbInsertNode',
|
||||
'DbUpdateNode',
|
||||
'DbQueryNode',
|
||||
'DbDeleteNode',
|
||||
'DbSqlNode',
|
||||
# 对话流节点
|
||||
'QuestionNode',
|
||||
'ChoiceNode',
|
||||
'MessageNode',
|
||||
'ConfirmNode',
|
||||
'IntentNode',
|
||||
# Snowflake Cortex 节点
|
||||
'SnowflakeCortexLLMNode',
|
||||
'SnowflakeCortexAnalystNode',
|
||||
# 循环节点
|
||||
'LoopNode',
|
||||
# 表单节点
|
||||
'FormBasicInfoNode',
|
||||
'FormDatabaseDesignNode',
|
||||
'FormDatabaseCreateNode',
|
||||
'FormUIDesignNode',
|
||||
'FormListDesignNode',
|
||||
'FormCreateNode',
|
||||
'FormPublishNode',
|
||||
# 应用节点
|
||||
'AppCreateNode',
|
||||
'AppDesignNode',
|
||||
'AppSettingsNode',
|
||||
'AppUpdateNode',
|
||||
# 仪表盘节点
|
||||
'DashboardBasicInfoNode',
|
||||
'DashboardDesignNode',
|
||||
'DashboardCreateNode',
|
||||
'DashboardPublishNode',
|
||||
# 系统总结节点
|
||||
'SystemSummaryNode',
|
||||
# 子流程节点
|
||||
'SubflowNode',
|
||||
# Text-to-SQL 节点
|
||||
'TextToSqlNode',
|
||||
# 表单数据节点
|
||||
'FormDataCreateNode',
|
||||
'FormDataReadNode',
|
||||
'FormDataUpdateNode',
|
||||
'FormDataDeleteNode',
|
||||
'FormDataListNode',
|
||||
'FormSchemaToLLMNode',
|
||||
# 知识库节点
|
||||
'KnowledgeRetrievalNode',
|
||||
]
|
||||
@@ -0,0 +1,182 @@
|
||||
"""
|
||||
代码执行节点
|
||||
"""
|
||||
import concurrent.futures
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Dict
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class CodeNode(BaseNode):
|
||||
"""
|
||||
代码执行节点
|
||||
|
||||
执行 Python 代码片段
|
||||
"""
|
||||
|
||||
node_type = 'code'
|
||||
node_name = '代码'
|
||||
node_category = 'logic'
|
||||
node_icon = 'code'
|
||||
node_description = '执行 Python 代码片段'
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'inputs',
|
||||
'type': 'object',
|
||||
'description': '输入变量',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'result',
|
||||
'type': 'any',
|
||||
'description': '执行结果',
|
||||
},
|
||||
]
|
||||
|
||||
# 安全的内置函数白名单
|
||||
SAFE_BUILTINS = {
|
||||
'abs', 'all', 'any', 'bool', 'dict', 'enumerate', 'filter',
|
||||
'float', 'int', 'len', 'list', 'map', 'max', 'min', 'range',
|
||||
'round', 'set', 'sorted', 'str', 'sum', 'tuple', 'zip',
|
||||
'True', 'False', 'None',
|
||||
}
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""执行代码（线程池隔离，防止死循环卡死主流程）"""
|
||||
timeout = self.config.get('timeout', 30)
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
|
||||
future = executor.submit(self._execute_sync, context)
|
||||
try:
|
||||
return future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
future.cancel()
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=f'代码执行超时（{timeout}秒），请检查是否存在死循环',
|
||||
elapsed_time=timeout * 1000,
|
||||
)
|
||||
|
||||
def _execute_sync(self, context: NodeContext) -> NodeResult:
|
||||
"""同步执行代码（在子线程中运行）"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
code = self.config.get('code', '')
|
||||
input_variables = self.config.get('inputs', []) or self.config.get('input_variables', [])
|
||||
output_variable = self.config.get('output_variable', 'result')
|
||||
|
||||
if not code:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='代码不能为空',
|
||||
)
|
||||
|
||||
safe_globals = {
|
||||
'__builtins__': {k: getattr(__builtins__, k) if hasattr(__builtins__, k) else __builtins__[k]
|
||||
for k in self.SAFE_BUILTINS if hasattr(__builtins__, k) or k in __builtins__},
|
||||
}
|
||||
|
||||
import json
|
||||
import re
|
||||
import math
|
||||
safe_globals['json'] = json
|
||||
safe_globals['re'] = re
|
||||
safe_globals['math'] = math
|
||||
|
||||
local_vars = {}
|
||||
for var_config in input_variables:
|
||||
if isinstance(var_config, dict):
|
||||
var_name = var_config.get('variable', '')
|
||||
default_value = var_config.get('default_value', None)
|
||||
if var_name:
|
||||
value = context.get_variable(var_name)
|
||||
if value is not None:
|
||||
local_vars[var_name] = value
|
||||
elif default_value:
|
||||
local_vars[var_name] = context.resolve_template(str(default_value))
|
||||
else:
|
||||
local_vars[var_name] = None
|
||||
elif isinstance(var_config, str):
|
||||
local_vars[var_config] = context.get_variable(var_config)
|
||||
|
||||
local_vars['user_input'] = context.user_input
|
||||
local_vars['variables'] = context.variables.copy()
|
||||
|
||||
exec(code, safe_globals, local_vars)
|
||||
|
||||
if 'main' in local_vars and callable(local_vars['main']):
|
||||
inputs_dict = {
|
||||
'user_input': context.user_input,
|
||||
**context.variables,
|
||||
**local_vars,
|
||||
}
|
||||
main_result = local_vars['main'](inputs_dict)
|
||||
if isinstance(main_result, dict):
|
||||
result = main_result.get('result', main_result)
|
||||
else:
|
||||
result = main_result
|
||||
else:
|
||||
result = local_vars.get('result', None)
|
||||
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=result,
|
||||
output_variables={output_variable: result},
|
||||
elapsed_time=elapsed_time,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'代码节点执行失败: {e}')
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
elapsed_time=elapsed_time,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'code': {
|
||||
'type': 'string',
|
||||
'title': '代码',
|
||||
'description': 'Python 代码，结果存储在 result 变量中',
|
||||
'format': 'code',
|
||||
'default': '# 在这里编写代码\n# 可用变量: user_input, variables\n# 将结果赋值给 result\n\nresult = user_input.upper()',
|
||||
},
|
||||
'input_variables': {
|
||||
'type': 'array',
|
||||
'title': '输入变量',
|
||||
'items': {'type': 'string'},
|
||||
'description': '需要传入代码的变量名列表',
|
||||
},
|
||||
'output_variable': {
|
||||
'type': 'string',
|
||||
'title': '输出变量名',
|
||||
'default': 'code_result',
|
||||
},
|
||||
'timeout': {
|
||||
'type': 'integer',
|
||||
'title': '超时时间（秒）',
|
||||
'default': 30,
|
||||
'minimum': 1,
|
||||
'maximum': 300,
|
||||
},
|
||||
},
|
||||
'required': ['code'],
|
||||
}
|
||||
@@ -0,0 +1,248 @@
|
||||
"""
|
||||
条件分支节点
|
||||
"""
|
||||
import logging
|
||||
import operator
|
||||
from typing import Any, Dict
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class ConditionNode(BaseNode):
|
||||
"""
|
||||
条件分支节点
|
||||
|
||||
根据条件判断选择不同的分支
|
||||
"""
|
||||
|
||||
node_type = 'condition'
|
||||
node_name = '条件分支'
|
||||
node_category = 'logic'
|
||||
node_icon = 'git-branch'
|
||||
node_description = '根据条件判断选择不同的执行分支'
|
||||
|
||||
supports_branches = True
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'value',
|
||||
'type': 'any',
|
||||
'description': '要判断的值',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'branch',
|
||||
'type': 'string',
|
||||
'description': '选中的分支',
|
||||
},
|
||||
]
|
||||
|
||||
# 支持的操作符
|
||||
OPERATORS = {
|
||||
'eq': operator.eq, # 等于
|
||||
'ne': operator.ne, # 不等于
|
||||
'gt': operator.gt, # 大于
|
||||
'gte': operator.ge, # 大于等于
|
||||
'lt': operator.lt, # 小于
|
||||
'lte': operator.le, # 小于等于
|
||||
'contains': lambda a, b: b in str(a), # 包含
|
||||
'not_contains': lambda a, b: b not in str(a), # 不包含
|
||||
'starts_with': lambda a, b: str(a).startswith(str(b)), # 开头是
|
||||
'ends_with': lambda a, b: str(a).endswith(str(b)), # 结尾是
|
||||
'is_empty': lambda a, _: not a, # 为空
|
||||
'is_not_empty': lambda a, _: bool(a), # 不为空
|
||||
}
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""执行条件判断"""
|
||||
try:
|
||||
branches = self.config.get('branches', [])
|
||||
|
||||
# 遍历所有 IF 分支
|
||||
for branch in branches:
|
||||
branch_id = branch.get('id')
|
||||
conditions = branch.get('conditions', [])
|
||||
|
||||
# 评估该分支的所有条件
|
||||
if self._evaluate_branch(conditions, context):
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=branch_id,
|
||||
next_node_id=branch_id,
|
||||
metadata={'matched_branch': branch_id},
|
||||
)
|
||||
|
||||
# 没有匹配的条件，走 ELSE 分支
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output='else',
|
||||
next_node_id='else',
|
||||
metadata={'matched_branch': 'else'},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'条件节点执行失败: {e}')
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
)
|
||||
|
||||
def _evaluate_branch(self, conditions: list, context: NodeContext) -> bool:
|
||||
"""评估分支的一组条件 (AND 关系)"""
|
||||
if not conditions:
|
||||
return True # 无条件默认为真
|
||||
|
||||
for condition in conditions:
|
||||
if not self._evaluate_condition(condition, context):
|
||||
return False
|
||||
return True
|
||||
|
||||
def _resolve_variable(self, variable: str, context: NodeContext) -> Any:
|
||||
"""
|
||||
解析变量引用
|
||||
|
||||
支持格式：
|
||||
- {{variable_name}} - 直接变量引用
|
||||
- {{NodeID.key}} - 节点输出引用
|
||||
- 普通字符串 - 直接返回
|
||||
"""
|
||||
if not isinstance(variable, str):
|
||||
return variable
|
||||
|
||||
# 检查是否是 {{...}} 格式
|
||||
if not (variable.startswith('{{') and variable.endswith('}}')):
|
||||
return variable
|
||||
|
||||
content = variable[2:-2].strip()
|
||||
|
||||
logger.info(f'解析变量: {variable}, content={content}, 上下文变量keys={list(context.variables.keys())}')
|
||||
|
||||
# 尝试解析 NodeID.key 格式
|
||||
if '.' in content:
|
||||
parts = content.split('.', 1)
|
||||
node_id = parts[0]
|
||||
var_key = parts[1]
|
||||
|
||||
logger.info(f'解析节点变量: node_id={node_id}, var_key={var_key}')
|
||||
|
||||
# 先尝试从节点输出命名空间获取
|
||||
node_outputs = context.get_variable(f'_node_{node_id}')
|
||||
logger.info(f'节点输出 _node_{node_id}: {node_outputs}')
|
||||
if isinstance(node_outputs, dict) and var_key in node_outputs:
|
||||
return node_outputs[var_key]
|
||||
|
||||
# 回退：尝试直接从变量获取
|
||||
node_data = context.get_variable(node_id)
|
||||
logger.info(f'直接变量 {node_id}: {node_data}')
|
||||
if isinstance(node_data, dict) and var_key in node_data:
|
||||
return node_data[var_key]
|
||||
|
||||
# 直接变量引用
|
||||
value = context.get_variable(content)
|
||||
if value is not None:
|
||||
return value
|
||||
|
||||
return None
|
||||
|
||||
def _evaluate_condition(self, condition: Dict, context: NodeContext) -> bool:
|
||||
"""评估单个条件"""
|
||||
variable = condition.get('variable', '')
|
||||
op_name = condition.get('operator', 'equals')
|
||||
value = condition.get('value', '')
|
||||
|
||||
# 解析变量和值
|
||||
left_val = self._resolve_variable(variable, context)
|
||||
right_val = self._resolve_variable(value, context)
|
||||
|
||||
logger.info(f'条件评估: variable={variable}, left_val={left_val}, op={op_name}, right_val={right_val}')
|
||||
|
||||
# 映射操作符
|
||||
op_mapping = {
|
||||
'equals': 'eq',
|
||||
'not_equals': 'ne',
|
||||
}
|
||||
op_key = op_mapping.get(op_name, op_name)
|
||||
op_func = self.OPERATORS.get(op_key, operator.eq)
|
||||
|
||||
try:
|
||||
# 特殊处理空值判断，不需要右值
|
||||
if op_key in ['is_empty', 'is_not_empty']:
|
||||
return op_func(left_val, None)
|
||||
|
||||
# 布尔值比较：将字符串 "true"/"false" 转换为布尔值
|
||||
if isinstance(left_val, bool) and isinstance(right_val, str):
|
||||
if right_val.lower() in ['true', '1', 'yes']:
|
||||
right_val = True
|
||||
elif right_val.lower() in ['false', '0', 'no']:
|
||||
right_val = False
|
||||
elif isinstance(right_val, bool) and isinstance(left_val, str):
|
||||
if left_val.lower() in ['true', '1', 'yes']:
|
||||
left_val = True
|
||||
elif left_val.lower() in ['false', '0', 'no']:
|
||||
left_val = False
|
||||
|
||||
# 尝试转换类型以进行比较 (如数字)
|
||||
if isinstance(left_val, (int, float)) and isinstance(right_val, str):
|
||||
try:
|
||||
if '.' in right_val:
|
||||
right_val = float(right_val)
|
||||
else:
|
||||
right_val = int(right_val)
|
||||
except:
|
||||
pass
|
||||
|
||||
# 字符串比较忽略类型差异
|
||||
if op_key in ['contains', 'not_contains', 'starts_with', 'ends_with']:
|
||||
return op_func(str(left_val), str(right_val))
|
||||
|
||||
return op_func(left_val, right_val)
|
||||
except Exception as e:
|
||||
logger.warning(f'条件评估失败: {e}')
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'conditions': {
|
||||
'type': 'array',
|
||||
'title': '条件列表',
|
||||
'items': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'variable': {
|
||||
'type': 'string',
|
||||
'title': '变量名',
|
||||
},
|
||||
'operator': {
|
||||
'type': 'string',
|
||||
'title': '操作符',
|
||||
'enum': list(cls.OPERATORS.keys()),
|
||||
'default': 'eq',
|
||||
},
|
||||
'value': {
|
||||
'type': 'string',
|
||||
'title': '比较值',
|
||||
},
|
||||
'branch_id': {
|
||||
'type': 'string',
|
||||
'title': '分支 ID',
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
'default_branch': {
|
||||
'type': 'string',
|
||||
'title': '默认分支',
|
||||
'description': '没有条件匹配时执行的分支',
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,794 @@
|
||||
"""
|
||||
数据库操作节点
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from datetime import date, datetime
|
||||
from decimal import Decimal
|
||||
from typing import Any, Dict, List, Optional
|
||||
from uuid import UUID
|
||||
|
||||
from sqlalchemy import text
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
from ..utils.db_execution import (
|
||||
DbTarget,
|
||||
build_sql_param_dict,
|
||||
build_where_clause_platform,
|
||||
build_where_clause_raw,
|
||||
default_connection_write_warnings,
|
||||
format_limit_clause,
|
||||
format_select_sql,
|
||||
merge_result_metadata,
|
||||
normalize_return_fields,
|
||||
quote_table_for_target,
|
||||
resolve_db_target,
|
||||
resolve_handler_schema_name,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def serialize_value(value: Any) -> Any:
|
||||
"""将数据库值转换为可 JSON 序列化的格式"""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, datetime):
|
||||
return value.isoformat()
|
||||
if isinstance(value, date):
|
||||
return value.isoformat()
|
||||
if isinstance(value, Decimal):
|
||||
return float(value)
|
||||
if isinstance(value, UUID):
|
||||
return str(value)
|
||||
if isinstance(value, bytes):
|
||||
return value.decode('utf-8', errors='replace')
|
||||
if isinstance(value, (list, tuple)):
|
||||
return [serialize_value(v) for v in value]
|
||||
if isinstance(value, dict):
|
||||
return {k: serialize_value(v) for k, v in value.items()}
|
||||
return value
|
||||
|
||||
|
||||
def serialize_row(row: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""序列化数据库行"""
|
||||
return {k: serialize_value(v) for k, v in row.items()}
|
||||
|
||||
|
||||
def prepare_value_for_db(value: Any) -> Any:
|
||||
"""将值转换为数据库可接受的格式（dict/list 转为 JSON 字符串）"""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, (dict, list)):
|
||||
return json.dumps(value, ensure_ascii=False)
|
||||
if isinstance(value, datetime):
|
||||
return value.isoformat()
|
||||
if isinstance(value, date):
|
||||
return value.isoformat()
|
||||
if isinstance(value, UUID):
|
||||
return str(value)
|
||||
return value
|
||||
|
||||
|
||||
ALLOWED_TABLES = []
|
||||
|
||||
PROTECTED_FIELDS = ['password', 'token', 'secret', 'api_key', 'private_key']
|
||||
|
||||
|
||||
class BaseDatabaseNode(BaseNode):
|
||||
"""
|
||||
数据库操作节点基类
|
||||
|
||||
default 连接走平台 AsyncSession；第三方连接走 AsyncDatabaseManagerService。
|
||||
"""
|
||||
|
||||
node_type = 'database'
|
||||
node_name = '数据库操作'
|
||||
node_category = 'data'
|
||||
node_icon = 'database'
|
||||
node_description = '对数据库进行增删改查操作'
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'data',
|
||||
'type': 'object',
|
||||
'description': '要操作的数据',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'result',
|
||||
'type': 'object',
|
||||
'description': '操作结果',
|
||||
},
|
||||
{
|
||||
'name': 'affected_rows',
|
||||
'type': 'number',
|
||||
'description': '影响的行数',
|
||||
},
|
||||
]
|
||||
|
||||
def _get_db_target(self) -> DbTarget:
|
||||
return resolve_db_target(self.config.get('db_config'))
|
||||
|
||||
def _build_full_table_name(self, table: str, target: Optional[DbTarget] = None) -> str:
|
||||
"""构建完整的表名（平台 PG 路径）"""
|
||||
db_config = self.config.get('db_config', {})
|
||||
schema = db_config.get('schema', '')
|
||||
if schema:
|
||||
return f'"{schema}"."{table}"'
|
||||
return f'"{table}"'
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
import asyncio
|
||||
loop = asyncio.get_event_loop()
|
||||
if loop.is_running():
|
||||
import concurrent.futures
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future = executor.submit(asyncio.run, self.execute_async(context))
|
||||
return future.result()
|
||||
return loop.run_until_complete(self.execute_async(context))
|
||||
|
||||
async def execute_async(self, context: NodeContext) -> NodeResult:
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
operation = self.config.get('operation', 'select').lower()
|
||||
table = self.config.get('table', '')
|
||||
output_variable = self.config.get('output_variable', 'db_result')
|
||||
frontend_max_rows = int(self.config.get('frontend_max_rows', 100))
|
||||
target = self._get_db_target()
|
||||
|
||||
if not table:
|
||||
raise ValueError('未指定目标表')
|
||||
|
||||
if ALLOWED_TABLES and table not in ALLOWED_TABLES:
|
||||
is_allowed = any(allowed == '*' or table == allowed for allowed in ALLOWED_TABLES)
|
||||
if not is_allowed:
|
||||
raise ValueError(f'表 {table} 不在允许操作的白名单中')
|
||||
|
||||
if operation == 'insert':
|
||||
result = await self._execute_insert(table, context, target)
|
||||
elif operation == 'update':
|
||||
result = await self._execute_update(table, context, target)
|
||||
elif operation == 'upsert':
|
||||
result = await self._execute_upsert(table, context, target)
|
||||
elif operation == 'select':
|
||||
result = await self._execute_select(table, context, target)
|
||||
elif operation == 'delete':
|
||||
result = await self._execute_delete(table, context, target)
|
||||
else:
|
||||
raise ValueError(f'不支持的操作类型: {operation}')
|
||||
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
full_data = result.get('data')
|
||||
affected_rows = result.get('affected_rows', 0)
|
||||
|
||||
output_variables = {
|
||||
output_variable: full_data,
|
||||
f'{output_variable}_count': affected_rows,
|
||||
}
|
||||
|
||||
frontend_output_variables = output_variables
|
||||
frontend_output = full_data
|
||||
|
||||
if operation == 'select' and isinstance(full_data, list) and len(full_data) > frontend_max_rows:
|
||||
truncated = full_data[:frontend_max_rows]
|
||||
frontend_output = truncated
|
||||
frontend_output_variables = {
|
||||
output_variable: truncated,
|
||||
f'{output_variable}_count': affected_rows,
|
||||
f'{output_variable}_total': len(full_data),
|
||||
}
|
||||
|
||||
warnings = default_connection_write_warnings(operation, target)
|
||||
metadata = merge_result_metadata(
|
||||
{'frontend_output_variables': frontend_output_variables},
|
||||
warnings,
|
||||
)
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=frontend_output,
|
||||
output_variables=output_variables,
|
||||
metadata=metadata,
|
||||
elapsed_time=elapsed_time,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'数据库节点执行失败: {e}')
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
elapsed_time=int((time.time() - start_time) * 1000),
|
||||
)
|
||||
|
||||
async def _create_db_service(self, context: NodeContext, target: DbTarget):
|
||||
from core.database_manager.service import AsyncDatabaseManagerService
|
||||
|
||||
try:
|
||||
return await AsyncDatabaseManagerService.create(target.db_name, context.db_session)
|
||||
except Exception as e:
|
||||
raise ValueError(f'无法连接数据库 {target.db_name}: {e}') from e
|
||||
|
||||
def _resolve_field_mapping(self, context: NodeContext) -> Dict[str, Any]:
|
||||
field_mapping = self.config.get('field_mapping', {})
|
||||
resolved = {}
|
||||
|
||||
for field, value in field_mapping.items():
|
||||
if field.lower() in PROTECTED_FIELDS:
|
||||
logger.warning(f'跳过保护字段: {field}')
|
||||
continue
|
||||
|
||||
if isinstance(value, str):
|
||||
resolved_value = context.resolve_template(value)
|
||||
try:
|
||||
resolved[field] = json.loads(resolved_value)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
resolved[field] = resolved_value
|
||||
else:
|
||||
resolved[field] = value
|
||||
|
||||
return resolved
|
||||
|
||||
def _resolve_conditions(self, context: NodeContext) -> List[Dict[str, Any]]:
|
||||
conditions = self.config.get('where_conditions', [])
|
||||
resolved = []
|
||||
|
||||
for condition in conditions:
|
||||
field = condition.get('field', '')
|
||||
operator = condition.get('operator', '=')
|
||||
value = condition.get('value', '')
|
||||
|
||||
if isinstance(value, str):
|
||||
resolved_value = context.resolve_template(value)
|
||||
try:
|
||||
resolved_value = json.loads(resolved_value)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
else:
|
||||
resolved_value = value
|
||||
|
||||
resolved.append({
|
||||
'field': field,
|
||||
'operator': operator,
|
||||
'value': resolved_value,
|
||||
})
|
||||
|
||||
return resolved
|
||||
|
||||
async def _execute_insert(
|
||||
self,
|
||||
table: str,
|
||||
context: NodeContext,
|
||||
target: DbTarget,
|
||||
) -> Dict[str, Any]:
|
||||
data = self._resolve_field_mapping(context)
|
||||
if not data:
|
||||
raise ValueError('没有要插入的数据')
|
||||
if 'id' not in data:
|
||||
data['id'] = str(uuid.uuid4())
|
||||
|
||||
if target.is_external:
|
||||
db_service = await self._create_db_service(context, target)
|
||||
schema_name = await resolve_handler_schema_name(db_service, target)
|
||||
payload = {k: prepare_value_for_db(v) for k, v in data.items()}
|
||||
result = await db_service.insert_data(table, payload, schema_name)
|
||||
if not result.get('success'):
|
||||
raise ValueError(result.get('message', '插入失败'))
|
||||
return {
|
||||
'data': {'id': data['id'], **data},
|
||||
'affected_rows': result.get('affected_rows', 1),
|
||||
}
|
||||
|
||||
db = context.db_session
|
||||
if not db:
|
||||
raise ValueError('数据库会话不可用，请确保工作流已配置数据库连接')
|
||||
|
||||
full_table_name = self._build_full_table_name(table)
|
||||
fields = list(data.keys())
|
||||
params = {f: prepare_value_for_db(v) for f, v in data.items()}
|
||||
placeholders = ', '.join([f':{f}' for f in fields])
|
||||
field_names = ', '.join([f'"{f}"' for f in fields])
|
||||
sql = f'INSERT INTO {full_table_name} ({field_names}) VALUES ({placeholders})'
|
||||
await db.execute(text(sql), params)
|
||||
return {
|
||||
'data': {'id': data['id'], **data},
|
||||
'affected_rows': 1,
|
||||
}
|
||||
|
||||
async def _execute_update(
|
||||
self,
|
||||
table: str,
|
||||
context: NodeContext,
|
||||
target: DbTarget,
|
||||
) -> Dict[str, Any]:
|
||||
data = self._resolve_field_mapping(context)
|
||||
conditions = self._resolve_conditions(context)
|
||||
|
||||
if not data:
|
||||
raise ValueError('没有要更新的数据')
|
||||
if not conditions:
|
||||
raise ValueError('UPDATE 操作必须指定条件，防止误更新全表')
|
||||
|
||||
if target.is_external:
|
||||
db_service = await self._create_db_service(context, target)
|
||||
schema_name = await resolve_handler_schema_name(db_service, target)
|
||||
where_raw = build_where_clause_raw(conditions, db_service.db_type)
|
||||
payload = {k: prepare_value_for_db(v) for k, v in data.items()}
|
||||
result = await db_service.update_data(table, payload, where_raw, schema_name)
|
||||
if not result.get('success'):
|
||||
raise ValueError(result.get('message', '更新失败'))
|
||||
affected_rows = result.get('affected_rows', 0)
|
||||
return {
|
||||
'data': {'updated': True, 'affected_rows': affected_rows, **data},
|
||||
'affected_rows': affected_rows,
|
||||
}
|
||||
|
||||
db = context.db_session
|
||||
if not db:
|
||||
raise ValueError('数据库会话不可用')
|
||||
|
||||
set_clauses = []
|
||||
params = {}
|
||||
for field, value in data.items():
|
||||
param_name = f's_{field}'
|
||||
set_clauses.append(f'"{field}" = :{param_name}')
|
||||
params[param_name] = prepare_value_for_db(value)
|
||||
|
||||
where_clause, where_params = build_where_clause_platform(conditions)
|
||||
params.update(where_params)
|
||||
full_table_name = self._build_full_table_name(table)
|
||||
sql = f'UPDATE {full_table_name} SET {", ".join(set_clauses)} {where_clause}'
|
||||
result = await db.execute(text(sql), params)
|
||||
affected_rows = result.rowcount
|
||||
return {
|
||||
'data': {'updated': True, 'affected_rows': affected_rows, **data},
|
||||
'affected_rows': affected_rows,
|
||||
}
|
||||
|
||||
async def _execute_upsert(
|
||||
self,
|
||||
table: str,
|
||||
context: NodeContext,
|
||||
target: DbTarget,
|
||||
) -> Dict[str, Any]:
|
||||
data = self._resolve_field_mapping(context)
|
||||
conditions = self._resolve_conditions(context)
|
||||
if not data:
|
||||
raise ValueError('没有要操作的数据')
|
||||
|
||||
if conditions:
|
||||
if target.is_external:
|
||||
db_service = await self._create_db_service(context, target)
|
||||
schema_name = await resolve_handler_schema_name(db_service, target)
|
||||
where_raw = build_where_clause_raw(conditions, db_service.db_type)
|
||||
full_table = quote_table_for_target(table, target)
|
||||
check_sql = f'SELECT 1 FROM {full_table}'
|
||||
if where_raw:
|
||||
check_sql += f' WHERE {where_raw}'
|
||||
check_sql += format_limit_clause(db_service.db_type, 1)
|
||||
check_result = await db_service.execute_sql(check_sql, is_query=True)
|
||||
rows = check_result.get('rows') or check_result.get('data') or []
|
||||
if rows:
|
||||
return await self._execute_update(table, context, target)
|
||||
else:
|
||||
db = context.db_session
|
||||
if not db:
|
||||
raise ValueError('数据库会话不可用')
|
||||
full_table_name = self._build_full_table_name(table)
|
||||
where_clause, where_params = build_where_clause_platform(conditions)
|
||||
check_sql = f'SELECT id FROM {full_table_name} {where_clause} LIMIT 1'
|
||||
result = await db.execute(text(check_sql), where_params)
|
||||
if result.fetchone():
|
||||
return await self._execute_update(table, context, target)
|
||||
|
||||
return await self._execute_insert(table, context, target)
|
||||
|
||||
async def _execute_select(
|
||||
self,
|
||||
table: str,
|
||||
context: NodeContext,
|
||||
target: DbTarget,
|
||||
) -> Dict[str, Any]:
|
||||
conditions = self._resolve_conditions(context)
|
||||
return_fields = self.config.get('return_fields', ['*'])
|
||||
limit = self.config.get('limit', 100)
|
||||
order_by = self.config.get('order_by', '')
|
||||
|
||||
if target.is_external:
|
||||
db_service = await self._create_db_service(context, target)
|
||||
sql = format_select_sql(
|
||||
table,
|
||||
target,
|
||||
return_fields=return_fields,
|
||||
conditions=conditions,
|
||||
order_by=order_by,
|
||||
limit=int(limit),
|
||||
)
|
||||
result_data = await db_service.execute_sql(sql, is_query=True)
|
||||
if result_data.get('success') is False:
|
||||
raise ValueError(result_data.get('message') or '查询失败')
|
||||
rows = result_data.get('rows') or result_data.get('data') or []
|
||||
result_data_list = []
|
||||
for row in rows:
|
||||
if isinstance(row, dict):
|
||||
result_data_list.append(serialize_row(row))
|
||||
else:
|
||||
result_data_list.append(serialize_row(dict(row)))
|
||||
return {
|
||||
'data': result_data_list,
|
||||
'affected_rows': len(result_data_list),
|
||||
}
|
||||
|
||||
db = context.db_session
|
||||
if not db:
|
||||
raise ValueError('数据库会话不可用')
|
||||
|
||||
normalized_fields = normalize_return_fields(return_fields)
|
||||
if normalized_fields == '*':
|
||||
field_list = '*'
|
||||
else:
|
||||
field_list = ', '.join([f'"{f}"' for f in normalized_fields])
|
||||
|
||||
where_clause, params = build_where_clause_platform(conditions)
|
||||
full_table_name = self._build_full_table_name(table)
|
||||
sql = f'SELECT {field_list} FROM {full_table_name} {where_clause}'
|
||||
if order_by:
|
||||
sql += f' ORDER BY {order_by}'
|
||||
sql += f' LIMIT {int(limit)}'
|
||||
|
||||
result = await db.execute(text(sql), params)
|
||||
columns = result.keys()
|
||||
rows = result.fetchall()
|
||||
result_data = [serialize_row(dict(zip(columns, row))) for row in rows]
|
||||
return {
|
||||
'data': result_data,
|
||||
'affected_rows': len(result_data),
|
||||
}
|
||||
|
||||
async def _execute_delete(
|
||||
self,
|
||||
table: str,
|
||||
context: NodeContext,
|
||||
target: DbTarget,
|
||||
) -> Dict[str, Any]:
|
||||
conditions = self._resolve_conditions(context)
|
||||
if not conditions:
|
||||
raise ValueError('DELETE 操作必须指定条件，防止误删全表')
|
||||
|
||||
if target.is_external:
|
||||
db_service = await self._create_db_service(context, target)
|
||||
schema_name = await resolve_handler_schema_name(db_service, target)
|
||||
where_raw = build_where_clause_raw(conditions, db_service.db_type)
|
||||
result = await db_service.delete_data(table, where_raw, schema_name)
|
||||
if not result.get('success'):
|
||||
raise ValueError(result.get('message', '删除失败'))
|
||||
affected_rows = result.get('affected_rows', 0)
|
||||
return {
|
||||
'data': {'deleted': True, 'affected_rows': affected_rows},
|
||||
'affected_rows': affected_rows,
|
||||
}
|
||||
|
||||
db = context.db_session
|
||||
if not db:
|
||||
raise ValueError('数据库会话不可用')
|
||||
|
||||
where_clause, params = build_where_clause_platform(conditions)
|
||||
full_table_name = self._build_full_table_name(table)
|
||||
sql = f'DELETE FROM {full_table_name} {where_clause}'
|
||||
result = await db.execute(text(sql), params)
|
||||
affected_rows = result.rowcount
|
||||
return {
|
||||
'data': {'deleted': True, 'affected_rows': affected_rows},
|
||||
'affected_rows': affected_rows,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'operation': {
|
||||
'type': 'string',
|
||||
'title': '操作类型',
|
||||
'enum': ['insert', 'update', 'upsert', 'select', 'delete'],
|
||||
'enumNames': ['插入', '更新', '插入或更新', '查询', '删除'],
|
||||
'default': 'insert',
|
||||
},
|
||||
'table': {
|
||||
'type': 'string',
|
||||
'title': '目标表',
|
||||
'description': '数据库表名',
|
||||
},
|
||||
'field_mapping': {
|
||||
'type': 'object',
|
||||
'title': '字段映射',
|
||||
'description': '数据库字段与变量的映射关系',
|
||||
'additionalProperties': {'type': 'string'},
|
||||
},
|
||||
'where_conditions': {
|
||||
'type': 'array',
|
||||
'title': '条件',
|
||||
'description': '查询/更新/删除的条件',
|
||||
'items': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'field': {'type': 'string', 'title': '字段'},
|
||||
'operator': {
|
||||
'type': 'string',
|
||||
'title': '操作符',
|
||||
'enum': ['=', '!=', '>', '>=', '<', '<=', 'like', 'in', 'is_null', 'is_not_null'],
|
||||
'default': '=',
|
||||
},
|
||||
'value': {'type': 'string', 'title': '值'},
|
||||
},
|
||||
},
|
||||
},
|
||||
'return_fields': {
|
||||
'type': 'array',
|
||||
'title': '返回字段',
|
||||
'description': '查询时返回的字段列表',
|
||||
'items': {'type': 'string'},
|
||||
'default': ['*'],
|
||||
},
|
||||
'limit': {
|
||||
'type': 'integer',
|
||||
'title': '限制条数',
|
||||
'description': 'SQL 查询时的最大返回条数',
|
||||
'default': 100,
|
||||
},
|
||||
'frontend_max_rows': {
|
||||
'type': 'integer',
|
||||
'title': '前端返回最大条数',
|
||||
'description': '前端 SSE 事件中返回的最大数据条数（默认100），超过此值仅截断前端传输，后续节点仍可获取全量数据',
|
||||
'default': 100,
|
||||
'minimum': 1,
|
||||
'maximum': 10000,
|
||||
},
|
||||
'order_by': {
|
||||
'type': 'string',
|
||||
'title': '排序',
|
||||
'description': '排序字段，如 created_at DESC',
|
||||
},
|
||||
'output_variable': {
|
||||
'type': 'string',
|
||||
'title': '输出变量名',
|
||||
'default': 'db_result',
|
||||
},
|
||||
},
|
||||
'required': ['operation', 'table'],
|
||||
}
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class DbInsertNode(BaseDatabaseNode):
|
||||
node_type = 'db_insert'
|
||||
node_name = 'DB 插入'
|
||||
node_icon = 'database-zap'
|
||||
node_description = '向数据库插入数据'
|
||||
|
||||
def __init__(self, config: Dict[str, Any] = None):
|
||||
super().__init__(config)
|
||||
if config:
|
||||
self.config['operation'] = 'upsert' if config.get('upsert') else 'insert'
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class DbUpdateNode(BaseDatabaseNode):
|
||||
node_type = 'db_update'
|
||||
node_name = 'DB 更新'
|
||||
node_icon = 'database-backup'
|
||||
node_description = '更新数据库记录'
|
||||
|
||||
def __init__(self, config: Dict[str, Any] = None):
|
||||
super().__init__(config)
|
||||
if config:
|
||||
self.config['operation'] = 'update'
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class DbQueryNode(BaseDatabaseNode):
|
||||
node_type = 'db_query'
|
||||
node_name = 'DB 查询'
|
||||
node_icon = 'search'
|
||||
node_description = '从数据库查询数据'
|
||||
|
||||
def __init__(self, config: Dict[str, Any] = None):
|
||||
super().__init__(config)
|
||||
if config:
|
||||
self.config['operation'] = 'select'
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class DbDeleteNode(BaseDatabaseNode):
|
||||
node_type = 'db_delete'
|
||||
node_name = 'DB 删除'
|
||||
node_icon = 'trash-2'
|
||||
node_description = '从数据库删除数据'
|
||||
|
||||
def __init__(self, config: Dict[str, Any] = None):
|
||||
super().__init__(config)
|
||||
if config:
|
||||
self.config['operation'] = 'delete'
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class DbSqlNode(BaseNode):
|
||||
"""自定义 SQL 执行节点"""
|
||||
|
||||
node_type = 'db_sql'
|
||||
node_name = 'SQL 执行'
|
||||
node_category = 'data'
|
||||
node_icon = 'database'
|
||||
node_description = '执行自定义 SQL 语句'
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'data',
|
||||
'type': 'object',
|
||||
'description': '输入数据',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'result',
|
||||
'type': 'any',
|
||||
'description': 'SQL 执行结果',
|
||||
},
|
||||
]
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
import asyncio
|
||||
loop = asyncio.get_event_loop()
|
||||
if loop.is_running():
|
||||
import concurrent.futures
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future = executor.submit(asyncio.run, self.execute_async(context))
|
||||
return future.result()
|
||||
return loop.run_until_complete(self.execute_async(context))
|
||||
|
||||
async def execute_async(self, context: NodeContext) -> NodeResult:
|
||||
start_time = time.time()
|
||||
|
||||
sql_type = self.config.get('sql_type', 'query')
|
||||
sql = self.config.get('sql', '')
|
||||
target = resolve_db_target(self.config.get('db_config'))
|
||||
output_variable = self.config.get('output_variable', 'sql_result')
|
||||
is_query = sql_type == 'query'
|
||||
|
||||
if not sql:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='SQL 语句不能为空',
|
||||
elapsed_time=int((time.time() - start_time) * 1000),
|
||||
)
|
||||
|
||||
try:
|
||||
resolved_sql = context.resolve_template(sql)
|
||||
param_dict = build_sql_param_dict(self.config.get('params'), context)
|
||||
logger.info(
|
||||
'执行 SQL [%s]: %s, params=%s',
|
||||
target.db_name,
|
||||
resolved_sql,
|
||||
list(param_dict.keys()),
|
||||
)
|
||||
|
||||
operation = 'query' if is_query else 'execute'
|
||||
warnings = default_connection_write_warnings(operation, target)
|
||||
|
||||
if not target.is_external:
|
||||
db = context.db_session
|
||||
if not db:
|
||||
raise ValueError('数据库会话不可用')
|
||||
result = await db.execute(text(resolved_sql), param_dict)
|
||||
if is_query:
|
||||
rows = result.mappings().all()
|
||||
output_result = [serialize_row(dict(row)) for row in rows]
|
||||
row_count = len(output_result)
|
||||
else:
|
||||
row_count = max(result.rowcount or 0, 0)
|
||||
output_result = row_count
|
||||
else:
|
||||
from core.database_manager.service import AsyncDatabaseManagerService
|
||||
from utils.sql_param_compile import compile_sql_with_named_params
|
||||
|
||||
try:
|
||||
db_service = await AsyncDatabaseManagerService.create(
|
||||
target.db_name,
|
||||
context.db_session,
|
||||
)
|
||||
except Exception as e:
|
||||
raise ValueError(f'无法连接数据库 {target.db_name}: {e}') from e
|
||||
|
||||
executable_sql = compile_sql_with_named_params(
|
||||
resolved_sql,
|
||||
param_dict,
|
||||
target.db_type,
|
||||
)
|
||||
result_data = await db_service.execute_sql(executable_sql, is_query=is_query)
|
||||
if result_data.get('success') is False:
|
||||
raise Exception(result_data.get('message') or 'SQL 执行失败')
|
||||
|
||||
if is_query:
|
||||
rows = result_data.get('rows') or result_data.get('data') or []
|
||||
output_result = []
|
||||
for row in rows:
|
||||
if isinstance(row, dict):
|
||||
output_result.append(serialize_row(row))
|
||||
else:
|
||||
output_result.append(serialize_row(dict(row)))
|
||||
row_count = len(output_result)
|
||||
else:
|
||||
output_result = result_data.get('affected_rows', 0)
|
||||
row_count = output_result
|
||||
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
metadata = merge_result_metadata({}, warnings)
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=output_result,
|
||||
output_variables={
|
||||
output_variable: output_result,
|
||||
f'{output_variable}_count': row_count,
|
||||
},
|
||||
metadata=metadata,
|
||||
elapsed_time=elapsed_time,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error('SQL 执行失败: %s', e)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=f'SQL 执行失败: {str(e)}',
|
||||
elapsed_time=int((time.time() - start_time) * 1000),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'sql_type': {
|
||||
'type': 'string',
|
||||
'title': '执行类型',
|
||||
'enum': ['query', 'execute'],
|
||||
'default': 'query',
|
||||
'description': 'query: 查询返回结果, execute: 执行不返回结果',
|
||||
},
|
||||
'sql': {
|
||||
'type': 'string',
|
||||
'title': 'SQL 语句',
|
||||
'description': '要执行的 SQL 语句，使用 :param_name 作为命名参数占位符',
|
||||
},
|
||||
'params': {
|
||||
'type': 'array',
|
||||
'title': '参数列表',
|
||||
'description': 'SQL 命名参数，与 SQL 中 :param_name 对应',
|
||||
'items': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'name': {'type': 'string', 'title': '参数名'},
|
||||
'type': {
|
||||
'type': 'string',
|
||||
'enum': ['string', 'integer', 'float', 'boolean', 'date', 'datetime'],
|
||||
'default': 'string',
|
||||
},
|
||||
'value': {'type': 'string', 'title': '参数值'},
|
||||
},
|
||||
},
|
||||
},
|
||||
'output_variable': {
|
||||
'type': 'string',
|
||||
'title': '输出变量名',
|
||||
'default': 'sql_result',
|
||||
},
|
||||
},
|
||||
'required': ['sql'],
|
||||
}
|
||||
@@ -0,0 +1,672 @@
|
||||
"""
|
||||
对话流节点
|
||||
|
||||
用于对话流模式的智能体，支持与用户的交互式对话
|
||||
"""
|
||||
import logging
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class QuestionNode(BaseNode):
|
||||
"""
|
||||
问答节点
|
||||
|
||||
向用户提出问题，等待用户输入回答
|
||||
支持输入验证和默认值
|
||||
"""
|
||||
|
||||
node_type = 'question'
|
||||
node_name = '问答节点'
|
||||
node_category = 'dialog'
|
||||
|
||||
inputs = {
|
||||
'question': {
|
||||
'type': 'string',
|
||||
'description': '要向用户提出的问题',
|
||||
'required': True,
|
||||
},
|
||||
'variable_name': {
|
||||
'type': 'string',
|
||||
'description': '存储用户回答的变量名',
|
||||
'required': True,
|
||||
},
|
||||
'input_type': {
|
||||
'type': 'string',
|
||||
'description': '输入类型：text/number/email/phone/date',
|
||||
'default': 'text',
|
||||
},
|
||||
'placeholder': {
|
||||
'type': 'string',
|
||||
'description': '输入框占位文本',
|
||||
'default': '',
|
||||
},
|
||||
'default_value': {
|
||||
'type': 'string',
|
||||
'description': '默认值',
|
||||
'default': '',
|
||||
},
|
||||
'required': {
|
||||
'type': 'boolean',
|
||||
'description': '是否必填',
|
||||
'default': True,
|
||||
},
|
||||
'validation_regex': {
|
||||
'type': 'string',
|
||||
'description': '验证正则表达式',
|
||||
'default': '',
|
||||
},
|
||||
'validation_message': {
|
||||
'type': 'string',
|
||||
'description': '验证失败提示',
|
||||
'default': '输入格式不正确',
|
||||
},
|
||||
'render_input': {
|
||||
'type': 'boolean',
|
||||
'description': '是否渲染输入框，默认不渲染（用户直接在聊天框输入）',
|
||||
'default': False,
|
||||
},
|
||||
}
|
||||
|
||||
outputs = {
|
||||
'answer': {
|
||||
'type': 'string',
|
||||
'description': '用户的回答',
|
||||
},
|
||||
'is_valid': {
|
||||
'type': 'boolean',
|
||||
'description': '回答是否有效',
|
||||
},
|
||||
}
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""
|
||||
执行问答节点
|
||||
|
||||
这个节点会暂停工作流执行，等待用户输入
|
||||
"""
|
||||
question = self.config.get('question', '')
|
||||
variable_name = self.config.get('variable_name', 'user_input')
|
||||
input_type = self.config.get('input_type', 'text')
|
||||
placeholder = self.config.get('placeholder', '')
|
||||
default_value = self.config.get('default_value', '')
|
||||
required = self.config.get('required', True)
|
||||
validation_regex = self.config.get('validation_regex', '')
|
||||
validation_message = self.config.get('validation_message', '输入格式不正确')
|
||||
render_input = self.config.get('render_input', False)
|
||||
|
||||
# 解析模板变量
|
||||
question = context.resolve_template(question)
|
||||
placeholder = context.resolve_template(placeholder)
|
||||
default_value = context.resolve_template(default_value)
|
||||
|
||||
# 检查是否已有用户输入（续流时）
|
||||
user_input = context.variables.get('__user_input__')
|
||||
|
||||
if user_input is not None:
|
||||
# 用户已输入，验证并继续
|
||||
import re
|
||||
|
||||
# 将非字符串输入转换为字符串
|
||||
if not isinstance(user_input, str):
|
||||
user_input = str(user_input)
|
||||
|
||||
is_valid = True
|
||||
|
||||
# 必填验证
|
||||
if required and not user_input.strip():
|
||||
is_valid = False
|
||||
|
||||
# 正则验证
|
||||
if is_valid and validation_regex:
|
||||
if not re.match(validation_regex, user_input):
|
||||
is_valid = False
|
||||
|
||||
# 类型验证
|
||||
if is_valid and input_type == 'number':
|
||||
try:
|
||||
user_input = float(user_input)
|
||||
except ValueError:
|
||||
is_valid = False
|
||||
elif is_valid and input_type == 'email':
|
||||
email_regex = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
|
||||
if not re.match(email_regex, user_input):
|
||||
is_valid = False
|
||||
elif is_valid and input_type == 'phone':
|
||||
phone_regex = r'^1[3-9]\d{9}$'
|
||||
if not re.match(phone_regex, user_input):
|
||||
is_valid = False
|
||||
|
||||
if is_valid:
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={
|
||||
'answer': user_input,
|
||||
'is_valid': True,
|
||||
},
|
||||
output_variables={
|
||||
variable_name: user_input,
|
||||
},
|
||||
)
|
||||
else:
|
||||
# 验证失败，重新等待输入
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={
|
||||
'answer': '',
|
||||
'is_valid': False,
|
||||
},
|
||||
waiting_for_input=True,
|
||||
waiting_config={
|
||||
'type': 'question',
|
||||
'question': question,
|
||||
'input_type': input_type,
|
||||
'placeholder': placeholder,
|
||||
'default_value': default_value,
|
||||
'required': required,
|
||||
'error_message': validation_message,
|
||||
'render_input': render_input,
|
||||
},
|
||||
)
|
||||
else:
|
||||
# 首次执行，等待用户输入
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={
|
||||
'answer': '',
|
||||
'is_valid': False,
|
||||
},
|
||||
waiting_for_input=True,
|
||||
waiting_config={
|
||||
'type': 'question',
|
||||
'question': question,
|
||||
'input_type': input_type,
|
||||
'placeholder': placeholder,
|
||||
'default_value': default_value,
|
||||
'required': required,
|
||||
'render_input': render_input,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class ChoiceNode(BaseNode):
|
||||
"""
|
||||
选项节点
|
||||
|
||||
向用户展示多个选项，用户选择后继续执行
|
||||
支持单选和多选
|
||||
"""
|
||||
|
||||
node_type = 'choice'
|
||||
node_name = '选项节点'
|
||||
node_category = 'dialog'
|
||||
|
||||
inputs = {
|
||||
'question': {
|
||||
'type': 'string',
|
||||
'description': '问题或提示文本',
|
||||
'required': True,
|
||||
},
|
||||
'variable_name': {
|
||||
'type': 'string',
|
||||
'description': '存储用户选择的变量名',
|
||||
'required': True,
|
||||
},
|
||||
'options': {
|
||||
'type': 'array',
|
||||
'description': '选项列表',
|
||||
'required': True,
|
||||
'items': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'value': {'type': 'string', 'description': '选项值'},
|
||||
'label': {'type': 'string', 'description': '显示文本'},
|
||||
'description': {'type': 'string', 'description': '选项描述'},
|
||||
},
|
||||
},
|
||||
},
|
||||
'multiple': {
|
||||
'type': 'boolean',
|
||||
'description': '是否多选',
|
||||
'default': False,
|
||||
},
|
||||
'min_select': {
|
||||
'type': 'integer',
|
||||
'description': '最少选择数量',
|
||||
'default': 1,
|
||||
},
|
||||
'max_select': {
|
||||
'type': 'integer',
|
||||
'description': '最多选择数量',
|
||||
'default': 1,
|
||||
},
|
||||
}
|
||||
|
||||
outputs = {
|
||||
'selected': {
|
||||
'type': 'any',
|
||||
'description': '用户选择的值（单选为字符串，多选为数组）',
|
||||
},
|
||||
'selected_labels': {
|
||||
'type': 'any',
|
||||
'description': '用户选择的显示文本',
|
||||
},
|
||||
}
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""
|
||||
执行选项节点
|
||||
"""
|
||||
question = self.config.get('question', '')
|
||||
variable_name = self.config.get('variable_name', 'user_choice')
|
||||
options = self.config.get('options', [])
|
||||
multiple = self.config.get('multiple', False)
|
||||
min_select = self.config.get('min_select', 1)
|
||||
max_select = self.config.get('max_select', 1)
|
||||
|
||||
# 解析模板变量
|
||||
question = context.resolve_template(question)
|
||||
|
||||
# 检查是否已有用户选择
|
||||
user_selection = context.variables.get('__user_input__')
|
||||
|
||||
if user_selection is not None:
|
||||
# 用户已选择
|
||||
if multiple:
|
||||
# 多选：user_selection 应该是数组
|
||||
if isinstance(user_selection, str):
|
||||
user_selection = [user_selection]
|
||||
elif not isinstance(user_selection, list):
|
||||
# 如果不是字符串也不是列表（比如布尔值），转换为字符串后放入列表
|
||||
user_selection = [str(user_selection)]
|
||||
|
||||
# 验证所有选项值是否有效
|
||||
valid_values = [opt.get('value') for opt in options]
|
||||
invalid_selections = [val for val in user_selection if val not in valid_values]
|
||||
|
||||
if invalid_selections:
|
||||
# 有无效选项，重新显示选择界面
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={'selected': [], 'selected_labels': []},
|
||||
waiting_for_input=True,
|
||||
waiting_config={
|
||||
'type': 'choice',
|
||||
'question': question,
|
||||
'options': options,
|
||||
'multiple': multiple,
|
||||
'min_select': min_select,
|
||||
'max_select': max_select,
|
||||
'error_message': '请从选项中选择',
|
||||
},
|
||||
)
|
||||
|
||||
# 验证选择数量
|
||||
if len(user_selection) < min_select or len(user_selection) > max_select:
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={'selected': [], 'selected_labels': []},
|
||||
waiting_for_input=True,
|
||||
waiting_config={
|
||||
'type': 'choice',
|
||||
'question': question,
|
||||
'options': options,
|
||||
'multiple': multiple,
|
||||
'min_select': min_select,
|
||||
'max_select': max_select,
|
||||
'error_message': f'请选择 {min_select}-{max_select} 个选项',
|
||||
},
|
||||
)
|
||||
|
||||
# 获取选中的标签
|
||||
selected_labels = []
|
||||
for opt in options:
|
||||
if opt.get('value') in user_selection:
|
||||
selected_labels.append(opt.get('label', opt.get('value')))
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={
|
||||
'selected': user_selection,
|
||||
'selected_labels': selected_labels,
|
||||
},
|
||||
output_variables={
|
||||
variable_name: user_selection,
|
||||
f'{variable_name}_labels': selected_labels,
|
||||
},
|
||||
)
|
||||
else:
|
||||
# 单选：验证选项值是否有效
|
||||
valid_values = [opt.get('value') for opt in options]
|
||||
if user_selection not in valid_values:
|
||||
# 无效选项，重新显示选择界面
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={'selected': None, 'selected_labels': None},
|
||||
waiting_for_input=True,
|
||||
waiting_config={
|
||||
'type': 'choice',
|
||||
'question': question,
|
||||
'options': options,
|
||||
'multiple': multiple,
|
||||
'min_select': min_select,
|
||||
'max_select': max_select,
|
||||
'error_message': '请从选项中选择',
|
||||
},
|
||||
)
|
||||
|
||||
selected_label = ''
|
||||
for opt in options:
|
||||
if opt.get('value') == user_selection:
|
||||
selected_label = opt.get('label', opt.get('value'))
|
||||
break
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={
|
||||
'selected': user_selection,
|
||||
'selected_labels': selected_label,
|
||||
},
|
||||
output_variables={
|
||||
variable_name: user_selection,
|
||||
f'{variable_name}_label': selected_label,
|
||||
},
|
||||
)
|
||||
else:
|
||||
# 首次执行，等待用户选择
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={
|
||||
'selected': None,
|
||||
'selected_labels': None,
|
||||
},
|
||||
waiting_for_input=True,
|
||||
waiting_config={
|
||||
'type': 'choice',
|
||||
'question': question,
|
||||
'options': options,
|
||||
'multiple': multiple,
|
||||
'min_select': min_select,
|
||||
'max_select': max_select,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class MessageNode(BaseNode):
|
||||
"""
|
||||
消息节点
|
||||
|
||||
向用户发送消息，不等待回复
|
||||
"""
|
||||
|
||||
node_type = 'message'
|
||||
node_name = '消息节点'
|
||||
node_category = 'dialog'
|
||||
|
||||
inputs = {
|
||||
'content': {
|
||||
'type': 'string',
|
||||
'description': '消息内容',
|
||||
'required': True,
|
||||
},
|
||||
'message_type': {
|
||||
'type': 'string',
|
||||
'description': '消息类型：text/markdown/html',
|
||||
'default': 'text',
|
||||
},
|
||||
}
|
||||
|
||||
outputs = {
|
||||
'sent': {
|
||||
'type': 'boolean',
|
||||
'description': '是否发送成功',
|
||||
},
|
||||
}
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""
|
||||
执行消息节点
|
||||
"""
|
||||
import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
content = self.config.get('content', '')
|
||||
message_type = self.config.get('message_type', 'text')
|
||||
|
||||
logger.info(f'[MessageNode] Original content: {content}')
|
||||
logger.info(f'[MessageNode] Context variables: item={context.get_variable("item")}, index={context.get_variable("index")}')
|
||||
|
||||
# 解析模板变量
|
||||
content = context.resolve_template(content)
|
||||
|
||||
logger.info(f'[MessageNode] Resolved content: {content}')
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={
|
||||
'sent': True,
|
||||
'content': content, # 也在 output 中返回消息内容
|
||||
},
|
||||
# 发送消息事件
|
||||
events=[{
|
||||
'type': 'message',
|
||||
'content': content,
|
||||
'message_type': message_type,
|
||||
}],
|
||||
)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class ConfirmNode(BaseNode):
|
||||
"""
|
||||
确认节点
|
||||
|
||||
向用户展示确认对话框，等待用户确认或取消
|
||||
"""
|
||||
|
||||
node_type = 'confirm'
|
||||
node_name = '确认节点'
|
||||
node_category = 'dialog'
|
||||
|
||||
# 扩展的确认关键词列表（覆盖常见表达）
|
||||
CONFIRM_KEYWORDS = {
|
||||
# 英文
|
||||
'true', 'yes', 'ok', 'okay', 'sure', 'confirm', 'confirmed', 'agree', 'accept', 'y',
|
||||
# 中文
|
||||
'是', '是的', '对', '对的', '好', '好的', '行', '行的', '可以', '没问题', '没有问题',
|
||||
'同意', '确认', '确定', '嗯', '嗯嗯', '好吧', '可', '中', '成', '得', '要', '要的',
|
||||
'继续', '执行', '进行', '开始', '去吧', '做吧', '干吧',
|
||||
# 数字
|
||||
'1',
|
||||
}
|
||||
|
||||
# 扩展的取消关键词列表
|
||||
CANCEL_KEYWORDS = {
|
||||
# 英文
|
||||
'false', 'no', 'cancel', 'reject', 'decline', 'deny', 'n', 'nope',
|
||||
# 中文
|
||||
'否', '不', '不是', '不行', '不可以', '不要', '不用', '取消', '拒绝', '算了',
|
||||
'停止', '终止', '放弃', '别', '别了', '不了', '不用了', '不需要',
|
||||
# 数字
|
||||
'0',
|
||||
}
|
||||
|
||||
inputs = {
|
||||
'title': {
|
||||
'type': 'string',
|
||||
'description': '确认框标题',
|
||||
'default': '确认',
|
||||
},
|
||||
'content': {
|
||||
'type': 'string',
|
||||
'description': '确认内容',
|
||||
'required': True,
|
||||
},
|
||||
'confirm_text': {
|
||||
'type': 'string',
|
||||
'description': '确认按钮文本',
|
||||
'default': '确认',
|
||||
},
|
||||
'cancel_text': {
|
||||
'type': 'string',
|
||||
'description': '取消按钮文本',
|
||||
'default': '取消',
|
||||
},
|
||||
'variable_name': {
|
||||
'type': 'string',
|
||||
'description': '存储结果的变量名',
|
||||
'default': 'confirmed',
|
||||
},
|
||||
'use_llm_intent': {
|
||||
'type': 'boolean',
|
||||
'description': '使用 LLM 进行意图识别（当关键词匹配失败时）',
|
||||
'default': False,
|
||||
},
|
||||
'llm_model_id': {
|
||||
'type': 'string',
|
||||
'description': '用于意图识别的 LLM 模型 ID',
|
||||
'default': '',
|
||||
},
|
||||
}
|
||||
|
||||
outputs = {
|
||||
'confirmed': {
|
||||
'type': 'boolean',
|
||||
'description': '用户是否确认',
|
||||
},
|
||||
}
|
||||
|
||||
def _match_keywords(self, user_input) -> tuple[bool, bool]:
|
||||
"""
|
||||
使用关键词匹配判断用户意图
|
||||
|
||||
Returns:
|
||||
(matched, confirmed): matched 表示是否匹配到关键词，confirmed 表示是否确认
|
||||
"""
|
||||
# 布尔值特殊处理（必须在字符串处理之前）
|
||||
if user_input is True or user_input == True:
|
||||
return True, True
|
||||
if user_input is False or user_input == False:
|
||||
return True, False
|
||||
|
||||
# 确保是字符串
|
||||
if not isinstance(user_input, str):
|
||||
user_input = str(user_input)
|
||||
|
||||
# 标准化输入：去除空格、转小写
|
||||
normalized = user_input.strip().lower()
|
||||
|
||||
# 精确匹配
|
||||
if normalized in self.CONFIRM_KEYWORDS:
|
||||
return True, True
|
||||
if normalized in self.CANCEL_KEYWORDS:
|
||||
return True, False
|
||||
|
||||
return False, False
|
||||
|
||||
def _llm_intent_recognition(self, user_input: str, context_content: str, model_id: str) -> bool:
|
||||
"""
|
||||
使用 LLM 进行意图识别
|
||||
|
||||
Returns:
|
||||
confirmed: 用户是否确认
|
||||
"""
|
||||
try:
|
||||
from ai_platform.services.llm_service import LLMService
|
||||
|
||||
system_prompt = """你是一个意图识别助手。用户正在回应一个确认请求，你需要判断用户的回复是"确认"还是"取消"。
|
||||
|
||||
规则：
|
||||
1. 如果用户表达同意、肯定、愿意继续的意思，返回 "confirm"
|
||||
2. 如果用户表达拒绝、否定、不愿意继续的意思，返回 "cancel"
|
||||
3. 如果无法判断，默认返回 "cancel"
|
||||
|
||||
只返回 "confirm" 或 "cancel"，不要返回其他内容。"""
|
||||
|
||||
user_prompt = f"""确认请求内容：{context_content}
|
||||
|
||||
用户回复：{user_input}
|
||||
|
||||
请判断用户意图："""
|
||||
|
||||
llm_service = LLMService()
|
||||
response = llm_service.chat(
|
||||
model_id=model_id,
|
||||
messages=[
|
||||
{'role': 'system', 'content': system_prompt},
|
||||
{'role': 'user', 'content': user_prompt},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=10,
|
||||
)
|
||||
|
||||
result = response.content.strip().lower()
|
||||
return result == 'confirm'
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f'LLM 意图识别失败，回退到默认行为: {e}')
|
||||
return False
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""
|
||||
执行确认节点
|
||||
"""
|
||||
title = self.config.get('title', '确认')
|
||||
content = self.config.get('content', '')
|
||||
confirm_text = self.config.get('confirm_text', '确认')
|
||||
cancel_text = self.config.get('cancel_text', '取消')
|
||||
variable_name = self.config.get('variable_name', 'confirmed')
|
||||
use_llm_intent = self.config.get('use_llm_intent', False)
|
||||
llm_model_id = self.config.get('llm_model_id', '')
|
||||
|
||||
# 解析模板变量
|
||||
title = context.resolve_template(title)
|
||||
content = context.resolve_template(content)
|
||||
|
||||
# 检查是否已有用户选择
|
||||
user_input = context.variables.get('__user_input__')
|
||||
|
||||
if user_input is not None:
|
||||
# 首先尝试关键词匹配
|
||||
matched, confirmed = self._match_keywords(user_input)
|
||||
|
||||
if not matched and use_llm_intent and llm_model_id:
|
||||
# 关键词未匹配，使用 LLM 意图识别
|
||||
logger.info(f'关键词未匹配，使用 LLM 意图识别: {user_input}')
|
||||
confirmed = self._llm_intent_recognition(str(user_input), content, llm_model_id)
|
||||
elif not matched:
|
||||
# 关键词未匹配且未启用 LLM，默认为取消
|
||||
logger.info(f'关键词未匹配，默认取消: {user_input}')
|
||||
confirmed = False
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={
|
||||
'confirmed': confirmed,
|
||||
},
|
||||
output_variables={
|
||||
variable_name: confirmed,
|
||||
},
|
||||
)
|
||||
else:
|
||||
# 等待用户确认
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={
|
||||
'confirmed': False,
|
||||
},
|
||||
waiting_for_input=True,
|
||||
waiting_config={
|
||||
'type': 'confirm',
|
||||
'title': title,
|
||||
'content': content,
|
||||
'confirm_text': confirm_text,
|
||||
'cancel_text': cancel_text,
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,98 @@
|
||||
"""
|
||||
结束节点
|
||||
"""
|
||||
from typing import Any, Dict
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class EndNode(BaseNode):
|
||||
"""
|
||||
结束节点
|
||||
|
||||
工作流的出口节点，输出最终结果
|
||||
"""
|
||||
|
||||
node_type = 'end'
|
||||
node_name = '结束'
|
||||
node_category = 'basic'
|
||||
node_icon = 'stop-circle'
|
||||
node_description = '工作流的结束节点，输出最终结果'
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'output',
|
||||
'type': 'any',
|
||||
'description': '输出内容',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = []
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""执行结束节点"""
|
||||
# 获取输出内容配置（前端保存的是字符串模板，如 "{{llm_response}}"）
|
||||
output_template = self.config.get('output', '')
|
||||
|
||||
if output_template and isinstance(output_template, str) and output_template.strip():
|
||||
# 如果配置了输出模板，解析模板中的变量
|
||||
output = context.resolve_template(output_template)
|
||||
else:
|
||||
# 没有配置输出模板，不输出任何内容
|
||||
output = None
|
||||
|
||||
# 处理输出变量（用于结构化输出）
|
||||
# 只有当 output 不为 None 时才处理结构化输出
|
||||
outputs_config = self.config.get('outputs', [])
|
||||
if output is not None and outputs_config and isinstance(outputs_config, list):
|
||||
# 如果定义了输出变量，构建结构化输出
|
||||
structured_output = {}
|
||||
for out_var in outputs_config:
|
||||
if isinstance(out_var, dict):
|
||||
var_name = out_var.get('variable', '')
|
||||
if var_name:
|
||||
# 从上下文获取变量值
|
||||
structured_output[var_name] = context.get_variable(var_name, None)
|
||||
|
||||
# 如果有结构化输出，合并到结果中
|
||||
if structured_output:
|
||||
if isinstance(output, dict):
|
||||
output = {**output, **structured_output}
|
||||
else:
|
||||
output = {
|
||||
'result': output,
|
||||
**structured_output
|
||||
}
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=output,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'output': {
|
||||
'type': 'object',
|
||||
'title': '输出配置',
|
||||
'properties': {
|
||||
'type': {
|
||||
'type': 'string',
|
||||
'title': '输出类型',
|
||||
'enum': ['variable', 'template', 'previous'],
|
||||
'default': 'previous',
|
||||
},
|
||||
'value': {
|
||||
'type': 'string',
|
||||
'title': '输出值',
|
||||
'description': '变量名或模板',
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
"""
|
||||
HTTP 请求节点
|
||||
"""
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Dict
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class HttpNode(BaseNode):
|
||||
"""
|
||||
HTTP 请求节点
|
||||
|
||||
发送 HTTP 请求并获取响应
|
||||
"""
|
||||
|
||||
node_type = 'http'
|
||||
node_name = 'HTTP 请求'
|
||||
node_category = 'data'
|
||||
node_icon = 'globe'
|
||||
node_description = '发送 HTTP 请求并获取响应'
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'url',
|
||||
'type': 'string',
|
||||
'description': '请求 URL',
|
||||
},
|
||||
{
|
||||
'name': 'body',
|
||||
'type': 'object',
|
||||
'description': '请求体',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'response',
|
||||
'type': 'object',
|
||||
'description': '响应数据',
|
||||
},
|
||||
{
|
||||
'name': 'status_code',
|
||||
'type': 'number',
|
||||
'description': '状态码',
|
||||
},
|
||||
]
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""执行 HTTP 请求"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
method = self.config.get('method', 'GET').upper()
|
||||
url = self.config.get('url', '')
|
||||
headers = self.config.get('headers', {})
|
||||
params = self.config.get('params', {})
|
||||
body = self.config.get('body', {})
|
||||
timeout = self.config.get('timeout', 30)
|
||||
output_variable = self.config.get('output_variable', 'http_response')
|
||||
|
||||
# 解析模板变量
|
||||
url = context.resolve_template(url)
|
||||
|
||||
# 解析 headers 中的变量
|
||||
resolved_headers = {}
|
||||
for key, value in headers.items():
|
||||
resolved_headers[key] = context.resolve_template(str(value))
|
||||
|
||||
# 解析 params 中的变量
|
||||
resolved_params = {}
|
||||
for key, value in params.items():
|
||||
resolved_params[key] = context.resolve_template(str(value))
|
||||
|
||||
# 解析 body 中的变量
|
||||
resolved_body = self._resolve_body(body, context)
|
||||
|
||||
# 发送请求
|
||||
with httpx.Client(timeout=timeout) as client:
|
||||
if method in ('GET', 'DELETE'):
|
||||
response = client.request(
|
||||
method=method,
|
||||
url=url,
|
||||
headers=resolved_headers,
|
||||
params=resolved_params,
|
||||
)
|
||||
else:
|
||||
response = client.request(
|
||||
method=method,
|
||||
url=url,
|
||||
headers=resolved_headers,
|
||||
params=resolved_params,
|
||||
json=resolved_body,
|
||||
)
|
||||
|
||||
# 解析响应
|
||||
try:
|
||||
response_data = response.json()
|
||||
except Exception:
|
||||
response_data = response.text
|
||||
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
|
||||
return NodeResult(
|
||||
success=response.is_success,
|
||||
output=response_data,
|
||||
output_variables={
|
||||
output_variable: response_data,
|
||||
f'{output_variable}_status': response.status_code,
|
||||
},
|
||||
elapsed_time=elapsed_time,
|
||||
metadata={
|
||||
'status_code': response.status_code,
|
||||
'headers': dict(response.headers),
|
||||
},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'HTTP 节点执行失败: {e}')
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
elapsed_time=elapsed_time,
|
||||
)
|
||||
|
||||
def _resolve_body(self, body: Any, context: NodeContext) -> Any:
|
||||
"""递归解析 body 中的变量"""
|
||||
if isinstance(body, str):
|
||||
return context.resolve_template(body)
|
||||
elif isinstance(body, dict):
|
||||
return {k: self._resolve_body(v, context) for k, v in body.items()}
|
||||
elif isinstance(body, list):
|
||||
return [self._resolve_body(item, context) for item in body]
|
||||
return body
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'method': {
|
||||
'type': 'string',
|
||||
'title': '请求方法',
|
||||
'enum': ['GET', 'POST', 'PUT', 'DELETE', 'PATCH'],
|
||||
'default': 'GET',
|
||||
},
|
||||
'url': {
|
||||
'type': 'string',
|
||||
'title': 'URL',
|
||||
'description': '支持变量引用，如 {{api_url}}',
|
||||
},
|
||||
'headers': {
|
||||
'type': 'object',
|
||||
'title': '请求头',
|
||||
'additionalProperties': {'type': 'string'},
|
||||
},
|
||||
'params': {
|
||||
'type': 'object',
|
||||
'title': 'URL 参数',
|
||||
'additionalProperties': {'type': 'string'},
|
||||
},
|
||||
'body': {
|
||||
'type': 'object',
|
||||
'title': '请求体',
|
||||
'description': 'POST/PUT/PATCH 请求的 JSON 数据',
|
||||
},
|
||||
'timeout': {
|
||||
'type': 'integer',
|
||||
'title': '超时时间（秒）',
|
||||
'default': 30,
|
||||
},
|
||||
'output_variable': {
|
||||
'type': 'string',
|
||||
'title': '输出变量名',
|
||||
'default': 'http_response',
|
||||
},
|
||||
},
|
||||
'required': ['url'],
|
||||
}
|
||||
@@ -0,0 +1,443 @@
|
||||
"""
|
||||
意图识别节点
|
||||
|
||||
使用 LLM 进行智能意图分类，根据用户输入自动路由到对应分支
|
||||
支持原生 Function Calling（更准确）和文本解析（回退方案）
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class IntentNode(BaseNode):
|
||||
"""
|
||||
意图识别节点
|
||||
|
||||
使用 LLM 分析用户输入，识别用户意图，并路由到对应分支
|
||||
支持定义多个意图，每个意图可配置名称、描述和示例
|
||||
"""
|
||||
|
||||
node_type = 'intent'
|
||||
node_name = '意图识别'
|
||||
node_category = 'logic'
|
||||
node_icon = 'brain'
|
||||
node_description = '使用 AI 识别用户意图，自动路由到对应分支'
|
||||
|
||||
supports_branches = True
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'user_input',
|
||||
'type': 'string',
|
||||
'description': '用户输入文本',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'intent',
|
||||
'type': 'string',
|
||||
'description': '识别到的意图名称',
|
||||
},
|
||||
{
|
||||
'name': 'confidence',
|
||||
'type': 'number',
|
||||
'description': '置信度（0-1）',
|
||||
},
|
||||
]
|
||||
|
||||
# 默认的意图识别 prompt 模板
|
||||
DEFAULT_SYSTEM_PROMPT = """你是一个意图分类器。根据用户输入，判断用户的意图属于以下哪个类别。
|
||||
|
||||
可选的意图类别：
|
||||
{intents_description}
|
||||
|
||||
请严格按照以下 JSON 格式输出，不要输出其他任何内容：
|
||||
{{"intent": "意图名称", "confidence": 0.95}}
|
||||
|
||||
注意：
|
||||
1. intent 必须是上述意图类别中的一个名称，如果都不匹配则输出 "other"
|
||||
2. confidence 是你对这个分类的置信度，范围 0-1
|
||||
3. 只输出 JSON，不要有任何解释"""
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""执行意图识别（同步方法，通过运行异步方法实现）"""
|
||||
import asyncio
|
||||
|
||||
# 在同步方法中运行异步代码
|
||||
loop = asyncio.get_event_loop()
|
||||
if loop.is_running():
|
||||
# 如果已有事件循环在运行，创建新任务
|
||||
import concurrent.futures
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future = executor.submit(asyncio.run, self.execute_async(context))
|
||||
return future.result()
|
||||
else:
|
||||
return loop.run_until_complete(self.execute_async(context))
|
||||
|
||||
async def execute_async(self, context: NodeContext) -> NodeResult:
|
||||
"""异步执行意图识别"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
from ai_platform.services.llm_service import LLMService
|
||||
from ai_platform.models import LLMModel
|
||||
|
||||
# 获取配置
|
||||
model_id = self.config.get('model_id', '')
|
||||
intents = self.config.get('intents', [])
|
||||
input_variable = self.config.get('input_variable', 'user_input')
|
||||
confidence_threshold = self.config.get('confidence_threshold', 0.6)
|
||||
|
||||
# 获取用户输入
|
||||
user_input = context.get_variable(input_variable)
|
||||
if not user_input:
|
||||
# 尝试从 __user_input__ 获取
|
||||
user_input = context.get_variable('__user_input__')
|
||||
if not user_input:
|
||||
user_input = context.user_input
|
||||
|
||||
if not user_input:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='未获取到用户输入',
|
||||
)
|
||||
|
||||
if not intents:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='请至少配置一个意图',
|
||||
)
|
||||
|
||||
if not model_id:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='请选择用于意图识别的模型',
|
||||
)
|
||||
|
||||
# 检查模型是否支持 Function Calling
|
||||
supports_fc = self._check_model_supports_function_call(model_id)
|
||||
|
||||
llm_service = LLMService(context.db_session)
|
||||
|
||||
if supports_fc:
|
||||
# 使用 Function Calling 方式（更准确）
|
||||
intent_name, confidence, total_tokens = await self._execute_with_function_calling_async(
|
||||
llm_service, model_id, user_input, intents
|
||||
)
|
||||
else:
|
||||
# 回退到文本解析方式
|
||||
intent_name, confidence, total_tokens = await self._execute_with_text_parsing_async(
|
||||
llm_service, model_id, user_input, intents
|
||||
)
|
||||
|
||||
logger.info(f'意图识别结果: intent_name={intent_name}, confidence={confidence}, use_fc={supports_fc}')
|
||||
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
|
||||
# 判断是否达到置信度阈值
|
||||
if confidence < confidence_threshold:
|
||||
logger.info(f'置信度 {confidence} 低于阈值 {confidence_threshold}，走 other 分支')
|
||||
intent_name = 'other'
|
||||
|
||||
# 查找对应的分支 ID
|
||||
next_node_id = self._get_branch_id(intent_name, intents)
|
||||
logger.info(f'意图识别结果: intent={intent_name}, next_node_id={next_node_id}')
|
||||
|
||||
# 设置输出变量
|
||||
output_var = self.config.get('output_variable', 'intent_result')
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=intent_name,
|
||||
next_node_id=next_node_id,
|
||||
output_variables={
|
||||
output_var: intent_name,
|
||||
f'{output_var}_confidence': confidence,
|
||||
f'{output_var}_input': user_input,
|
||||
},
|
||||
tokens_used=total_tokens,
|
||||
elapsed_time=elapsed_time,
|
||||
metadata={
|
||||
'intent': intent_name,
|
||||
'confidence': confidence,
|
||||
'matched_branch': next_node_id,
|
||||
},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'意图识别节点执行失败: {e}')
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
elapsed_time=elapsed_time,
|
||||
)
|
||||
|
||||
def _check_model_supports_function_call(self, model_id) -> bool:
|
||||
"""检查模型是否支持 Function Calling
|
||||
|
||||
TODO: 此方法需要改造为异步版本，当前使用同步查询作为临时方案
|
||||
"""
|
||||
# 临时方案：默认返回True，让调用方尝试使用Function Calling
|
||||
# 实际使用时应通过上下文传入模型信息
|
||||
return True
|
||||
|
||||
async def _execute_with_function_calling_async(
|
||||
self,
|
||||
llm_service,
|
||||
model_id: str,
|
||||
user_input: str,
|
||||
intents: List[Dict],
|
||||
) -> tuple:
|
||||
"""使用 Function Calling 执行意图识别（更准确，异步版本）"""
|
||||
# 构建意图名称列表（用于 enum）
|
||||
intent_names = [intent.get('name') for intent in intents] + ['other']
|
||||
|
||||
# 构建意图描述（用于 LLM 理解）
|
||||
intents_description = self._build_intents_description(intents)
|
||||
|
||||
# 构建 Function Calling 工具定义
|
||||
tools = [{
|
||||
'name': 'classify_intent',
|
||||
'description': f'根据用户输入对意图进行分类。可选的意图类别：\n{intents_description}',
|
||||
'parameters': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'intent': {
|
||||
'type': 'string',
|
||||
'enum': intent_names,
|
||||
'description': '识别到的意图名称',
|
||||
},
|
||||
'confidence': {
|
||||
'type': 'number',
|
||||
'description': '置信度，范围 0-1',
|
||||
},
|
||||
},
|
||||
'required': ['intent', 'confidence'],
|
||||
},
|
||||
}]
|
||||
|
||||
# 简化的系统提示词
|
||||
system_prompt = "你是一个意图分类器。分析用户输入，调用 classify_intent 函数返回分类结果。"
|
||||
|
||||
messages = [
|
||||
{'role': 'system', 'content': system_prompt},
|
||||
{'role': 'user', 'content': user_input},
|
||||
]
|
||||
|
||||
# 调用 LLM（带 Function Calling，异步）
|
||||
response = await llm_service.chat_async(
|
||||
model_id=model_id,
|
||||
messages=messages,
|
||||
temperature=0.1,
|
||||
max_tokens=100,
|
||||
tools=tools,
|
||||
tool_choice='required', # 强制使用 Function Calling
|
||||
)
|
||||
|
||||
# 从 tool_calls 中提取结果
|
||||
if response.has_tool_calls and response.tool_calls:
|
||||
tool_call = response.tool_calls[0]
|
||||
args = tool_call.arguments
|
||||
intent_name = args.get('intent', 'other')
|
||||
confidence = float(args.get('confidence', 0.8))
|
||||
|
||||
# 验证意图名称
|
||||
if intent_name not in intent_names:
|
||||
intent_name = 'other'
|
||||
confidence = 0.5
|
||||
|
||||
return intent_name, confidence, response.total_tokens
|
||||
else:
|
||||
# Function Calling 失败，返回默认值
|
||||
logger.warning('Function Calling 未返回 tool_calls，回退到 other')
|
||||
return 'other', 0.5, response.total_tokens
|
||||
|
||||
async def _execute_with_text_parsing_async(
|
||||
self,
|
||||
llm_service,
|
||||
model_id: str,
|
||||
user_input: str,
|
||||
intents: List[Dict],
|
||||
) -> tuple:
|
||||
"""使用文本解析执行意图识别（回退方案，异步版本）"""
|
||||
# 构建意图描述
|
||||
intents_description = self._build_intents_description(intents)
|
||||
|
||||
# 构建 prompt
|
||||
system_prompt = self.DEFAULT_SYSTEM_PROMPT.format(
|
||||
intents_description=intents_description
|
||||
)
|
||||
|
||||
messages = [
|
||||
{'role': 'system', 'content': system_prompt},
|
||||
{'role': 'user', 'content': user_input},
|
||||
]
|
||||
|
||||
# 调用 LLM（异步）
|
||||
response = await llm_service.chat_async(
|
||||
model_id=model_id,
|
||||
messages=messages,
|
||||
temperature=0.1,
|
||||
max_tokens=100,
|
||||
)
|
||||
|
||||
# 解析响应
|
||||
intent_name, confidence = self._parse_response(response.content, intents)
|
||||
|
||||
return intent_name, confidence, response.total_tokens
|
||||
|
||||
def _build_intents_description(self, intents: List[Dict]) -> str:
|
||||
"""构建意图描述文本"""
|
||||
lines = []
|
||||
for i, intent in enumerate(intents, 1):
|
||||
name = intent.get('name', '')
|
||||
description = intent.get('description', '')
|
||||
examples = intent.get('examples', [])
|
||||
|
||||
line = f"{i}. {name}"
|
||||
if description:
|
||||
line += f" - {description}"
|
||||
|
||||
lines.append(line)
|
||||
|
||||
# 添加示例
|
||||
if examples:
|
||||
for example in examples[:3]: # 最多3个示例
|
||||
lines.append(f" 示例: \"{example}\"")
|
||||
|
||||
# 添加 other 选项
|
||||
lines.append(f"{len(intents) + 1}. other - 以上都不匹配时选择此项")
|
||||
|
||||
return '\n'.join(lines)
|
||||
|
||||
def _parse_response(self, response: str, intents: List[Dict]) -> tuple:
|
||||
"""解析 LLM 响应"""
|
||||
logger.info(f'开始解析 LLM 响应: {response}')
|
||||
try:
|
||||
# 尝试提取 JSON
|
||||
response = response.strip()
|
||||
|
||||
# 处理可能的 markdown 代码块
|
||||
if response.startswith('```'):
|
||||
lines = response.split('\n')
|
||||
json_lines = []
|
||||
in_json = False
|
||||
for line in lines:
|
||||
if line.startswith('```') and not in_json:
|
||||
in_json = True
|
||||
continue
|
||||
elif line.startswith('```') and in_json:
|
||||
break
|
||||
elif in_json:
|
||||
json_lines.append(line)
|
||||
response = '\n'.join(json_lines)
|
||||
|
||||
# 解析 JSON
|
||||
result = json.loads(response)
|
||||
intent_name = result.get('intent', 'other')
|
||||
confidence = float(result.get('confidence', 0.5))
|
||||
|
||||
# 验证意图名称是否有效
|
||||
valid_names = [intent.get('name') for intent in intents] + ['other']
|
||||
if intent_name not in valid_names:
|
||||
# 尝试模糊匹配
|
||||
intent_name_lower = intent_name.lower()
|
||||
for valid_name in valid_names:
|
||||
if valid_name.lower() == intent_name_lower:
|
||||
intent_name = valid_name
|
||||
break
|
||||
else:
|
||||
intent_name = 'other'
|
||||
confidence = 0.5
|
||||
|
||||
return intent_name, confidence
|
||||
|
||||
except (json.JSONDecodeError, ValueError, KeyError) as e:
|
||||
logger.warning(f'解析意图识别响应失败: {e}, response: {response}')
|
||||
return 'other', 0.5
|
||||
|
||||
def _get_branch_id(self, intent_name: str, intents: List[Dict]) -> str:
|
||||
"""获取意图对应的分支 ID"""
|
||||
logger.info(f'查找意图分支: intent_name={intent_name}, intents={intents}')
|
||||
for intent in intents:
|
||||
if intent.get('name') == intent_name:
|
||||
branch_id = intent.get('branch_id', intent_name)
|
||||
logger.info(f'找到匹配意图: {intent}, 返回 branch_id={branch_id}')
|
||||
return branch_id
|
||||
logger.info(f'未找到匹配意图，返回 other')
|
||||
return 'other'
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'model_id': {
|
||||
'type': 'string',
|
||||
'title': '模型',
|
||||
'description': '选择用于意图识别的 LLM 模型',
|
||||
},
|
||||
'input_variable': {
|
||||
'type': 'string',
|
||||
'title': '输入变量',
|
||||
'description': '包含用户输入的变量名',
|
||||
'default': 'user_input',
|
||||
},
|
||||
'intents': {
|
||||
'type': 'array',
|
||||
'title': '意图列表',
|
||||
'description': '定义要识别的意图',
|
||||
'items': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'name': {
|
||||
'type': 'string',
|
||||
'title': '意图名称',
|
||||
'description': '唯一标识，如 consult, complaint',
|
||||
},
|
||||
'description': {
|
||||
'type': 'string',
|
||||
'title': '意图描述',
|
||||
'description': '描述这个意图的含义',
|
||||
},
|
||||
'examples': {
|
||||
'type': 'array',
|
||||
'title': '示例',
|
||||
'description': '用户可能的输入示例',
|
||||
'items': {'type': 'string'},
|
||||
},
|
||||
'branch_id': {
|
||||
'type': 'string',
|
||||
'title': '分支 ID',
|
||||
'description': '匹配时跳转的分支',
|
||||
},
|
||||
},
|
||||
'required': ['name'],
|
||||
},
|
||||
},
|
||||
'confidence_threshold': {
|
||||
'type': 'number',
|
||||
'title': '置信度阈值',
|
||||
'description': '低于此阈值将走 other 分支',
|
||||
'default': 0.6,
|
||||
'minimum': 0,
|
||||
'maximum': 1,
|
||||
},
|
||||
'output_variable': {
|
||||
'type': 'string',
|
||||
'title': '输出变量名',
|
||||
'default': 'intent_result',
|
||||
},
|
||||
},
|
||||
'required': ['model_id', 'intents'],
|
||||
}
|
||||
@@ -0,0 +1,238 @@
|
||||
"""
|
||||
知识库检索节点
|
||||
|
||||
在 AI 工作流中检索知识库，返回与查询最相关的文档分段
|
||||
支持向量检索、全文检索和混合检索
|
||||
"""
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class KnowledgeRetrievalNode(BaseNode):
|
||||
"""
|
||||
知识库检索节点
|
||||
|
||||
从指定知识库中检索与查询文本最相关的文档分段，
|
||||
输出可直接作为 LLM 节点的上下文使用
|
||||
"""
|
||||
|
||||
node_type = 'knowledge_retrieval'
|
||||
node_name = '知识库检索'
|
||||
node_category = 'knowledge'
|
||||
node_icon = 'BookOpen'
|
||||
node_description = '从知识库中检索相关文档，为 LLM 提供上下文'
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'query',
|
||||
'type': 'string',
|
||||
'description': '检索查询文本',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'results',
|
||||
'type': 'array',
|
||||
'description': '检索结果列表',
|
||||
},
|
||||
{
|
||||
'name': 'context',
|
||||
'type': 'string',
|
||||
'description': '拼接后的上下文文本',
|
||||
},
|
||||
]
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""同步执行（不支持，需要异步）"""
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='知识库检索节点必须异步执行',
|
||||
)
|
||||
|
||||
async def execute_async(self, context: NodeContext) -> NodeResult:
|
||||
"""异步执行知识库检索"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 读取配置
|
||||
knowledge_base_ids = self.config.get('knowledge_base_ids', [])
|
||||
query_template = self.config.get('query', '')
|
||||
top_k = self.config.get('top_k', 5)
|
||||
score_threshold = self.config.get('score_threshold', 0.5)
|
||||
retrieval_mode = self.config.get('retrieval_mode', None)
|
||||
rerank_enabled = self.config.get('rerank_enabled', None)
|
||||
rerank_model_id = self.config.get('rerank_model_id', None)
|
||||
output_variable = self.config.get('output_variable', 'knowledge_results')
|
||||
context_variable = self.config.get('context_variable', 'knowledge_context')
|
||||
context_template = self.config.get('context_template', '')
|
||||
|
||||
if not knowledge_base_ids:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='未配置知识库',
|
||||
elapsed_time=int((time.time() - start_time) * 1000),
|
||||
)
|
||||
|
||||
# 解析查询模板中的变量
|
||||
query = context.resolve_template(query_template) if query_template else context.user_input
|
||||
if not query or not query.strip():
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='检索查询文本为空',
|
||||
elapsed_time=int((time.time() - start_time) * 1000),
|
||||
)
|
||||
|
||||
# 调用检索服务
|
||||
db = context.db_session
|
||||
if not db:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='数据库会话不可用',
|
||||
elapsed_time=int((time.time() - start_time) * 1000),
|
||||
)
|
||||
|
||||
from ai_platform.knowledge.services.retrieval_service import RetrievalService
|
||||
|
||||
service = RetrievalService(db)
|
||||
results = await service.retrieve(
|
||||
query=query,
|
||||
knowledge_base_ids=knowledge_base_ids,
|
||||
top_k=top_k,
|
||||
score_threshold=score_threshold,
|
||||
retrieval_mode=retrieval_mode,
|
||||
rerank_enabled=rerank_enabled,
|
||||
rerank_model_id=rerank_model_id,
|
||||
)
|
||||
|
||||
# 构建结果列表
|
||||
result_list = []
|
||||
for r in results:
|
||||
result_list.append({
|
||||
'segment_id': r.segment_id,
|
||||
'document_id': r.document_id,
|
||||
'document_name': r.document_name,
|
||||
'knowledge_base_id': r.knowledge_base_id,
|
||||
'knowledge_base_name': r.knowledge_base_name,
|
||||
'content': r.content,
|
||||
'score': r.score,
|
||||
'token_count': r.token_count,
|
||||
})
|
||||
|
||||
# 构建上下文文本
|
||||
if context_template:
|
||||
# 自定义模板
|
||||
context_text = context.resolve_template(context_template)
|
||||
else:
|
||||
# 默认：拼接所有检索结果内容
|
||||
context_parts = []
|
||||
for i, r in enumerate(result_list, 1):
|
||||
context_parts.append(
|
||||
f"[{i}] (来源: {r['document_name']}, 相似度: {r['score']:.2f})\n{r['content']}"
|
||||
)
|
||||
context_text = '\n\n'.join(context_parts)
|
||||
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={
|
||||
'results': result_list,
|
||||
'context': context_text,
|
||||
'total': len(result_list),
|
||||
'query': query,
|
||||
},
|
||||
output_variables={
|
||||
output_variable: result_list,
|
||||
context_variable: context_text,
|
||||
f'{output_variable}_total': len(result_list),
|
||||
},
|
||||
elapsed_time=elapsed_time,
|
||||
metadata={
|
||||
'query': query,
|
||||
'knowledge_base_ids': knowledge_base_ids,
|
||||
'top_k': top_k,
|
||||
'score_threshold': score_threshold,
|
||||
'retrieval_mode': retrieval_mode,
|
||||
'rerank_enabled': rerank_enabled,
|
||||
'result_count': len(result_list),
|
||||
},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'知识库检索节点执行失败: {e}')
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=f'检索失败: {str(e)}',
|
||||
elapsed_time=elapsed_time,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'knowledge_base_ids': {
|
||||
'type': 'array',
|
||||
'title': '知识库',
|
||||
'description': '选择要检索的知识库',
|
||||
'items': {'type': 'string'},
|
||||
},
|
||||
'query': {
|
||||
'type': 'string',
|
||||
'title': '查询文本',
|
||||
'description': '支持变量引用，如 {{user_input}}，留空则使用用户输入',
|
||||
},
|
||||
'top_k': {
|
||||
'type': 'integer',
|
||||
'title': '返回数量',
|
||||
'default': 5,
|
||||
'minimum': 1,
|
||||
'maximum': 20,
|
||||
},
|
||||
'score_threshold': {
|
||||
'type': 'number',
|
||||
'title': '相似度阈值',
|
||||
'default': 0.5,
|
||||
'minimum': 0,
|
||||
'maximum': 1,
|
||||
},
|
||||
'retrieval_mode': {
|
||||
'type': 'string',
|
||||
'title': '检索模式',
|
||||
'description': '留空则使用知识库默认配置',
|
||||
'enum': ['vector', 'fulltext', 'hybrid'],
|
||||
},
|
||||
'rerank_enabled': {
|
||||
'type': 'boolean',
|
||||
'title': '启用重排序',
|
||||
'description': '不设置则使用知识库默认配置',
|
||||
'default': None,
|
||||
},
|
||||
'rerank_model_id': {
|
||||
'type': 'string',
|
||||
'title': '重排序模型',
|
||||
'description': '不设置则使用知识库默认配置',
|
||||
},
|
||||
'output_variable': {
|
||||
'type': 'string',
|
||||
'title': '结果变量名',
|
||||
'default': 'knowledge_results',
|
||||
},
|
||||
'context_variable': {
|
||||
'type': 'string',
|
||||
'title': '上下文变量名',
|
||||
'default': 'knowledge_context',
|
||||
},
|
||||
},
|
||||
'required': ['knowledge_base_ids'],
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,405 @@
|
||||
"""
|
||||
循环节点
|
||||
|
||||
支持两种循环模式：
|
||||
1. for_each - 遍历数组，对每个元素执行循环体
|
||||
2. while - 条件循环，满足条件时持续执行
|
||||
"""
|
||||
import ast
|
||||
import json
|
||||
import logging
|
||||
import operator
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class LoopNode(BaseNode):
|
||||
"""
|
||||
循环节点
|
||||
|
||||
支持 for_each 和 while 两种循环模式
|
||||
"""
|
||||
|
||||
node_type = 'loop'
|
||||
node_name = '循环'
|
||||
node_category = 'logic'
|
||||
node_icon = 'repeat'
|
||||
node_description = '循环执行一组节点，支持遍历数组或条件循环'
|
||||
|
||||
supports_branches = True
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'items',
|
||||
'type': 'array',
|
||||
'description': '要遍历的数组（for_each 模式）',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'results',
|
||||
'type': 'array',
|
||||
'description': '每次循环的结果数组',
|
||||
},
|
||||
{
|
||||
'name': 'current_item',
|
||||
'type': 'any',
|
||||
'description': '当前循环项（循环体内可用）',
|
||||
},
|
||||
{
|
||||
'name': 'current_index',
|
||||
'type': 'number',
|
||||
'description': '当前循环索引（从 0 开始）',
|
||||
},
|
||||
]
|
||||
|
||||
# 支持的操作符（用于 while 条件判断）
|
||||
OPERATORS = {
|
||||
'eq': operator.eq, # 等于
|
||||
'ne': operator.ne, # 不等于
|
||||
'gt': operator.gt, # 大于
|
||||
'gte': operator.ge, # 大于等于
|
||||
'lt': operator.lt, # 小于
|
||||
'lte': operator.le, # 小于等于
|
||||
'is_empty': lambda a, _: not a, # 为空
|
||||
'is_not_empty': lambda a, _: bool(a), # 不为空
|
||||
'is_true': lambda a, _: bool(a), # 为真
|
||||
'is_false': lambda a, _: not bool(a), # 为假
|
||||
}
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""
|
||||
执行循环节点
|
||||
|
||||
循环节点本身只负责初始化循环状态，实际的循环执行由 WorkflowService 处理
|
||||
"""
|
||||
try:
|
||||
loop_mode = self.config.get('loop_mode', 'for_each')
|
||||
max_iterations = self.config.get('max_iterations', 100)
|
||||
|
||||
if loop_mode == 'for_each':
|
||||
return self._init_for_each(context, max_iterations)
|
||||
elif loop_mode == 'while':
|
||||
return self._init_while(context, max_iterations)
|
||||
else:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=f'不支持的循环模式: {loop_mode}',
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'循环节点执行失败: {e}')
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
)
|
||||
|
||||
def _init_for_each(self, context: NodeContext, max_iterations: int) -> NodeResult:
|
||||
"""初始化 for_each 循环"""
|
||||
items_var = self.config.get('items_variable', '')
|
||||
|
||||
logger.info(f'[LoopNode] 初始化 for_each 循环，items_variable={items_var}')
|
||||
|
||||
# 解析数组变量
|
||||
items = self._resolve_variable(items_var, context)
|
||||
|
||||
logger.info(f'[LoopNode] 解析后的 items 类型: {type(items).__name__}, 值: {items}')
|
||||
|
||||
# 如果是字符串，尝试解析为 JSON 或 Python 字面量
|
||||
if isinstance(items, str):
|
||||
try:
|
||||
items = json.loads(items)
|
||||
logger.info(f'[LoopNode] JSON 解析成功，items={items}')
|
||||
except json.JSONDecodeError:
|
||||
# 尝试使用 ast.literal_eval 解析 Python 字面量（如 str() 输出的列表）
|
||||
try:
|
||||
items = ast.literal_eval(items)
|
||||
logger.info(f'[LoopNode] Python 字面量解析成功，items={items}')
|
||||
except (ValueError, SyntaxError):
|
||||
# 尝试按逗号分割（仅当不是列表/字典格式时）
|
||||
if not (items.strip().startswith('[') or items.strip().startswith('{')):
|
||||
items = [item.strip() for item in items.split(',') if item.strip()]
|
||||
logger.info(f'[LoopNode] 按逗号分割，items={items}')
|
||||
else:
|
||||
logger.error(f'[LoopNode] 无法解析 items 字符串: {items[:200]}...')
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=f'无法解析 items 变量，格式不正确',
|
||||
)
|
||||
|
||||
if not isinstance(items, (list, tuple)):
|
||||
error_msg = f'items 必须是数组，当前类型: {type(items).__name__}, 值: {items}, items_variable={items_var}'
|
||||
logger.error(f'[LoopNode] {error_msg}')
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=error_msg,
|
||||
)
|
||||
|
||||
# 检查数组是否为空
|
||||
if len(items) == 0:
|
||||
logger.warning(f'[LoopNode] items 数组为空，items_variable={items_var}')
|
||||
# 返回成功但不执行循环体
|
||||
|
||||
# 限制最大迭代次数
|
||||
if len(items) > max_iterations:
|
||||
logger.warning(f'数组长度 {len(items)} 超过最大迭代次数 {max_iterations}，将被截断')
|
||||
items = list(items)[:max_iterations]
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={
|
||||
'items': list(items),
|
||||
'total_count': len(items),
|
||||
},
|
||||
output_variables={
|
||||
'_loop_items': list(items),
|
||||
'_loop_index': 0,
|
||||
'_loop_total': len(items),
|
||||
'_loop_mode': 'for_each',
|
||||
'_loop_results': [],
|
||||
},
|
||||
metadata={
|
||||
'is_loop': True,
|
||||
'loop_mode': 'for_each',
|
||||
'total_iterations': len(items),
|
||||
'max_iterations': max_iterations,
|
||||
},
|
||||
)
|
||||
|
||||
def _init_while(self, context: NodeContext, max_iterations: int) -> NodeResult:
|
||||
"""初始化 while 循环"""
|
||||
# 检查初始条件
|
||||
condition_met = self._evaluate_condition(context)
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={
|
||||
'condition_met': condition_met,
|
||||
},
|
||||
output_variables={
|
||||
'_loop_index': 0,
|
||||
'_loop_mode': 'while',
|
||||
'_loop_condition_met': condition_met,
|
||||
'_loop_results': [],
|
||||
'_loop_max_iterations': max_iterations,
|
||||
},
|
||||
metadata={
|
||||
'is_loop': True,
|
||||
'loop_mode': 'while',
|
||||
'condition_met': condition_met,
|
||||
'max_iterations': max_iterations,
|
||||
},
|
||||
)
|
||||
|
||||
def check_continue(self, context: NodeContext) -> bool:
|
||||
"""
|
||||
检查是否继续循环
|
||||
|
||||
由 WorkflowService 在每次循环迭代后调用
|
||||
"""
|
||||
loop_mode = context.get_variable('_loop_mode')
|
||||
current_index = context.get_variable('_loop_index', 0)
|
||||
max_iterations = self.config.get('max_iterations', 100)
|
||||
|
||||
# 检查最大迭代次数
|
||||
if current_index >= max_iterations:
|
||||
logger.warning(f'达到最大迭代次数 {max_iterations}，停止循环')
|
||||
return False
|
||||
|
||||
if loop_mode == 'for_each':
|
||||
total = context.get_variable('_loop_total', 0)
|
||||
return current_index < total
|
||||
elif loop_mode == 'while':
|
||||
return self._evaluate_condition(context)
|
||||
|
||||
return False
|
||||
|
||||
def get_current_item(self, context: NodeContext) -> Any:
|
||||
"""获取当前循环项（for_each 模式）"""
|
||||
items = context.get_variable('_loop_items', [])
|
||||
index = context.get_variable('_loop_index', 0)
|
||||
|
||||
if 0 <= index < len(items):
|
||||
return items[index]
|
||||
return None
|
||||
|
||||
def increment_index(self, context: NodeContext) -> int:
|
||||
"""增加循环索引"""
|
||||
current_index = context.get_variable('_loop_index', 0)
|
||||
new_index = current_index + 1
|
||||
context.set_variable('_loop_index', new_index)
|
||||
return new_index
|
||||
|
||||
def add_result(self, context: NodeContext, result: Any) -> None:
|
||||
"""添加循环结果"""
|
||||
results = context.get_variable('_loop_results', [])
|
||||
results.append(result)
|
||||
context.set_variable('_loop_results', results)
|
||||
|
||||
def _resolve_variable(self, variable: str, context: NodeContext) -> Any:
|
||||
"""解析变量引用，支持多层路径访问"""
|
||||
if not isinstance(variable, str):
|
||||
return variable
|
||||
|
||||
# 如果是空字符串，返回 None
|
||||
if not variable.strip():
|
||||
return None
|
||||
|
||||
# 检查是否是 {{...}} 格式的变量引用
|
||||
if not (variable.startswith('{{') and variable.endswith('}}')):
|
||||
# 不是变量引用，直接返回原始字符串值（可能是硬编码的 JSON 数组或逗号分隔的值）
|
||||
return variable
|
||||
|
||||
# 使用 resolve_template 解析变量，它支持多层路径访问（如 node.key.subkey）
|
||||
resolved = context.resolve_template(variable)
|
||||
|
||||
# 如果解析结果与原始变量相同，说明变量不存在或解析失败
|
||||
if resolved == variable:
|
||||
return None
|
||||
|
||||
# 如果解析结果是字符串，尝试解析为 JSON
|
||||
if isinstance(resolved, str):
|
||||
try:
|
||||
import json
|
||||
return json.loads(resolved)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return resolved
|
||||
|
||||
return resolved
|
||||
|
||||
def _evaluate_condition(self, context: NodeContext) -> bool:
|
||||
"""评估 while 条件"""
|
||||
conditions = self.config.get('conditions', [])
|
||||
logic = self.config.get('condition_logic', 'and') # and 或 or
|
||||
|
||||
if not conditions:
|
||||
return False # 无条件默认不继续
|
||||
|
||||
results = []
|
||||
for condition in conditions:
|
||||
result = self._evaluate_single_condition(condition, context)
|
||||
results.append(result)
|
||||
|
||||
if logic == 'or':
|
||||
return any(results)
|
||||
else: # and
|
||||
return all(results)
|
||||
|
||||
def _evaluate_single_condition(self, condition: Dict, context: NodeContext) -> bool:
|
||||
"""评估单个条件"""
|
||||
variable = condition.get('variable', '')
|
||||
op_name = condition.get('operator', 'eq')
|
||||
value = condition.get('value', '')
|
||||
|
||||
# 解析变量和值
|
||||
left_val = self._resolve_variable(variable, context)
|
||||
right_val = self._resolve_variable(value, context) if value else None
|
||||
|
||||
# 获取操作符函数
|
||||
op_func = self.OPERATORS.get(op_name, operator.eq)
|
||||
|
||||
try:
|
||||
# 特殊处理不需要右值的操作符
|
||||
if op_name in ['is_empty', 'is_not_empty', 'is_true', 'is_false']:
|
||||
return op_func(left_val, None)
|
||||
|
||||
# 尝试类型转换
|
||||
if isinstance(left_val, (int, float)) and isinstance(right_val, str):
|
||||
try:
|
||||
if '.' in right_val:
|
||||
right_val = float(right_val)
|
||||
else:
|
||||
right_val = int(right_val)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
return op_func(left_val, right_val)
|
||||
except Exception as e:
|
||||
logger.warning(f'条件评估失败: {e}')
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'loop_mode': {
|
||||
'type': 'string',
|
||||
'title': '循环模式',
|
||||
'enum': ['for_each', 'while'],
|
||||
'enumNames': ['遍历数组 (For Each)', '条件循环 (While)'],
|
||||
'default': 'for_each',
|
||||
},
|
||||
'items_variable': {
|
||||
'type': 'string',
|
||||
'title': '数组变量',
|
||||
'description': '要遍历的数组变量名（for_each 模式）',
|
||||
},
|
||||
'item_variable_name': {
|
||||
'type': 'string',
|
||||
'title': '循环项变量名',
|
||||
'description': '当前循环项的变量名，默认为 item',
|
||||
'default': 'item',
|
||||
},
|
||||
'index_variable_name': {
|
||||
'type': 'string',
|
||||
'title': '索引变量名',
|
||||
'description': '当前索引的变量名，默认为 index',
|
||||
'default': 'index',
|
||||
},
|
||||
'conditions': {
|
||||
'type': 'array',
|
||||
'title': '循环条件',
|
||||
'description': 'while 模式的循环条件',
|
||||
'items': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'variable': {
|
||||
'type': 'string',
|
||||
'title': '变量',
|
||||
},
|
||||
'operator': {
|
||||
'type': 'string',
|
||||
'title': '操作符',
|
||||
'enum': list(cls.OPERATORS.keys()),
|
||||
'default': 'eq',
|
||||
},
|
||||
'value': {
|
||||
'type': 'string',
|
||||
'title': '比较值',
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
'condition_logic': {
|
||||
'type': 'string',
|
||||
'title': '条件逻辑',
|
||||
'enum': ['and', 'or'],
|
||||
'enumNames': ['全部满足 (AND)', '任一满足 (OR)'],
|
||||
'default': 'and',
|
||||
},
|
||||
'max_iterations': {
|
||||
'type': 'integer',
|
||||
'title': '最大迭代次数',
|
||||
'description': '防止无限循环，默认 100 次',
|
||||
'default': 100,
|
||||
'minimum': 1,
|
||||
'maximum': 10000,
|
||||
},
|
||||
'output_variable': {
|
||||
'type': 'string',
|
||||
'title': '输出变量名',
|
||||
'description': '存储所有循环结果的变量名',
|
||||
'default': 'loop_results',
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
"""
|
||||
合并节点
|
||||
"""
|
||||
import logging
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class MergeNode(BaseNode):
|
||||
"""
|
||||
合并节点
|
||||
|
||||
等待所有并行分支执行完成后，合并结果继续执行
|
||||
"""
|
||||
|
||||
node_type = 'merge'
|
||||
node_name = '合并'
|
||||
node_category = 'logic'
|
||||
node_icon = 'git-merge'
|
||||
node_description = '等待所有并行分支完成后合并结果'
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'branch_results',
|
||||
'type': 'array',
|
||||
'description': '各分支的执行结果',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'merged_result',
|
||||
'type': 'object',
|
||||
'description': '合并后的结果',
|
||||
},
|
||||
]
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""
|
||||
执行合并节点
|
||||
|
||||
从上下文中获取所有并行分支的结果并合并
|
||||
"""
|
||||
try:
|
||||
# 获取合并模式
|
||||
merge_mode = self.config.get('merge_mode', 'object')
|
||||
output_variable = self.config.get('output_variable', 'merged_result')
|
||||
|
||||
# 从上下文获取并行分支结果
|
||||
# 并行执行时，每个分支的结果会存储在 _parallel_results 中
|
||||
parallel_results = context.get_variable('_parallel_results', {})
|
||||
|
||||
if merge_mode == 'object':
|
||||
# 对象模式：将各分支结果合并为一个对象
|
||||
merged = {}
|
||||
for branch_id, result in parallel_results.items():
|
||||
merged[branch_id] = result
|
||||
elif merge_mode == 'array':
|
||||
# 数组模式：将各分支结果合并为数组
|
||||
merged = list(parallel_results.values())
|
||||
elif merge_mode == 'first':
|
||||
# 取第一个完成的结果
|
||||
merged = list(parallel_results.values())[0] if parallel_results else None
|
||||
elif merge_mode == 'concat':
|
||||
# 字符串拼接模式
|
||||
separator = self.config.get('separator', '\n')
|
||||
merged = separator.join(str(v) for v in parallel_results.values())
|
||||
else:
|
||||
merged = parallel_results
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=merged,
|
||||
output_variables={
|
||||
output_variable: merged,
|
||||
'branch_count': len(parallel_results),
|
||||
},
|
||||
metadata={
|
||||
'merge_mode': merge_mode,
|
||||
'branch_ids': list(parallel_results.keys()),
|
||||
},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'合并节点执行失败: {e}')
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'merge_mode': {
|
||||
'type': 'string',
|
||||
'title': '合并模式',
|
||||
'description': '如何合并各分支的结果',
|
||||
'enum': ['object', 'array', 'first', 'concat'],
|
||||
'enumNames': ['对象（按分支ID）', '数组', '取第一个', '字符串拼接'],
|
||||
'default': 'object',
|
||||
},
|
||||
'separator': {
|
||||
'type': 'string',
|
||||
'title': '分隔符',
|
||||
'description': '字符串拼接模式的分隔符',
|
||||
'default': '\n',
|
||||
},
|
||||
'output_variable': {
|
||||
'type': 'string',
|
||||
'title': '输出变量名',
|
||||
'default': 'merged_result',
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
"""
|
||||
并行分支节点
|
||||
"""
|
||||
import logging
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class ParallelNode(BaseNode):
|
||||
"""
|
||||
并行分支节点
|
||||
|
||||
将工作流分成多个并行分支同时执行
|
||||
"""
|
||||
|
||||
node_type = 'parallel'
|
||||
node_name = '并行分支'
|
||||
node_category = 'logic'
|
||||
node_icon = 'git-fork'
|
||||
node_description = '将工作流分成多个并行分支同时执行'
|
||||
|
||||
supports_branches = True
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'input',
|
||||
'type': 'any',
|
||||
'description': '输入数据',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'branches',
|
||||
'type': 'array',
|
||||
'description': '并行分支 ID 列表',
|
||||
},
|
||||
]
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""
|
||||
执行并行分支节点
|
||||
|
||||
返回所有需要并行执行的分支 ID 列表
|
||||
"""
|
||||
try:
|
||||
# 获取配置的分支
|
||||
branches = self.config.get('branches', [])
|
||||
|
||||
if not branches:
|
||||
# 如果没有配置分支，返回默认的两个分支
|
||||
branches = [
|
||||
{'id': 'branch_1', 'name': '分支 1'},
|
||||
{'id': 'branch_2', 'name': '分支 2'},
|
||||
]
|
||||
|
||||
branch_ids = [b.get('id') for b in branches if b.get('id')]
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=branch_ids,
|
||||
output_variables={
|
||||
'parallel_branches': branch_ids,
|
||||
},
|
||||
metadata={
|
||||
'is_parallel': True,
|
||||
'branch_count': len(branch_ids),
|
||||
'branches': branches,
|
||||
},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'并行分支节点执行失败: {e}')
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'branches': {
|
||||
'type': 'array',
|
||||
'title': '并行分支',
|
||||
'description': '定义并行执行的分支',
|
||||
'items': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'id': {
|
||||
'type': 'string',
|
||||
'title': '分支 ID',
|
||||
},
|
||||
'name': {
|
||||
'type': 'string',
|
||||
'title': '分支名称',
|
||||
},
|
||||
},
|
||||
'required': ['id'],
|
||||
},
|
||||
'default': [
|
||||
{'id': 'branch_1', 'name': '分支 1'},
|
||||
{'id': 'branch_2', 'name': '分支 2'},
|
||||
],
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,834 @@
|
||||
"""
|
||||
Snowflake Cortex AI 节点
|
||||
|
||||
包含两个节点：
|
||||
1. SnowflakeCortexLLMNode - Cortex LLM Functions (COMPLETE, SUMMARIZE, TRANSLATE 等)
|
||||
2. SnowflakeCortexAnalystNode - Cortex Analyst (自然语言查询数据)
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SnowflakeConnectionMixin:
|
||||
"""Snowflake 连接混入类"""
|
||||
|
||||
# 支持的认证方式
|
||||
AUTH_TYPES = [
|
||||
('password', '用户名/密码'),
|
||||
('key_pair', '密钥对 (Key Pair)'),
|
||||
('externalbrowser', '外部浏览器 SSO'),
|
||||
]
|
||||
|
||||
def _get_connection(self, config: Dict, context: NodeContext):
|
||||
"""
|
||||
获取 Snowflake 连接
|
||||
|
||||
支持多种认证方式：
|
||||
- password: 用户名密码认证
|
||||
- key_pair: RSA 密钥对认证（推荐生产环境使用）
|
||||
- externalbrowser: 外部浏览器 SSO 认证
|
||||
|
||||
Args:
|
||||
config: 节点配置
|
||||
context: 执行上下文
|
||||
|
||||
Returns:
|
||||
snowflake.connector.connection
|
||||
"""
|
||||
try:
|
||||
import snowflake.connector
|
||||
except ImportError:
|
||||
raise ImportError('请安装 snowflake-connector-python: pip install snowflake-connector-python')
|
||||
|
||||
# 从配置获取连接信息
|
||||
connection_config = config.get('connection', {})
|
||||
|
||||
# 认证方式，默认为密码认证
|
||||
auth_type = connection_config.get('auth_type', 'password')
|
||||
|
||||
# 支持变量引用 - 基础参数
|
||||
account = context.resolve_template(connection_config.get('account', ''))
|
||||
user = context.resolve_template(connection_config.get('user', ''))
|
||||
warehouse = context.resolve_template(connection_config.get('warehouse', ''))
|
||||
database = context.resolve_template(connection_config.get('database', ''))
|
||||
schema = context.resolve_template(connection_config.get('schema', 'PUBLIC'))
|
||||
role = context.resolve_template(connection_config.get('role', ''))
|
||||
|
||||
# 验证必填参数
|
||||
if not account:
|
||||
raise ValueError('Snowflake 连接缺少 account 配置')
|
||||
if not user:
|
||||
raise ValueError('Snowflake 连接缺少 user 配置')
|
||||
if not warehouse:
|
||||
raise ValueError('Snowflake 连接缺少 warehouse 配置')
|
||||
if not database:
|
||||
raise ValueError('Snowflake 连接缺少 database 配置')
|
||||
|
||||
conn_params = {
|
||||
'account': account,
|
||||
'user': user,
|
||||
'warehouse': warehouse,
|
||||
'database': database,
|
||||
'schema': schema,
|
||||
}
|
||||
|
||||
if role:
|
||||
conn_params['role'] = role
|
||||
|
||||
# 根据认证方式设置认证参数
|
||||
if auth_type == 'password':
|
||||
password = context.resolve_template(connection_config.get('password', ''))
|
||||
if not password:
|
||||
raise ValueError('密码认证方式需要配置 password')
|
||||
conn_params['password'] = password
|
||||
|
||||
elif auth_type == 'key_pair':
|
||||
# 密钥对认证
|
||||
private_key = context.resolve_template(connection_config.get('private_key', ''))
|
||||
private_key_path = context.resolve_template(connection_config.get('private_key_path', ''))
|
||||
private_key_passphrase = context.resolve_template(connection_config.get('private_key_passphrase', ''))
|
||||
|
||||
if private_key:
|
||||
# 直接使用私钥内容
|
||||
conn_params['private_key'] = self._load_private_key_from_string(
|
||||
private_key, private_key_passphrase
|
||||
)
|
||||
elif private_key_path:
|
||||
# 从文件路径加载私钥
|
||||
conn_params['private_key'] = self._load_private_key_from_file(
|
||||
private_key_path, private_key_passphrase
|
||||
)
|
||||
else:
|
||||
raise ValueError('密钥对认证需要配置 private_key 或 private_key_path')
|
||||
|
||||
elif auth_type == 'externalbrowser':
|
||||
# 外部浏览器 SSO 认证
|
||||
conn_params['authenticator'] = 'externalbrowser'
|
||||
|
||||
else:
|
||||
raise ValueError(f'不支持的认证方式: {auth_type}')
|
||||
|
||||
logger.info(f'Snowflake 连接: account={account}, user={user}, auth_type={auth_type}')
|
||||
|
||||
return snowflake.connector.connect(**conn_params)
|
||||
|
||||
def _load_private_key_from_string(self, private_key_str: str, passphrase: str = '') -> bytes:
|
||||
"""从字符串加载 RSA 私钥"""
|
||||
from cryptography.hazmat.backends import default_backend
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
|
||||
# 处理可能的转义换行符
|
||||
private_key_str = private_key_str.replace('\\n', '\n')
|
||||
|
||||
p_key = serialization.load_pem_private_key(
|
||||
private_key_str.encode('utf-8'),
|
||||
password=passphrase.encode('utf-8') if passphrase else None,
|
||||
backend=default_backend()
|
||||
)
|
||||
|
||||
return p_key.private_bytes(
|
||||
encoding=serialization.Encoding.DER,
|
||||
format=serialization.PrivateFormat.PKCS8,
|
||||
encryption_algorithm=serialization.NoEncryption()
|
||||
)
|
||||
|
||||
def _load_private_key_from_file(self, file_path: str, passphrase: str = '') -> bytes:
|
||||
"""从文件加载 RSA 私钥"""
|
||||
from cryptography.hazmat.backends import default_backend
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
|
||||
with open(file_path, 'rb') as key_file:
|
||||
p_key = serialization.load_pem_private_key(
|
||||
key_file.read(),
|
||||
password=passphrase.encode('utf-8') if passphrase else None,
|
||||
backend=default_backend()
|
||||
)
|
||||
|
||||
return p_key.private_bytes(
|
||||
encoding=serialization.Encoding.DER,
|
||||
format=serialization.PrivateFormat.PKCS8,
|
||||
encryption_algorithm=serialization.NoEncryption()
|
||||
)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class SnowflakeCortexLLMNode(BaseNode, SnowflakeConnectionMixin):
|
||||
"""
|
||||
Snowflake Cortex LLM Functions 节点
|
||||
|
||||
支持的功能：
|
||||
- COMPLETE: LLM 文本生成
|
||||
- SUMMARIZE: 文本摘要
|
||||
- TRANSLATE: 翻译
|
||||
- SENTIMENT: 情感分析
|
||||
- EXTRACT_ANSWER: 问答提取
|
||||
"""
|
||||
|
||||
node_type = 'snowflake_cortex_llm'
|
||||
node_name = 'Snowflake Cortex LLM'
|
||||
node_category = 'ai'
|
||||
node_icon = 'snowflake'
|
||||
node_description = 'Snowflake Cortex AI 函数（LLM 生成、摘要、翻译、情感分析等）'
|
||||
|
||||
# 支持的 Cortex 模型
|
||||
CORTEX_MODELS = [
|
||||
'mistral-large',
|
||||
'mistral-large2',
|
||||
'mistral-7b',
|
||||
'mixtral-8x7b',
|
||||
'llama3-8b',
|
||||
'llama3-70b',
|
||||
'llama3.1-8b',
|
||||
'llama3.1-70b',
|
||||
'llama3.1-405b',
|
||||
'llama3.2-1b',
|
||||
'llama3.2-3b',
|
||||
'snowflake-arctic',
|
||||
'reka-core',
|
||||
'reka-flash',
|
||||
'jamba-instruct',
|
||||
'jamba-1.5-mini',
|
||||
'jamba-1.5-large',
|
||||
'gemma-7b',
|
||||
]
|
||||
|
||||
# 支持的功能类型
|
||||
FUNCTION_TYPES = [
|
||||
('complete', 'LLM 生成 (COMPLETE)'),
|
||||
('summarize', '文本摘要 (SUMMARIZE)'),
|
||||
('translate', '翻译 (TRANSLATE)'),
|
||||
('sentiment', '情感分析 (SENTIMENT)'),
|
||||
('extract_answer', '问答提取 (EXTRACT_ANSWER)'),
|
||||
]
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'text',
|
||||
'type': 'string',
|
||||
'description': '输入文本',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'result',
|
||||
'type': 'string',
|
||||
'description': 'Cortex 输出结果',
|
||||
},
|
||||
]
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""执行 Cortex LLM 函数"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
function_type = self.config.get('function', 'complete')
|
||||
|
||||
# 根据功能类型执行
|
||||
if function_type == 'complete':
|
||||
result = self._execute_complete(context)
|
||||
elif function_type == 'summarize':
|
||||
result = self._execute_summarize(context)
|
||||
elif function_type == 'translate':
|
||||
result = self._execute_translate(context)
|
||||
elif function_type == 'sentiment':
|
||||
result = self._execute_sentiment(context)
|
||||
elif function_type == 'extract_answer':
|
||||
result = self._execute_extract_answer(context)
|
||||
else:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=f'不支持的功能类型: {function_type}',
|
||||
)
|
||||
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
output_var = self.config.get('output_variable', 'cortex_result')
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=result,
|
||||
output_variables={output_var: result},
|
||||
elapsed_time=elapsed_time,
|
||||
metadata={'function': function_type},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'Snowflake Cortex LLM 节点执行失败: {e}')
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
elapsed_time=elapsed_time,
|
||||
)
|
||||
|
||||
def _execute_complete(self, context: NodeContext) -> str:
|
||||
"""
|
||||
执行 COMPLETE 函数 - 使用 REST API 方式
|
||||
|
||||
API 端点: /api/v2/cortex/inference:complete
|
||||
"""
|
||||
import requests
|
||||
|
||||
config = self.config.get('complete', {})
|
||||
model = config.get('model', 'mistral-large')
|
||||
prompt = context.resolve_template(config.get('prompt', ''))
|
||||
system_prompt = context.resolve_template(config.get('system_prompt', ''))
|
||||
temperature = config.get('temperature', 0.7)
|
||||
max_tokens = config.get('max_tokens', 1024)
|
||||
|
||||
if not prompt:
|
||||
raise ValueError('COMPLETE 功能需要配置 prompt')
|
||||
|
||||
connection_config = self.config.get('connection', {})
|
||||
account = context.resolve_template(connection_config.get('account', ''))
|
||||
|
||||
# 获取连接以获取 session token
|
||||
conn = self._get_connection(self.config, context)
|
||||
|
||||
try:
|
||||
# 从连接中获取 REST session token
|
||||
rest_token = conn.rest.token
|
||||
|
||||
# 构建 API URL
|
||||
if '.snowflakecomputing.com' in account:
|
||||
host = account
|
||||
else:
|
||||
host = conn.host if hasattr(conn, 'host') else f'{account}.snowflakecomputing.com'
|
||||
|
||||
api_url = f"https://{host}/api/v2/cortex/inference:complete"
|
||||
|
||||
# 构建请求头
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
'Accept': 'application/json',
|
||||
'Authorization': f'Snowflake Token="{rest_token}"',
|
||||
}
|
||||
|
||||
# 构建消息
|
||||
messages = []
|
||||
if system_prompt:
|
||||
messages.append({'role': 'system', 'content': system_prompt})
|
||||
messages.append({'role': 'user', 'content': prompt})
|
||||
|
||||
# 构建请求体
|
||||
payload = {
|
||||
'model': model,
|
||||
'messages': messages,
|
||||
'temperature': temperature,
|
||||
'max_tokens': max_tokens,
|
||||
}
|
||||
|
||||
logger.info(f'调用 Cortex LLM API: {api_url}')
|
||||
logger.info(f'Model: {model}, Temperature: {temperature}, Max Tokens: {max_tokens}')
|
||||
|
||||
response = requests.post(
|
||||
api_url,
|
||||
json=payload,
|
||||
headers=headers,
|
||||
timeout=120,
|
||||
)
|
||||
|
||||
logger.info(f'Cortex LLM 响应状态: {response.status_code}')
|
||||
|
||||
if response.status_code != 200:
|
||||
error_text = response.text
|
||||
logger.error(f'Cortex LLM API 错误: {error_text}')
|
||||
raise ValueError(f'Cortex LLM API 调用失败: {response.status_code} - {error_text}')
|
||||
|
||||
# 解析响应 - REST API 返回 SSE 格式
|
||||
result_content = ''
|
||||
response_text = response.text
|
||||
|
||||
# 解析 SSE 格式的响应
|
||||
for line in response_text.split('\n'):
|
||||
if line.startswith('data:'):
|
||||
data_str = line[5:].strip()
|
||||
if data_str:
|
||||
try:
|
||||
data = json.loads(data_str)
|
||||
if 'choices' in data and len(data['choices']) > 0:
|
||||
choice = data['choices'][0]
|
||||
if 'delta' in choice and 'content' in choice['delta']:
|
||||
result_content += choice['delta']['content']
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
return result_content if result_content else response_text
|
||||
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _execute_summarize(self, context: NodeContext) -> str:
|
||||
"""执行 SUMMARIZE 函数"""
|
||||
config = self.config.get('summarize', {})
|
||||
text = context.resolve_template(config.get('text', ''))
|
||||
|
||||
if not text:
|
||||
raise ValueError('SUMMARIZE 功能需要配置 text')
|
||||
|
||||
conn = self._get_connection(self.config, context)
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
sql = "SELECT SNOWFLAKE.CORTEX.SUMMARIZE(%s)"
|
||||
cursor.execute(sql, (text,))
|
||||
row = cursor.fetchone()
|
||||
return row[0] if row else ''
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _execute_translate(self, context: NodeContext) -> str:
|
||||
"""执行 TRANSLATE 函数"""
|
||||
config = self.config.get('translate', {})
|
||||
text = context.resolve_template(config.get('text', ''))
|
||||
source_language = config.get('source_language', 'en')
|
||||
target_language = config.get('target_language', 'zh')
|
||||
|
||||
if not text:
|
||||
raise ValueError('TRANSLATE 功能需要配置 text')
|
||||
|
||||
conn = self._get_connection(self.config, context)
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
sql = "SELECT SNOWFLAKE.CORTEX.TRANSLATE(%s, %s, %s)"
|
||||
cursor.execute(sql, (text, source_language, target_language))
|
||||
row = cursor.fetchone()
|
||||
return row[0] if row else ''
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _execute_sentiment(self, context: NodeContext) -> str:
|
||||
"""执行 SENTIMENT 函数"""
|
||||
config = self.config.get('sentiment', {})
|
||||
text = context.resolve_template(config.get('text', ''))
|
||||
|
||||
if not text:
|
||||
raise ValueError('SENTIMENT 功能需要配置 text')
|
||||
|
||||
conn = self._get_connection(self.config, context)
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
sql = "SELECT SNOWFLAKE.CORTEX.SENTIMENT(%s)"
|
||||
cursor.execute(sql, (text,))
|
||||
row = cursor.fetchone()
|
||||
# SENTIMENT 返回 -1 到 1 的数值
|
||||
result = row[0] if row else 0
|
||||
return str(result)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _execute_extract_answer(self, context: NodeContext) -> str:
|
||||
"""执行 EXTRACT_ANSWER 函数"""
|
||||
config = self.config.get('extract_answer', {})
|
||||
document = context.resolve_template(config.get('document', ''))
|
||||
question = context.resolve_template(config.get('question', ''))
|
||||
|
||||
if not document or not question:
|
||||
raise ValueError('EXTRACT_ANSWER 功能需要配置 document 和 question')
|
||||
|
||||
conn = self._get_connection(self.config, context)
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
sql = "SELECT SNOWFLAKE.CORTEX.EXTRACT_ANSWER(%s, %s)"
|
||||
cursor.execute(sql, (document, question))
|
||||
row = cursor.fetchone()
|
||||
result = row[0] if row else ''
|
||||
|
||||
# 解析 JSON 结果
|
||||
if isinstance(result, str):
|
||||
try:
|
||||
parsed = json.loads(result)
|
||||
if isinstance(parsed, list) and len(parsed) > 0:
|
||||
# 返回第一个答案
|
||||
return parsed[0].get('answer', result)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
return result
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'connection': {
|
||||
'type': 'object',
|
||||
'title': 'Snowflake 连接',
|
||||
'properties': {
|
||||
'account': {'type': 'string', 'title': 'Account'},
|
||||
'user': {'type': 'string', 'title': 'User'},
|
||||
'password': {'type': 'string', 'title': 'Password', 'format': 'password'},
|
||||
'warehouse': {'type': 'string', 'title': 'Warehouse'},
|
||||
'database': {'type': 'string', 'title': 'Database'},
|
||||
'schema': {'type': 'string', 'title': 'Schema', 'default': 'PUBLIC'},
|
||||
'role': {'type': 'string', 'title': 'Role'},
|
||||
},
|
||||
'required': ['account', 'user', 'password', 'warehouse', 'database'],
|
||||
},
|
||||
'function': {
|
||||
'type': 'string',
|
||||
'title': '功能类型',
|
||||
'enum': ['complete', 'summarize', 'translate', 'sentiment', 'extract_answer'],
|
||||
'default': 'complete',
|
||||
},
|
||||
'complete': {
|
||||
'type': 'object',
|
||||
'title': 'COMPLETE 配置',
|
||||
'properties': {
|
||||
'model': {'type': 'string', 'title': '模型', 'default': 'mistral-large'},
|
||||
'prompt': {'type': 'string', 'title': '提示词', 'format': 'textarea'},
|
||||
'system_prompt': {'type': 'string', 'title': '系统提示词', 'format': 'textarea'},
|
||||
'temperature': {'type': 'number', 'title': '温度', 'default': 0.7},
|
||||
'max_tokens': {'type': 'integer', 'title': '最大 Token', 'default': 1024},
|
||||
},
|
||||
},
|
||||
'summarize': {
|
||||
'type': 'object',
|
||||
'title': 'SUMMARIZE 配置',
|
||||
'properties': {
|
||||
'text': {'type': 'string', 'title': '待摘要文本', 'format': 'textarea'},
|
||||
},
|
||||
},
|
||||
'translate': {
|
||||
'type': 'object',
|
||||
'title': 'TRANSLATE 配置',
|
||||
'properties': {
|
||||
'text': {'type': 'string', 'title': '待翻译文本', 'format': 'textarea'},
|
||||
'source_language': {'type': 'string', 'title': '源语言', 'default': 'en'},
|
||||
'target_language': {'type': 'string', 'title': '目标语言', 'default': 'zh'},
|
||||
},
|
||||
},
|
||||
'sentiment': {
|
||||
'type': 'object',
|
||||
'title': 'SENTIMENT 配置',
|
||||
'properties': {
|
||||
'text': {'type': 'string', 'title': '待分析文本', 'format': 'textarea'},
|
||||
},
|
||||
},
|
||||
'extract_answer': {
|
||||
'type': 'object',
|
||||
'title': 'EXTRACT_ANSWER 配置',
|
||||
'properties': {
|
||||
'document': {'type': 'string', 'title': '文档内容', 'format': 'textarea'},
|
||||
'question': {'type': 'string', 'title': '问题'},
|
||||
},
|
||||
},
|
||||
'output_variable': {
|
||||
'type': 'string',
|
||||
'title': '输出变量名',
|
||||
'default': 'cortex_result',
|
||||
},
|
||||
},
|
||||
'required': ['connection', 'function'],
|
||||
}
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class SnowflakeCortexAnalystNode(BaseNode, SnowflakeConnectionMixin):
|
||||
"""
|
||||
Snowflake Cortex Analyst 节点
|
||||
|
||||
使用自然语言查询数据，自动生成 SQL 并返回结果
|
||||
"""
|
||||
|
||||
node_type = 'snowflake_cortex_analyst'
|
||||
node_name = 'Snowflake Cortex Analyst'
|
||||
node_category = 'ai'
|
||||
node_icon = 'snowflake'
|
||||
node_description = 'Snowflake Cortex Analyst - 自然语言查询数据'
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'question',
|
||||
'type': 'string',
|
||||
'description': '自然语言问题',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'result',
|
||||
'type': 'object',
|
||||
'description': '查询结果',
|
||||
},
|
||||
{
|
||||
'name': 'sql',
|
||||
'type': 'string',
|
||||
'description': '生成的 SQL',
|
||||
},
|
||||
]
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""执行 Cortex Analyst 查询"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 获取配置
|
||||
question = context.resolve_template(self.config.get('question', ''))
|
||||
semantic_model_file = self.config.get('semantic_model_file', '')
|
||||
|
||||
if not question:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='请配置查询问题',
|
||||
)
|
||||
|
||||
if not semantic_model_file:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='请配置语义模型文件路径',
|
||||
)
|
||||
|
||||
# 调用 Cortex Analyst API
|
||||
result = self._call_cortex_analyst(context, question, semantic_model_file)
|
||||
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
output_var = self.config.get('output_variable', 'analyst_result')
|
||||
sql_var = self.config.get('sql_variable', 'analyst_sql')
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=result,
|
||||
output_variables={
|
||||
output_var: result.get('data', []),
|
||||
sql_var: result.get('sql', ''),
|
||||
f'{output_var}_raw': result,
|
||||
},
|
||||
elapsed_time=elapsed_time,
|
||||
metadata={'question': question},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'Snowflake Cortex Analyst 节点执行失败: {e}')
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
elapsed_time=elapsed_time,
|
||||
)
|
||||
|
||||
def _call_cortex_analyst(
|
||||
self,
|
||||
context: NodeContext,
|
||||
question: str,
|
||||
semantic_model: str,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
调用 Cortex Analyst API
|
||||
|
||||
Cortex Analyst 必须通过 REST API 调用，不支持 SQL 函数方式。
|
||||
|
||||
支持两种语义模型来源：
|
||||
1. Semantic View: database.schema.semantic_view_name (如 AI_TEST.PUBLIC.CAR_AI_TEST)
|
||||
2. Stage 文件: @database.schema.stage/file.yaml
|
||||
|
||||
Args:
|
||||
context: 节点上下文
|
||||
question: 用户问题
|
||||
semantic_model: 语义模型路径（Semantic View 或 Stage 文件）
|
||||
"""
|
||||
import requests
|
||||
|
||||
connection_config = self.config.get('connection', {})
|
||||
account = context.resolve_template(connection_config.get('account', ''))
|
||||
database = context.resolve_template(connection_config.get('database', ''))
|
||||
schema = context.resolve_template(connection_config.get('schema', 'PUBLIC'))
|
||||
|
||||
# 获取连接以获取 session token
|
||||
conn = self._get_connection(self.config, context)
|
||||
|
||||
try:
|
||||
# 从连接中获取 REST session token
|
||||
# Snowflake Python Connector 可以提供 REST session token
|
||||
rest_token = conn.rest.token
|
||||
master_token = conn.rest.master_token
|
||||
|
||||
# 构建 API URL
|
||||
# 处理 account 格式：可能是 xxx.snowflakecomputing.com 或 account_identifier
|
||||
if '.snowflakecomputing.com' in account:
|
||||
host = account
|
||||
else:
|
||||
# 从连接获取实际 host
|
||||
host = conn.host if hasattr(conn, 'host') else f'{account}.snowflakecomputing.com'
|
||||
|
||||
base_url = f"https://{host}"
|
||||
api_url = f"{base_url}/api/v2/cortex/analyst/message"
|
||||
|
||||
# 构建请求头
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
'Accept': 'application/json',
|
||||
'Authorization': f'Snowflake Token="{rest_token}"',
|
||||
}
|
||||
|
||||
# 判断是 Semantic View 还是 Stage 文件
|
||||
is_stage_file = semantic_model.startswith('@')
|
||||
|
||||
# 构建请求体
|
||||
messages = [
|
||||
{
|
||||
'role': 'user',
|
||||
'content': [{'type': 'text', 'text': question}]
|
||||
}
|
||||
]
|
||||
|
||||
if is_stage_file:
|
||||
# Stage 文件格式: @database.schema.stage/file.yaml
|
||||
payload = {
|
||||
'messages': messages,
|
||||
'semantic_model_file': semantic_model,
|
||||
}
|
||||
else:
|
||||
# Semantic View 格式: database.schema.view_name
|
||||
# 使用 semantic_models 数组包含 semantic_view
|
||||
payload = {
|
||||
'messages': messages,
|
||||
'semantic_models': [
|
||||
{'semantic_view': semantic_model}
|
||||
],
|
||||
}
|
||||
|
||||
logger.info(f'调用 Cortex Analyst API: {api_url}')
|
||||
logger.info(f'Payload: {payload}')
|
||||
|
||||
response = requests.post(
|
||||
api_url,
|
||||
json=payload,
|
||||
headers=headers,
|
||||
timeout=60,
|
||||
)
|
||||
|
||||
logger.info(f'Cortex Analyst 响应状态: {response.status_code}')
|
||||
|
||||
if response.status_code != 200:
|
||||
error_text = response.text
|
||||
logger.error(f'Cortex Analyst API 错误: {error_text}')
|
||||
raise ValueError(f'Cortex Analyst API 调用失败: {response.status_code} - {error_text}')
|
||||
|
||||
result = response.json()
|
||||
logger.info(f'Cortex Analyst 返回: {result}')
|
||||
|
||||
# 解析响应，提取 SQL 和执行结果
|
||||
return self._parse_analyst_response(conn, result)
|
||||
|
||||
except requests.RequestException as e:
|
||||
logger.exception(f'Cortex Analyst REST API 调用失败: {e}')
|
||||
raise ValueError(f'Cortex Analyst API 调用失败: {e}')
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _parse_analyst_response(self, conn, result: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""解析 Cortex Analyst 响应并执行生成的 SQL"""
|
||||
generated_sql = None
|
||||
answer_text = None
|
||||
|
||||
# 尝试从响应中提取 SQL
|
||||
if isinstance(result, dict):
|
||||
# 格式1: message.content 数组
|
||||
if 'message' in result and 'content' in result['message']:
|
||||
content = result['message']['content']
|
||||
if isinstance(content, list):
|
||||
for item in content:
|
||||
if isinstance(item, dict):
|
||||
if item.get('type') == 'sql':
|
||||
generated_sql = item.get('statement', item.get('text', ''))
|
||||
elif item.get('type') == 'text':
|
||||
answer_text = item.get('text', '')
|
||||
elif isinstance(content, str):
|
||||
answer_text = content
|
||||
|
||||
# 格式2: 直接 sql 字段
|
||||
elif 'sql' in result:
|
||||
generated_sql = result['sql']
|
||||
|
||||
# 格式3: choices 数组
|
||||
elif 'choices' in result and len(result['choices']) > 0:
|
||||
choice = result['choices'][0]
|
||||
if 'message' in choice and 'content' in choice['message']:
|
||||
content = choice['message']['content']
|
||||
if isinstance(content, list):
|
||||
for item in content:
|
||||
if isinstance(item, dict):
|
||||
if item.get('type') == 'sql':
|
||||
generated_sql = item.get('statement', '')
|
||||
elif item.get('type') == 'text':
|
||||
answer_text = item.get('text', '')
|
||||
|
||||
# 如果有生成的 SQL，执行它获取数据
|
||||
if generated_sql:
|
||||
logger.info(f'执行生成的 SQL: {generated_sql}')
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute(generated_sql)
|
||||
columns = [desc[0] for desc in cursor.description]
|
||||
rows = cursor.fetchall()
|
||||
data = [dict(zip(columns, r)) for r in rows]
|
||||
|
||||
return {
|
||||
'sql': generated_sql,
|
||||
'data': data,
|
||||
'columns': columns,
|
||||
'row_count': len(data),
|
||||
'answer': answer_text,
|
||||
}
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
# 没有 SQL，返回文本回答
|
||||
return {
|
||||
'sql': '',
|
||||
'data': [],
|
||||
'columns': [],
|
||||
'row_count': 0,
|
||||
'answer': answer_text or str(result),
|
||||
'raw_response': result,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'connection': {
|
||||
'type': 'object',
|
||||
'title': 'Snowflake 连接',
|
||||
'properties': {
|
||||
'account': {'type': 'string', 'title': 'Account'},
|
||||
'user': {'type': 'string', 'title': 'User'},
|
||||
'password': {'type': 'string', 'title': 'Password', 'format': 'password'},
|
||||
'warehouse': {'type': 'string', 'title': 'Warehouse'},
|
||||
'database': {'type': 'string', 'title': 'Database'},
|
||||
'schema': {'type': 'string', 'title': 'Schema', 'default': 'PUBLIC'},
|
||||
'role': {'type': 'string', 'title': 'Role'},
|
||||
},
|
||||
'required': ['account', 'user', 'password', 'warehouse', 'database'],
|
||||
},
|
||||
'question': {
|
||||
'type': 'string',
|
||||
'title': '查询问题',
|
||||
'description': '用自然语言描述你想查询的数据',
|
||||
'format': 'textarea',
|
||||
},
|
||||
'semantic_model_file': {
|
||||
'type': 'string',
|
||||
'title': '语义模型文件',
|
||||
'description': '语义模型文件路径，如 @my_db.my_schema.my_stage/semantic_model.yaml',
|
||||
},
|
||||
'output_variable': {
|
||||
'type': 'string',
|
||||
'title': '输出变量名',
|
||||
'default': 'analyst_result',
|
||||
},
|
||||
},
|
||||
'required': ['connection', 'question', 'semantic_model_file'],
|
||||
}
|
||||
@@ -0,0 +1,155 @@
|
||||
"""
|
||||
开始节点
|
||||
"""
|
||||
import logging
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class StartNode(BaseNode):
|
||||
"""
|
||||
开始节点
|
||||
|
||||
工作流的入口节点，接收用户输入
|
||||
"""
|
||||
|
||||
node_type = 'start'
|
||||
node_name = '开始'
|
||||
node_category = 'basic'
|
||||
node_icon = 'play-circle'
|
||||
node_description = '工作流的开始节点，接收用户输入'
|
||||
|
||||
inputs = []
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'user_input',
|
||||
'type': 'string',
|
||||
'description': '用户输入',
|
||||
},
|
||||
{
|
||||
'name': 'application_id',
|
||||
'type': 'string',
|
||||
'description': '子应用ID（在子应用模式下自动注入）',
|
||||
},
|
||||
{
|
||||
'name': 'application_code',
|
||||
'type': 'string',
|
||||
'description': '子应用编码（在子应用模式下自动注入）',
|
||||
},
|
||||
{
|
||||
'name': 'form_code',
|
||||
'type': 'string',
|
||||
'description': '表单编码（从表单列表调用时自动注入）',
|
||||
},
|
||||
]
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""执行开始节点"""
|
||||
output_variables = {}
|
||||
|
||||
# 调试日志
|
||||
logger.info(f"StartNode execute - context.variables: {list(context.variables.keys())}")
|
||||
logger.info(f"StartNode execute - application_id in variables: {'application_id' in context.variables}")
|
||||
if 'application_id' in context.variables:
|
||||
logger.info(f"StartNode execute - application_id value: {context.variables['application_id']}")
|
||||
|
||||
# 首先将 user_input 作为开始节点的输出变量
|
||||
# 这样下游节点可以通过 {{start-xxx.user_input}} 引用
|
||||
if context.user_input:
|
||||
output_variables['user_input'] = context.user_input
|
||||
elif 'user_input' in context.variables:
|
||||
output_variables['user_input'] = context.variables['user_input']
|
||||
|
||||
# 将 application_id 作为开始节点的输出变量
|
||||
# 这样下游节点可以通过 {{start-xxx.application_id}} 引用
|
||||
# 主应用模式下为空字符串，子应用模式下为实际的应用ID
|
||||
output_variables['application_id'] = context.variables.get('application_id', 'main')
|
||||
|
||||
# 将 application_code 作为开始节点的输出变量
|
||||
# 这样下游节点可以通过 {{start-xxx.application_code}} 引用
|
||||
# 主应用模式下为空字符串，子应用模式下为实际的应用编码
|
||||
output_variables['application_code'] = context.variables.get('application_code', '')
|
||||
|
||||
# 将 form_code 作为开始节点的输出变量
|
||||
# 这样下游节点可以通过 {{start-xxx.form_code}} 引用
|
||||
# 从表单列表调用时会自动注入，否则为空字符串
|
||||
output_variables['form_code'] = context.variables.get('form_code', '')
|
||||
|
||||
logger.info(f"StartNode execute - output_variables: {output_variables}")
|
||||
|
||||
# 处理前端定义的变量（variables 数组）
|
||||
# 格式: [{ variable: 'name', type: 'string', label: '名称', required: true, default_value: '' }]
|
||||
variables = self.config.get('variables', [])
|
||||
for var in variables:
|
||||
var_name = var.get('variable', '')
|
||||
if not var_name:
|
||||
continue
|
||||
|
||||
default_value = var.get('default_value', '')
|
||||
var_type = var.get('type', 'string')
|
||||
|
||||
# 如果 context 中已有该变量（从 inputs 传入），使用传入的值
|
||||
# 否则使用默认值
|
||||
if var_name in context.variables and context.variables[var_name]:
|
||||
value = context.variables[var_name]
|
||||
else:
|
||||
value = default_value
|
||||
|
||||
# 类型转换
|
||||
if var_type == 'number' and value:
|
||||
try:
|
||||
value = float(value) if '.' in str(value) else int(value)
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
elif var_type == 'boolean':
|
||||
if isinstance(value, str):
|
||||
value = value.lower() in ('true', '1', 'yes')
|
||||
|
||||
context.set_variable(var_name, value)
|
||||
output_variables[var_name] = value
|
||||
|
||||
# 处理自定义输入变量（兼容旧格式 input_variables）
|
||||
input_variables = self.config.get('input_variables', [])
|
||||
for var in input_variables:
|
||||
var_name = var.get('name', '')
|
||||
var_value = var.get('default', '')
|
||||
if var_name:
|
||||
# 如果 context 中没有该变量，使用默认值
|
||||
if var_name not in context.variables or not context.variables[var_name]:
|
||||
context.set_variable(var_name, var_value)
|
||||
output_variables[var_name] = context.get_variable(var_name)
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=output_variables,
|
||||
output_variables=output_variables,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'input_variables': {
|
||||
'type': 'array',
|
||||
'title': '输入变量',
|
||||
'items': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'name': {'type': 'string', 'title': '变量名'},
|
||||
'type': {'type': 'string', 'title': '类型', 'enum': ['string', 'number', 'boolean']},
|
||||
'description': {'type': 'string', 'title': '描述'},
|
||||
'default': {'type': 'string', 'title': '默认值'},
|
||||
'required': {'type': 'boolean', 'title': '是否必填'},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,84 @@
|
||||
"""
|
||||
模板渲染节点
|
||||
"""
|
||||
import logging
|
||||
from typing import Any, Dict
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class TemplateNode(BaseNode):
|
||||
"""
|
||||
模板渲染节点
|
||||
|
||||
使用变量渲染模板字符串
|
||||
"""
|
||||
|
||||
node_type = 'template'
|
||||
node_name = '模板'
|
||||
node_category = 'data'
|
||||
node_icon = 'file-text'
|
||||
node_description = '使用变量渲染模板字符串'
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'template',
|
||||
'type': 'string',
|
||||
'description': '模板字符串',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'result',
|
||||
'type': 'string',
|
||||
'description': '渲染结果',
|
||||
},
|
||||
]
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""执行模板渲染"""
|
||||
try:
|
||||
template = self.config.get('template', '')
|
||||
output_variable = self.config.get('output_variable', 'template_result')
|
||||
|
||||
# 渲染模板
|
||||
result = context.resolve_template(template)
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=result,
|
||||
output_variables={output_variable: result},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'模板节点执行失败: {e}')
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'template': {
|
||||
'type': 'string',
|
||||
'title': '模板',
|
||||
'description': '支持变量引用，如 {{variable_name}}',
|
||||
'format': 'textarea',
|
||||
},
|
||||
'output_variable': {
|
||||
'type': 'string',
|
||||
'title': '输出变量名',
|
||||
'default': 'template_result',
|
||||
},
|
||||
},
|
||||
'required': ['template'],
|
||||
}
|
||||
@@ -0,0 +1,947 @@
|
||||
"""
|
||||
Text-to-SQL 节点
|
||||
|
||||
将自然语言转换为 SQL 查询并执行，支持流式输出和图表推荐
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Generator, List, Optional, Tuple
|
||||
|
||||
import sqlparse
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextToSqlStreamEvent:
|
||||
"""Text-to-SQL 流式输出事件"""
|
||||
event_type: str = '' # thought, sql_chunk, sql_complete, executing, result
|
||||
content: str = ''
|
||||
is_finished: bool = False
|
||||
data: Any = None
|
||||
|
||||
|
||||
# Text-to-SQL Function Calling 工具定义
|
||||
TEXT_TO_SQL_TOOL = {
|
||||
'name': 'generate_sql',
|
||||
'description': '根据用户的自然语言问题生成 SQL 查询语句',
|
||||
'parameters': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'thought': {
|
||||
'type': 'string',
|
||||
'description': '分析思路，简要描述如何理解用户问题并设计 SQL',
|
||||
},
|
||||
'sql': {
|
||||
'type': 'string',
|
||||
'description': '生成的 SQL 查询语句，必须是有效的 SELECT 语句，必须格式化，如果表结构中包含 schema 信息，SQL 中的表名必须使用完整格式：schema.table_name',
|
||||
},
|
||||
},
|
||||
'required': ['thought', 'sql'],
|
||||
},
|
||||
}
|
||||
|
||||
# Text-to-SQL Function Calling 版 System Prompt（简化版）
|
||||
TEXT_TO_SQL_SYSTEM_PROMPT_FC = """你是一个专业的数据分析助手，擅长将自然语言转换为 SQL 查询。
|
||||
|
||||
## 数据库信息
|
||||
数据库类型: {db_type}
|
||||
当前日期: {current_date}
|
||||
|
||||
## 可用的表结构
|
||||
{schema_context}
|
||||
|
||||
## 任务要求
|
||||
1. 根据用户问题生成准确的 SQL 查询
|
||||
2. 只生成 SELECT 语句，禁止生成 INSERT/UPDATE/DELETE/DROP 等修改语句
|
||||
3. 使用表的注释理解业务含义
|
||||
4. 添加必要的 WHERE 条件和 ORDER BY
|
||||
5. 如果表结构中包含 schema 信息，SQL 中的表名必须使用完整格式：schema.table_name
|
||||
|
||||
## 注意事项
|
||||
- SQL 必须是有效的 {db_type} 语法
|
||||
- 避免使用 SELECT *，明确指定需要的字段
|
||||
- 对于大数据量查询，添加 LIMIT 限制
|
||||
- 使用 DISTINCT 时，ORDER BY 的列必须出现在 SELECT 列表中
|
||||
|
||||
请使用 generate_sql 函数返回结果。
|
||||
"""
|
||||
|
||||
# Text-to-SQL 专用 System Prompt（原有版本，作为 fallback）
|
||||
TEXT_TO_SQL_SYSTEM_PROMPT = """你是一个专业的数据分析助手，擅长将自然语言转换为 SQL 查询。
|
||||
|
||||
## 数据库信息
|
||||
数据库类型: {db_type}
|
||||
当前日期: {current_date}
|
||||
|
||||
## 可用的表结构
|
||||
{schema_context}
|
||||
|
||||
## 任务要求
|
||||
1. 根据用户问题生成准确的 SQL 查询
|
||||
2. 只生成 SELECT 语句，禁止生成 INSERT/UPDATE/DELETE/DROP 等修改语句
|
||||
3. 使用表的注释理解业务含义
|
||||
4. 添加必要的 WHERE 条件和 ORDER BY
|
||||
5. 如果表结构中包含 schema 信息，SQL 中的表名必须使用完整格式：schema.table_name
|
||||
6. 如果需要，推荐最合适的图表类型
|
||||
|
||||
## 输出格式
|
||||
严格按照以下 JSON 格式输出，不要输出其他内容：
|
||||
{{
|
||||
"thought": "你的分析思路（简洁描述）",
|
||||
"sql": "生成的 SQL 语句",
|
||||
"chart_type": "bar|line|pie|scatter|radar|table|none",
|
||||
"chart_config": {{
|
||||
"x_field": "X轴字段名（用于 bar/line/scatter）",
|
||||
"y_field": "Y轴字段名（用于 scatter）",
|
||||
"series_fields": ["系列字段1", "系列字段2"],
|
||||
"name_field": "名称字段（用于 pie）",
|
||||
"value_field": "数值字段（用于 pie/gauge）",
|
||||
"title": "图表标题"
|
||||
}}
|
||||
}}
|
||||
|
||||
## 图表类型选择指南
|
||||
- bar（柱状图）: 分类对比，如各部门销售额
|
||||
- line（折线图）: 时间趋势，如每日销售额变化
|
||||
- pie（饼图）: 占比分析，如各类别占比
|
||||
- scatter（散点图）: 分布分析，如价格与销量关系
|
||||
- radar（雷达图）: 多维对比，如产品多指标评分
|
||||
- table（表格）: 详细数据展示
|
||||
- none（无图表）: 不适合可视化的数据
|
||||
|
||||
## 注意事项
|
||||
- SQL 必须是有效的 {db_type} 语法
|
||||
- 避免使用 SELECT *，明确指定需要的字段
|
||||
- 对于大数据量查询，添加 LIMIT 限制
|
||||
- 时间字段使用标准格式
|
||||
- 使用 DISTINCT 时，ORDER BY 的列必须出现在 SELECT 列表中
|
||||
- 如需去重并排序，考虑使用子查询或窗口函数
|
||||
"""
|
||||
|
||||
|
||||
def serialize_db_value(value: Any) -> Any:
|
||||
"""将数据库值转换为可 JSON 序列化的格式"""
|
||||
from datetime import date, datetime
|
||||
from decimal import Decimal
|
||||
from uuid import UUID
|
||||
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, datetime):
|
||||
return value.isoformat()
|
||||
if isinstance(value, date):
|
||||
return value.isoformat()
|
||||
if isinstance(value, Decimal):
|
||||
return float(value)
|
||||
if isinstance(value, UUID):
|
||||
return str(value)
|
||||
if isinstance(value, bytes):
|
||||
return value.decode('utf-8', errors='replace')
|
||||
if isinstance(value, (list, tuple)):
|
||||
return [serialize_db_value(v) for v in value]
|
||||
if isinstance(value, dict):
|
||||
return {k: serialize_db_value(v) for k, v in value.items()}
|
||||
return value
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class TextToSqlNode(BaseNode):
|
||||
"""
|
||||
Text-to-SQL 节点
|
||||
|
||||
将自然语言转换为 SQL 查询并执行
|
||||
"""
|
||||
|
||||
node_type = 'text_to_sql'
|
||||
node_name = 'Text-to-SQL'
|
||||
node_category = 'data'
|
||||
node_icon = 'database-zap'
|
||||
node_description = '将自然语言转换为 SQL 查询并执行，支持图表推荐'
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'user_question',
|
||||
'type': 'string',
|
||||
'description': '用户的自然语言问题',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'text_to_sql_result_sql',
|
||||
'type': 'string',
|
||||
'description': '生成的 SQL 语句',
|
||||
},
|
||||
{
|
||||
'name': 'text_to_sql_result_thought',
|
||||
'type': 'string',
|
||||
'description': 'SQL 生成思路',
|
||||
},
|
||||
]
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""执行 Text-to-SQL（同步方法）"""
|
||||
import asyncio
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
if loop.is_running():
|
||||
import concurrent.futures
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future = executor.submit(asyncio.run, self.execute_async(context))
|
||||
return future.result()
|
||||
else:
|
||||
return loop.run_until_complete(self.execute_async(context))
|
||||
|
||||
async def execute_async(self, context: NodeContext) -> NodeResult:
|
||||
"""异步执行 Text-to-SQL"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 获取配置
|
||||
user_question = self.config.get('user_question', '')
|
||||
db_config = self.config.get('db_config') or {}
|
||||
db_connection = db_config.get('dbName', 'default')
|
||||
schema_name = db_config.get('schema', 'public')
|
||||
model_id = self.config.get('model_id', '')
|
||||
selected_tables = self.config.get('selected_tables', [])
|
||||
table_relations = self.config.get('table_relations', []) # 手动指定的表关系
|
||||
output_variable = self.config.get('output_variable', 'text_to_sql_result')
|
||||
include_relations = self.config.get('include_table_relations', True)
|
||||
|
||||
# 解析变量
|
||||
if user_question:
|
||||
user_question = context.resolve_template(user_question)
|
||||
|
||||
if not user_question:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='请输入要查询的问题',
|
||||
)
|
||||
|
||||
if not model_id:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='请选择 LLM 模型',
|
||||
)
|
||||
|
||||
# Step 1: 获取数据库 Schema
|
||||
|
||||
schema_context = await self._get_schema_context(
|
||||
db_connection,
|
||||
schema_name,
|
||||
selected_tables,
|
||||
include_relations,
|
||||
table_relations
|
||||
)
|
||||
if not schema_context:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=f'无法获取数据库 {db_connection} 的表结构信息',
|
||||
)
|
||||
|
||||
# Step 2: 调用 LLM 生成 SQL
|
||||
from datetime import datetime
|
||||
from ai_platform.services.llm_service import LLMService
|
||||
|
||||
db_type = await self._get_db_type(db_connection)
|
||||
llm_service = LLMService(context.db_session)
|
||||
llm_result = None
|
||||
use_function_calling = self.config.get('use_function_calling', True)
|
||||
|
||||
# 优先尝试 Function Calling
|
||||
if use_function_calling:
|
||||
try:
|
||||
system_prompt = TEXT_TO_SQL_SYSTEM_PROMPT_FC.format(
|
||||
db_type=db_type,
|
||||
current_date=datetime.now().strftime('%Y-%m-%d'),
|
||||
schema_context=schema_context,
|
||||
)
|
||||
|
||||
messages = [
|
||||
{'role': 'system', 'content': system_prompt},
|
||||
{'role': 'user', 'content': f'用户问题: {user_question}'},
|
||||
]
|
||||
|
||||
response = await llm_service.chat_async(
|
||||
model_id=model_id,
|
||||
messages=messages,
|
||||
temperature=0.1,
|
||||
max_tokens=2048,
|
||||
tools=[TEXT_TO_SQL_TOOL],
|
||||
tool_choice='required',
|
||||
)
|
||||
|
||||
# 解析 Function Calling 响应
|
||||
if response.tool_calls and len(response.tool_calls) > 0:
|
||||
tool_call = response.tool_calls[0]
|
||||
if tool_call.name == 'generate_sql':
|
||||
import json
|
||||
arguments = tool_call.arguments
|
||||
if isinstance(arguments, str):
|
||||
llm_result = json.loads(arguments)
|
||||
else:
|
||||
llm_result = arguments
|
||||
logger.info(f'Function Calling 成功生成 SQL')
|
||||
except Exception as fc_error:
|
||||
logger.warning(f'Function Calling 失败，回退到 JSON 解析模式: {fc_error}')
|
||||
llm_result = None
|
||||
|
||||
# Fallback: 使用 JSON 解析方式
|
||||
if not llm_result:
|
||||
system_prompt = TEXT_TO_SQL_SYSTEM_PROMPT.format(
|
||||
db_type=db_type,
|
||||
current_date=datetime.now().strftime('%Y-%m-%d'),
|
||||
schema_context=schema_context,
|
||||
)
|
||||
|
||||
messages = [
|
||||
{'role': 'system', 'content': system_prompt},
|
||||
{'role': 'user', 'content': f'用户问题: {user_question}'},
|
||||
]
|
||||
|
||||
response = await llm_service.chat_async(
|
||||
model_id=model_id,
|
||||
messages=messages,
|
||||
temperature=0.1,
|
||||
max_tokens=2048,
|
||||
)
|
||||
|
||||
# 解析 LLM 响应
|
||||
llm_result = self._parse_llm_response(response.content)
|
||||
if not llm_result:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='LLM 响应解析失败，请重试',
|
||||
metadata={'raw_response': response.content},
|
||||
)
|
||||
|
||||
sql = llm_result.get('sql', '')
|
||||
thought = llm_result.get('thought', '')
|
||||
|
||||
# Step 3: 验证 SQL 安全性
|
||||
if not self._validate_sql(sql):
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='生成的 SQL 不安全或不是 SELECT 语句',
|
||||
metadata={'sql': sql},
|
||||
)
|
||||
|
||||
# Step 4: 格式化 SQL
|
||||
sql = self._format_sql(sql)
|
||||
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
|
||||
# 构建输出
|
||||
result_data = {
|
||||
'sql': sql,
|
||||
'thought': thought,
|
||||
}
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=result_data,
|
||||
output_variables={
|
||||
f'{output_variable}_sql': sql,
|
||||
f'{output_variable}_thought': thought,
|
||||
},
|
||||
tokens_used=response.total_tokens,
|
||||
elapsed_time=elapsed_time,
|
||||
metadata={
|
||||
'model': response.model,
|
||||
'thought': thought,
|
||||
'suggested_next_node': 'db_sql',
|
||||
},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'Text-to-SQL 节点执行失败: {e}')
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
elapsed_time=elapsed_time,
|
||||
)
|
||||
|
||||
def execute_stream(self, context: NodeContext) -> Generator[TextToSqlStreamEvent, None, NodeResult]:
|
||||
"""
|
||||
流式执行 Text-to-SQL
|
||||
|
||||
Yields:
|
||||
TextToSqlStreamEvent: 流式输出事件
|
||||
|
||||
Returns:
|
||||
NodeResult: 最终执行结果
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 获取配置
|
||||
user_question = self.config.get('user_question', '')
|
||||
db_config = self.config.get('db_config') or {}
|
||||
db_connection = db_config.get('dbName', 'default')
|
||||
schema_name = db_config.get('schema', 'public')
|
||||
model_id = self.config.get('model_id', '')
|
||||
selected_tables = self.config.get('selected_tables', [])
|
||||
output_variable = self.config.get('output_variable', 'text_to_sql_result')
|
||||
include_relations = self.config.get('include_table_relations', True)
|
||||
|
||||
# 解析变量
|
||||
if user_question:
|
||||
user_question = context.resolve_template(user_question)
|
||||
|
||||
if not user_question:
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='error',
|
||||
content='请输入要查询的问题',
|
||||
is_finished=True,
|
||||
)
|
||||
return NodeResult(success=False, error='请输入要查询的问题')
|
||||
|
||||
if not model_id:
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='error',
|
||||
content='请选择 LLM 模型',
|
||||
is_finished=True,
|
||||
)
|
||||
return NodeResult(success=False, error='请选择 LLM 模型')
|
||||
|
||||
# Step 1: 获取数据库 Schema
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='thought',
|
||||
content='正在获取数据库表结构...',
|
||||
)
|
||||
|
||||
table_relations = self.config.get('table_relations', []) # 手动指定的表关系
|
||||
|
||||
import asyncio
|
||||
schema_context = asyncio.get_event_loop().run_until_complete(
|
||||
self._get_schema_context(
|
||||
db_connection,
|
||||
schema_name,
|
||||
selected_tables,
|
||||
include_relations,
|
||||
table_relations
|
||||
)
|
||||
)
|
||||
|
||||
if not schema_context:
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='error',
|
||||
content=f'无法获取数据库 {db_connection} 的表结构信息',
|
||||
is_finished=True,
|
||||
)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=f'无法获取数据库 {db_connection} 的表结构信息',
|
||||
)
|
||||
|
||||
# Step 2: 调用 LLM 生成 SQL（流式）
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='thought',
|
||||
content='正在分析问题并生成 SQL...',
|
||||
)
|
||||
|
||||
from datetime import datetime
|
||||
from ai_platform.services.llm_service import LLMService
|
||||
|
||||
db_type = asyncio.get_event_loop().run_until_complete(
|
||||
self._get_db_type(db_connection)
|
||||
)
|
||||
|
||||
llm_service = LLMService()
|
||||
accumulated_content = ''
|
||||
total_tokens = 0
|
||||
llm_result = None
|
||||
use_function_calling = self.config.get('use_function_calling', True)
|
||||
|
||||
# 优先尝试 Function Calling
|
||||
if use_function_calling:
|
||||
try:
|
||||
llm_result, total_tokens = yield from self._call_llm_with_function_calling(
|
||||
llm_service, model_id, db_type, schema_context, user_question
|
||||
)
|
||||
logger.info(f'Function Calling 成功: {llm_result}')
|
||||
except Exception as fc_error:
|
||||
logger.warning(f'Function Calling 失败，回退到 JSON 解析模式: {fc_error}')
|
||||
llm_result = None
|
||||
|
||||
# Fallback: 使用原有的 JSON 解析方式
|
||||
if not llm_result:
|
||||
system_prompt = TEXT_TO_SQL_SYSTEM_PROMPT.format(
|
||||
db_type=db_type,
|
||||
current_date=datetime.now().strftime('%Y-%m-%d'),
|
||||
schema_context=schema_context,
|
||||
)
|
||||
|
||||
messages = [
|
||||
{'role': 'system', 'content': system_prompt},
|
||||
{'role': 'user', 'content': f'用户问题: {user_question}'},
|
||||
]
|
||||
|
||||
for chunk in llm_service.chat_stream_sync(
|
||||
model_id=model_id,
|
||||
messages=messages,
|
||||
temperature=0.1,
|
||||
max_tokens=2048,
|
||||
):
|
||||
accumulated_content += chunk.content
|
||||
|
||||
if chunk.is_finished and chunk.total_tokens > 0:
|
||||
total_tokens = chunk.total_tokens
|
||||
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='sql_chunk',
|
||||
content=chunk.content,
|
||||
is_finished=chunk.is_finished,
|
||||
)
|
||||
|
||||
# 解析 LLM 响应
|
||||
llm_result = self._parse_llm_response(accumulated_content)
|
||||
if not llm_result:
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='error',
|
||||
content='LLM 响应解析失败',
|
||||
is_finished=True,
|
||||
)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='LLM 响应解析失败',
|
||||
metadata={'raw_response': accumulated_content},
|
||||
)
|
||||
|
||||
sql = llm_result.get('sql', '')
|
||||
thought = llm_result.get('thought', '')
|
||||
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='sql_complete',
|
||||
content=sql,
|
||||
data={'thought': thought},
|
||||
)
|
||||
|
||||
# Step 3: 验证 SQL 安全性
|
||||
if not self._validate_sql(sql):
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='error',
|
||||
content='生成的 SQL 不安全或不是 SELECT 语句',
|
||||
is_finished=True,
|
||||
)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='生成的 SQL 不安全或不是 SELECT 语句',
|
||||
metadata={'sql': sql},
|
||||
)
|
||||
|
||||
# Step 4: 格式化 SQL
|
||||
sql = self._format_sql(sql)
|
||||
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
|
||||
# 构建输出
|
||||
result_data = {
|
||||
'sql': sql,
|
||||
'thought': thought,
|
||||
}
|
||||
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='result',
|
||||
content='SQL 生成完成',
|
||||
is_finished=True,
|
||||
data=result_data,
|
||||
)
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=result_data,
|
||||
output_variables={
|
||||
f'{output_variable}_sql': sql,
|
||||
f'{output_variable}_thought': thought,
|
||||
},
|
||||
tokens_used=total_tokens,
|
||||
elapsed_time=elapsed_time,
|
||||
metadata={
|
||||
'thought': thought,
|
||||
'suggested_next_node': 'db_sql',
|
||||
},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'Text-to-SQL 流式执行失败: {e}')
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='error',
|
||||
content=str(e),
|
||||
is_finished=True,
|
||||
)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
elapsed_time=int((time.time() - start_time) * 1000),
|
||||
)
|
||||
|
||||
async def _get_schema_context(
|
||||
self,
|
||||
db_connection: str,
|
||||
schema_name: str = 'public',
|
||||
selected_tables: List[str] = None,
|
||||
include_relations: bool = True,
|
||||
manual_relations: List[Dict[str, Any]] = None
|
||||
) -> str:
|
||||
"""获取数据库 Schema 上下文"""
|
||||
try:
|
||||
from core.database_manager.service import AsyncDatabaseManagerService
|
||||
|
||||
# 创建数据库服务实例
|
||||
db_service = await AsyncDatabaseManagerService.create(db_connection)
|
||||
schema_name = self._resolve_schema_name(db_service, schema_name)
|
||||
|
||||
# 获取表列表
|
||||
if selected_tables and len(selected_tables) > 0:
|
||||
tables = selected_tables
|
||||
else:
|
||||
# 获取所有表
|
||||
tables_info = await db_service.get_tables(schema_name=schema_name)
|
||||
tables = [
|
||||
t.get('table_name') or t.get('name')
|
||||
for t in tables_info[:20]
|
||||
if t.get('table_name') or t.get('name')
|
||||
] # 限制最多 20 个表
|
||||
|
||||
schema_parts = []
|
||||
table_relations = []
|
||||
|
||||
for table_name in tables:
|
||||
try:
|
||||
columns = await db_service.get_table_columns(table_name, schema_name)
|
||||
|
||||
col_desc = []
|
||||
for col in columns:
|
||||
col_name = col.get('column_name') or col.get('name', '')
|
||||
col_type = col.get('data_type') or col.get('type', '')
|
||||
col_str = f" - {col_name} ({col_type})"
|
||||
if col.get('description') or col.get('comment'):
|
||||
col_str += f" -- {col.get('description') or col.get('comment')}"
|
||||
if col.get('is_primary_key'):
|
||||
col_str += " [PK]"
|
||||
if col.get('is_foreign_key'):
|
||||
col_str += " [FK]"
|
||||
col_desc.append(col_str)
|
||||
|
||||
full_table_name = self._format_schema_table_name(
|
||||
db_service, schema_name, table_name
|
||||
)
|
||||
table_schema = f"表: {full_table_name}\n" + "\n".join(col_desc)
|
||||
schema_parts.append(table_schema)
|
||||
|
||||
# 获取外键关系（从数据库）
|
||||
if include_relations:
|
||||
try:
|
||||
constraints = await db_service.get_table_constraints(table_name, schema_name)
|
||||
for constraint in constraints:
|
||||
constraint_type = (
|
||||
constraint.get('constraint_type')
|
||||
or constraint.get('type', '')
|
||||
)
|
||||
if constraint_type != 'FOREIGN KEY':
|
||||
continue
|
||||
source_full_name = self._format_schema_table_name(
|
||||
db_service, schema_name, table_name
|
||||
)
|
||||
foreign_table = (
|
||||
constraint.get('referenced_table')
|
||||
or constraint.get('foreign_table')
|
||||
or ''
|
||||
)
|
||||
target_full_name = self._format_schema_table_name(
|
||||
db_service, schema_name, foreign_table
|
||||
)
|
||||
source_columns = (
|
||||
constraint.get('columns')
|
||||
or constraint.get('column')
|
||||
or ''
|
||||
)
|
||||
foreign_columns = (
|
||||
constraint.get('referenced_columns')
|
||||
or constraint.get('foreign_column')
|
||||
or ''
|
||||
)
|
||||
relation = (
|
||||
f" {source_full_name}.{source_columns} → "
|
||||
f"{target_full_name}.{foreign_columns}"
|
||||
)
|
||||
table_relations.append(relation)
|
||||
except Exception as e:
|
||||
logger.debug(f'获取表 {table_name} 外键关系失败: {e}')
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f'获取表 {table_name} 结构失败: {e}')
|
||||
|
||||
# 添加手动指定的表关系（逻辑外键）
|
||||
if manual_relations:
|
||||
for rel in manual_relations:
|
||||
source_table = rel.get('sourceTable', '')
|
||||
source_field = rel.get('sourceField', '')
|
||||
target_table = rel.get('targetTable', '')
|
||||
target_field = rel.get('targetField', '')
|
||||
if source_table and source_field and target_table and target_field:
|
||||
relation = f" {source_table}.{source_field} → {target_table}.{target_field} (逻辑外键)"
|
||||
table_relations.append(relation)
|
||||
|
||||
# 组装最终的 Schema 上下文
|
||||
result = "\n\n".join(schema_parts)
|
||||
|
||||
if table_relations:
|
||||
result += "\n\n## 表关系（外键）\n" + "\n".join(table_relations)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'获取 Schema 上下文失败: {e}')
|
||||
return ''
|
||||
|
||||
@staticmethod
|
||||
def _resolve_schema_name(db_service, schema_name: str) -> str:
|
||||
"""按连接类型解析有效 schema，避免非 PG 库误用 public。"""
|
||||
if db_service._uses_schema_layer():
|
||||
if not schema_name or (
|
||||
schema_name == 'public' and db_service.db_type != 'postgresql'
|
||||
):
|
||||
return db_service._default_schema()
|
||||
return schema_name
|
||||
return schema_name or ''
|
||||
|
||||
@staticmethod
|
||||
def _format_schema_table_name(db_service, schema_name: str, table_name: str) -> str:
|
||||
if not table_name:
|
||||
return table_name
|
||||
if db_service._uses_schema_layer() and schema_name:
|
||||
return f"{schema_name}.{table_name}"
|
||||
return table_name
|
||||
|
||||
async def _get_db_type(self, db_connection: str) -> str:
|
||||
"""获取数据库类型"""
|
||||
try:
|
||||
from core.database_manager.service import AsyncDatabaseManagerService
|
||||
db_service = await AsyncDatabaseManagerService.create(db_connection)
|
||||
return db_service.db_type.upper()
|
||||
except Exception:
|
||||
return 'PostgreSQL'
|
||||
|
||||
def _call_llm_with_function_calling(
|
||||
self,
|
||||
llm_service,
|
||||
model_id: str,
|
||||
db_type: str,
|
||||
schema_context: str,
|
||||
user_question: str
|
||||
) -> Generator[TextToSqlStreamEvent, None, Tuple[Dict[str, Any], int]]:
|
||||
"""
|
||||
使用 Function Calling 调用 LLM 生成 SQL
|
||||
|
||||
Args:
|
||||
llm_service: LLM 服务实例
|
||||
model_id: 模型 ID
|
||||
db_type: 数据库类型
|
||||
schema_context: Schema 上下文
|
||||
user_question: 用户问题
|
||||
|
||||
Yields:
|
||||
TextToSqlStreamEvent: 流式输出事件
|
||||
|
||||
Returns:
|
||||
tuple: (llm_result, total_tokens)
|
||||
"""
|
||||
from datetime import datetime
|
||||
|
||||
system_prompt = TEXT_TO_SQL_SYSTEM_PROMPT_FC.format(
|
||||
db_type=db_type,
|
||||
current_date=datetime.now().strftime('%Y-%m-%d'),
|
||||
schema_context=schema_context,
|
||||
)
|
||||
|
||||
messages = [
|
||||
{'role': 'system', 'content': system_prompt},
|
||||
{'role': 'user', 'content': f'用户问题: {user_question}'},
|
||||
]
|
||||
|
||||
# 定义工具
|
||||
tools = [TEXT_TO_SQL_TOOL]
|
||||
|
||||
# 累积工具调用参数
|
||||
accumulated_arguments = ''
|
||||
total_tokens = 0
|
||||
tool_call_received = False
|
||||
|
||||
for chunk in llm_service.chat_stream_sync(
|
||||
model_id=model_id,
|
||||
messages=messages,
|
||||
temperature=0.1,
|
||||
max_tokens=2048,
|
||||
tools=tools,
|
||||
tool_choice='required', # 强制使用工具
|
||||
):
|
||||
# 处理工具调用增量
|
||||
if chunk.tool_call_delta:
|
||||
delta_args = chunk.tool_call_delta.get('arguments', '')
|
||||
if delta_args:
|
||||
accumulated_arguments += delta_args
|
||||
# 流式输出工具调用参数
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='sql_chunk',
|
||||
content=delta_args,
|
||||
is_finished=False,
|
||||
)
|
||||
|
||||
# 处理完整的工具调用
|
||||
if chunk.tool_calls:
|
||||
tool_call_received = True
|
||||
for tool_call in chunk.tool_calls:
|
||||
if tool_call.name == 'generate_sql':
|
||||
# 解析工具调用参数
|
||||
arguments = tool_call.arguments
|
||||
if isinstance(arguments, str):
|
||||
arguments = json.loads(arguments)
|
||||
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='sql_chunk',
|
||||
content='',
|
||||
is_finished=True,
|
||||
)
|
||||
|
||||
return (arguments, total_tokens)
|
||||
|
||||
# 记录 token 使用
|
||||
if chunk.is_finished and chunk.total_tokens > 0:
|
||||
total_tokens = chunk.total_tokens
|
||||
|
||||
# 如果没有收到工具调用，尝试从累积的参数中解析
|
||||
if accumulated_arguments and not tool_call_received:
|
||||
try:
|
||||
arguments = json.loads(accumulated_arguments)
|
||||
yield TextToSqlStreamEvent(
|
||||
event_type='sql_chunk',
|
||||
content='',
|
||||
is_finished=True,
|
||||
)
|
||||
return (arguments, total_tokens)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Function Calling 失败
|
||||
raise ValueError('Function Calling 未返回有效的工具调用')
|
||||
|
||||
def _parse_llm_response(self, content: str) -> Optional[Dict[str, Any]]:
|
||||
"""解析 LLM 响应（Fallback 方式）"""
|
||||
try:
|
||||
# 尝试直接解析 JSON
|
||||
return json.loads(content)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# 尝试从 markdown 代码块中提取 JSON
|
||||
import re
|
||||
json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', content)
|
||||
if json_match:
|
||||
try:
|
||||
return json.loads(json_match.group(1))
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# 尝试提取 {...} 部分
|
||||
brace_match = re.search(r'\{[\s\S]*\}', content)
|
||||
if brace_match:
|
||||
try:
|
||||
return json.loads(brace_match.group(0))
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
logger.warning(f'无法解析 LLM 响应: {content[:500]}')
|
||||
return None
|
||||
|
||||
def _validate_sql(self, sql: str) -> bool:
|
||||
"""验证 SQL 安全性"""
|
||||
if not sql:
|
||||
logger.warning("SQL 验证失败: SQL 为空")
|
||||
return False
|
||||
|
||||
# 去除前后空白和可能的 Markdown 代码块标记
|
||||
sql_cleaned = sql.strip()
|
||||
if sql_cleaned.startswith('```'):
|
||||
# 移除 Markdown 代码块
|
||||
lines = sql_cleaned.split('\n')
|
||||
# 移除第一行（```sql 或 ```）
|
||||
if lines:
|
||||
lines = lines[1:]
|
||||
# 移除最后一行（```）
|
||||
if lines and lines[-1].strip() == '```':
|
||||
lines = lines[:-1]
|
||||
sql_cleaned = '\n'.join(lines).strip()
|
||||
|
||||
sql_upper = sql_cleaned.upper().strip()
|
||||
|
||||
# 只允许 SELECT 语句（也允许 WITH ... SELECT 即 CTE）
|
||||
if not sql_upper.startswith('SELECT') and not sql_upper.startswith('WITH'):
|
||||
logger.warning(f"SQL 验证失败: 不是 SELECT/WITH 语句, SQL 开头: {sql_upper[:50]}")
|
||||
return False
|
||||
|
||||
# 禁止危险关键字（作为独立语句，不在子查询或 CTE 中）
|
||||
dangerous_keywords = [
|
||||
'INSERT', 'UPDATE', 'DELETE', 'DROP', 'TRUNCATE',
|
||||
'ALTER', 'CREATE', 'GRANT', 'REVOKE', 'EXEC', 'EXECUTE',
|
||||
]
|
||||
|
||||
for keyword in dangerous_keywords:
|
||||
# 检查是否作为独立关键字出现（前后有空格或在开头/结尾）
|
||||
if re.search(rf'\b{keyword}\b', sql_upper):
|
||||
logger.warning(f"SQL 验证失败: 包含危险关键字 {keyword}, SQL: {sql_cleaned[:200]}")
|
||||
return False
|
||||
|
||||
# 禁止多语句（分号后还有内容）
|
||||
if ';' in sql_cleaned:
|
||||
parts = sql_cleaned.split(';')
|
||||
if any(p.strip() for p in parts[1:]):
|
||||
logger.warning(f"SQL 验证失败: 包含多条语句")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def _format_sql(self, sql: str) -> str:
|
||||
"""
|
||||
格式化 SQL 语句
|
||||
|
||||
Args:
|
||||
sql: 原始 SQL 语句
|
||||
|
||||
Returns:
|
||||
格式化后的 SQL 语句
|
||||
"""
|
||||
try:
|
||||
# 去除可能的 Markdown 代码块标记
|
||||
sql_cleaned = sql.strip()
|
||||
if sql_cleaned.startswith('```'):
|
||||
lines = sql_cleaned.split('\n')
|
||||
if lines:
|
||||
lines = lines[1:]
|
||||
if lines and lines[-1].strip() == '```':
|
||||
lines = lines[:-1]
|
||||
sql_cleaned = '\n'.join(lines).strip()
|
||||
|
||||
# 使用 sqlparse 格式化
|
||||
formatted_sql = sqlparse.format(
|
||||
sql_cleaned,
|
||||
reindent=True, # 重新缩进
|
||||
keyword_case='upper', # 关键字大写
|
||||
identifier_case='lower', # 标识符小写
|
||||
strip_comments=False, # 保留注释
|
||||
use_space_around_operators=True, # 操作符周围加空格
|
||||
)
|
||||
|
||||
return formatted_sql.strip()
|
||||
except Exception as e:
|
||||
logger.warning(f"SQL 格式化失败: {e}, 返回原始 SQL")
|
||||
return sql
|
||||
@@ -0,0 +1,151 @@
|
||||
"""
|
||||
变量赋值节点
|
||||
"""
|
||||
import logging
|
||||
from typing import Any, Dict
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class VariableNode(BaseNode):
|
||||
"""
|
||||
变量赋值节点
|
||||
|
||||
设置或修改变量的值
|
||||
"""
|
||||
|
||||
node_type = 'variable'
|
||||
node_name = '变量'
|
||||
node_category = 'data'
|
||||
node_icon = 'variable'
|
||||
node_description = '设置或修改变量的值'
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'value',
|
||||
'type': 'any',
|
||||
'description': '要设置的值',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'value',
|
||||
'type': 'any',
|
||||
'description': '设置后的值',
|
||||
},
|
||||
]
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""执行变量赋值"""
|
||||
try:
|
||||
assignments = self.config.get('assignments', [])
|
||||
output_variables = {}
|
||||
|
||||
for assignment in assignments:
|
||||
var_name = assignment.get('name', '')
|
||||
value_type = assignment.get('type', 'static')
|
||||
value = assignment.get('value', '')
|
||||
|
||||
if not var_name:
|
||||
continue
|
||||
|
||||
# 根据类型处理值
|
||||
if value_type == 'static':
|
||||
# 静态值
|
||||
final_value = value
|
||||
elif value_type == 'variable':
|
||||
# 从其他变量获取
|
||||
final_value = context.get_variable(value, '')
|
||||
elif value_type == 'template':
|
||||
# 模板渲染
|
||||
final_value = context.resolve_template(value)
|
||||
elif value_type == 'json':
|
||||
# JSON 解析
|
||||
import json
|
||||
final_value = json.loads(value)
|
||||
elif value_type == 'expression':
|
||||
# 简单表达式（仅支持基本运算）
|
||||
final_value = self._evaluate_expression(value, context)
|
||||
else:
|
||||
final_value = value
|
||||
|
||||
# 设置变量
|
||||
context.set_variable(var_name, final_value)
|
||||
output_variables[var_name] = final_value
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output=output_variables,
|
||||
output_variables=output_variables,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'变量节点执行失败: {e}')
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
)
|
||||
|
||||
def _evaluate_expression(self, expression: str, context: NodeContext) -> Any:
|
||||
"""
|
||||
评估简单表达式
|
||||
|
||||
仅支持基本的数学运算和字符串操作
|
||||
"""
|
||||
# 替换变量引用
|
||||
resolved = context.resolve_template(expression)
|
||||
|
||||
# 安全的评估环境
|
||||
safe_dict = {
|
||||
'abs': abs,
|
||||
'int': int,
|
||||
'float': float,
|
||||
'str': str,
|
||||
'len': len,
|
||||
'min': min,
|
||||
'max': max,
|
||||
'sum': sum,
|
||||
'round': round,
|
||||
}
|
||||
|
||||
try:
|
||||
return eval(resolved, {"__builtins__": {}}, safe_dict)
|
||||
except Exception:
|
||||
return resolved
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'assignments': {
|
||||
'type': 'array',
|
||||
'title': '变量赋值',
|
||||
'items': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'name': {
|
||||
'type': 'string',
|
||||
'title': '变量名',
|
||||
},
|
||||
'type': {
|
||||
'type': 'string',
|
||||
'title': '值类型',
|
||||
'enum': ['static', 'variable', 'template', 'json', 'expression'],
|
||||
'default': 'static',
|
||||
},
|
||||
'value': {
|
||||
'type': 'string',
|
||||
'title': '值',
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,155 @@
|
||||
"""
|
||||
节点注册中心
|
||||
"""
|
||||
import logging
|
||||
from typing import Dict, List, Optional, Type
|
||||
|
||||
from .base import BaseNode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class NodeRegistry:
|
||||
"""
|
||||
节点注册中心
|
||||
|
||||
管理所有工作流节点的注册和获取
|
||||
"""
|
||||
|
||||
_nodes: Dict[str, Type[BaseNode]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, node_class: Type[BaseNode]) -> Type[BaseNode]:
|
||||
"""
|
||||
注册节点（可作为装饰器使用）
|
||||
|
||||
Args:
|
||||
node_class: 节点类
|
||||
|
||||
Returns:
|
||||
节点类
|
||||
"""
|
||||
node_type = node_class.node_type
|
||||
if not node_type:
|
||||
raise ValueError(f'Node {node_class.__name__} must have a node_type')
|
||||
|
||||
cls._nodes[node_type] = node_class
|
||||
logger.info(f'Registered AI workflow node: {node_type}')
|
||||
return node_class
|
||||
|
||||
@classmethod
|
||||
def get(cls, node_type: str) -> Optional[Type[BaseNode]]:
|
||||
"""
|
||||
获取节点类
|
||||
|
||||
Args:
|
||||
node_type: 节点类型
|
||||
|
||||
Returns:
|
||||
节点类或 None
|
||||
"""
|
||||
return cls._nodes.get(node_type)
|
||||
|
||||
@classmethod
|
||||
def create_instance(cls, node_type: str, config: Dict = None) -> Optional[BaseNode]:
|
||||
"""
|
||||
创建节点实例
|
||||
|
||||
Args:
|
||||
node_type: 节点类型
|
||||
config: 节点配置
|
||||
|
||||
Returns:
|
||||
节点实例或 None
|
||||
"""
|
||||
node_class = cls.get(node_type)
|
||||
if not node_class:
|
||||
logger.warning(f'Unknown node type: {node_type}')
|
||||
return None
|
||||
|
||||
return node_class(config=config)
|
||||
|
||||
@classmethod
|
||||
def get_all_schemas(cls) -> List[Dict]:
|
||||
"""
|
||||
获取所有节点 Schema
|
||||
|
||||
Returns:
|
||||
节点 Schema 列表
|
||||
"""
|
||||
return [node.get_schema() for node in cls._nodes.values()]
|
||||
|
||||
@classmethod
|
||||
def get_schemas_by_category(cls) -> Dict[str, List[Dict]]:
|
||||
"""
|
||||
按分类获取节点 Schema
|
||||
|
||||
Returns:
|
||||
按分类分组的节点 Schema
|
||||
"""
|
||||
result = {}
|
||||
for node in cls._nodes.values():
|
||||
category = node.node_category
|
||||
if category not in result:
|
||||
result[category] = []
|
||||
result[category].append(node.get_schema())
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def get_all_types(cls) -> List[str]:
|
||||
"""
|
||||
获取所有已注册的节点类型
|
||||
|
||||
Returns:
|
||||
节点类型列表
|
||||
"""
|
||||
return list(cls._nodes.keys())
|
||||
|
||||
|
||||
# 自动加载所有内置节点
|
||||
def _load_builtin_nodes():
|
||||
"""加载所有内置节点"""
|
||||
required_modules = [
|
||||
'start_node',
|
||||
'end_node',
|
||||
'condition_node',
|
||||
'template_node',
|
||||
'parallel_node',
|
||||
'merge_node',
|
||||
]
|
||||
optional_modules = [
|
||||
'llm_node',
|
||||
'code_node',
|
||||
'http_node',
|
||||
'variable_node',
|
||||
'database_node',
|
||||
'dialog_nodes',
|
||||
'intent_node',
|
||||
'loop_node',
|
||||
'subflow_node',
|
||||
'text_to_sql_node',
|
||||
'snowflake_cortex_node',
|
||||
'knowledge_retrieval_node',
|
||||
]
|
||||
|
||||
from importlib import import_module
|
||||
|
||||
for module_name in required_modules:
|
||||
import_module(f'{__package__}.builtin.{module_name}')
|
||||
|
||||
for module_name in optional_modules:
|
||||
try:
|
||||
import_module(f'{__package__}.builtin.{module_name}')
|
||||
except ImportError as e:
|
||||
logger.warning(
|
||||
'Skipped optional AI workflow node module %s because dependencies are unavailable: %s',
|
||||
module_name,
|
||||
e,
|
||||
)
|
||||
|
||||
|
||||
# 延迟加载
|
||||
try:
|
||||
_load_builtin_nodes()
|
||||
except ImportError as e:
|
||||
logger.warning(f'Failed to load some builtin nodes: {e}')
|
||||
@@ -0,0 +1 @@
|
||||
"""AI workflow node utilities."""
|
||||
@@ -0,0 +1,66 @@
|
||||
"""
|
||||
节点配置解析工具
|
||||
|
||||
从节点配置或上下文变量中解析 object / list，避免 dict 被 stringify 后无法反序列化。
|
||||
"""
|
||||
import ast
|
||||
import json
|
||||
import re
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from ai_platform.nodes.base import NodeContext
|
||||
|
||||
_VAR_REF_PATTERN = re.compile(r'^\{\{\s*([^}]+)\s*\}\}$')
|
||||
|
||||
|
||||
def _parse_structured_string(value: str) -> Any:
|
||||
text = value.strip()
|
||||
if not text:
|
||||
return None
|
||||
try:
|
||||
return json.loads(text)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
try:
|
||||
return ast.literal_eval(text)
|
||||
except (ValueError, SyntaxError):
|
||||
pass
|
||||
return value
|
||||
|
||||
|
||||
def resolve_config_value(context: NodeContext, raw: Any) -> Any:
|
||||
"""解析配置值：支持 dict/list 直传、{{var}} 引用、JSON/Python 字面量字符串。"""
|
||||
if raw is None or raw == '':
|
||||
return None
|
||||
|
||||
if isinstance(raw, (dict, list)):
|
||||
return raw
|
||||
|
||||
if not isinstance(raw, str):
|
||||
return raw
|
||||
|
||||
stripped = raw.strip()
|
||||
var_match = _VAR_REF_PATTERN.match(stripped)
|
||||
if var_match:
|
||||
var_name = var_match.group(1).strip()
|
||||
if var_name in context.variables:
|
||||
return context.variables[var_name]
|
||||
|
||||
resolved = context.resolve_template(raw)
|
||||
if isinstance(resolved, (dict, list)):
|
||||
return resolved
|
||||
|
||||
if isinstance(resolved, str):
|
||||
return _parse_structured_string(resolved)
|
||||
|
||||
return resolved
|
||||
|
||||
|
||||
def resolve_object_config(context: NodeContext, raw: Any) -> Optional[dict]:
|
||||
value = resolve_config_value(context, raw)
|
||||
return value if isinstance(value, dict) else None
|
||||
|
||||
|
||||
def resolve_list_config(context: NodeContext, raw: Any) -> Optional[List[Any]]:
|
||||
value = resolve_config_value(context, raw)
|
||||
return value if isinstance(value, list) else None
|
||||
@@ -0,0 +1,361 @@
|
||||
"""
|
||||
工作流数据库节点：连接解析与方言 SQL 构建
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from datetime import date, datetime
|
||||
from decimal import Decimal
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from uuid import UUID
|
||||
|
||||
from core.database_manager.sql_utils import quote_identifier, quote_table
|
||||
|
||||
DEFAULT_CONNECTION_CODE = "default"
|
||||
DEFAULT_CONNECTION_WRITE_WARNING = "default_connection_write"
|
||||
|
||||
OPERATOR_MAP = {
|
||||
"=": "=",
|
||||
"!=": "!=",
|
||||
">": ">",
|
||||
">=": ">=",
|
||||
"<": "<",
|
||||
"<=": "<=",
|
||||
"like": "LIKE",
|
||||
"in": "IN",
|
||||
"is_null": "IS NULL",
|
||||
"is_not_null": "IS NOT NULL",
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class DbTarget:
|
||||
"""工作流节点数据库目标"""
|
||||
|
||||
db_name: str = DEFAULT_CONNECTION_CODE
|
||||
db_type: str = "postgresql"
|
||||
database: str = ""
|
||||
schema: str = ""
|
||||
|
||||
@property
|
||||
def is_external(self) -> bool:
|
||||
return is_external_connection(self.db_name)
|
||||
|
||||
|
||||
def is_external_connection(db_name: Optional[str]) -> bool:
|
||||
code = (db_name or DEFAULT_CONNECTION_CODE).strip() or DEFAULT_CONNECTION_CODE
|
||||
return code != DEFAULT_CONNECTION_CODE
|
||||
|
||||
|
||||
def resolve_db_target(config: Optional[Dict[str, Any]]) -> DbTarget:
|
||||
"""从节点 db_config 解析连接目标"""
|
||||
db_config = config or {}
|
||||
db_name = (db_config.get("dbName") or DEFAULT_CONNECTION_CODE).strip() or DEFAULT_CONNECTION_CODE
|
||||
db_type = (db_config.get("dbType") or "postgresql").lower()
|
||||
database = (db_config.get("database") or "").strip()
|
||||
schema = (db_config.get("schema") or "").strip()
|
||||
return DbTarget(
|
||||
db_name=db_name,
|
||||
db_type=db_type,
|
||||
database=database,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
|
||||
def default_connection_write_warnings(operation: str, target: DbTarget) -> List[str]:
|
||||
"""default 连接写操作返回运行时警告码"""
|
||||
if target.is_external:
|
||||
return []
|
||||
write_ops = {"insert", "update", "delete", "upsert", "execute"}
|
||||
if operation.lower() in write_ops:
|
||||
return [DEFAULT_CONNECTION_WRITE_WARNING]
|
||||
return []
|
||||
|
||||
|
||||
def resolve_schema_for_handler(db_type: str, schema: str, default_schema: str = "") -> str:
|
||||
"""解析 handler 使用的 schema/database 参数"""
|
||||
db = (db_type or "postgresql").lower()
|
||||
if db == "mysql":
|
||||
return schema or default_schema or ""
|
||||
if not schema or (schema == "public" and db != "postgresql"):
|
||||
return default_schema or schema or ""
|
||||
return schema
|
||||
|
||||
|
||||
async def resolve_handler_schema_name(db_service, target: DbTarget) -> str:
|
||||
"""结合 AsyncDatabaseManagerService 解析 schema"""
|
||||
db_type = (db_service.db_type or "postgresql").lower()
|
||||
if db_type == "mysql":
|
||||
default_schema = db_service._default_schema() if hasattr(db_service, "_default_schema") else ""
|
||||
return target.database or target.schema or default_schema or ""
|
||||
default_schema = ""
|
||||
if hasattr(db_service, "_default_schema"):
|
||||
default_schema = db_service._default_schema() or ""
|
||||
return resolve_schema_for_handler(db_type, target.schema, default_schema)
|
||||
|
||||
|
||||
def format_sql_literal(value: Any, db_type: str) -> str:
|
||||
"""将 Python 值格式化为 SQL 字面量（用于 handler raw WHERE）"""
|
||||
if value is None:
|
||||
return "NULL"
|
||||
if isinstance(value, bool):
|
||||
if (db_type or "").lower() == "postgresql":
|
||||
return "TRUE" if value else "FALSE"
|
||||
return "1" if value else "0"
|
||||
if isinstance(value, (int, float, Decimal)):
|
||||
return str(value)
|
||||
if isinstance(value, (datetime, date)):
|
||||
return f"'{value.isoformat()}'"
|
||||
if isinstance(value, UUID):
|
||||
return f"'{value}'"
|
||||
if isinstance(value, (list, tuple)):
|
||||
inner = ", ".join(format_sql_literal(v, db_type) for v in value)
|
||||
return f"({inner})"
|
||||
if isinstance(value, (dict, list)):
|
||||
encoded = json.dumps(value, ensure_ascii=False)
|
||||
escaped = encoded.replace("'", "''")
|
||||
return f"'{escaped}'"
|
||||
escaped = str(value).replace("'", "''")
|
||||
return f"'{escaped}'"
|
||||
|
||||
|
||||
def build_where_clause_raw(
|
||||
conditions: List[Dict[str, Any]],
|
||||
db_type: str,
|
||||
) -> str:
|
||||
"""
|
||||
构建 raw WHERE 子句（不含 WHERE 关键字），供 database_manager handler 使用。
|
||||
"""
|
||||
if not conditions:
|
||||
return ""
|
||||
|
||||
clauses: List[str] = []
|
||||
for condition in conditions:
|
||||
field = condition.get("field", "")
|
||||
op = (condition.get("operator") or "=").lower()
|
||||
value = condition.get("value")
|
||||
sql_op = OPERATOR_MAP.get(op, "=")
|
||||
quoted_field = quote_identifier(field, db_type)
|
||||
|
||||
if op in ("is_null", "is_not_null"):
|
||||
clauses.append(f"{quoted_field} {sql_op}")
|
||||
elif op == "in":
|
||||
if isinstance(value, list):
|
||||
literals = ", ".join(format_sql_literal(v, db_type) for v in value)
|
||||
clauses.append(f"{quoted_field} IN ({literals})")
|
||||
else:
|
||||
clauses.append(f"{quoted_field} IN ({format_sql_literal(value, db_type)})")
|
||||
elif op == "like":
|
||||
clauses.append(f"{quoted_field} LIKE {format_sql_literal(value, db_type)}")
|
||||
else:
|
||||
clauses.append(f"{quoted_field} {sql_op} {format_sql_literal(value, db_type)}")
|
||||
|
||||
return " AND ".join(clauses)
|
||||
|
||||
|
||||
def build_where_clause_platform(
|
||||
conditions: List[Dict[str, Any]],
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
"""构建平台 PostgreSQL WHERE 子句（SQLAlchemy 命名参数）"""
|
||||
if not conditions:
|
||||
return "", {}
|
||||
|
||||
clauses: List[str] = []
|
||||
params: Dict[str, Any] = {}
|
||||
param_index = 0
|
||||
|
||||
for condition in conditions:
|
||||
field = condition["field"]
|
||||
op = condition["operator"].lower()
|
||||
value = condition["value"]
|
||||
sql_op = OPERATOR_MAP.get(op, "=")
|
||||
|
||||
if op in ("is_null", "is_not_null"):
|
||||
clauses.append(f'"{field}" {sql_op}')
|
||||
elif op == "in":
|
||||
if isinstance(value, list):
|
||||
param_placeholders = []
|
||||
for v in value:
|
||||
param_name = f"p{param_index}"
|
||||
param_placeholders.append(f":{param_name}")
|
||||
params[param_name] = v
|
||||
param_index += 1
|
||||
clauses.append(f'"{field}" IN ({", ".join(param_placeholders)})')
|
||||
else:
|
||||
param_name = f"p{param_index}"
|
||||
clauses.append(f'"{field}" {sql_op} :{param_name}')
|
||||
params[param_name] = value
|
||||
param_index += 1
|
||||
else:
|
||||
param_name = f"p{param_index}"
|
||||
clauses.append(f'"{field}" {sql_op} :{param_name}')
|
||||
params[param_name] = value
|
||||
param_index += 1
|
||||
|
||||
where_clause = " AND ".join(clauses)
|
||||
return f"WHERE {where_clause}", params
|
||||
|
||||
|
||||
def quote_table_for_target(table: str, target: DbTarget) -> str:
|
||||
"""按目标连接类型构建完整表名"""
|
||||
schema = target.schema or None
|
||||
if target.db_type == "mysql" and not schema and target.database:
|
||||
schema = target.database
|
||||
return quote_table(schema, table, target.db_type)
|
||||
|
||||
|
||||
def format_select_sql(
|
||||
table: str,
|
||||
target: DbTarget,
|
||||
return_fields: Any = None,
|
||||
conditions: Optional[List[Dict[str, Any]]] = None,
|
||||
order_by: str = "",
|
||||
limit: int = 100,
|
||||
) -> str:
|
||||
"""按 dialect 生成 SELECT SQL"""
|
||||
normalized_fields = normalize_return_fields(return_fields)
|
||||
if normalized_fields == "*":
|
||||
field_list = "*"
|
||||
else:
|
||||
field_list = ", ".join(
|
||||
quote_identifier(f, target.db_type) for f in normalized_fields
|
||||
)
|
||||
|
||||
full_table = quote_table_for_target(table, target)
|
||||
where_raw = build_where_clause_raw(conditions or [], target.db_type)
|
||||
where_part = f" WHERE {where_raw}" if where_raw else ""
|
||||
|
||||
sql = f"SELECT {field_list} FROM {full_table}{where_part}"
|
||||
|
||||
db = (target.db_type or "postgresql").lower()
|
||||
if order_by:
|
||||
sql += f" ORDER BY {order_by}"
|
||||
elif db == "sqlserver":
|
||||
sql += " ORDER BY (SELECT NULL)"
|
||||
|
||||
if db == "sqlserver":
|
||||
sql += f" OFFSET 0 ROWS FETCH NEXT {int(limit)} ROWS ONLY"
|
||||
elif db == "oracle":
|
||||
sql += f" FETCH FIRST {int(limit)} ROWS ONLY"
|
||||
else:
|
||||
sql += f" LIMIT {int(limit)}"
|
||||
|
||||
return sql
|
||||
|
||||
|
||||
def format_limit_clause(db_type: str, limit: int = 1) -> str:
|
||||
db = (db_type or "postgresql").lower()
|
||||
if db == "sqlserver":
|
||||
return f" ORDER BY (SELECT NULL) OFFSET 0 ROWS FETCH NEXT {int(limit)} ROWS ONLY"
|
||||
if db == "oracle":
|
||||
return f" FETCH FIRST {int(limit)} ROWS ONLY"
|
||||
return f" LIMIT {int(limit)}"
|
||||
|
||||
|
||||
def normalize_return_fields(return_fields: Any) -> Any:
|
||||
"""将 return_fields 规范化为 * 或字段名列表"""
|
||||
if return_fields in (None, "", "*", ["*"]):
|
||||
return "*"
|
||||
if isinstance(return_fields, str):
|
||||
stripped = return_fields.strip()
|
||||
if not stripped or stripped == "*":
|
||||
return "*"
|
||||
return [field.strip() for field in stripped.split(",") if field.strip()]
|
||||
if isinstance(return_fields, list):
|
||||
if not return_fields or "*" in return_fields:
|
||||
return "*"
|
||||
return return_fields
|
||||
return "*"
|
||||
|
||||
|
||||
def merge_result_metadata(
|
||||
base_metadata: Optional[Dict[str, Any]],
|
||||
warnings: Optional[List[str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
metadata = dict(base_metadata or {})
|
||||
if warnings:
|
||||
existing = list(metadata.get("warnings") or [])
|
||||
for code in warnings:
|
||||
if code not in existing:
|
||||
existing.append(code)
|
||||
metadata["warnings"] = existing
|
||||
return metadata
|
||||
|
||||
|
||||
def convert_sql_param_type(value: Any, param_type: str) -> Any:
|
||||
"""将参数值转换为指定 SQL 参数类型（与数据源 _convert_param_type 对齐）"""
|
||||
if value is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
if param_type == "integer":
|
||||
return int(value)
|
||||
if param_type == "float":
|
||||
return float(value)
|
||||
if param_type == "boolean":
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
return str(value).lower() in ("true", "1", "yes")
|
||||
if param_type == "date":
|
||||
if isinstance(value, date):
|
||||
return value
|
||||
if isinstance(value, datetime):
|
||||
return value.date()
|
||||
return date.fromisoformat(str(value).strip()[:10])
|
||||
if param_type == "datetime":
|
||||
if isinstance(value, datetime):
|
||||
return value
|
||||
return datetime.fromisoformat(str(value).strip())
|
||||
return str(value)
|
||||
except (ValueError, TypeError):
|
||||
return value
|
||||
|
||||
|
||||
def resolve_sql_param_value(raw: Any, context: Any) -> Any:
|
||||
"""解析单个 SQL 参数值：模板变量 + JSON 字面量"""
|
||||
if raw is None:
|
||||
return None
|
||||
if isinstance(raw, str):
|
||||
resolved = context.resolve_template(raw)
|
||||
if resolved == "":
|
||||
return None
|
||||
try:
|
||||
return json.loads(resolved)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return resolved
|
||||
return raw
|
||||
|
||||
|
||||
def build_sql_param_dict(
|
||||
param_defs: Optional[List[Dict[str, Any]]],
|
||||
context: Any,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
从节点 params 配置构建 SQLAlchemy 命名参数字典。
|
||||
|
||||
参数定义字段：name, type, value（或 default）, required
|
||||
SQL 中使用 :name 占位符，与数据源一致。
|
||||
"""
|
||||
result: Dict[str, Any] = {}
|
||||
for param in param_defs or []:
|
||||
name = (param.get("name") or "").strip()
|
||||
if not name:
|
||||
continue
|
||||
|
||||
param_type = param.get("type") or "string"
|
||||
required = bool(param.get("required", False))
|
||||
raw = param.get("value")
|
||||
if raw is None or raw == "":
|
||||
raw = param.get("default")
|
||||
|
||||
if raw is None or raw == "":
|
||||
if required:
|
||||
raise ValueError(f"缺少必填 SQL 参数: {name}")
|
||||
result[name] = None
|
||||
continue
|
||||
|
||||
value = resolve_sql_param_value(raw, context)
|
||||
result[name] = convert_sql_param_type(value, param_type)
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,177 @@
|
||||
"""Tests for workflow database execution utilities."""
|
||||
import unittest
|
||||
|
||||
from ai_platform.nodes.utils.db_execution import (
|
||||
DbTarget,
|
||||
build_sql_param_dict,
|
||||
build_where_clause_platform,
|
||||
build_where_clause_raw,
|
||||
convert_sql_param_type,
|
||||
default_connection_write_warnings,
|
||||
format_select_sql,
|
||||
format_limit_clause,
|
||||
merge_result_metadata,
|
||||
normalize_return_fields,
|
||||
resolve_db_target,
|
||||
resolve_handler_schema_name,
|
||||
resolve_schema_for_handler,
|
||||
resolve_sql_param_value,
|
||||
)
|
||||
|
||||
|
||||
class DbExecutionTestCase(unittest.TestCase):
|
||||
def test_resolve_db_target_default(self):
|
||||
target = resolve_db_target({})
|
||||
self.assertEqual(target.db_name, "default")
|
||||
self.assertFalse(target.is_external)
|
||||
|
||||
def test_resolve_db_target_external(self):
|
||||
target = resolve_db_target(
|
||||
{
|
||||
"dbName": "erp_mysql",
|
||||
"dbType": "mysql",
|
||||
"database": "sales",
|
||||
"schema": "",
|
||||
}
|
||||
)
|
||||
self.assertEqual(target.db_name, "erp_mysql")
|
||||
self.assertTrue(target.is_external)
|
||||
|
||||
def test_default_connection_write_warnings(self):
|
||||
target = resolve_db_target({"dbName": "default"})
|
||||
self.assertEqual(
|
||||
default_connection_write_warnings("insert", target),
|
||||
["default_connection_write"],
|
||||
)
|
||||
self.assertEqual(default_connection_write_warnings("select", target), [])
|
||||
|
||||
def test_build_where_clause_raw_postgresql(self):
|
||||
where = build_where_clause_raw(
|
||||
[
|
||||
{"field": "status", "operator": "=", "value": "active"},
|
||||
{"field": "age", "operator": ">", "value": 18},
|
||||
],
|
||||
"postgresql",
|
||||
)
|
||||
self.assertIn('"status" = \'active\'', where)
|
||||
self.assertIn('"age" > 18', where)
|
||||
|
||||
def test_build_where_clause_raw_mysql_like(self):
|
||||
where = build_where_clause_raw(
|
||||
[{"field": "name", "operator": "like", "value": "张%"}],
|
||||
"mysql",
|
||||
)
|
||||
self.assertIn("`name` LIKE '张%'", where)
|
||||
|
||||
def test_build_where_clause_platform_named_params(self):
|
||||
clause, params = build_where_clause_platform(
|
||||
[{"field": "id", "operator": "=", "value": "1"}]
|
||||
)
|
||||
self.assertTrue(clause.startswith("WHERE"))
|
||||
self.assertEqual(params["p0"], "1")
|
||||
|
||||
def test_format_select_sql_postgresql(self):
|
||||
target = DbTarget(db_name="erp", db_type="postgresql", schema="public")
|
||||
sql = format_select_sql(
|
||||
"users",
|
||||
target,
|
||||
return_fields=["id", "name"],
|
||||
conditions=[{"field": "status", "operator": "=", "value": 1}],
|
||||
limit=10,
|
||||
)
|
||||
self.assertTrue(sql.startswith("SELECT"))
|
||||
self.assertIn('"public"."users"', sql)
|
||||
self.assertIn("LIMIT 10", sql)
|
||||
|
||||
def test_resolve_schema_for_handler_mysql(self):
|
||||
self.assertEqual(resolve_schema_for_handler("mysql", "", "sales_db"), "sales_db")
|
||||
|
||||
def test_merge_result_metadata(self):
|
||||
metadata = merge_result_metadata({}, ["default_connection_write"])
|
||||
self.assertEqual(metadata["warnings"], ["default_connection_write"])
|
||||
|
||||
def test_normalize_return_fields_comma_string(self):
|
||||
self.assertEqual(normalize_return_fields("id, name"), ["id", "name"])
|
||||
self.assertEqual(normalize_return_fields("*"), "*")
|
||||
|
||||
def test_format_select_sql_sqlserver_order_by(self):
|
||||
target = DbTarget(db_name="erp", db_type="sqlserver", schema="dbo")
|
||||
sql = format_select_sql("users", target, limit=10)
|
||||
self.assertIn("ORDER BY (SELECT NULL)", sql)
|
||||
self.assertIn("OFFSET 0 ROWS FETCH NEXT 10 ROWS ONLY", sql)
|
||||
|
||||
def test_format_limit_clause_sqlserver(self):
|
||||
clause = format_limit_clause("sqlserver", 1)
|
||||
self.assertIn("ORDER BY (SELECT NULL)", clause)
|
||||
self.assertIn("OFFSET 0 ROWS FETCH NEXT 1 ROWS ONLY", clause)
|
||||
|
||||
def test_resolve_handler_schema_mysql_prefers_database(self):
|
||||
class FakeService:
|
||||
db_type = "mysql"
|
||||
|
||||
def _default_schema(self):
|
||||
return "fallback"
|
||||
|
||||
target = DbTarget(db_name="erp", db_type="mysql", database="sales_db", schema="")
|
||||
import asyncio
|
||||
|
||||
schema = asyncio.run(resolve_handler_schema_name(FakeService(), target))
|
||||
self.assertEqual(schema, "sales_db")
|
||||
|
||||
def test_convert_sql_param_type(self):
|
||||
self.assertEqual(convert_sql_param_type("42", "integer"), 42)
|
||||
self.assertEqual(convert_sql_param_type("3.14", "float"), 3.14)
|
||||
self.assertTrue(convert_sql_param_type("true", "boolean"))
|
||||
self.assertEqual(convert_sql_param_type("2024-01-15", "date").isoformat(), "2024-01-15")
|
||||
|
||||
def test_build_sql_param_dict_basic(self):
|
||||
class FakeContext:
|
||||
def resolve_template(self, raw: str) -> str:
|
||||
return raw.replace("{{user_id}}", "99")
|
||||
|
||||
params = build_sql_param_dict(
|
||||
[
|
||||
{"name": "status", "type": "string", "value": "active"},
|
||||
{"name": "limit", "type": "integer", "value": "10"},
|
||||
{"name": "user_id", "type": "integer", "value": "{{user_id}}"},
|
||||
],
|
||||
FakeContext(),
|
||||
)
|
||||
self.assertEqual(params["status"], "active")
|
||||
self.assertEqual(params["limit"], 10)
|
||||
self.assertEqual(params["user_id"], 99)
|
||||
|
||||
def test_build_sql_param_dict_required_missing(self):
|
||||
class FakeContext:
|
||||
def resolve_template(self, raw: str) -> str:
|
||||
return raw
|
||||
|
||||
with self.assertRaises(ValueError) as ctx:
|
||||
build_sql_param_dict(
|
||||
[{"name": "id", "type": "integer", "required": True}],
|
||||
FakeContext(),
|
||||
)
|
||||
self.assertIn("id", str(ctx.exception))
|
||||
|
||||
def test_build_sql_param_dict_uses_default(self):
|
||||
class FakeContext:
|
||||
def resolve_template(self, raw: str) -> str:
|
||||
return raw
|
||||
|
||||
params = build_sql_param_dict(
|
||||
[{"name": "offset", "type": "integer", "default": "0"}],
|
||||
FakeContext(),
|
||||
)
|
||||
self.assertEqual(params["offset"], 0)
|
||||
|
||||
def test_resolve_sql_param_value_json(self):
|
||||
class FakeContext:
|
||||
def resolve_template(self, raw: str) -> str:
|
||||
return '["a", "b"]'
|
||||
|
||||
value = resolve_sql_param_value("{{ids}}", FakeContext())
|
||||
self.assertEqual(value, ["a", "b"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,24 @@
|
||||
"""
|
||||
LLM 提供商适配器
|
||||
"""
|
||||
from .base import (
|
||||
BaseLLMProvider,
|
||||
LLMResponse,
|
||||
LLMMessage,
|
||||
LLMConfig,
|
||||
LLMStreamChunk,
|
||||
ToolCall,
|
||||
ToolDefinition,
|
||||
)
|
||||
from .registry import ProviderRegistry
|
||||
|
||||
__all__ = [
|
||||
'BaseLLMProvider',
|
||||
'LLMResponse',
|
||||
'LLMMessage',
|
||||
'LLMConfig',
|
||||
'LLMStreamChunk',
|
||||
'ToolCall',
|
||||
'ToolDefinition',
|
||||
'ProviderRegistry',
|
||||
]
|
||||
@@ -0,0 +1,360 @@
|
||||
"""
|
||||
LLM 提供商基类
|
||||
"""
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, AsyncGenerator, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMMessage:
|
||||
"""LLM 消息
|
||||
|
||||
支持多模态内容格式：
|
||||
- 纯文本: content 为字符串
|
||||
- 多模态: content 为列表 [{"type": "text", "text": "..."}, {"type": "image_url", "image_url": {"url": "..."}}]
|
||||
"""
|
||||
role: str # system, user, assistant, tool
|
||||
content: Any # str 或 List[Dict] (多模态内容)
|
||||
name: Optional[str] = None
|
||||
# Function Calling 相关
|
||||
tool_calls: Optional[List[Dict[str, Any]]] = None # assistant 消息中的工具调用
|
||||
tool_call_id: Optional[str] = None # tool 消息中的工具调用 ID
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result = {'role': self.role, 'content': self.content}
|
||||
if self.name:
|
||||
result['name'] = self.name
|
||||
if self.tool_calls:
|
||||
result['tool_calls'] = self.tool_calls
|
||||
if self.tool_call_id:
|
||||
result['tool_call_id'] = self.tool_call_id
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def create_multimodal_content(text: str, attachments: List[Dict] = None) -> Any:
|
||||
"""
|
||||
创建多模态内容
|
||||
|
||||
Args:
|
||||
text: 文本内容
|
||||
attachments: 附件列表 [{"type": "image/file", "url": "...", "mime_type": "...", "base64": "..."}]
|
||||
|
||||
Returns:
|
||||
str 或 List[Dict] 格式的内容
|
||||
|
||||
Note:
|
||||
对于图片类型，优先使用 base64 格式（因为 LLM 无法访问内部 URL）
|
||||
base64 格式: data:<mime_type>;base64,<base64_content>
|
||||
"""
|
||||
if not attachments:
|
||||
return text
|
||||
|
||||
content = []
|
||||
|
||||
# 添加文本
|
||||
if text:
|
||||
content.append({"type": "text", "text": text})
|
||||
|
||||
# 添加附件
|
||||
for att in attachments:
|
||||
att_type = att.get('type', 'file')
|
||||
mime_type = att.get('mime_type', '')
|
||||
url = att.get('url', '')
|
||||
base64_content = att.get('base64', '')
|
||||
|
||||
if att_type == 'image' or mime_type.startswith('image/'):
|
||||
# 图片类型
|
||||
text_content = att.get('text_content', '')
|
||||
|
||||
if base64_content:
|
||||
# 多模态模型：使用 OpenAI 多模态格式
|
||||
# 使用 data URL 格式
|
||||
image_url = f"data:{mime_type};base64,{base64_content}"
|
||||
content.append({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url}
|
||||
})
|
||||
elif text_content:
|
||||
# 非多模态模型：使用 OCR 识别的文本内容
|
||||
file_name = att.get('name', 'unknown')
|
||||
content.append({
|
||||
"type": "text",
|
||||
"text": f"\n{text_content}"
|
||||
})
|
||||
elif url:
|
||||
# 回退到 URL（仅适用于公网可访问的 URL）
|
||||
content.append({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": url}
|
||||
})
|
||||
else:
|
||||
# 其他文件类型 - 提取文本内容或显示文件信息
|
||||
file_name = att.get('name', 'unknown')
|
||||
text_content = att.get('text_content', '')
|
||||
|
||||
if text_content:
|
||||
# 有提取的文本内容，将其添加到消息中
|
||||
content.append({
|
||||
"type": "text",
|
||||
"text": f"\n--- 文件: {file_name} ---\n{text_content}\n--- 文件结束 ---"
|
||||
})
|
||||
else:
|
||||
# 无法提取内容，只显示文件信息
|
||||
file_size = att.get('size', 0)
|
||||
size_str = f"{file_size / 1024:.1f}KB" if file_size < 1024 * 1024 else f"{file_size / 1024 / 1024:.1f}MB"
|
||||
content.append({
|
||||
"type": "text",
|
||||
"text": f"\n[附件: {file_name} ({size_str})，该文件类型暂不支持内容提取]"
|
||||
})
|
||||
|
||||
return content if len(content) > 1 else (content[0].get('text', '') if content else text)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCall:
|
||||
"""工具调用"""
|
||||
id: str # 工具调用 ID
|
||||
name: str # 工具/函数名称
|
||||
arguments: Dict[str, Any] # 参数
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
import json
|
||||
return {
|
||||
'id': self.id,
|
||||
'type': 'function',
|
||||
'function': {
|
||||
'name': self.name,
|
||||
'arguments': json.dumps(self.arguments, ensure_ascii=False) if isinstance(self.arguments, dict) else self.arguments,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMResponse:
|
||||
"""LLM 响应"""
|
||||
content: str
|
||||
model: str = ''
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
finish_reason: str = ''
|
||||
raw_response: Optional[Dict] = None
|
||||
# Function Calling 相关
|
||||
tool_calls: Optional[List[ToolCall]] = None # 工具调用列表
|
||||
|
||||
@property
|
||||
def is_complete(self) -> bool:
|
||||
return self.finish_reason in ('stop', 'end_turn', 'length')
|
||||
|
||||
@property
|
||||
def has_tool_calls(self) -> bool:
|
||||
"""是否包含工具调用"""
|
||||
return bool(self.tool_calls)
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMStreamChunk:
|
||||
"""LLM 流式响应块"""
|
||||
content: str = ''
|
||||
is_finished: bool = False
|
||||
finish_reason: str = ''
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
# Function Calling 相关
|
||||
tool_calls: Optional[List[ToolCall]] = None # 工具调用（流式时可能分多次返回）
|
||||
tool_call_delta: Optional[Dict[str, Any]] = None # 工具调用增量
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolDefinition:
|
||||
"""工具定义（用于 Function Calling）"""
|
||||
name: str # 工具名称
|
||||
description: str # 工具描述
|
||||
parameters: Dict[str, Any] # 参数 Schema (JSON Schema 格式)
|
||||
|
||||
def to_openai_format(self) -> Dict[str, Any]:
|
||||
"""转换为 OpenAI 格式"""
|
||||
return {
|
||||
'type': 'function',
|
||||
'function': {
|
||||
'name': self.name,
|
||||
'description': self.description,
|
||||
'parameters': self.parameters,
|
||||
}
|
||||
}
|
||||
|
||||
def to_claude_format(self) -> Dict[str, Any]:
|
||||
"""转换为 Claude 格式"""
|
||||
return {
|
||||
'name': self.name,
|
||||
'description': self.description,
|
||||
'input_schema': self.parameters,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMConfig:
|
||||
"""LLM 配置"""
|
||||
model: str
|
||||
temperature: float = 0.7
|
||||
top_p: float = 1.0
|
||||
max_tokens: int = 2048
|
||||
stop: Optional[List[str]] = None
|
||||
presence_penalty: float = 0.0
|
||||
frequency_penalty: float = 0.0
|
||||
extra_params: Dict[str, Any] = field(default_factory=dict)
|
||||
# Function Calling 相关
|
||||
tools: Optional[List[ToolDefinition]] = None # 工具定义列表
|
||||
tool_choice: str = 'auto' # 工具选择策略: auto, none, required, 或具体工具名
|
||||
|
||||
|
||||
class BaseLLMProvider(ABC):
|
||||
"""
|
||||
LLM 提供商基类
|
||||
|
||||
所有提供商适配器必须继承此类并实现抽象方法
|
||||
"""
|
||||
|
||||
# 提供商类型标识
|
||||
provider_type: str = ''
|
||||
|
||||
# 提供商显示名称
|
||||
provider_name: str = ''
|
||||
|
||||
# 支持的模型类型
|
||||
supported_model_types: List[str] = ['chat']
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str = '',
|
||||
api_base: str = '',
|
||||
**kwargs
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.api_base = api_base
|
||||
self.extra_config = kwargs
|
||||
|
||||
@abstractmethod
|
||||
def chat(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> LLMResponse:
|
||||
"""
|
||||
同步对话
|
||||
|
||||
Args:
|
||||
messages: 消息列表
|
||||
config: 配置
|
||||
|
||||
Returns:
|
||||
LLMResponse
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def chat_async(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> LLMResponse:
|
||||
"""
|
||||
异步对话
|
||||
|
||||
Args:
|
||||
messages: 消息列表
|
||||
config: 配置
|
||||
|
||||
Returns:
|
||||
LLMResponse
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def chat_stream(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> AsyncGenerator[LLMStreamChunk, None]:
|
||||
"""
|
||||
异步流式对话
|
||||
|
||||
Args:
|
||||
messages: 消息列表
|
||||
config: 配置
|
||||
|
||||
Yields:
|
||||
LLMStreamChunk
|
||||
"""
|
||||
pass
|
||||
|
||||
def chat_stream_sync(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
):
|
||||
"""
|
||||
同步流式对话（生成器）
|
||||
|
||||
Args:
|
||||
messages: 消息列表
|
||||
config: 配置
|
||||
|
||||
Yields:
|
||||
LLMStreamChunk
|
||||
"""
|
||||
# 默认实现：调用同步 chat 并返回单个结果
|
||||
response = self.chat(messages, config)
|
||||
yield LLMStreamChunk(
|
||||
content=response.content,
|
||||
is_finished=True,
|
||||
finish_reason=response.finish_reason,
|
||||
prompt_tokens=response.prompt_tokens,
|
||||
completion_tokens=response.completion_tokens,
|
||||
total_tokens=response.total_tokens,
|
||||
)
|
||||
|
||||
def validate_config(self) -> bool:
|
||||
"""
|
||||
验证配置是否有效
|
||||
|
||||
Returns:
|
||||
是否有效
|
||||
"""
|
||||
return bool(self.api_key)
|
||||
|
||||
def get_available_models(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
获取可用模型列表
|
||||
|
||||
Returns:
|
||||
模型列表
|
||||
"""
|
||||
return []
|
||||
|
||||
async def fetch_models_from_api(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
从提供商 API 在线拉取最新模型列表
|
||||
|
||||
子类可覆盖此方法实现具体的 API 调用逻辑。
|
||||
默认返回空列表，表示该提供商不支持在线拉取。
|
||||
|
||||
Returns:
|
||||
模型列表，格式与 get_default_models 一致
|
||||
"""
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def get_default_models(cls) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
获取默认模型列表（用于初始化）
|
||||
|
||||
Returns:
|
||||
默认模型列表
|
||||
"""
|
||||
return []
|
||||
@@ -0,0 +1,528 @@
|
||||
"""
|
||||
Anthropic Claude 提供商适配器
|
||||
"""
|
||||
import logging
|
||||
from typing import AsyncGenerator, Dict, List, Any
|
||||
|
||||
from .base import BaseLLMProvider, LLMConfig, LLMMessage, LLMResponse, LLMStreamChunk, ToolCall
|
||||
from .registry import ProviderRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ProviderRegistry.register
|
||||
class ClaudeProvider(BaseLLMProvider):
|
||||
"""
|
||||
Anthropic Claude 提供商适配器
|
||||
|
||||
支持 Claude 3 系列模型
|
||||
"""
|
||||
|
||||
provider_type = 'claude'
|
||||
provider_name = 'Anthropic Claude'
|
||||
supported_model_types = ['chat']
|
||||
|
||||
DEFAULT_API_BASE = 'https://api.anthropic.com'
|
||||
|
||||
def __init__(self, api_key: str = '', api_base: str = '', **kwargs):
|
||||
super().__init__(api_key, api_base, **kwargs)
|
||||
self.api_base = api_base or self.DEFAULT_API_BASE
|
||||
self._client = None
|
||||
self._async_client = None
|
||||
|
||||
def _get_client(self):
|
||||
"""获取同步客户端"""
|
||||
if self._client is None:
|
||||
try:
|
||||
import anthropic
|
||||
self._client = anthropic.Anthropic(
|
||||
api_key=self.api_key,
|
||||
base_url=self.api_base if self.api_base != self.DEFAULT_API_BASE else None,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError('请安装 anthropic 库: pip install anthropic')
|
||||
return self._client
|
||||
|
||||
def _get_async_client(self):
|
||||
"""获取异步客户端"""
|
||||
if self._async_client is None:
|
||||
try:
|
||||
import anthropic
|
||||
self._async_client = anthropic.AsyncAnthropic(
|
||||
api_key=self.api_key,
|
||||
base_url=self.api_base if self.api_base != self.DEFAULT_API_BASE else None,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError('请安装 anthropic 库: pip install anthropic')
|
||||
return self._async_client
|
||||
|
||||
def _convert_messages(self, messages: List[LLMMessage]) -> tuple:
|
||||
"""
|
||||
转换消息格式
|
||||
|
||||
Claude API 需要将 system 消息单独传递
|
||||
"""
|
||||
system_prompt = ''
|
||||
converted_messages = []
|
||||
|
||||
for msg in messages:
|
||||
if msg.role == 'system':
|
||||
system_prompt = msg.content
|
||||
elif msg.role == 'tool':
|
||||
# Claude 的工具结果消息格式
|
||||
converted_messages.append({
|
||||
'role': 'user',
|
||||
'content': [{
|
||||
'type': 'tool_result',
|
||||
'tool_use_id': msg.tool_call_id,
|
||||
'content': msg.content,
|
||||
}],
|
||||
})
|
||||
elif msg.role == 'assistant' and msg.tool_calls:
|
||||
# 包含工具调用的助手消息
|
||||
content = []
|
||||
if msg.content:
|
||||
content.append({'type': 'text', 'text': msg.content})
|
||||
for tc in msg.tool_calls:
|
||||
content.append({
|
||||
'type': 'tool_use',
|
||||
'id': tc.get('id', ''),
|
||||
'name': tc.get('function', {}).get('name', ''),
|
||||
'input': tc.get('function', {}).get('arguments', {}),
|
||||
})
|
||||
converted_messages.append({
|
||||
'role': 'assistant',
|
||||
'content': content,
|
||||
})
|
||||
else:
|
||||
converted_messages.append({
|
||||
'role': msg.role,
|
||||
'content': msg.content,
|
||||
})
|
||||
|
||||
return system_prompt, converted_messages
|
||||
|
||||
def chat(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> LLMResponse:
|
||||
"""同步对话"""
|
||||
client = self._get_client()
|
||||
system_prompt, converted_messages = self._convert_messages(messages)
|
||||
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': converted_messages,
|
||||
'max_tokens': config.max_tokens,
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
}
|
||||
|
||||
if system_prompt:
|
||||
kwargs['system'] = system_prompt
|
||||
|
||||
if config.stop:
|
||||
kwargs['stop_sequences'] = config.stop
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_claude_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = {'type': 'any'}
|
||||
elif config.tool_choice == 'none':
|
||||
# Claude 不支持 none，不传递 tools 即可
|
||||
del kwargs['tools']
|
||||
elif config.tool_choice != 'auto':
|
||||
# 指定具体工具
|
||||
kwargs['tool_choice'] = {'type': 'tool', 'name': config.tool_choice}
|
||||
else:
|
||||
kwargs['tool_choice'] = {'type': 'auto'}
|
||||
|
||||
response = client.messages.create(**kwargs)
|
||||
|
||||
# 解析响应内容和工具调用
|
||||
content = ''
|
||||
tool_calls = []
|
||||
|
||||
for block in response.content:
|
||||
if block.type == 'text':
|
||||
content = block.text
|
||||
elif block.type == 'tool_use':
|
||||
tool_calls.append(ToolCall(
|
||||
id=block.id,
|
||||
name=block.name,
|
||||
arguments=block.input if isinstance(block.input, dict) else {},
|
||||
))
|
||||
|
||||
return LLMResponse(
|
||||
content=content,
|
||||
model=response.model,
|
||||
prompt_tokens=response.usage.input_tokens,
|
||||
completion_tokens=response.usage.output_tokens,
|
||||
total_tokens=response.usage.input_tokens + response.usage.output_tokens,
|
||||
finish_reason=response.stop_reason or '',
|
||||
raw_response=response.model_dump(),
|
||||
tool_calls=tool_calls if tool_calls else None,
|
||||
)
|
||||
|
||||
async def chat_async(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> LLMResponse:
|
||||
"""异步对话"""
|
||||
client = self._get_async_client()
|
||||
system_prompt, converted_messages = self._convert_messages(messages)
|
||||
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': converted_messages,
|
||||
'max_tokens': config.max_tokens,
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
}
|
||||
|
||||
if system_prompt:
|
||||
kwargs['system'] = system_prompt
|
||||
|
||||
if config.stop:
|
||||
kwargs['stop_sequences'] = config.stop
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_claude_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = {'type': 'any'}
|
||||
elif config.tool_choice == 'none':
|
||||
del kwargs['tools']
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'tool', 'name': config.tool_choice}
|
||||
else:
|
||||
kwargs['tool_choice'] = {'type': 'auto'}
|
||||
|
||||
response = await client.messages.create(**kwargs)
|
||||
|
||||
# 解析响应内容和工具调用
|
||||
content = ''
|
||||
tool_calls = []
|
||||
|
||||
for block in response.content:
|
||||
if block.type == 'text':
|
||||
content = block.text
|
||||
elif block.type == 'tool_use':
|
||||
tool_calls.append(ToolCall(
|
||||
id=block.id,
|
||||
name=block.name,
|
||||
arguments=block.input if isinstance(block.input, dict) else {},
|
||||
))
|
||||
|
||||
return LLMResponse(
|
||||
content=content,
|
||||
model=response.model,
|
||||
prompt_tokens=response.usage.input_tokens,
|
||||
completion_tokens=response.usage.output_tokens,
|
||||
total_tokens=response.usage.input_tokens + response.usage.output_tokens,
|
||||
finish_reason=response.stop_reason or '',
|
||||
raw_response=response.model_dump(),
|
||||
tool_calls=tool_calls if tool_calls else None,
|
||||
)
|
||||
|
||||
async def chat_stream(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> AsyncGenerator[LLMStreamChunk, None]:
|
||||
"""流式对话"""
|
||||
client = self._get_async_client()
|
||||
system_prompt, converted_messages = self._convert_messages(messages)
|
||||
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': converted_messages,
|
||||
'max_tokens': config.max_tokens,
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
}
|
||||
|
||||
if system_prompt:
|
||||
kwargs['system'] = system_prompt
|
||||
|
||||
if config.stop:
|
||||
kwargs['stop_sequences'] = config.stop
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_claude_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = {'type': 'any'}
|
||||
elif config.tool_choice == 'none':
|
||||
del kwargs['tools']
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'tool', 'name': config.tool_choice}
|
||||
else:
|
||||
kwargs['tool_choice'] = {'type': 'auto'}
|
||||
|
||||
# 用于累积工具调用
|
||||
tool_calls_accumulator = [] # List of {id, name, input_json}
|
||||
current_tool_call = None
|
||||
|
||||
async with client.messages.stream(**kwargs) as stream:
|
||||
async for event in stream:
|
||||
if hasattr(event, 'type'):
|
||||
if event.type == 'content_block_start':
|
||||
if hasattr(event, 'content_block') and event.content_block.type == 'tool_use':
|
||||
current_tool_call = {
|
||||
'id': event.content_block.id,
|
||||
'name': event.content_block.name,
|
||||
'input_json': '',
|
||||
}
|
||||
elif event.type == 'content_block_delta':
|
||||
if hasattr(event, 'delta'):
|
||||
if event.delta.type == 'text_delta':
|
||||
yield LLMStreamChunk(
|
||||
content=event.delta.text,
|
||||
is_finished=False,
|
||||
)
|
||||
elif event.delta.type == 'input_json_delta' and current_tool_call:
|
||||
current_tool_call['input_json'] += event.delta.partial_json
|
||||
elif event.type == 'content_block_stop':
|
||||
if current_tool_call:
|
||||
tool_calls_accumulator.append(current_tool_call)
|
||||
current_tool_call = None
|
||||
|
||||
# 获取最终的 usage 信息
|
||||
final_message = await stream.get_final_message()
|
||||
|
||||
# 解析工具调用
|
||||
tool_calls = None
|
||||
if tool_calls_accumulator:
|
||||
tool_calls = self._parse_accumulated_tool_calls(tool_calls_accumulator)
|
||||
|
||||
yield LLMStreamChunk(
|
||||
content='',
|
||||
is_finished=True,
|
||||
finish_reason=final_message.stop_reason or '',
|
||||
prompt_tokens=final_message.usage.input_tokens,
|
||||
completion_tokens=final_message.usage.output_tokens,
|
||||
total_tokens=final_message.usage.input_tokens + final_message.usage.output_tokens,
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
|
||||
def _parse_accumulated_tool_calls(self, accumulator: List[Dict]) -> List[ToolCall]:
|
||||
"""解析累积的工具调用"""
|
||||
import json
|
||||
result = []
|
||||
for tc in accumulator:
|
||||
try:
|
||||
arguments = json.loads(tc['input_json']) if tc['input_json'] else {}
|
||||
except json.JSONDecodeError:
|
||||
arguments = {'raw': tc['input_json']}
|
||||
|
||||
result.append(ToolCall(
|
||||
id=tc['id'],
|
||||
name=tc['name'],
|
||||
arguments=arguments,
|
||||
))
|
||||
return result
|
||||
|
||||
def chat_stream_sync(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
):
|
||||
"""同步流式对话"""
|
||||
client = self._get_client()
|
||||
system_prompt, converted_messages = self._convert_messages(messages)
|
||||
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': converted_messages,
|
||||
'max_tokens': config.max_tokens,
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
}
|
||||
|
||||
if system_prompt:
|
||||
kwargs['system'] = system_prompt
|
||||
|
||||
if config.stop:
|
||||
kwargs['stop_sequences'] = config.stop
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_claude_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = {'type': 'any'}
|
||||
elif config.tool_choice == 'none':
|
||||
del kwargs['tools']
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'tool', 'name': config.tool_choice}
|
||||
else:
|
||||
kwargs['tool_choice'] = {'type': 'auto'}
|
||||
|
||||
# 用于累积工具调用
|
||||
tool_calls_accumulator = []
|
||||
current_tool_call = None
|
||||
|
||||
with client.messages.stream(**kwargs) as stream:
|
||||
for event in stream:
|
||||
if hasattr(event, 'type'):
|
||||
if event.type == 'content_block_start':
|
||||
if hasattr(event, 'content_block') and event.content_block.type == 'tool_use':
|
||||
current_tool_call = {
|
||||
'id': event.content_block.id,
|
||||
'name': event.content_block.name,
|
||||
'input_json': '',
|
||||
}
|
||||
elif event.type == 'content_block_delta':
|
||||
if hasattr(event, 'delta'):
|
||||
if event.delta.type == 'text_delta':
|
||||
yield LLMStreamChunk(
|
||||
content=event.delta.text,
|
||||
is_finished=False,
|
||||
)
|
||||
elif event.delta.type == 'input_json_delta' and current_tool_call:
|
||||
current_tool_call['input_json'] += event.delta.partial_json
|
||||
elif event.type == 'content_block_stop':
|
||||
if current_tool_call:
|
||||
tool_calls_accumulator.append(current_tool_call)
|
||||
current_tool_call = None
|
||||
|
||||
# 获取最终的 usage 信息
|
||||
final_message = stream.get_final_message()
|
||||
|
||||
# 解析工具调用
|
||||
tool_calls = None
|
||||
if tool_calls_accumulator:
|
||||
tool_calls = self._parse_accumulated_tool_calls(tool_calls_accumulator)
|
||||
|
||||
yield LLMStreamChunk(
|
||||
content='',
|
||||
is_finished=True,
|
||||
finish_reason=final_message.stop_reason or '',
|
||||
prompt_tokens=final_message.usage.input_tokens,
|
||||
completion_tokens=final_message.usage.output_tokens,
|
||||
total_tokens=final_message.usage.input_tokens + final_message.usage.output_tokens,
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
|
||||
async def fetch_models_from_api(self) -> List[Dict[str, Any]]:
|
||||
"""通过 Anthropic /v1/models 端点在线拉取模型列表"""
|
||||
import httpx
|
||||
|
||||
base_url = (self.api_base or self.DEFAULT_API_BASE).rstrip('/')
|
||||
headers = {
|
||||
'x-api-key': self.api_key,
|
||||
'anthropic-version': '2023-06-01',
|
||||
}
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=15) as client:
|
||||
resp = await client.get(f'{base_url}/v1/models', headers=headers)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
model_list = data.get('data', [])
|
||||
if not model_list:
|
||||
return []
|
||||
|
||||
results = []
|
||||
for item in model_list:
|
||||
model_id = item.get('id', '')
|
||||
display = item.get('display_name', model_id)
|
||||
if not model_id:
|
||||
continue
|
||||
|
||||
results.append({
|
||||
'model_name': model_id,
|
||||
'display_name': display,
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 4096,
|
||||
'context_window': 200000,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
})
|
||||
|
||||
results.sort(key=lambda m: m['model_name'])
|
||||
logger.info(f'从 Anthropic ({base_url}) 拉取到 {len(results)} 个模型')
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f'在线拉取 Anthropic 模型列表失败 ({base_url}): {e}')
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def get_default_models(cls) -> List[Dict[str, Any]]:
|
||||
"""获取默认模型列表"""
|
||||
return [
|
||||
# ---- Chat 模型 ----
|
||||
{
|
||||
'model_name': 'claude-sonnet-4-20250514',
|
||||
'display_name': 'Claude Sonnet 4',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 16384,
|
||||
'context_window': 200000,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.003,
|
||||
'output_price': 0.015,
|
||||
},
|
||||
{
|
||||
'model_name': 'claude-3-7-sonnet-20250219',
|
||||
'display_name': 'Claude 3.7 Sonnet',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 16384,
|
||||
'context_window': 200000,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.003,
|
||||
'output_price': 0.015,
|
||||
},
|
||||
{
|
||||
'model_name': 'claude-3-5-sonnet-20241022',
|
||||
'display_name': 'Claude 3.5 Sonnet',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 200000,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.003,
|
||||
'output_price': 0.015,
|
||||
},
|
||||
{
|
||||
'model_name': 'claude-3-5-haiku-20241022',
|
||||
'display_name': 'Claude 3.5 Haiku',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 200000,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.001,
|
||||
'output_price': 0.005,
|
||||
},
|
||||
{
|
||||
'model_name': 'claude-3-opus-20240229',
|
||||
'display_name': 'Claude 3 Opus',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 4096,
|
||||
'context_window': 200000,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.015,
|
||||
'output_price': 0.075,
|
||||
},
|
||||
{
|
||||
'model_name': 'claude-3-haiku-20240307',
|
||||
'display_name': 'Claude 3 Haiku',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 4096,
|
||||
'context_window': 200000,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.00025,
|
||||
'output_price': 0.00125,
|
||||
},
|
||||
]
|
||||
@@ -0,0 +1,56 @@
|
||||
"""
|
||||
DeepSeek 提供商适配器
|
||||
|
||||
兼容 OpenAI API，继承 OpenAI 适配器
|
||||
"""
|
||||
import logging
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from .openai_provider import OpenAIProvider
|
||||
from .registry import ProviderRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ProviderRegistry.register
|
||||
class DeepSeekProvider(OpenAIProvider):
|
||||
"""
|
||||
DeepSeek 提供商适配器
|
||||
|
||||
兼容 OpenAI API 格式，支持 DeepSeek-V3、DeepSeek-R1 等模型
|
||||
"""
|
||||
|
||||
provider_type = 'deepseek'
|
||||
provider_name = 'DeepSeek'
|
||||
supported_model_types = ['chat', 'completion']
|
||||
|
||||
DEFAULT_API_BASE = 'https://api.deepseek.com/v1'
|
||||
|
||||
@classmethod
|
||||
def get_default_models(cls) -> List[Dict[str, Any]]:
|
||||
"""获取默认模型列表"""
|
||||
return [
|
||||
# ---- Chat 模型 ----
|
||||
{
|
||||
'model_name': 'deepseek-chat',
|
||||
'display_name': 'DeepSeek V3',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 65536,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.0014,
|
||||
'output_price': 0.0028,
|
||||
},
|
||||
{
|
||||
'model_name': 'deepseek-reasoner',
|
||||
'display_name': 'DeepSeek R1',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 65536,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': False,
|
||||
'input_price': 0.004,
|
||||
'output_price': 0.016,
|
||||
},
|
||||
]
|
||||
@@ -0,0 +1,551 @@
|
||||
"""
|
||||
Ollama 本地模型提供商适配器
|
||||
"""
|
||||
import logging
|
||||
from typing import AsyncGenerator, Dict, List, Any
|
||||
|
||||
from .base import BaseLLMProvider, LLMConfig, LLMMessage, LLMResponse, LLMStreamChunk, ToolCall
|
||||
from .registry import ProviderRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ProviderRegistry.register
|
||||
class OllamaProvider(BaseLLMProvider):
|
||||
"""
|
||||
Ollama 本地模型提供商适配器
|
||||
|
||||
支持本地部署的开源模型
|
||||
"""
|
||||
|
||||
provider_type = 'ollama'
|
||||
provider_name = 'Ollama (本地)'
|
||||
supported_model_types = ['chat', 'embedding']
|
||||
|
||||
DEFAULT_API_BASE = 'http://localhost:11434/v1'
|
||||
|
||||
def __init__(self, api_key: str = '', api_base: str = '', **kwargs):
|
||||
super().__init__(api_key, api_base, **kwargs)
|
||||
# Ollama 使用 ollama_host 参数
|
||||
ollama_host = kwargs.get('ollama_host', '')
|
||||
if ollama_host:
|
||||
self.api_base = f'{ollama_host.rstrip("/")}/v1'
|
||||
else:
|
||||
self.api_base = api_base or self.DEFAULT_API_BASE
|
||||
self._client = None
|
||||
self._async_client = None
|
||||
|
||||
def _get_client(self):
|
||||
"""获取同步客户端"""
|
||||
if self._client is None:
|
||||
try:
|
||||
from openai import OpenAI
|
||||
self._client = OpenAI(
|
||||
api_key='ollama', # Ollama 不需要 API Key
|
||||
base_url=self.api_base,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError('请安装 openai 库: pip install openai')
|
||||
return self._client
|
||||
|
||||
def _get_async_client(self):
|
||||
"""获取异步客户端"""
|
||||
if self._async_client is None:
|
||||
try:
|
||||
from openai import AsyncOpenAI
|
||||
self._async_client = AsyncOpenAI(
|
||||
api_key='ollama', # Ollama 不需要 API Key
|
||||
base_url=self.api_base,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError('请安装 openai 库: pip install openai')
|
||||
return self._async_client
|
||||
|
||||
def validate_config(self) -> bool:
|
||||
"""Ollama 不需要 API Key"""
|
||||
return True
|
||||
|
||||
def chat(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> LLMResponse:
|
||||
"""同步对话"""
|
||||
client = self._get_client()
|
||||
|
||||
# 构建请求参数
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': [m.to_dict() for m in messages],
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
'max_tokens': config.max_tokens,
|
||||
'stop': config.stop,
|
||||
}
|
||||
|
||||
# 添加 Function Calling 参数（Ollama 部分模型支持）
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = 'required'
|
||||
elif config.tool_choice == 'none':
|
||||
kwargs['tool_choice'] = 'none'
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||||
else:
|
||||
kwargs['tool_choice'] = 'auto'
|
||||
|
||||
response = client.chat.completions.create(**kwargs)
|
||||
|
||||
choice = response.choices[0]
|
||||
usage = response.usage
|
||||
|
||||
# 解析工具调用
|
||||
tool_calls = None
|
||||
if choice.message.tool_calls:
|
||||
tool_calls = self._parse_tool_calls(choice.message.tool_calls)
|
||||
|
||||
return LLMResponse(
|
||||
content=choice.message.content or '',
|
||||
model=response.model,
|
||||
prompt_tokens=usage.prompt_tokens if usage else 0,
|
||||
completion_tokens=usage.completion_tokens if usage else 0,
|
||||
total_tokens=usage.total_tokens if usage else 0,
|
||||
finish_reason=choice.finish_reason or '',
|
||||
raw_response=response.model_dump(),
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
|
||||
def _parse_tool_calls(self, tool_calls) -> List[ToolCall]:
|
||||
"""解析工具调用"""
|
||||
import json
|
||||
result = []
|
||||
for tc in tool_calls:
|
||||
try:
|
||||
arguments = json.loads(tc.function.arguments) if tc.function.arguments else {}
|
||||
except json.JSONDecodeError:
|
||||
arguments = {'raw': tc.function.arguments}
|
||||
|
||||
result.append(ToolCall(
|
||||
id=tc.id or f'call_{len(result)}',
|
||||
name=tc.function.name,
|
||||
arguments=arguments,
|
||||
))
|
||||
return result
|
||||
|
||||
async def chat_async(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> LLMResponse:
|
||||
"""异步对话"""
|
||||
client = self._get_async_client()
|
||||
|
||||
# 构建请求参数
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': [m.to_dict() for m in messages],
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
'max_tokens': config.max_tokens,
|
||||
'stop': config.stop,
|
||||
}
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = 'required'
|
||||
elif config.tool_choice == 'none':
|
||||
kwargs['tool_choice'] = 'none'
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||||
else:
|
||||
kwargs['tool_choice'] = 'auto'
|
||||
|
||||
response = await client.chat.completions.create(**kwargs)
|
||||
|
||||
choice = response.choices[0]
|
||||
usage = response.usage
|
||||
|
||||
# 解析工具调用
|
||||
tool_calls = None
|
||||
if choice.message.tool_calls:
|
||||
tool_calls = self._parse_tool_calls(choice.message.tool_calls)
|
||||
|
||||
return LLMResponse(
|
||||
content=choice.message.content or '',
|
||||
model=response.model,
|
||||
prompt_tokens=usage.prompt_tokens if usage else 0,
|
||||
completion_tokens=usage.completion_tokens if usage else 0,
|
||||
total_tokens=usage.total_tokens if usage else 0,
|
||||
finish_reason=choice.finish_reason or '',
|
||||
raw_response=response.model_dump(),
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
|
||||
async def chat_stream(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> AsyncGenerator[LLMStreamChunk, None]:
|
||||
"""流式对话"""
|
||||
client = self._get_async_client()
|
||||
|
||||
# 构建请求参数
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': [m.to_dict() for m in messages],
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
'max_tokens': config.max_tokens,
|
||||
'stop': config.stop,
|
||||
'stream': True,
|
||||
}
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = 'required'
|
||||
elif config.tool_choice == 'none':
|
||||
kwargs['tool_choice'] = 'none'
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||||
else:
|
||||
kwargs['tool_choice'] = 'auto'
|
||||
|
||||
stream = await client.chat.completions.create(**kwargs)
|
||||
|
||||
# 用于累积工具调用
|
||||
tool_call_accumulator = {}
|
||||
|
||||
async for chunk in stream:
|
||||
if chunk.choices:
|
||||
choice = chunk.choices[0]
|
||||
delta = choice.delta
|
||||
|
||||
# 处理工具调用增量
|
||||
tool_call_delta = None
|
||||
if hasattr(delta, 'tool_calls') and delta.tool_calls:
|
||||
for tc_delta in delta.tool_calls:
|
||||
tc_index = tc_delta.index if hasattr(tc_delta, 'index') else 0
|
||||
if tc_index not in tool_call_accumulator:
|
||||
tool_call_accumulator[tc_index] = {
|
||||
'id': getattr(tc_delta, 'id', '') or f'call_{tc_index}',
|
||||
'name': '',
|
||||
'arguments': '',
|
||||
}
|
||||
if hasattr(tc_delta, 'id') and tc_delta.id:
|
||||
tool_call_accumulator[tc_index]['id'] = tc_delta.id
|
||||
if tc_delta.function:
|
||||
if tc_delta.function.name:
|
||||
tool_call_accumulator[tc_index]['name'] = tc_delta.function.name
|
||||
if tc_delta.function.arguments:
|
||||
tool_call_accumulator[tc_index]['arguments'] += tc_delta.function.arguments
|
||||
tool_call_delta = tool_call_accumulator[tc_index]
|
||||
|
||||
# 检查是否完成
|
||||
tool_calls = None
|
||||
if choice.finish_reason and tool_call_accumulator:
|
||||
tool_calls = self._parse_accumulated_tool_calls(tool_call_accumulator)
|
||||
|
||||
yield LLMStreamChunk(
|
||||
content=delta.content or '',
|
||||
is_finished=choice.finish_reason is not None,
|
||||
finish_reason=choice.finish_reason or '',
|
||||
tool_calls=tool_calls,
|
||||
tool_call_delta=tool_call_delta,
|
||||
)
|
||||
|
||||
def _parse_accumulated_tool_calls(self, accumulator: Dict) -> List[ToolCall]:
|
||||
"""解析累积的工具调用"""
|
||||
import json
|
||||
result = []
|
||||
for idx in sorted(accumulator.keys()):
|
||||
tc = accumulator[idx]
|
||||
try:
|
||||
arguments = json.loads(tc['arguments']) if tc['arguments'] else {}
|
||||
except json.JSONDecodeError:
|
||||
arguments = {'raw': tc['arguments']}
|
||||
|
||||
result.append(ToolCall(
|
||||
id=tc['id'],
|
||||
name=tc['name'],
|
||||
arguments=arguments,
|
||||
))
|
||||
return result
|
||||
|
||||
def chat_stream_sync(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
):
|
||||
"""同步流式对话"""
|
||||
client = self._get_client()
|
||||
|
||||
# 构建请求参数
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': [m.to_dict() for m in messages],
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
'max_tokens': config.max_tokens,
|
||||
'stop': config.stop,
|
||||
'stream': True,
|
||||
}
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = 'required'
|
||||
elif config.tool_choice == 'none':
|
||||
kwargs['tool_choice'] = 'none'
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||||
else:
|
||||
kwargs['tool_choice'] = 'auto'
|
||||
|
||||
stream = client.chat.completions.create(**kwargs)
|
||||
|
||||
# 用于累积工具调用
|
||||
tool_call_accumulator = {}
|
||||
|
||||
for chunk in stream:
|
||||
if chunk.choices:
|
||||
choice = chunk.choices[0]
|
||||
delta = choice.delta
|
||||
|
||||
# 处理工具调用增量
|
||||
tool_call_delta = None
|
||||
if hasattr(delta, 'tool_calls') and delta.tool_calls:
|
||||
for tc_delta in delta.tool_calls:
|
||||
tc_index = tc_delta.index if hasattr(tc_delta, 'index') else 0
|
||||
if tc_index not in tool_call_accumulator:
|
||||
tool_call_accumulator[tc_index] = {
|
||||
'id': getattr(tc_delta, 'id', '') or f'call_{tc_index}',
|
||||
'name': '',
|
||||
'arguments': '',
|
||||
}
|
||||
if hasattr(tc_delta, 'id') and tc_delta.id:
|
||||
tool_call_accumulator[tc_index]['id'] = tc_delta.id
|
||||
if tc_delta.function:
|
||||
if tc_delta.function.name:
|
||||
tool_call_accumulator[tc_index]['name'] = tc_delta.function.name
|
||||
if tc_delta.function.arguments:
|
||||
tool_call_accumulator[tc_index]['arguments'] += tc_delta.function.arguments
|
||||
tool_call_delta = tool_call_accumulator[tc_index]
|
||||
|
||||
# 检查是否完成
|
||||
tool_calls = None
|
||||
if choice.finish_reason and tool_call_accumulator:
|
||||
tool_calls = self._parse_accumulated_tool_calls(tool_call_accumulator)
|
||||
|
||||
yield LLMStreamChunk(
|
||||
content=delta.content or '',
|
||||
is_finished=choice.finish_reason is not None,
|
||||
finish_reason=choice.finish_reason or '',
|
||||
tool_calls=tool_calls,
|
||||
tool_call_delta=tool_call_delta,
|
||||
)
|
||||
|
||||
_EMBEDDING_KEYWORDS = ('embed', 'bge', 'nomic-embed', 'mxbai-embed')
|
||||
_RERANK_KEYWORDS = ('rerank', 'reranker')
|
||||
|
||||
def get_available_models(self) -> List[Dict[str, Any]]:
|
||||
"""获取本地可用的模型列表"""
|
||||
try:
|
||||
import httpx
|
||||
base_url = self.api_base.replace('/v1', '')
|
||||
response = httpx.get(f'{base_url}/api/tags', timeout=10)
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
models = []
|
||||
for model in data.get('models', []):
|
||||
models.append({
|
||||
'model_name': model.get('name', ''),
|
||||
'display_name': model.get('name', ''),
|
||||
'size': model.get('size', 0),
|
||||
'modified_at': model.get('modified_at', ''),
|
||||
})
|
||||
return models
|
||||
except Exception as e:
|
||||
logger.warning(f'获取 Ollama 模型列表失败: {e}')
|
||||
return []
|
||||
|
||||
async def fetch_models_from_api(self) -> List[Dict[str, Any]]:
|
||||
"""通过 /api/tags 端点拉取 Ollama 本地已安装的模型"""
|
||||
import httpx
|
||||
|
||||
base_url = (self.api_base or 'http://localhost:11434').rstrip('/')
|
||||
# Ollama 的 /api/tags 在根路径，去掉 /v1
|
||||
if base_url.endswith('/v1'):
|
||||
base_url = base_url[:-3]
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=10) as client:
|
||||
resp = await client.get(f'{base_url}/api/tags')
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
results = []
|
||||
for item in data.get('models', []):
|
||||
name = item.get('name', '')
|
||||
if not name:
|
||||
continue
|
||||
|
||||
name_lower = name.lower()
|
||||
if any(kw in name_lower for kw in self._RERANK_KEYWORDS):
|
||||
model_type = 'rerank'
|
||||
elif any(kw in name_lower for kw in self._EMBEDDING_KEYWORDS):
|
||||
model_type = 'embedding'
|
||||
else:
|
||||
model_type = 'chat'
|
||||
|
||||
# 从 size 推算显示名称
|
||||
size_gb = round(item.get('size', 0) / (1024 ** 3), 1)
|
||||
display = f'{name} ({size_gb}GB)' if size_gb > 0 else name
|
||||
|
||||
results.append({
|
||||
'model_name': name,
|
||||
'display_name': display,
|
||||
'model_type': model_type,
|
||||
'max_tokens': 4096,
|
||||
'context_window': 4096,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': model_type == 'chat',
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
})
|
||||
|
||||
type_order = {'chat': 0, 'embedding': 1, 'rerank': 2}
|
||||
results.sort(key=lambda m: (type_order.get(m['model_type'], 9), m['model_name']))
|
||||
logger.info(f'从 Ollama ({base_url}) 拉取到 {len(results)} 个本地模型')
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f'在线拉取 Ollama 模型列表失败 ({base_url}): {e}')
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def get_default_models(cls) -> List[Dict[str, Any]]:
|
||||
"""获取默认模型列表（常用的开源模型）"""
|
||||
return [
|
||||
# ---- Chat 模型 ----
|
||||
{
|
||||
'model_name': 'qwen3:32b',
|
||||
'display_name': 'Qwen3 32B',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 131072,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'qwen3:8b',
|
||||
'display_name': 'Qwen3 8B',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 131072,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'llama3.3',
|
||||
'display_name': 'Llama 3.3 70B',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 4096,
|
||||
'context_window': 128000,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'llama3.2',
|
||||
'display_name': 'Llama 3.2',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 4096,
|
||||
'context_window': 128000,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'deepseek-r1',
|
||||
'display_name': 'DeepSeek R1',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 4096,
|
||||
'context_window': 64000,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': False,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'mistral',
|
||||
'display_name': 'Mistral',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 4096,
|
||||
'context_window': 32000,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'gemma3',
|
||||
'display_name': 'Gemma 3',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 128000,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
# ---- Embedding 模型 ----
|
||||
{
|
||||
'model_name': 'nomic-embed-text',
|
||||
'display_name': 'Nomic Embed Text',
|
||||
'model_type': 'embedding',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 8192,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'bge-m3',
|
||||
'display_name': 'BGE-M3',
|
||||
'model_type': 'embedding',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 8192,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'mxbai-embed-large',
|
||||
'display_name': 'MxBai Embed Large',
|
||||
'model_type': 'embedding',
|
||||
'max_tokens': 512,
|
||||
'context_window': 512,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
# ---- Rerank 模型 ----
|
||||
{
|
||||
'model_name': 'bge-reranker-v2-m3',
|
||||
'display_name': 'BGE Reranker V2 M3',
|
||||
'model_type': 'rerank',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 8192,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
]
|
||||
@@ -0,0 +1,565 @@
|
||||
"""
|
||||
OpenAI 提供商适配器
|
||||
"""
|
||||
import logging
|
||||
from typing import AsyncGenerator, Dict, List, Any
|
||||
|
||||
from .base import BaseLLMProvider, LLMConfig, LLMMessage, LLMResponse, LLMStreamChunk, ToolCall
|
||||
from .registry import ProviderRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ProviderRegistry.register
|
||||
class OpenAIProvider(BaseLLMProvider):
|
||||
"""
|
||||
OpenAI 提供商适配器
|
||||
|
||||
支持 GPT-3.5, GPT-4 等模型
|
||||
"""
|
||||
|
||||
provider_type = 'openai'
|
||||
provider_name = 'OpenAI'
|
||||
supported_model_types = ['chat', 'completion', 'embedding']
|
||||
|
||||
DEFAULT_API_BASE = 'https://api.openai.com/v1'
|
||||
|
||||
def __init__(self, api_key: str = '', api_base: str = '', **kwargs):
|
||||
super().__init__(api_key, api_base, **kwargs)
|
||||
self.api_base = api_base or self.DEFAULT_API_BASE
|
||||
self._client = None
|
||||
self._async_client = None
|
||||
|
||||
def _get_client(self):
|
||||
"""获取同步客户端"""
|
||||
if self._client is None:
|
||||
try:
|
||||
from openai import OpenAI
|
||||
self._client = OpenAI(
|
||||
api_key=self.api_key,
|
||||
base_url=self.api_base,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError('请安装 openai 库: pip install openai')
|
||||
return self._client
|
||||
|
||||
def _get_async_client(self):
|
||||
"""获取异步客户端"""
|
||||
if self._async_client is None:
|
||||
try:
|
||||
from openai import AsyncOpenAI
|
||||
self._async_client = AsyncOpenAI(
|
||||
api_key=self.api_key,
|
||||
base_url=self.api_base,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError('请安装 openai 库: pip install openai')
|
||||
return self._async_client
|
||||
|
||||
def chat(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> LLMResponse:
|
||||
"""同步对话"""
|
||||
client = self._get_client()
|
||||
|
||||
# 构建请求参数
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': [m.to_dict() for m in messages],
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
'max_tokens': config.max_tokens,
|
||||
'stop': config.stop,
|
||||
'presence_penalty': config.presence_penalty,
|
||||
'frequency_penalty': config.frequency_penalty,
|
||||
**config.extra_params,
|
||||
}
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = 'required'
|
||||
elif config.tool_choice == 'none':
|
||||
kwargs['tool_choice'] = 'none'
|
||||
elif config.tool_choice != 'auto':
|
||||
# 指定具体工具
|
||||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||||
else:
|
||||
kwargs['tool_choice'] = 'auto'
|
||||
|
||||
response = client.chat.completions.create(**kwargs)
|
||||
|
||||
choice = response.choices[0]
|
||||
usage = response.usage
|
||||
|
||||
# 解析工具调用
|
||||
tool_calls = None
|
||||
if choice.message.tool_calls:
|
||||
tool_calls = self._parse_tool_calls(choice.message.tool_calls)
|
||||
|
||||
return LLMResponse(
|
||||
content=choice.message.content or '',
|
||||
model=response.model,
|
||||
prompt_tokens=usage.prompt_tokens if usage else 0,
|
||||
completion_tokens=usage.completion_tokens if usage else 0,
|
||||
total_tokens=usage.total_tokens if usage else 0,
|
||||
finish_reason=choice.finish_reason or '',
|
||||
raw_response=response.model_dump(),
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
|
||||
def _parse_tool_calls(self, tool_calls) -> List[ToolCall]:
|
||||
"""解析工具调用"""
|
||||
import json
|
||||
result = []
|
||||
for tc in tool_calls:
|
||||
try:
|
||||
arguments = json.loads(tc.function.arguments) if tc.function.arguments else {}
|
||||
except json.JSONDecodeError:
|
||||
arguments = {'raw': tc.function.arguments}
|
||||
|
||||
result.append(ToolCall(
|
||||
id=tc.id,
|
||||
name=tc.function.name,
|
||||
arguments=arguments,
|
||||
))
|
||||
return result
|
||||
|
||||
async def chat_async(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> LLMResponse:
|
||||
"""异步对话"""
|
||||
client = self._get_async_client()
|
||||
|
||||
# 构建请求参数
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': [m.to_dict() for m in messages],
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
'max_tokens': config.max_tokens,
|
||||
'stop': config.stop,
|
||||
'presence_penalty': config.presence_penalty,
|
||||
'frequency_penalty': config.frequency_penalty,
|
||||
**config.extra_params,
|
||||
}
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = 'required'
|
||||
elif config.tool_choice == 'none':
|
||||
kwargs['tool_choice'] = 'none'
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||||
else:
|
||||
kwargs['tool_choice'] = 'auto'
|
||||
|
||||
response = await client.chat.completions.create(**kwargs)
|
||||
|
||||
choice = response.choices[0]
|
||||
usage = response.usage
|
||||
|
||||
# 解析工具调用
|
||||
tool_calls = None
|
||||
if choice.message.tool_calls:
|
||||
tool_calls = self._parse_tool_calls(choice.message.tool_calls)
|
||||
|
||||
return LLMResponse(
|
||||
content=choice.message.content or '',
|
||||
model=response.model,
|
||||
prompt_tokens=usage.prompt_tokens if usage else 0,
|
||||
completion_tokens=usage.completion_tokens if usage else 0,
|
||||
total_tokens=usage.total_tokens if usage else 0,
|
||||
finish_reason=choice.finish_reason or '',
|
||||
raw_response=response.model_dump(),
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
|
||||
async def chat_stream(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> AsyncGenerator[LLMStreamChunk, None]:
|
||||
"""流式对话"""
|
||||
client = self._get_async_client()
|
||||
|
||||
# 构建请求参数
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': [m.to_dict() for m in messages],
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
'max_tokens': config.max_tokens,
|
||||
'stop': config.stop,
|
||||
'presence_penalty': config.presence_penalty,
|
||||
'frequency_penalty': config.frequency_penalty,
|
||||
'stream': True,
|
||||
'stream_options': {"include_usage": True},
|
||||
**config.extra_params,
|
||||
}
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = 'required'
|
||||
elif config.tool_choice == 'none':
|
||||
kwargs['tool_choice'] = 'none'
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||||
else:
|
||||
kwargs['tool_choice'] = 'auto'
|
||||
|
||||
stream = await client.chat.completions.create(**kwargs)
|
||||
|
||||
# 用于累积工具调用
|
||||
tool_call_accumulator = {} # id -> {name, arguments}
|
||||
|
||||
async for chunk in stream:
|
||||
if chunk.choices:
|
||||
choice = chunk.choices[0]
|
||||
delta = choice.delta
|
||||
|
||||
# 处理工具调用增量
|
||||
tool_call_delta = None
|
||||
if hasattr(delta, 'tool_calls') and delta.tool_calls:
|
||||
for tc_delta in delta.tool_calls:
|
||||
tc_index = tc_delta.index
|
||||
if tc_index not in tool_call_accumulator:
|
||||
tool_call_accumulator[tc_index] = {
|
||||
'id': tc_delta.id or '',
|
||||
'name': '',
|
||||
'arguments': '',
|
||||
}
|
||||
if tc_delta.id:
|
||||
tool_call_accumulator[tc_index]['id'] = tc_delta.id
|
||||
if tc_delta.function:
|
||||
if tc_delta.function.name:
|
||||
tool_call_accumulator[tc_index]['name'] = tc_delta.function.name
|
||||
if tc_delta.function.arguments:
|
||||
tool_call_accumulator[tc_index]['arguments'] += tc_delta.function.arguments
|
||||
tool_call_delta = tool_call_accumulator[tc_index]
|
||||
|
||||
# 检查是否完成（包含工具调用）
|
||||
tool_calls = None
|
||||
if choice.finish_reason == 'tool_calls' or (choice.finish_reason and tool_call_accumulator):
|
||||
tool_calls = self._parse_accumulated_tool_calls(tool_call_accumulator)
|
||||
|
||||
yield LLMStreamChunk(
|
||||
content=delta.content or '',
|
||||
is_finished=choice.finish_reason is not None,
|
||||
finish_reason=choice.finish_reason or '',
|
||||
tool_calls=tool_calls,
|
||||
tool_call_delta=tool_call_delta,
|
||||
)
|
||||
|
||||
# 最后一个 chunk 包含 usage 信息
|
||||
if chunk.usage:
|
||||
yield LLMStreamChunk(
|
||||
content='',
|
||||
is_finished=True,
|
||||
prompt_tokens=chunk.usage.prompt_tokens,
|
||||
completion_tokens=chunk.usage.completion_tokens,
|
||||
total_tokens=chunk.usage.total_tokens,
|
||||
)
|
||||
|
||||
def _parse_accumulated_tool_calls(self, accumulator: Dict) -> List[ToolCall]:
|
||||
"""解析累积的工具调用"""
|
||||
import json
|
||||
result = []
|
||||
for idx in sorted(accumulator.keys()):
|
||||
tc = accumulator[idx]
|
||||
try:
|
||||
arguments = json.loads(tc['arguments']) if tc['arguments'] else {}
|
||||
except json.JSONDecodeError:
|
||||
arguments = {'raw': tc['arguments']}
|
||||
|
||||
result.append(ToolCall(
|
||||
id=tc['id'],
|
||||
name=tc['name'],
|
||||
arguments=arguments,
|
||||
))
|
||||
return result
|
||||
|
||||
def chat_stream_sync(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
):
|
||||
"""同步流式对话"""
|
||||
client = self._get_client()
|
||||
|
||||
# 构建请求参数
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': [m.to_dict() for m in messages],
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
'max_tokens': config.max_tokens,
|
||||
'stop': config.stop,
|
||||
'presence_penalty': config.presence_penalty,
|
||||
'frequency_penalty': config.frequency_penalty,
|
||||
'stream': True,
|
||||
'stream_options': {"include_usage": True},
|
||||
**config.extra_params,
|
||||
}
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = 'required'
|
||||
elif config.tool_choice == 'none':
|
||||
kwargs['tool_choice'] = 'none'
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||||
else:
|
||||
kwargs['tool_choice'] = 'auto'
|
||||
|
||||
stream = client.chat.completions.create(**kwargs)
|
||||
|
||||
# 用于累积工具调用
|
||||
tool_call_accumulator = {}
|
||||
|
||||
for chunk in stream:
|
||||
if chunk.choices:
|
||||
choice = chunk.choices[0]
|
||||
delta = choice.delta
|
||||
|
||||
# 处理工具调用增量
|
||||
tool_call_delta = None
|
||||
if hasattr(delta, 'tool_calls') and delta.tool_calls:
|
||||
for tc_delta in delta.tool_calls:
|
||||
tc_index = tc_delta.index
|
||||
if tc_index not in tool_call_accumulator:
|
||||
tool_call_accumulator[tc_index] = {
|
||||
'id': tc_delta.id or '',
|
||||
'name': '',
|
||||
'arguments': '',
|
||||
}
|
||||
if tc_delta.id:
|
||||
tool_call_accumulator[tc_index]['id'] = tc_delta.id
|
||||
if tc_delta.function:
|
||||
if tc_delta.function.name:
|
||||
tool_call_accumulator[tc_index]['name'] = tc_delta.function.name
|
||||
if tc_delta.function.arguments:
|
||||
tool_call_accumulator[tc_index]['arguments'] += tc_delta.function.arguments
|
||||
tool_call_delta = tool_call_accumulator[tc_index]
|
||||
|
||||
# 检查是否完成（包含工具调用）
|
||||
tool_calls = None
|
||||
if choice.finish_reason == 'tool_calls' or (choice.finish_reason and tool_call_accumulator):
|
||||
tool_calls = self._parse_accumulated_tool_calls(tool_call_accumulator)
|
||||
|
||||
yield LLMStreamChunk(
|
||||
content=delta.content or '',
|
||||
is_finished=choice.finish_reason is not None,
|
||||
finish_reason=choice.finish_reason or '',
|
||||
tool_calls=tool_calls,
|
||||
tool_call_delta=tool_call_delta,
|
||||
)
|
||||
|
||||
# 最后一个 chunk 包含 usage 信息
|
||||
if chunk.usage:
|
||||
yield LLMStreamChunk(
|
||||
content='',
|
||||
is_finished=True,
|
||||
prompt_tokens=chunk.usage.prompt_tokens,
|
||||
completion_tokens=chunk.usage.completion_tokens,
|
||||
total_tokens=chunk.usage.total_tokens,
|
||||
)
|
||||
|
||||
# 模型类型推断关键词
|
||||
_EMBEDDING_KEYWORDS = ('embed', 'bge', 'nomic-embed', 'mxbai-embed', 'jina-embedding')
|
||||
_RERANK_KEYWORDS = ('rerank', 'reranker')
|
||||
_SKIP_KEYWORDS = ('dall-e', 'tts', 'whisper', 'audio', 'moderation', 'davinci', 'babbage')
|
||||
|
||||
async def fetch_models_from_api(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
通过 /v1/models API 在线拉取最新模型列表
|
||||
|
||||
适用于所有 OpenAI 兼容的提供商
|
||||
"""
|
||||
import httpx
|
||||
|
||||
base_url = (self.api_base or self.DEFAULT_API_BASE).rstrip('/')
|
||||
headers = {}
|
||||
if self.api_key:
|
||||
headers['Authorization'] = f'Bearer {self.api_key}'
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=15) as client:
|
||||
resp = await client.get(f'{base_url}/models', headers=headers)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
model_list = data.get('data', [])
|
||||
if not model_list:
|
||||
return []
|
||||
|
||||
results = []
|
||||
for item in model_list:
|
||||
model_id = item.get('id', '')
|
||||
if not model_id:
|
||||
continue
|
||||
|
||||
model_lower = model_id.lower()
|
||||
|
||||
# 跳过非文本模型
|
||||
if any(kw in model_lower for kw in self._SKIP_KEYWORDS):
|
||||
continue
|
||||
|
||||
# 推断模型类型
|
||||
if any(kw in model_lower for kw in self._RERANK_KEYWORDS):
|
||||
model_type = 'rerank'
|
||||
elif any(kw in model_lower for kw in self._EMBEDDING_KEYWORDS):
|
||||
model_type = 'embedding'
|
||||
else:
|
||||
model_type = 'chat'
|
||||
|
||||
results.append({
|
||||
'model_name': model_id,
|
||||
'display_name': model_id,
|
||||
'model_type': model_type,
|
||||
'max_tokens': 4096,
|
||||
'context_window': 4096,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': model_type == 'chat',
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
})
|
||||
|
||||
# 补充默认列表中的 embedding/rerank 模型（API 可能不返回这些类型）
|
||||
fetched_names = {m['model_name'] for m in results}
|
||||
for dm in self.__class__.get_default_models():
|
||||
if dm['model_type'] in ('embedding', 'rerank') and dm['model_name'] not in fetched_names:
|
||||
results.append(dm)
|
||||
|
||||
# 按类型排序：chat → embedding → rerank
|
||||
type_order = {'chat': 0, 'embedding': 1, 'rerank': 2}
|
||||
results.sort(key=lambda m: (type_order.get(m['model_type'], 9), m['model_name']))
|
||||
logger.info(f'从 {base_url} 拉取到 {len(results)} 个模型')
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f'在线拉取模型列表失败 ({base_url}): {e}')
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def get_default_models(cls) -> List[Dict[str, Any]]:
|
||||
"""获取默认模型列表"""
|
||||
return [
|
||||
# ---- Chat 模型 ----
|
||||
{
|
||||
'model_name': 'o4-mini',
|
||||
'display_name': 'o4-mini',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 16384,
|
||||
'context_window': 200000,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.0011,
|
||||
'output_price': 0.0044,
|
||||
},
|
||||
{
|
||||
'model_name': 'o3-mini',
|
||||
'display_name': 'o3-mini',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 16384,
|
||||
'context_window': 200000,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.0011,
|
||||
'output_price': 0.0044,
|
||||
},
|
||||
{
|
||||
'model_name': 'gpt-4.1',
|
||||
'display_name': 'GPT-4.1',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 32768,
|
||||
'context_window': 1047576,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.002,
|
||||
'output_price': 0.008,
|
||||
},
|
||||
{
|
||||
'model_name': 'gpt-4.1-mini',
|
||||
'display_name': 'GPT-4.1 Mini',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 32768,
|
||||
'context_window': 1047576,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.0004,
|
||||
'output_price': 0.0016,
|
||||
},
|
||||
{
|
||||
'model_name': 'gpt-4.1-nano',
|
||||
'display_name': 'GPT-4.1 Nano',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 32768,
|
||||
'context_window': 1047576,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.0001,
|
||||
'output_price': 0.0004,
|
||||
},
|
||||
{
|
||||
'model_name': 'gpt-4o',
|
||||
'display_name': 'GPT-4o',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 16384,
|
||||
'context_window': 128000,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.0025,
|
||||
'output_price': 0.01,
|
||||
},
|
||||
{
|
||||
'model_name': 'gpt-4o-mini',
|
||||
'display_name': 'GPT-4o Mini',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 16384,
|
||||
'context_window': 128000,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.00015,
|
||||
'output_price': 0.0006,
|
||||
},
|
||||
# ---- Embedding 模型 ----
|
||||
{
|
||||
'model_name': 'text-embedding-3-small',
|
||||
'display_name': 'Text Embedding 3 Small',
|
||||
'model_type': 'embedding',
|
||||
'max_tokens': 8191,
|
||||
'context_window': 8191,
|
||||
'input_price': 0.00002,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'text-embedding-3-large',
|
||||
'display_name': 'Text Embedding 3 Large',
|
||||
'model_type': 'embedding',
|
||||
'max_tokens': 8191,
|
||||
'context_window': 8191,
|
||||
'input_price': 0.00013,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'text-embedding-ada-002',
|
||||
'display_name': 'Text Embedding Ada 002',
|
||||
'model_type': 'embedding',
|
||||
'max_tokens': 8191,
|
||||
'context_window': 8191,
|
||||
'input_price': 0.0001,
|
||||
'output_price': 0,
|
||||
},
|
||||
]
|
||||
@@ -0,0 +1,555 @@
|
||||
"""
|
||||
阿里通义千问提供商适配器
|
||||
"""
|
||||
import logging
|
||||
from typing import AsyncGenerator, Dict, List, Any
|
||||
|
||||
from .base import BaseLLMProvider, LLMConfig, LLMMessage, LLMResponse, LLMStreamChunk, ToolCall
|
||||
from .registry import ProviderRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ProviderRegistry.register
|
||||
class QwenProvider(BaseLLMProvider):
|
||||
"""
|
||||
阿里通义千问提供商适配器
|
||||
|
||||
使用 OpenAI 兼容接口
|
||||
"""
|
||||
|
||||
provider_type = 'qwen'
|
||||
provider_name = '通义千问'
|
||||
supported_model_types = ['chat', 'embedding']
|
||||
|
||||
DEFAULT_API_BASE = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
|
||||
|
||||
def __init__(self, api_key: str = '', api_base: str = '', **kwargs):
|
||||
super().__init__(api_key, api_base, **kwargs)
|
||||
self.api_base = api_base or self.DEFAULT_API_BASE
|
||||
self._client = None
|
||||
self._async_client = None
|
||||
|
||||
def _get_client(self):
|
||||
"""获取同步客户端"""
|
||||
if self._client is None:
|
||||
try:
|
||||
from openai import OpenAI
|
||||
self._client = OpenAI(
|
||||
api_key=self.api_key,
|
||||
base_url=self.api_base,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError('请安装 openai 库: pip install openai')
|
||||
return self._client
|
||||
|
||||
def _get_async_client(self):
|
||||
"""获取异步客户端"""
|
||||
if self._async_client is None:
|
||||
try:
|
||||
from openai import AsyncOpenAI
|
||||
self._async_client = AsyncOpenAI(
|
||||
api_key=self.api_key,
|
||||
base_url=self.api_base,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError('请安装 openai 库: pip install openai')
|
||||
return self._async_client
|
||||
|
||||
def chat(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> LLMResponse:
|
||||
"""同步对话"""
|
||||
client = self._get_client()
|
||||
|
||||
# 构建请求参数
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': [m.to_dict() for m in messages],
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
'max_tokens': config.max_tokens,
|
||||
'stop': config.stop,
|
||||
**config.extra_params,
|
||||
}
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = 'required'
|
||||
elif config.tool_choice == 'none':
|
||||
kwargs['tool_choice'] = 'none'
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||||
else:
|
||||
kwargs['tool_choice'] = 'auto'
|
||||
|
||||
response = client.chat.completions.create(**kwargs)
|
||||
|
||||
choice = response.choices[0]
|
||||
usage = response.usage
|
||||
|
||||
# 解析工具调用
|
||||
tool_calls = None
|
||||
if choice.message.tool_calls:
|
||||
tool_calls = self._parse_tool_calls(choice.message.tool_calls)
|
||||
|
||||
return LLMResponse(
|
||||
content=choice.message.content or '',
|
||||
model=response.model,
|
||||
prompt_tokens=usage.prompt_tokens if usage else 0,
|
||||
completion_tokens=usage.completion_tokens if usage else 0,
|
||||
total_tokens=usage.total_tokens if usage else 0,
|
||||
finish_reason=choice.finish_reason or '',
|
||||
raw_response=response.model_dump(),
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
|
||||
def _parse_tool_calls(self, tool_calls) -> List[ToolCall]:
|
||||
"""解析工具调用"""
|
||||
import json
|
||||
result = []
|
||||
for tc in tool_calls:
|
||||
try:
|
||||
arguments = json.loads(tc.function.arguments) if tc.function.arguments else {}
|
||||
except json.JSONDecodeError:
|
||||
arguments = {'raw': tc.function.arguments}
|
||||
|
||||
result.append(ToolCall(
|
||||
id=tc.id,
|
||||
name=tc.function.name,
|
||||
arguments=arguments,
|
||||
))
|
||||
return result
|
||||
|
||||
async def chat_async(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> LLMResponse:
|
||||
"""异步对话"""
|
||||
client = self._get_async_client()
|
||||
|
||||
# 构建请求参数
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': [m.to_dict() for m in messages],
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
'max_tokens': config.max_tokens,
|
||||
'stop': config.stop,
|
||||
**config.extra_params,
|
||||
}
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = 'required'
|
||||
elif config.tool_choice == 'none':
|
||||
kwargs['tool_choice'] = 'none'
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||||
else:
|
||||
kwargs['tool_choice'] = 'auto'
|
||||
|
||||
response = await client.chat.completions.create(**kwargs)
|
||||
|
||||
choice = response.choices[0]
|
||||
usage = response.usage
|
||||
|
||||
# 解析工具调用
|
||||
tool_calls = None
|
||||
if choice.message.tool_calls:
|
||||
tool_calls = self._parse_tool_calls(choice.message.tool_calls)
|
||||
|
||||
return LLMResponse(
|
||||
content=choice.message.content or '',
|
||||
model=response.model,
|
||||
prompt_tokens=usage.prompt_tokens if usage else 0,
|
||||
completion_tokens=usage.completion_tokens if usage else 0,
|
||||
total_tokens=usage.total_tokens if usage else 0,
|
||||
finish_reason=choice.finish_reason or '',
|
||||
raw_response=response.model_dump(),
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
|
||||
async def chat_stream(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
) -> AsyncGenerator[LLMStreamChunk, None]:
|
||||
"""流式对话"""
|
||||
client = self._get_async_client()
|
||||
|
||||
# 构建请求参数
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': [m.to_dict() for m in messages],
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
'max_tokens': config.max_tokens,
|
||||
'stop': config.stop,
|
||||
'stream': True,
|
||||
'stream_options': {"include_usage": True},
|
||||
**config.extra_params,
|
||||
}
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = 'required'
|
||||
elif config.tool_choice == 'none':
|
||||
kwargs['tool_choice'] = 'none'
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||||
else:
|
||||
kwargs['tool_choice'] = 'auto'
|
||||
|
||||
stream = await client.chat.completions.create(**kwargs)
|
||||
|
||||
# 用于累积工具调用
|
||||
tool_call_accumulator = {}
|
||||
|
||||
async for chunk in stream:
|
||||
if chunk.choices:
|
||||
choice = chunk.choices[0]
|
||||
delta = choice.delta
|
||||
|
||||
# 处理工具调用增量
|
||||
tool_call_delta = None
|
||||
if hasattr(delta, 'tool_calls') and delta.tool_calls:
|
||||
for tc_delta in delta.tool_calls:
|
||||
tc_index = tc_delta.index
|
||||
if tc_index not in tool_call_accumulator:
|
||||
tool_call_accumulator[tc_index] = {
|
||||
'id': tc_delta.id or '',
|
||||
'name': '',
|
||||
'arguments': '',
|
||||
}
|
||||
if tc_delta.id:
|
||||
tool_call_accumulator[tc_index]['id'] = tc_delta.id
|
||||
if tc_delta.function:
|
||||
if tc_delta.function.name:
|
||||
tool_call_accumulator[tc_index]['name'] = tc_delta.function.name
|
||||
if tc_delta.function.arguments:
|
||||
tool_call_accumulator[tc_index]['arguments'] += tc_delta.function.arguments
|
||||
tool_call_delta = tool_call_accumulator[tc_index]
|
||||
|
||||
# 检查是否完成
|
||||
tool_calls = None
|
||||
if choice.finish_reason == 'tool_calls' or (choice.finish_reason and tool_call_accumulator):
|
||||
tool_calls = self._parse_accumulated_tool_calls(tool_call_accumulator)
|
||||
|
||||
yield LLMStreamChunk(
|
||||
content=delta.content or '',
|
||||
is_finished=choice.finish_reason is not None,
|
||||
finish_reason=choice.finish_reason or '',
|
||||
tool_calls=tool_calls,
|
||||
tool_call_delta=tool_call_delta,
|
||||
)
|
||||
|
||||
if chunk.usage:
|
||||
yield LLMStreamChunk(
|
||||
content='',
|
||||
is_finished=True,
|
||||
prompt_tokens=chunk.usage.prompt_tokens,
|
||||
completion_tokens=chunk.usage.completion_tokens,
|
||||
total_tokens=chunk.usage.total_tokens,
|
||||
)
|
||||
|
||||
def _parse_accumulated_tool_calls(self, accumulator: Dict) -> List[ToolCall]:
|
||||
"""解析累积的工具调用"""
|
||||
import json
|
||||
result = []
|
||||
for idx in sorted(accumulator.keys()):
|
||||
tc = accumulator[idx]
|
||||
try:
|
||||
arguments = json.loads(tc['arguments']) if tc['arguments'] else {}
|
||||
except json.JSONDecodeError:
|
||||
arguments = {'raw': tc['arguments']}
|
||||
|
||||
result.append(ToolCall(
|
||||
id=tc['id'],
|
||||
name=tc['name'],
|
||||
arguments=arguments,
|
||||
))
|
||||
return result
|
||||
|
||||
def chat_stream_sync(
|
||||
self,
|
||||
messages: List[LLMMessage],
|
||||
config: LLMConfig,
|
||||
):
|
||||
"""同步流式对话"""
|
||||
client = self._get_client()
|
||||
|
||||
# 构建请求参数
|
||||
kwargs = {
|
||||
'model': config.model,
|
||||
'messages': [m.to_dict() for m in messages],
|
||||
'temperature': config.temperature,
|
||||
'top_p': config.top_p,
|
||||
'max_tokens': config.max_tokens,
|
||||
'stop': config.stop,
|
||||
'stream': True,
|
||||
'stream_options': {"include_usage": True},
|
||||
**config.extra_params,
|
||||
}
|
||||
|
||||
# 添加 Function Calling 参数
|
||||
if config.tools:
|
||||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||||
if config.tool_choice == 'required':
|
||||
kwargs['tool_choice'] = 'required'
|
||||
elif config.tool_choice == 'none':
|
||||
kwargs['tool_choice'] = 'none'
|
||||
elif config.tool_choice != 'auto':
|
||||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||||
else:
|
||||
kwargs['tool_choice'] = 'auto'
|
||||
|
||||
stream = client.chat.completions.create(**kwargs)
|
||||
|
||||
# 用于累积工具调用
|
||||
tool_call_accumulator = {}
|
||||
|
||||
for chunk in stream:
|
||||
if chunk.choices:
|
||||
choice = chunk.choices[0]
|
||||
delta = choice.delta
|
||||
|
||||
# 处理工具调用增量
|
||||
tool_call_delta = None
|
||||
if hasattr(delta, 'tool_calls') and delta.tool_calls:
|
||||
for tc_delta in delta.tool_calls:
|
||||
tc_index = tc_delta.index
|
||||
if tc_index not in tool_call_accumulator:
|
||||
tool_call_accumulator[tc_index] = {
|
||||
'id': tc_delta.id or '',
|
||||
'name': '',
|
||||
'arguments': '',
|
||||
}
|
||||
if tc_delta.id:
|
||||
tool_call_accumulator[tc_index]['id'] = tc_delta.id
|
||||
if tc_delta.function:
|
||||
if tc_delta.function.name:
|
||||
tool_call_accumulator[tc_index]['name'] = tc_delta.function.name
|
||||
if tc_delta.function.arguments:
|
||||
tool_call_accumulator[tc_index]['arguments'] += tc_delta.function.arguments
|
||||
tool_call_delta = tool_call_accumulator[tc_index]
|
||||
|
||||
# 检查是否完成
|
||||
tool_calls = None
|
||||
if choice.finish_reason == 'tool_calls' or (choice.finish_reason and tool_call_accumulator):
|
||||
tool_calls = self._parse_accumulated_tool_calls(tool_call_accumulator)
|
||||
|
||||
yield LLMStreamChunk(
|
||||
content=delta.content or '',
|
||||
is_finished=choice.finish_reason is not None,
|
||||
finish_reason=choice.finish_reason or '',
|
||||
tool_calls=tool_calls,
|
||||
tool_call_delta=tool_call_delta,
|
||||
)
|
||||
|
||||
if chunk.usage:
|
||||
yield LLMStreamChunk(
|
||||
content='',
|
||||
is_finished=True,
|
||||
prompt_tokens=chunk.usage.prompt_tokens,
|
||||
completion_tokens=chunk.usage.completion_tokens,
|
||||
total_tokens=chunk.usage.total_tokens,
|
||||
)
|
||||
|
||||
_EMBEDDING_KEYWORDS = ('embed', 'bge')
|
||||
_RERANK_KEYWORDS = ('rerank',)
|
||||
_SKIP_KEYWORDS = (
|
||||
'paraformer', 'sambert', 'wordart', 'wanx', 'cosyvoice',
|
||||
'qwen-vl-ocr', 'tts', 'asr', 'realtime',
|
||||
'tongyi-xiaomi', 'flux', 'stable-diffusion',
|
||||
'qwen-image-', 'z-image-', 'gui-plus',
|
||||
'qwen-mt-', 'livetranslate', 'deep-search',
|
||||
'flash-character',
|
||||
)
|
||||
|
||||
async def fetch_models_from_api(self) -> List[Dict[str, Any]]:
|
||||
"""通过 DashScope 兼容的 /v1/models 端点在线拉取模型列表"""
|
||||
import httpx
|
||||
|
||||
base_url = (self.api_base or self.DEFAULT_API_BASE).rstrip('/')
|
||||
headers = {}
|
||||
if self.api_key:
|
||||
headers['Authorization'] = f'Bearer {self.api_key}'
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=15) as client:
|
||||
resp = await client.get(f'{base_url}/models', headers=headers)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
model_list = data.get('data', [])
|
||||
if not model_list:
|
||||
return []
|
||||
|
||||
results = []
|
||||
for item in model_list:
|
||||
model_id = item.get('id', '')
|
||||
if not model_id:
|
||||
continue
|
||||
|
||||
model_lower = model_id.lower()
|
||||
|
||||
if any(kw in model_lower for kw in self._SKIP_KEYWORDS):
|
||||
continue
|
||||
|
||||
if any(kw in model_lower for kw in self._RERANK_KEYWORDS):
|
||||
model_type = 'rerank'
|
||||
elif any(kw in model_lower for kw in self._EMBEDDING_KEYWORDS):
|
||||
model_type = 'embedding'
|
||||
else:
|
||||
model_type = 'chat'
|
||||
|
||||
results.append({
|
||||
'model_name': model_id,
|
||||
'display_name': model_id,
|
||||
'model_type': model_type,
|
||||
'max_tokens': 4096,
|
||||
'context_window': 4096,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': model_type == 'chat',
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
})
|
||||
|
||||
# DashScope /v1/models 可能只返回 chat 模型，
|
||||
# 将默认列表中的 embedding/rerank 模型合并进去
|
||||
fetched_names = {m['model_name'] for m in results}
|
||||
for dm in self.__class__.get_default_models():
|
||||
if dm['model_type'] in ('embedding', 'rerank') and dm['model_name'] not in fetched_names:
|
||||
results.append(dm)
|
||||
|
||||
type_order = {'chat': 0, 'embedding': 1, 'rerank': 2}
|
||||
results.sort(key=lambda m: (type_order.get(m['model_type'], 9), m['model_name']))
|
||||
logger.info(f'从 {base_url} 拉取到 {len(results)} 个模型（含补充的 embedding/rerank）')
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f'在线拉取模型列表失败 ({base_url}): {e}')
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def get_default_models(cls) -> List[Dict[str, Any]]:
|
||||
"""获取默认模型列表"""
|
||||
return [
|
||||
# ---- Chat 模型 ----
|
||||
{
|
||||
'model_name': 'qwen-max',
|
||||
'display_name': '通义千问 Max',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 32768,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.02,
|
||||
'output_price': 0.06,
|
||||
},
|
||||
{
|
||||
'model_name': 'qwen-plus',
|
||||
'display_name': '通义千问 Plus',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 131072,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.0008,
|
||||
'output_price': 0.002,
|
||||
},
|
||||
{
|
||||
'model_name': 'qwen-turbo',
|
||||
'display_name': '通义千问 Turbo',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 131072,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.0003,
|
||||
'output_price': 0.0006,
|
||||
},
|
||||
{
|
||||
'model_name': 'qwen3-235b-a22b',
|
||||
'display_name': 'Qwen3 235B-A22B',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 131072,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.004,
|
||||
'output_price': 0.012,
|
||||
},
|
||||
{
|
||||
'model_name': 'qwen3-32b',
|
||||
'display_name': 'Qwen3 32B',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 131072,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.002,
|
||||
'output_price': 0.006,
|
||||
},
|
||||
{
|
||||
'model_name': 'qwen-vl-max',
|
||||
'display_name': '通义千问 VL Max',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 2048,
|
||||
'context_window': 32768,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': False,
|
||||
'input_price': 0.02,
|
||||
'output_price': 0.06,
|
||||
},
|
||||
{
|
||||
'model_name': 'qwen-vl-plus',
|
||||
'display_name': '通义千问 VL Plus',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 2048,
|
||||
'context_window': 32768,
|
||||
'supports_vision': True,
|
||||
'supports_function_call': False,
|
||||
'input_price': 0.008,
|
||||
'output_price': 0.02,
|
||||
},
|
||||
# ---- Embedding 模型 ----
|
||||
{
|
||||
'model_name': 'text-embedding-v3',
|
||||
'display_name': '通义文本向量 V3',
|
||||
'model_type': 'embedding',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 8192,
|
||||
'input_price': 0.0007,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'text-embedding-v2',
|
||||
'display_name': '通义文本向量 V2',
|
||||
'model_type': 'embedding',
|
||||
'max_tokens': 2048,
|
||||
'context_window': 2048,
|
||||
'input_price': 0.0007,
|
||||
'output_price': 0,
|
||||
},
|
||||
# ---- Rerank 模型 ----
|
||||
{
|
||||
'model_name': 'gte-rerank',
|
||||
'display_name': 'GTE Rerank',
|
||||
'model_type': 'rerank',
|
||||
'max_tokens': 4096,
|
||||
'context_window': 4096,
|
||||
'input_price': 0.001,
|
||||
'output_price': 0,
|
||||
},
|
||||
]
|
||||
@@ -0,0 +1,152 @@
|
||||
"""
|
||||
提供商注册中心
|
||||
"""
|
||||
import logging
|
||||
from typing import Dict, List, Optional, Type
|
||||
|
||||
from .base import BaseLLMProvider
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ProviderRegistry:
|
||||
"""
|
||||
提供商注册中心
|
||||
|
||||
管理所有 LLM 提供商适配器的注册和获取
|
||||
"""
|
||||
|
||||
_providers: Dict[str, Type[BaseLLMProvider]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, provider_class: Type[BaseLLMProvider]) -> Type[BaseLLMProvider]:
|
||||
"""
|
||||
注册提供商（可作为装饰器使用）
|
||||
|
||||
Args:
|
||||
provider_class: 提供商类
|
||||
|
||||
Returns:
|
||||
提供商类
|
||||
"""
|
||||
provider_type = provider_class.provider_type
|
||||
if not provider_type:
|
||||
raise ValueError(f'Provider {provider_class.__name__} must have a provider_type')
|
||||
|
||||
cls._providers[provider_type] = provider_class
|
||||
logger.info(f'Registered LLM provider: {provider_type}')
|
||||
return provider_class
|
||||
|
||||
@classmethod
|
||||
def get(cls, provider_type: str) -> Optional[Type[BaseLLMProvider]]:
|
||||
"""
|
||||
获取提供商类
|
||||
|
||||
Args:
|
||||
provider_type: 提供商类型
|
||||
|
||||
Returns:
|
||||
提供商类或 None
|
||||
"""
|
||||
return cls._providers.get(provider_type)
|
||||
|
||||
@classmethod
|
||||
def create_instance(
|
||||
cls,
|
||||
provider_type: str,
|
||||
api_key: str = '',
|
||||
api_base: str = '',
|
||||
**kwargs
|
||||
) -> Optional[BaseLLMProvider]:
|
||||
"""
|
||||
创建提供商实例
|
||||
|
||||
Args:
|
||||
provider_type: 提供商类型
|
||||
api_key: API Key
|
||||
api_base: API 地址
|
||||
**kwargs: 其他配置
|
||||
|
||||
Returns:
|
||||
提供商实例或 None
|
||||
"""
|
||||
provider_class = cls.get(provider_type)
|
||||
if not provider_class:
|
||||
logger.warning(f'Unknown provider type: {provider_type}')
|
||||
return None
|
||||
|
||||
return provider_class(api_key=api_key, api_base=api_base, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def get_all_types(cls) -> List[Dict[str, str]]:
|
||||
"""
|
||||
获取所有已注册的提供商类型
|
||||
|
||||
Returns:
|
||||
提供商类型列表
|
||||
"""
|
||||
return [
|
||||
{
|
||||
'type': provider_type,
|
||||
'name': provider_class.provider_name,
|
||||
}
|
||||
for provider_type, provider_class in cls._providers.items()
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def get_default_models(cls, provider_type: str) -> List[Dict]:
|
||||
"""
|
||||
获取提供商的默认模型列表
|
||||
|
||||
Args:
|
||||
provider_type: 提供商类型
|
||||
|
||||
Returns:
|
||||
默认模型列表
|
||||
"""
|
||||
provider_class = cls.get(provider_type)
|
||||
if not provider_class:
|
||||
return []
|
||||
return provider_class.get_default_models()
|
||||
|
||||
@classmethod
|
||||
async def fetch_models_from_api(
|
||||
cls,
|
||||
provider_type: str,
|
||||
api_key: str = '',
|
||||
api_base: str = '',
|
||||
**kwargs
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
通过 API 在线拉取提供商的最新模型列表
|
||||
|
||||
Args:
|
||||
provider_type: 提供商类型
|
||||
api_key: API Key
|
||||
api_base: API 地址
|
||||
|
||||
Returns:
|
||||
模型列表，失败时返回空列表
|
||||
"""
|
||||
instance = cls.create_instance(provider_type, api_key=api_key, api_base=api_base, **kwargs)
|
||||
if not instance:
|
||||
return []
|
||||
return await instance.fetch_models_from_api()
|
||||
|
||||
|
||||
# 自动加载所有提供商
|
||||
def _load_providers():
|
||||
"""加载所有提供商适配器"""
|
||||
from . import openai_provider
|
||||
from . import claude_provider
|
||||
from . import qwen_provider
|
||||
from . import ollama_provider
|
||||
from . import deepseek_provider
|
||||
from . import siliconflow_provider
|
||||
|
||||
|
||||
# 延迟加载
|
||||
try:
|
||||
_load_providers()
|
||||
except ImportError as e:
|
||||
logger.warning(f'Failed to load some providers: {e}')
|
||||
@@ -0,0 +1,147 @@
|
||||
"""
|
||||
硅基流动 (SiliconFlow) 提供商适配器
|
||||
|
||||
兼容 OpenAI API，继承 OpenAI 适配器
|
||||
"""
|
||||
import logging
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from .openai_provider import OpenAIProvider
|
||||
from .registry import ProviderRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ProviderRegistry.register
|
||||
class SiliconFlowProvider(OpenAIProvider):
|
||||
"""
|
||||
硅基流动 (SiliconFlow) 提供商适配器
|
||||
|
||||
兼容 OpenAI API 格式，支持多种开源模型、Embedding 和 Rerank
|
||||
"""
|
||||
|
||||
provider_type = 'siliconflow'
|
||||
provider_name = 'SiliconFlow'
|
||||
supported_model_types = ['chat', 'embedding', 'rerank']
|
||||
|
||||
DEFAULT_API_BASE = 'https://api.siliconflow.cn/v1'
|
||||
|
||||
@classmethod
|
||||
def get_default_models(cls) -> List[Dict[str, Any]]:
|
||||
"""获取默认模型列表"""
|
||||
return [
|
||||
# ---- Chat 模型 ----
|
||||
{
|
||||
'model_name': 'Qwen/Qwen3-235B-A22B',
|
||||
'display_name': 'Qwen3 235B-A22B',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 131072,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.004,
|
||||
'output_price': 0.012,
|
||||
},
|
||||
{
|
||||
'model_name': 'Qwen/Qwen3-32B',
|
||||
'display_name': 'Qwen3 32B',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 131072,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.002,
|
||||
'output_price': 0.006,
|
||||
},
|
||||
{
|
||||
'model_name': 'Qwen/Qwen3-8B',
|
||||
'display_name': 'Qwen3 8B',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 131072,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'deepseek-ai/DeepSeek-V3',
|
||||
'display_name': 'DeepSeek V3',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 65536,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': True,
|
||||
'input_price': 0.002,
|
||||
'output_price': 0.008,
|
||||
},
|
||||
{
|
||||
'model_name': 'deepseek-ai/DeepSeek-R1',
|
||||
'display_name': 'DeepSeek R1',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 65536,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': False,
|
||||
'input_price': 0.004,
|
||||
'output_price': 0.016,
|
||||
},
|
||||
{
|
||||
'model_name': 'Pro/deepseek-ai/DeepSeek-R1',
|
||||
'display_name': 'DeepSeek R1 (Pro)',
|
||||
'model_type': 'chat',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 65536,
|
||||
'supports_vision': False,
|
||||
'supports_function_call': False,
|
||||
'input_price': 0.008,
|
||||
'output_price': 0.032,
|
||||
},
|
||||
# ---- Embedding 模型 ----
|
||||
{
|
||||
'model_name': 'BAAI/bge-m3',
|
||||
'display_name': 'BGE-M3',
|
||||
'model_type': 'embedding',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 8192,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'BAAI/bge-large-zh-v1.5',
|
||||
'display_name': 'BGE Large ZH v1.5',
|
||||
'model_type': 'embedding',
|
||||
'max_tokens': 512,
|
||||
'context_window': 512,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'jinaai/jina-embeddings-v3',
|
||||
'display_name': 'Jina Embeddings V3',
|
||||
'model_type': 'embedding',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 8192,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
# ---- Rerank 模型 ----
|
||||
{
|
||||
'model_name': 'BAAI/bge-reranker-v2-m3',
|
||||
'display_name': 'BGE Reranker V2 M3',
|
||||
'model_type': 'rerank',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 8192,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
{
|
||||
'model_name': 'jinaai/jina-reranker-v2-base-multilingual',
|
||||
'display_name': 'Jina Reranker V2 Multilingual',
|
||||
'model_type': 'rerank',
|
||||
'max_tokens': 8192,
|
||||
'context_window': 8192,
|
||||
'input_price': 0,
|
||||
'output_price': 0,
|
||||
},
|
||||
]
|
||||
@@ -0,0 +1,23 @@
|
||||
from fastapi import APIRouter
|
||||
|
||||
from ai_platform.api.provider_api import router as provider_router
|
||||
from ai_platform.api.model_api import router as model_router
|
||||
from ai_platform.api.app_api import router as app_router
|
||||
from ai_platform.api.chat_api import router as chat_router
|
||||
from ai_platform.api.workflow_api import router as workflow_router
|
||||
from ai_platform.api.agent_api import router as agent_router
|
||||
from ai_platform.api.speech_api import router as speech_router
|
||||
from ai_platform.knowledge.api import router as knowledge_router
|
||||
|
||||
# 创建总路由
|
||||
router = APIRouter(tags=["AI平台"])
|
||||
|
||||
# 注册子路由
|
||||
router.include_router(provider_router)
|
||||
router.include_router(model_router)
|
||||
router.include_router(app_router)
|
||||
router.include_router(chat_router)
|
||||
router.include_router(workflow_router)
|
||||
router.include_router(agent_router)
|
||||
router.include_router(speech_router)
|
||||
router.include_router(knowledge_router)
|
||||
@@ -0,0 +1,9 @@
|
||||
"""
|
||||
AI 平台 Schema 定义
|
||||
"""
|
||||
from .provider_schema import *
|
||||
from .model_schema import *
|
||||
from .app_schema import *
|
||||
from .chat_schema import *
|
||||
from .workflow_schema import *
|
||||
from .agent_schema import *
|
||||
@@ -0,0 +1,312 @@
|
||||
"""
|
||||
智能体 Schema
|
||||
"""
|
||||
from typing import Optional, List, Any, Dict
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, Field, ConfigDict
|
||||
|
||||
from app.base_schema import CSTDatetime
|
||||
|
||||
|
||||
# ==================== Agent Schema ====================
|
||||
|
||||
class PersonaConfig(BaseModel):
|
||||
"""人设配置"""
|
||||
role: str = Field(default="", description="角色定位")
|
||||
personality: List[str] = Field(default_factory=list, description="性格特点")
|
||||
skills: List[str] = Field(default_factory=list, description="技能")
|
||||
constraints: List[str] = Field(default_factory=list, description="约束条件")
|
||||
background: str = Field(default="", description="背景介绍")
|
||||
examples: List[Dict[str, str]] = Field(default_factory=list, description="对话示例")
|
||||
|
||||
|
||||
class AgentCreate(BaseModel):
|
||||
"""创建智能体"""
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
application_id: Optional[str] = Field(None, description="所属应用ID")
|
||||
is_global: bool = Field(default=False, description="是否在子应用中可见")
|
||||
name: str = Field(..., max_length=100, description="智能体名称")
|
||||
code: str = Field(..., max_length=100, pattern=r"^[a-zA-Z][a-zA-Z0-9_]*$", description="智能体编码（字母开头，只能包含字母、数字和下划线）")
|
||||
description: str = Field(default="", description="描述")
|
||||
avatar: str = Field(default="", description="头像")
|
||||
mode: str = Field(default="autonomous", description="运行模式")
|
||||
persona: Dict[str, Any] = Field(default_factory=dict, description="人设配置")
|
||||
system_prompt: str = Field(default="", description="系统提示词")
|
||||
model_id: Optional[str] = Field(None, description="模型 ID")
|
||||
temperature: float = Field(default=0.7, description="温度")
|
||||
top_p: float = Field(default=1.0, description="top_p")
|
||||
max_tokens: int = Field(default=4096, description="最大 Token")
|
||||
max_iterations: int = Field(default=10, description="最大推理轮数")
|
||||
welcome_message: str = Field(default="", description="开场白")
|
||||
suggested_questions: List[str] = Field(default_factory=list, description="推荐问题")
|
||||
workflow_id: Optional[str] = Field(None, description="工作流 ID（对话流模式）")
|
||||
knowledge_base_ids: List[str] = Field(default_factory=list, description="关联的知识库ID列表")
|
||||
knowledge_config: Dict[str, Any] = Field(default_factory=dict, description="知识库检索配置")
|
||||
is_public: bool = Field(default=False, description="是否公开")
|
||||
|
||||
|
||||
class AgentUpdate(BaseModel):
|
||||
"""更新智能体"""
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
application_id: Optional[str] = None
|
||||
is_global: Optional[bool] = None
|
||||
name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
avatar: Optional[str] = None
|
||||
mode: Optional[str] = None
|
||||
persona: Optional[Dict[str, Any]] = None
|
||||
system_prompt: Optional[str] = None
|
||||
model_id: Optional[str] = None
|
||||
temperature: Optional[float] = None
|
||||
top_p: Optional[float] = None
|
||||
max_tokens: Optional[int] = None
|
||||
max_iterations: Optional[int] = None
|
||||
welcome_message: Optional[str] = None
|
||||
suggested_questions: Optional[List[str]] = None
|
||||
workflow_id: Optional[str] = None
|
||||
knowledge_base_ids: Optional[List[str]] = None
|
||||
knowledge_config: Optional[Dict[str, Any]] = None
|
||||
is_public: Optional[bool] = None
|
||||
enable_memory: Optional[bool] = None
|
||||
memory_window: Optional[int] = None
|
||||
enable_streaming: Optional[bool] = None
|
||||
|
||||
|
||||
class AgentResponse(BaseModel):
|
||||
"""智能体输出"""
|
||||
model_config = ConfigDict(from_attributes=True, protected_namespaces=())
|
||||
|
||||
id: str
|
||||
application_id: Optional[str] = None
|
||||
is_global: bool = False
|
||||
name: str
|
||||
code: str
|
||||
description: str = ""
|
||||
avatar: str = ""
|
||||
mode: str = "autonomous"
|
||||
status: str = "draft"
|
||||
persona: Dict[str, Any] = Field(default_factory=dict)
|
||||
system_prompt: str = ""
|
||||
model_id: Optional[str] = None
|
||||
model_name: str = ""
|
||||
temperature: float = 0.7
|
||||
top_p: float = 1.0
|
||||
max_tokens: int = 4096
|
||||
max_iterations: int = 10
|
||||
welcome_message: str = ""
|
||||
suggested_questions: List[str] = Field(default_factory=list)
|
||||
workflow_id: Optional[str] = None
|
||||
workflow_name: str = ""
|
||||
workflow_type: str = "general"
|
||||
is_public: bool = False
|
||||
enable_memory: bool = False
|
||||
memory_window: int = 10
|
||||
enable_streaming: bool = True
|
||||
knowledge_base_ids: List[str] = Field(default_factory=list)
|
||||
knowledge_config: Dict[str, Any] = Field(default_factory=dict)
|
||||
conversation_count: int = 0
|
||||
message_count: int = 0
|
||||
total_tokens: int = 0
|
||||
sort: int = 0
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
sys_update_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
|
||||
class AgentListResponse(BaseModel):
|
||||
"""智能体列表输出"""
|
||||
model_config = ConfigDict(from_attributes=True, protected_namespaces=())
|
||||
|
||||
id: str
|
||||
application_id: Optional[str] = None
|
||||
application_name: str = ""
|
||||
is_global: bool = False
|
||||
name: str
|
||||
code: str
|
||||
description: str = ""
|
||||
avatar: str = ""
|
||||
mode: str = "autonomous"
|
||||
status: str = "draft"
|
||||
model_name: str = ""
|
||||
is_public: bool = False
|
||||
conversation_count: int = 0
|
||||
message_count: int = 0
|
||||
has_menu: bool = False
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
|
||||
# ==================== 导入/导出 Schema ====================
|
||||
|
||||
class AgentImportCheckIn(BaseModel):
|
||||
"""智能体导入预检查请求"""
|
||||
code: str = Field(..., description="智能体编码")
|
||||
|
||||
|
||||
class AgentImportCheckOut(BaseModel):
|
||||
"""智能体导入预检查结果"""
|
||||
code_exists: bool = Field(..., description="智能体编码是否已存在")
|
||||
can_import: bool = Field(..., description="是否可以直接导入（编码不冲突）")
|
||||
|
||||
|
||||
class AgentImportIn(BaseModel):
|
||||
"""智能体配置导入"""
|
||||
application_id: Optional[str] = Field(None, description="所属应用ID")
|
||||
is_global: bool = Field(default=False, description="是否在子应用中可见")
|
||||
name: str = Field(..., description="智能体名称")
|
||||
code: str = Field(..., description="智能体编码")
|
||||
description: str = Field(default="", description="描述")
|
||||
avatar: str = Field(default="", description="头像")
|
||||
mode: str = Field(default="autonomous", description="运行模式")
|
||||
persona: Dict[str, Any] = Field(default_factory=dict, description="人设配置")
|
||||
system_prompt: str = Field(default="", description="系统提示词")
|
||||
model_name: str = Field(default="", description="模型名称（API model_name）")
|
||||
temperature: float = Field(default=0.7, description="温度")
|
||||
top_p: float = Field(default=1.0, description="top_p")
|
||||
max_tokens: int = Field(default=4096, description="最大 Token")
|
||||
max_iterations: int = Field(default=10, description="最大推理轮数")
|
||||
welcome_message: str = Field(default="", description="开场白")
|
||||
suggested_questions: List[str] = Field(default_factory=list, description="推荐问题")
|
||||
workflow_code: str = Field(default="", description="关联工作流编码")
|
||||
knowledge_base_codes: List[str] = Field(default_factory=list, description="知识库编码列表")
|
||||
knowledge_config: Dict[str, Any] = Field(default_factory=dict, description="知识库检索配置")
|
||||
enable_memory: bool = Field(default=False, description="是否启用对话记忆")
|
||||
memory_window: int = Field(default=10, description="记忆窗口大小")
|
||||
enable_streaming: bool = Field(default=True, description="是否启用流式输出")
|
||||
is_public: bool = Field(default=False, description="是否公开")
|
||||
|
||||
|
||||
# ==================== Conversation Schema ====================
|
||||
|
||||
class AgentConversationCreate(BaseModel):
|
||||
"""创建对话"""
|
||||
title: str = Field(default="", description="对话标题")
|
||||
|
||||
|
||||
class AgentConversationResponse(BaseModel):
|
||||
"""对话输出"""
|
||||
id: str
|
||||
agent_id: str
|
||||
agent_name: str = ""
|
||||
title: str = ""
|
||||
summary: str = ""
|
||||
message_count: int = 0
|
||||
total_tokens: int = 0
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
sys_update_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class AgentConversationListResponse(BaseModel):
|
||||
"""对话列表输出"""
|
||||
id: str
|
||||
agent_id: str
|
||||
agent_name: str = ""
|
||||
title: str = ""
|
||||
message_count: int = 0
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
# ==================== Message Schema ====================
|
||||
|
||||
class ReasoningStep(BaseModel):
|
||||
"""推理步骤"""
|
||||
type: str = Field(..., description="步骤类型: thought/action/observation")
|
||||
content: str = Field(default="", description="内容")
|
||||
tool: str = Field(default="", description="工具名称（action 类型）")
|
||||
params: Dict[str, Any] = Field(default_factory=dict, description="工具参数（action 类型）")
|
||||
timestamp: str = Field(default="", description="时间戳")
|
||||
|
||||
|
||||
class ToolCallRecord(BaseModel):
|
||||
"""工具调用记录"""
|
||||
tool: str = Field(..., description="工具名称")
|
||||
params: Dict[str, Any] = Field(default_factory=dict, description="参数")
|
||||
result: Any = Field(default=None, description="结果")
|
||||
elapsed_time: int = Field(default=0, description="耗时（毫秒）")
|
||||
status: str = Field(default="success", description="状态")
|
||||
|
||||
|
||||
class AttachmentResponse(BaseModel):
|
||||
"""附件输出
|
||||
|
||||
附件通过 file_id 关联文件管理系统(core_file_manager)
|
||||
前端通过 file_id 调用文件管理 API 获取访问 URL
|
||||
"""
|
||||
file_id: str = Field(..., description="文件管理系统中的文件 ID")
|
||||
type: str = Field(default="file", description="附件类型: image/file/audio/video")
|
||||
name: str = Field(default="", description="文件名")
|
||||
mime_type: str = Field(default="", description="MIME 类型")
|
||||
size: int = Field(default=0, description="文件大小（字节）")
|
||||
|
||||
|
||||
class AgentMessageResponse(BaseModel):
|
||||
"""消息输出"""
|
||||
id: str
|
||||
role: str
|
||||
content: str = ""
|
||||
attachments: List[AttachmentResponse] = Field(default_factory=list)
|
||||
status: str = "completed"
|
||||
reasoning_steps: List[ReasoningStep] = Field(default_factory=list)
|
||||
tool_calls: List[ToolCallRecord] = Field(default_factory=list)
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
elapsed_time: int = 0
|
||||
error_message: str = ""
|
||||
feedback: str = ""
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class AttachmentInput(BaseModel):
|
||||
"""附件输入
|
||||
|
||||
附件通过 file_id 关联文件管理系统(core_file_manager)
|
||||
"""
|
||||
file_id: str = Field(..., description="文件管理系统中的文件 ID")
|
||||
type: str = Field(default="file", description="附件类型: image/file/audio/video")
|
||||
|
||||
|
||||
class ChatInput(BaseModel):
|
||||
"""对话输入"""
|
||||
message: str = Field(..., description="用户消息")
|
||||
conversation_id: Optional[str] = Field(None, description="对话 ID（续聊时提供）")
|
||||
application_id: Optional[str] = Field(None, description="子应用 ID（在子应用下自动注入）")
|
||||
form_code: Optional[str] = Field(None, description="表单编码（从表单列表调用时自动注入）")
|
||||
attachments: Optional[List[AttachmentInput]] = Field(None, description="附件列表")
|
||||
|
||||
|
||||
class ChatStreamEvent(BaseModel):
|
||||
"""对话流式事件"""
|
||||
type: str = Field(..., description="事件类型")
|
||||
content: str = Field(default="", description="内容")
|
||||
tool: str = Field(default="", description="工具名称")
|
||||
params: Dict[str, Any] = Field(default_factory=dict, description="工具参数")
|
||||
message_id: str = Field(default="", description="消息 ID")
|
||||
conversation_id: str = Field(default="", description="对话 ID")
|
||||
tokens_used: int = Field(default=0, description="Token 使用量")
|
||||
elapsed_time: int = Field(default=0, description="耗时")
|
||||
|
||||
|
||||
# ==================== Feedback Schema ====================
|
||||
|
||||
class MessageFeedback(BaseModel):
|
||||
"""消息反馈"""
|
||||
feedback: str = Field(..., description="反馈类型: like/dislike")
|
||||
content: str = Field(default="", description="反馈内容")
|
||||
|
||||
|
||||
# ==================== Publish Schema ====================
|
||||
|
||||
class AgentPublishInput(BaseModel):
|
||||
"""发布智能体到菜单"""
|
||||
menu_name: str = Field(..., description="菜单名称")
|
||||
menu_parent_id: Optional[str] = Field(None, description="上级菜单 ID")
|
||||
menu_icon: str = Field(default="lucide:bot", description="菜单图标")
|
||||
menu_order: int = Field(default=0, description="菜单排序")
|
||||
@@ -0,0 +1,92 @@
|
||||
"""
|
||||
AI 应用 Schema
|
||||
"""
|
||||
from typing import Optional, List, Any
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, Field, ConfigDict
|
||||
|
||||
from app.base_schema import CSTDatetime
|
||||
|
||||
|
||||
class AppCreate(BaseModel):
|
||||
"""创建应用"""
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
name: str = Field(..., max_length=100, description="应用名称")
|
||||
code: str = Field(..., max_length=100, description="应用编码")
|
||||
description: str = Field(default="", description="描述")
|
||||
icon: str = Field(default="", description="图标")
|
||||
app_type: str = Field(default="chat", description="应用类型")
|
||||
model_id: Optional[str] = Field(None, description="默认模型 ID")
|
||||
system_prompt: str = Field(default="", description="系统提示词")
|
||||
temperature: float = Field(default=0.7, description="温度")
|
||||
top_p: float = Field(default=1.0, description="top_p")
|
||||
max_tokens: int = Field(default=2048, description="最大 Token")
|
||||
opening_statement: str = Field(default="", description="开场白")
|
||||
suggested_questions: List[str] = Field(default_factory=list, description="建议问题")
|
||||
is_public: bool = Field(default=False, description="是否公开")
|
||||
|
||||
|
||||
class AppUpdate(BaseModel):
|
||||
"""更新应用"""
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
icon: Optional[str] = None
|
||||
model_id: Optional[str] = None
|
||||
system_prompt: Optional[str] = None
|
||||
temperature: Optional[float] = None
|
||||
top_p: Optional[float] = None
|
||||
max_tokens: Optional[int] = None
|
||||
opening_statement: Optional[str] = None
|
||||
suggested_questions: Optional[List[str]] = None
|
||||
is_public: Optional[bool] = None
|
||||
workflow_definition: Optional[dict] = None
|
||||
|
||||
|
||||
class AppResponse(BaseModel):
|
||||
"""应用输出"""
|
||||
model_config = ConfigDict(from_attributes=True, protected_namespaces=())
|
||||
|
||||
id: str
|
||||
name: str
|
||||
code: str
|
||||
description: str = ""
|
||||
icon: str = ""
|
||||
app_type: str = "chat"
|
||||
status: str = "draft"
|
||||
model_id: Optional[str] = None
|
||||
model_name: str = ""
|
||||
system_prompt: str = ""
|
||||
temperature: float = 0.7
|
||||
top_p: float = 1.0
|
||||
max_tokens: int = 2048
|
||||
opening_statement: str = ""
|
||||
suggested_questions: List[Any] = Field(default_factory=list)
|
||||
workflow_definition: dict = Field(default_factory=dict)
|
||||
is_public: bool = False
|
||||
conversation_count: int = 0
|
||||
message_count: int = 0
|
||||
sort: int = 0
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
sys_update_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
|
||||
class AppListResponse(BaseModel):
|
||||
"""应用列表输出"""
|
||||
model_config = ConfigDict(from_attributes=True, protected_namespaces=())
|
||||
|
||||
id: str
|
||||
name: str
|
||||
code: str
|
||||
description: str = ""
|
||||
icon: str = ""
|
||||
app_type: str = "chat"
|
||||
status: str = "draft"
|
||||
model_name: str = ""
|
||||
is_public: bool = False
|
||||
conversation_count: int = 0
|
||||
message_count: int = 0
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
@@ -0,0 +1,76 @@
|
||||
"""
|
||||
对话 Schema
|
||||
"""
|
||||
from typing import Optional
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, Field, ConfigDict
|
||||
|
||||
from app.base_schema import CSTDatetime
|
||||
|
||||
|
||||
class ConversationCreate(BaseModel):
|
||||
"""创建对话"""
|
||||
app_id: str = Field(..., description="应用 ID")
|
||||
title: str = Field(default="", description="对话标题")
|
||||
|
||||
|
||||
class ConversationUpdate(BaseModel):
|
||||
"""更新对话"""
|
||||
title: Optional[str] = None
|
||||
is_pinned: Optional[bool] = None
|
||||
|
||||
|
||||
class ConversationResponse(BaseModel):
|
||||
"""对话输出"""
|
||||
id: str
|
||||
app_id: str
|
||||
app_name: str = ""
|
||||
title: str = ""
|
||||
message_count: int = 0
|
||||
total_tokens: int = 0
|
||||
is_pinned: bool = False
|
||||
sort: int = 0
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
sys_update_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class ConversationListResponse(BaseModel):
|
||||
"""对话列表输出"""
|
||||
id: str
|
||||
title: str = ""
|
||||
message_count: int = 0
|
||||
is_pinned: bool = False
|
||||
sys_update_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class MessageResponse(BaseModel):
|
||||
"""消息输出"""
|
||||
model_config = ConfigDict(from_attributes=True, protected_namespaces=())
|
||||
|
||||
id: str
|
||||
role: str
|
||||
content: str = ""
|
||||
status: str = "completed"
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
model_name: str = ""
|
||||
latency: int = 0
|
||||
error_message: str = ""
|
||||
feedback: str = ""
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
|
||||
class SendMessageInput(BaseModel):
|
||||
"""发送消息"""
|
||||
content: str = Field(..., description="消息内容")
|
||||
|
||||
|
||||
class MessageFeedbackInput(BaseModel):
|
||||
"""消息反馈"""
|
||||
feedback: str = Field(..., description="反馈: like/dislike")
|
||||
@@ -0,0 +1,101 @@
|
||||
"""
|
||||
LLM 模型 Schema
|
||||
"""
|
||||
from typing import Optional
|
||||
from datetime import datetime
|
||||
from decimal import Decimal
|
||||
|
||||
from pydantic import BaseModel, Field, ConfigDict
|
||||
|
||||
from app.base_schema import CSTDatetime
|
||||
|
||||
|
||||
class ModelCreate(BaseModel):
|
||||
"""创建模型"""
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
provider_id: str = Field(..., description="提供商 ID")
|
||||
model_name: str = Field(..., max_length=100, description="模型名称")
|
||||
display_name: str = Field(..., max_length=100, description="显示名称")
|
||||
model_type: str = Field(default="chat", description="模型类型")
|
||||
max_tokens: int = Field(default=4096, description="最大 Token")
|
||||
context_window: int = Field(default=4096, description="上下文窗口")
|
||||
default_temperature: float = Field(default=0.7, description="默认温度")
|
||||
default_top_p: float = Field(default=1.0, description="默认 top_p")
|
||||
input_price: Decimal = Field(default=0, description="输入价格")
|
||||
output_price: Decimal = Field(default=0, description="输出价格")
|
||||
supports_vision: bool = Field(default=False, description="支持视觉")
|
||||
supports_function_call: bool = Field(default=False, description="支持函数调用")
|
||||
supports_streaming: bool = Field(default=True, description="支持流式")
|
||||
is_active: bool = Field(default=True, description="是否启用")
|
||||
|
||||
|
||||
class ModelUpdate(BaseModel):
|
||||
"""更新模型"""
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
display_name: Optional[str] = None
|
||||
model_type: Optional[str] = None
|
||||
max_tokens: Optional[int] = None
|
||||
context_window: Optional[int] = None
|
||||
default_temperature: Optional[float] = None
|
||||
default_top_p: Optional[float] = None
|
||||
input_price: Optional[Decimal] = None
|
||||
output_price: Optional[Decimal] = None
|
||||
supports_vision: Optional[bool] = None
|
||||
supports_function_call: Optional[bool] = None
|
||||
supports_streaming: Optional[bool] = None
|
||||
is_active: Optional[bool] = None
|
||||
|
||||
|
||||
class ModelResponse(BaseModel):
|
||||
"""模型输出"""
|
||||
model_config = ConfigDict(from_attributes=True, protected_namespaces=())
|
||||
|
||||
id: str
|
||||
provider_id: str
|
||||
provider_name: str = ""
|
||||
model_name: str
|
||||
display_name: str
|
||||
model_type: str = "chat"
|
||||
max_tokens: int = 4096
|
||||
context_window: int = 4096
|
||||
default_temperature: float = 0.7
|
||||
default_top_p: float = 1.0
|
||||
input_price: Decimal = Decimal(0)
|
||||
output_price: Decimal = Decimal(0)
|
||||
supports_vision: bool = False
|
||||
supports_function_call: bool = False
|
||||
supports_streaming: bool = True
|
||||
is_active: bool = True
|
||||
sort: int = 0
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
sys_update_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
|
||||
class ModelListResponse(BaseModel):
|
||||
"""模型列表输出"""
|
||||
model_config = ConfigDict(from_attributes=True, protected_namespaces=())
|
||||
|
||||
id: str
|
||||
provider_id: str
|
||||
provider_name: str = ""
|
||||
model_name: str
|
||||
display_name: str
|
||||
model_type: str = "chat"
|
||||
is_active: bool = True
|
||||
|
||||
|
||||
class DefaultModelResponse(BaseModel):
|
||||
"""默认模型信息"""
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
model_name: str
|
||||
display_name: str
|
||||
model_type: str = "chat"
|
||||
max_tokens: int = 4096
|
||||
context_window: int = 4096
|
||||
supports_vision: bool = False
|
||||
supports_function_call: bool = False
|
||||
input_price: float = 0
|
||||
output_price: float = 0
|
||||
@@ -0,0 +1,70 @@
|
||||
"""
|
||||
LLM 提供商 Schema
|
||||
"""
|
||||
from typing import Optional
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, Field, ConfigDict
|
||||
|
||||
from app.base_schema import CSTDatetime
|
||||
|
||||
|
||||
class ProviderCreate(BaseModel):
|
||||
"""创建提供商"""
|
||||
name: str = Field(..., max_length=100, description="提供商名称")
|
||||
provider_type: str = Field(..., description="提供商类型")
|
||||
api_key: str = Field(default="", description="API Key")
|
||||
api_base: str = Field(default="", description="API 地址")
|
||||
api_version: str = Field(default="", description="API 版本")
|
||||
ollama_host: str = Field(default="http://localhost:11434", description="Ollama 地址")
|
||||
description: str = Field(default="", description="描述")
|
||||
is_active: bool = Field(default=True, description="是否启用")
|
||||
|
||||
|
||||
class ProviderUpdate(BaseModel):
|
||||
"""更新提供商"""
|
||||
name: Optional[str] = Field(None, max_length=100)
|
||||
api_key: Optional[str] = None
|
||||
api_base: Optional[str] = None
|
||||
api_version: Optional[str] = None
|
||||
ollama_host: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
is_active: Optional[bool] = None
|
||||
|
||||
|
||||
class ProviderResponse(BaseModel):
|
||||
"""提供商输出"""
|
||||
id: str
|
||||
name: str
|
||||
provider_type: str
|
||||
api_key_masked: str = ""
|
||||
api_base: str = ""
|
||||
api_version: str = ""
|
||||
ollama_host: str = ""
|
||||
description: str = ""
|
||||
is_active: bool = True
|
||||
quota_limit: int = 0
|
||||
quota_used: int = 0
|
||||
sort: int = 0
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
sys_update_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class ProviderListResponse(BaseModel):
|
||||
"""提供商列表输出"""
|
||||
id: str
|
||||
name: str
|
||||
provider_type: str
|
||||
is_active: bool = True
|
||||
description: str = ""
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class ProviderTypeResponse(BaseModel):
|
||||
"""提供商类型"""
|
||||
type: str
|
||||
name: str
|
||||
@@ -0,0 +1,170 @@
|
||||
"""
|
||||
AI 工作流 Schema
|
||||
"""
|
||||
from typing import Optional, List, Any
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, Field, ConfigDict
|
||||
|
||||
from app.base_schema import CSTDatetime
|
||||
|
||||
|
||||
class WorkflowCreate(BaseModel):
|
||||
"""创建工作流"""
|
||||
application_id: Optional[str] = Field(None, description="所属应用ID")
|
||||
is_global: bool = Field(default=False, description="是否在子应用中可见")
|
||||
name: str = Field(..., max_length=100, description="工作流名称")
|
||||
code: str = Field(..., max_length=100, pattern=r"^[a-zA-Z][a-zA-Z0-9_]*$", description="工作流编码（字母开头，只能包含字母、数字和下划线）")
|
||||
workflow_type: str = Field(default="general", description="工作流类型: general/application/form/report/data_process/automation")
|
||||
description: str = Field(default="", description="描述")
|
||||
definition: dict = Field(default_factory=dict, description="工作流定义")
|
||||
input_variables: List[dict] = Field(default_factory=list, description="输入变量")
|
||||
output_variables: List[dict] = Field(default_factory=list, description="输出变量")
|
||||
|
||||
|
||||
class WorkflowUpdate(BaseModel):
|
||||
"""更新工作流"""
|
||||
application_id: Optional[str] = None
|
||||
is_global: Optional[bool] = None
|
||||
name: Optional[str] = None
|
||||
code: Optional[str] = Field(None, max_length=100, pattern=r"^[a-zA-Z][a-zA-Z0-9_]*$", description="工作流编码（字母开头，只能包含字母、数字和下划线）")
|
||||
workflow_type: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
definition: Optional[dict] = None
|
||||
input_variables: Optional[List[dict]] = None
|
||||
output_variables: Optional[List[dict]] = None
|
||||
|
||||
|
||||
class WorkflowResponse(BaseModel):
|
||||
"""工作流输出"""
|
||||
id: str
|
||||
application_id: Optional[str] = None
|
||||
is_global: bool = False
|
||||
name: str
|
||||
code: str
|
||||
workflow_type: str = "general"
|
||||
description: str = ""
|
||||
status: str = "draft"
|
||||
version: int = 1
|
||||
published_version: Optional[int] = None
|
||||
published_at: Optional[CSTDatetime] = None
|
||||
published_definition: Optional[dict] = None
|
||||
definition: dict = Field(default_factory=dict)
|
||||
input_variables: List[dict] = Field(default_factory=list)
|
||||
output_variables: List[dict] = Field(default_factory=list)
|
||||
run_count: int = 0
|
||||
success_count: int = 0
|
||||
sort: int = 0
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
sys_update_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class WorkflowListResponse(BaseModel):
|
||||
"""工作流列表输出"""
|
||||
id: str
|
||||
application_id: Optional[str] = None
|
||||
application_name: str = ""
|
||||
is_global: bool = False
|
||||
name: str
|
||||
code: str
|
||||
workflow_type: str = "general"
|
||||
description: str = ""
|
||||
status: str = "draft"
|
||||
version: int = 1
|
||||
run_count: int = 0
|
||||
success_count: int = 0
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class WorkflowRunInput(BaseModel):
|
||||
"""运行工作流"""
|
||||
inputs: dict = Field(default_factory=dict, description="输入数据")
|
||||
use_draft: bool = Field(default=False, description="是否使用草稿版本（编辑器调试用）")
|
||||
|
||||
|
||||
class WorkflowResumeInput(BaseModel):
|
||||
"""恢复工作流执行"""
|
||||
user_input: Any = Field(..., description="用户输入的值")
|
||||
|
||||
|
||||
class WorkflowRunResponse(BaseModel):
|
||||
"""工作流运行记录输出"""
|
||||
id: str
|
||||
workflow_id: str
|
||||
workflow_name: str = ""
|
||||
status: str = "pending"
|
||||
trigger_type: str = "api"
|
||||
use_draft: bool = False
|
||||
workflow_version: Optional[int] = None
|
||||
definition_snapshot: dict = Field(default_factory=dict)
|
||||
inputs: dict = Field(default_factory=dict)
|
||||
outputs: dict = Field(default_factory=dict)
|
||||
execution_log: List[dict] = Field(default_factory=list)
|
||||
current_node_id: str = ""
|
||||
waiting_config: dict = Field(default_factory=dict)
|
||||
error_message: str = ""
|
||||
total_tokens: int = 0
|
||||
total_steps: int = 0
|
||||
elapsed_time: int = 0
|
||||
started_at: Optional[CSTDatetime] = None
|
||||
completed_at: Optional[CSTDatetime] = None
|
||||
sys_create_datetime: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class WorkflowRunListResponse(BaseModel):
|
||||
"""工作流运行记录列表输出"""
|
||||
id: str
|
||||
workflow_id: str
|
||||
workflow_name: str = ""
|
||||
status: str = "pending"
|
||||
trigger_type: str = "api"
|
||||
total_steps: int = 0
|
||||
total_tokens: int = 0
|
||||
elapsed_time: int = 0
|
||||
error_message: str = ""
|
||||
started_at: Optional[CSTDatetime] = None
|
||||
completed_at: Optional[CSTDatetime] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class WorkflowImportCheckIn(BaseModel):
|
||||
"""工作流导入预检查请求"""
|
||||
code: str = Field(..., description="工作流编码")
|
||||
|
||||
|
||||
class WorkflowImportCheckOut(BaseModel):
|
||||
"""工作流导入预检查结果"""
|
||||
code_exists: bool = Field(..., description="工作流编码是否已存在")
|
||||
can_import: bool = Field(..., description="是否可以直接导入（编码不冲突）")
|
||||
|
||||
|
||||
class WorkflowImportIn(BaseModel):
|
||||
"""工作流配置导入"""
|
||||
application_id: Optional[str] = Field(None, description="所属应用ID")
|
||||
is_global: bool = Field(default=False, description="是否在子应用中可见")
|
||||
name: str = Field(..., description="工作流名称")
|
||||
code: str = Field(..., description="工作流编码")
|
||||
workflow_type: str = Field(default="general", description="工作流类型")
|
||||
description: str = Field(default="", description="描述")
|
||||
definition: dict = Field(default_factory=dict, description="工作流定义")
|
||||
input_variables: List[dict] = Field(default_factory=list, description="输入变量")
|
||||
output_variables: List[dict] = Field(default_factory=list, description="输出变量")
|
||||
|
||||
|
||||
class NodeSchemaResponse(BaseModel):
|
||||
"""节点 Schema 输出"""
|
||||
type: str
|
||||
name: str
|
||||
category: str = ""
|
||||
icon: str = ""
|
||||
description: str = ""
|
||||
inputs: List[dict] = Field(default_factory=list)
|
||||
outputs: List[dict] = Field(default_factory=list)
|
||||
supports_branches: bool = False
|
||||
@@ -0,0 +1,14 @@
|
||||
"""
|
||||
AI 平台服务层
|
||||
"""
|
||||
from .llm_service import LLMService
|
||||
from .chat_service import ChatService
|
||||
from .workflow_service import AIWorkflowService
|
||||
from .agent_service import AgentService
|
||||
|
||||
__all__ = [
|
||||
'LLMService',
|
||||
'ChatService',
|
||||
'AIWorkflowService',
|
||||
'AgentService',
|
||||
]
|
||||
@@ -0,0 +1,229 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
智能体导入/导出
|
||||
"""
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from ai_platform.models import Agent, AIWorkflow, LLMModel
|
||||
from ai_platform.knowledge.models import KnowledgeBase
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AgentImportExportException(Exception):
|
||||
"""智能体导入导出异常"""
|
||||
pass
|
||||
|
||||
|
||||
async def export_config(db: AsyncSession, agent_id: str) -> Dict[str, Any]:
|
||||
"""导出智能体配置（外部引用使用 code / model_name）"""
|
||||
result = await db.execute(
|
||||
select(Agent).where(
|
||||
Agent.id == agent_id,
|
||||
Agent.is_deleted == False,
|
||||
)
|
||||
)
|
||||
agent = result.scalar_one_or_none()
|
||||
if not agent:
|
||||
raise AgentImportExportException("智能体不存在")
|
||||
|
||||
model_name = ""
|
||||
if agent.model_id:
|
||||
model_result = await db.execute(
|
||||
select(LLMModel.model_name).where(
|
||||
LLMModel.id == agent.model_id,
|
||||
LLMModel.is_deleted == False,
|
||||
)
|
||||
)
|
||||
row = model_result.first()
|
||||
if row:
|
||||
model_name = row[0] or ""
|
||||
|
||||
workflow_code = ""
|
||||
if agent.workflow_id:
|
||||
wf_result = await db.execute(
|
||||
select(AIWorkflow.code).where(
|
||||
AIWorkflow.id == agent.workflow_id,
|
||||
AIWorkflow.is_deleted == False,
|
||||
)
|
||||
)
|
||||
row = wf_result.first()
|
||||
if row:
|
||||
workflow_code = row[0] or ""
|
||||
|
||||
knowledge_base_codes: List[str] = []
|
||||
kb_ids = agent.knowledge_base_ids or []
|
||||
if kb_ids:
|
||||
kb_result = await db.execute(
|
||||
select(KnowledgeBase.code).where(
|
||||
KnowledgeBase.id.in_(kb_ids),
|
||||
KnowledgeBase.is_deleted == False,
|
||||
)
|
||||
)
|
||||
knowledge_base_codes = [row[0] for row in kb_result if row[0]]
|
||||
|
||||
return {
|
||||
"name": agent.name,
|
||||
"code": agent.code,
|
||||
"description": agent.description or "",
|
||||
"avatar": agent.avatar or "",
|
||||
"mode": agent.mode or "autonomous",
|
||||
"is_global": agent.is_global or False,
|
||||
"persona": agent.persona or {},
|
||||
"system_prompt": agent.system_prompt or "",
|
||||
"model_name": model_name,
|
||||
"temperature": agent.temperature if agent.temperature is not None else 0.7,
|
||||
"top_p": agent.top_p if agent.top_p is not None else 1.0,
|
||||
"max_tokens": agent.max_tokens or 4096,
|
||||
"max_iterations": agent.max_iterations or 10,
|
||||
"welcome_message": agent.welcome_message or "",
|
||||
"suggested_questions": agent.suggested_questions or [],
|
||||
"workflow_code": workflow_code,
|
||||
"knowledge_base_codes": knowledge_base_codes,
|
||||
"knowledge_config": agent.knowledge_config or {},
|
||||
"enable_memory": agent.enable_memory or False,
|
||||
"memory_window": agent.memory_window or 10,
|
||||
"enable_streaming": (
|
||||
agent.enable_streaming if agent.enable_streaming is not None else True
|
||||
),
|
||||
"is_public": agent.is_public or False,
|
||||
}
|
||||
|
||||
|
||||
async def check_import(db: AsyncSession, code: str) -> Dict[str, Any]:
|
||||
"""导入预检查：编码是否冲突"""
|
||||
code_exists = False
|
||||
if code:
|
||||
result = await db.execute(
|
||||
select(Agent).where(
|
||||
Agent.code == code,
|
||||
Agent.is_deleted == False,
|
||||
)
|
||||
)
|
||||
code_exists = result.scalar_one_or_none() is not None
|
||||
|
||||
return {
|
||||
"code_exists": code_exists,
|
||||
"can_import": not code_exists,
|
||||
}
|
||||
|
||||
|
||||
async def _resolve_model_id(db: AsyncSession, model_name: str) -> Optional[str]:
|
||||
if not model_name:
|
||||
return None
|
||||
result = await db.execute(
|
||||
select(LLMModel.id).where(
|
||||
LLMModel.model_name == model_name,
|
||||
LLMModel.is_deleted == False,
|
||||
LLMModel.is_active == True,
|
||||
).limit(1)
|
||||
)
|
||||
row = result.first()
|
||||
return row[0] if row else None
|
||||
|
||||
|
||||
async def _resolve_workflow_id(db: AsyncSession, workflow_code: str) -> Optional[str]:
|
||||
if not workflow_code:
|
||||
return None
|
||||
result = await db.execute(
|
||||
select(AIWorkflow.id).where(
|
||||
AIWorkflow.code == workflow_code,
|
||||
AIWorkflow.is_deleted == False,
|
||||
)
|
||||
)
|
||||
row = result.first()
|
||||
return row[0] if row else None
|
||||
|
||||
|
||||
async def _resolve_knowledge_base_ids(
|
||||
db: AsyncSession,
|
||||
codes: List[str],
|
||||
) -> tuple[List[str], List[str]]:
|
||||
if not codes:
|
||||
return [], []
|
||||
result = await db.execute(
|
||||
select(KnowledgeBase.id, KnowledgeBase.code).where(
|
||||
KnowledgeBase.code.in_(codes),
|
||||
KnowledgeBase.is_deleted == False,
|
||||
)
|
||||
)
|
||||
code_to_id = {row.code: row.id for row in result}
|
||||
resolved_ids = []
|
||||
unresolved: List[str] = []
|
||||
for code in codes:
|
||||
if code in code_to_id:
|
||||
resolved_ids.append(code_to_id[code])
|
||||
else:
|
||||
unresolved.append(code)
|
||||
return resolved_ids, unresolved
|
||||
|
||||
|
||||
async def import_config(db: AsyncSession, data: Dict[str, Any]) -> Agent:
|
||||
"""导入智能体配置"""
|
||||
required_fields = ["name", "code"]
|
||||
for field in required_fields:
|
||||
if not data.get(field):
|
||||
raise AgentImportExportException(f"缺少必要字段: {field}")
|
||||
|
||||
exists = await db.execute(
|
||||
select(Agent).where(
|
||||
Agent.code == data["code"],
|
||||
Agent.is_deleted == False,
|
||||
)
|
||||
)
|
||||
if exists.scalar_one_or_none():
|
||||
raise AgentImportExportException(f"智能体编码已存在: {data['code']}")
|
||||
|
||||
model_id = await _resolve_model_id(db, data.get("model_name") or "")
|
||||
if data.get("model_name") and not model_id:
|
||||
logger.warning("导入智能体时未找到模型: %s", data.get("model_name"))
|
||||
|
||||
workflow_id = await _resolve_workflow_id(db, data.get("workflow_code") or "")
|
||||
if data.get("workflow_code") and not workflow_id:
|
||||
logger.warning("导入智能体时未找到工作流: %s", data.get("workflow_code"))
|
||||
|
||||
kb_codes = data.get("knowledge_base_codes") or []
|
||||
knowledge_base_ids, unresolved_kb = await _resolve_knowledge_base_ids(db, kb_codes)
|
||||
if unresolved_kb:
|
||||
logger.warning(
|
||||
"导入智能体时部分知识库未找到: %s",
|
||||
", ".join(unresolved_kb),
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
application_id=data.get("application_id"),
|
||||
is_global=data.get("is_global", False),
|
||||
name=data["name"],
|
||||
code=data["code"],
|
||||
description=data.get("description", ""),
|
||||
avatar=data.get("avatar", ""),
|
||||
mode=data.get("mode", "autonomous"),
|
||||
status="draft",
|
||||
persona=data.get("persona") or {},
|
||||
system_prompt=data.get("system_prompt", ""),
|
||||
model_id=model_id,
|
||||
temperature=data.get("temperature", 0.7),
|
||||
top_p=data.get("top_p", 1.0),
|
||||
max_tokens=data.get("max_tokens", 4096),
|
||||
max_iterations=data.get("max_iterations", 10),
|
||||
welcome_message=data.get("welcome_message", ""),
|
||||
suggested_questions=data.get("suggested_questions") or [],
|
||||
workflow_id=workflow_id,
|
||||
knowledge_base_ids=knowledge_base_ids,
|
||||
knowledge_config=data.get("knowledge_config") or {},
|
||||
enable_memory=data.get("enable_memory", False),
|
||||
memory_window=data.get("memory_window", 10),
|
||||
enable_streaming=data.get("enable_streaming", True),
|
||||
is_public=data.get("is_public", False),
|
||||
)
|
||||
db.add(agent)
|
||||
await db.flush()
|
||||
await db.refresh(agent)
|
||||
|
||||
logger.info("智能体导入成功: %s", agent.code)
|
||||
return agent
|
||||
@@ -0,0 +1,172 @@
|
||||
"""
|
||||
Built-in agent Chinese display normalization.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from copy import deepcopy
|
||||
from typing import Any, Dict, Iterable, Optional
|
||||
|
||||
|
||||
class BuiltinAgentText:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
name: str,
|
||||
code_label: str,
|
||||
description: str,
|
||||
role: str = "",
|
||||
skills: Optional[list[str]] = None,
|
||||
constraints: Optional[list[str]] = None,
|
||||
background: str = "",
|
||||
keys: Iterable[str],
|
||||
description_keywords: Iterable[str] = (),
|
||||
persona_keywords: Iterable[str] = (),
|
||||
):
|
||||
self.name = name
|
||||
self.code_label = code_label
|
||||
self.description = description
|
||||
self.role = role
|
||||
self.skills = skills or []
|
||||
self.constraints = constraints or []
|
||||
self.background = background
|
||||
self.keys = list(keys)
|
||||
self.description_keywords = list(description_keywords)
|
||||
self.persona_keywords = list(persona_keywords)
|
||||
|
||||
|
||||
BUILTIN_AGENT_TEXTS = [
|
||||
BuiltinAgentText(
|
||||
name="项目经理",
|
||||
code_label="项目管理",
|
||||
description="协调交付状态、进度风险和跨角色协作，确保项目按计划推进。",
|
||||
keys=("project manager", "project_manager", "project-manager", "projectmanager"),
|
||||
description_keywords=("coordinates delivery status",),
|
||||
),
|
||||
BuiltinAgentText(
|
||||
name="Multica 产品经理",
|
||||
code_label="产品管理",
|
||||
description="负责产品需求分诊、优先级判断和方案澄清，推动产品决策落地。",
|
||||
role="负责 Multica 协作工作的产品需求分诊、范围澄清、验收标准和质量推进决策。",
|
||||
skills=["范围澄清", "验收标准设计", "优先级取舍分析", "干系人沟通"],
|
||||
constraints=[
|
||||
"需求必须能追溯到具体 issue 或用户请求。",
|
||||
"没有明确验收标准时，不推进下游工作。",
|
||||
"遇到未决产品决策时要明确指出，不能擅自假设。",
|
||||
],
|
||||
background="擅长把模糊的平台需求转化为责任清晰、可评审、可交付的 Multica issue。",
|
||||
keys=(
|
||||
"multica product manager",
|
||||
"multica_product_manager",
|
||||
"multica-product-manager",
|
||||
"product manager",
|
||||
"product_manager",
|
||||
"product-manager",
|
||||
),
|
||||
description_keywords=("owns product triage",),
|
||||
persona_keywords=(
|
||||
"product manager for multica",
|
||||
"scope clarification",
|
||||
"acceptance criteria design",
|
||||
"priority tradeoff analysis",
|
||||
"stakeholder communication",
|
||||
"keep requirements traceable",
|
||||
),
|
||||
),
|
||||
BuiltinAgentText(
|
||||
name="业务需求分析师",
|
||||
code_label="需求分析",
|
||||
description="将业务诉求转化为清晰需求、验收标准和可执行的交付范围。",
|
||||
keys=(
|
||||
"business requirements analyst",
|
||||
"business_requirements_analyst",
|
||||
"business-requirements-analyst",
|
||||
"business analyst",
|
||||
"business_analyst",
|
||||
"business-analyst",
|
||||
),
|
||||
description_keywords=("translates business needs",),
|
||||
),
|
||||
BuiltinAgentText(
|
||||
name="系统架构师",
|
||||
code_label="系统架构",
|
||||
description="设计系统边界、模块职责和技术方案，保障架构一致性与可演进性。",
|
||||
keys=("system architect", "system_architect", "system-architect"),
|
||||
),
|
||||
BuiltinAgentText(
|
||||
name="前端工程师",
|
||||
code_label="前端开发",
|
||||
description="负责前端页面、组件交互、路由状态和浏览器端体验实现。",
|
||||
keys=("frontend engineer", "frontend_engineer", "frontend-engineer"),
|
||||
),
|
||||
BuiltinAgentText(
|
||||
name="后端工程师",
|
||||
code_label="后端开发",
|
||||
description="负责后端接口、业务服务、数据模型和系统集成能力建设。",
|
||||
keys=("backend engineer", "backend_engineer", "backend-engineer"),
|
||||
),
|
||||
BuiltinAgentText(
|
||||
name="测试工程师",
|
||||
code_label="质量保障",
|
||||
description="负责测试策略、用例设计、缺陷验证和发布质量把关。",
|
||||
keys=("qa engineer", "qa_engineer", "qa-engineer"),
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _normalize_text(value: Any) -> str:
|
||||
return " ".join(str(value or "").strip().lower().replace("-", " ").replace("_", " ").split())
|
||||
|
||||
|
||||
def _iter_persona_values(persona: Any) -> Iterable[str]:
|
||||
if not isinstance(persona, dict):
|
||||
return []
|
||||
|
||||
values: list[str] = []
|
||||
for key in ("role", "background"):
|
||||
value = persona.get(key)
|
||||
if isinstance(value, str):
|
||||
values.append(value)
|
||||
for key in ("personality", "skills", "constraints"):
|
||||
values.extend(item for item in persona.get(key) or [] if isinstance(item, str))
|
||||
return values
|
||||
|
||||
|
||||
def find_builtin_agent_text(data: Dict[str, Any]) -> Optional[BuiltinAgentText]:
|
||||
name = _normalize_text(data.get("name"))
|
||||
code = _normalize_text(data.get("code"))
|
||||
description = _normalize_text(data.get("description"))
|
||||
persona_text = _normalize_text(" ".join(_iter_persona_values(data.get("persona"))))
|
||||
|
||||
for display in BUILTIN_AGENT_TEXTS:
|
||||
keys = {_normalize_text(key) for key in display.keys}
|
||||
if name in keys or code in keys:
|
||||
return display
|
||||
if any(keyword in description for keyword in display.description_keywords):
|
||||
return display
|
||||
if any(keyword in persona_text for keyword in display.persona_keywords):
|
||||
return display
|
||||
return None
|
||||
|
||||
|
||||
def normalize_builtin_agent_payload(data: Dict[str, Any], *, include_code_label: bool = False) -> Dict[str, Any]:
|
||||
display = find_builtin_agent_text(data)
|
||||
if not display:
|
||||
return data
|
||||
|
||||
normalized = deepcopy(data)
|
||||
normalized["name"] = display.name
|
||||
normalized["description"] = display.description
|
||||
if include_code_label:
|
||||
normalized["code_label"] = display.code_label
|
||||
|
||||
persona = deepcopy(normalized.get("persona") or {})
|
||||
if display.role:
|
||||
persona["role"] = display.role
|
||||
if display.skills:
|
||||
persona["skills"] = list(display.skills)
|
||||
if display.constraints:
|
||||
persona["constraints"] = list(display.constraints)
|
||||
if display.background:
|
||||
persona["background"] = display.background
|
||||
normalized["persona"] = persona
|
||||
return normalized
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user