From b536672000a56e73e9a6e5e25b685517963a0e32 Mon Sep 17 00:00:00 2001 From: ben Date: Wed, 12 Aug 2026 15:16:32 +0800 Subject: [PATCH] =?UTF-8?q?feat(m8):=20P2=20pgvector=20+=20Ollama=20?= =?UTF-8?q?=E5=90=91=E9=87=8F=E5=8C=96=E4=B8=8E=E8=AF=AD=E4=B9=89=E6=A3=80?= =?UTF-8?q?=E7=B4=A2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - PostgreSQL 切换 pgvector/pgvector:pg17 镜像;迁移 000024 建 vector 扩展、 knowledge_chunks.embedding vector(1024) + HNSW 余弦索引,retrieval_mode 放宽三态 - OllamaEmbedder 本地 bge-m3 批量嵌入,404 惰性 pull 重试,维度/超时校验,可整体关闭 - SemanticRetriever/HybridRetriever + NewRetriever 按 retrieval_mode 分发,缺 embedder 回退 FTS - 文档入库同步批量向量化;Ollama 故障降级入库 + embedding_failed 事件 - 修复 pgx CopyFrom 对 vector 列二进制编码误读:COPY 基础列后同事务 unnest 批量回填 - 修复降级路径 embeddings=nil 索引越界 panic(Add 与 Reprocess) - 知识库列表 vectorized_chunk_count + 前端三态检索模式选择与向量化覆盖率 - 单测 embedder/retrievers + 集成 TestKnowledgeVectorLifecycle 全绿 Co-Authored-By: Claude --- .env.example | 9 + .gitignore | 1 + README.md | 7 +- cmd/gateway-api/main.go | 18 +- deploy/PRODUCTION.md | 21 +- deploy/docker-compose.production.yml | 22 +- deploy/docker-compose.yml | 23 ++- deploy/production.env.example | 10 + docs/adr/0001-baseline-storage.md | 4 +- .../0002-vector-storage-and-object-store.md | 27 +++ docs/rewrite-progress.md | 11 + docs/旗舰版需求规划与完成情况.md | 26 +-- internal/platform/config/config.go | 43 +++- internal/workbench/admin_http.go | 2 +- internal/workbench/embedder.go | 189 ++++++++++++++++++ internal/workbench/embedder_test.go | 107 ++++++++++ internal/workbench/knowledge.go | 96 +++++++-- internal/workbench/retrievers.go | 159 +++++++++++++++ internal/workbench/retrievers_test.go | 108 ++++++++++ internal/workbench/service.go | 11 +- internal/workbench/types.go | 11 +- internal/workbench/vector_integration_test.go | 90 +++++++++ .../workbench/workbench_integration_test.go | 4 +- migrations/000024_pgvector_embeddings.sql | 21 ++ web/apps/admin/src/api/workbench.ts | 2 +- .../src/views/gateway/knowledge/index.vue | 4 +- 26 files changed, 966 insertions(+), 60 deletions(-) create mode 100644 docs/adr/0002-vector-storage-and-object-store.md create mode 100644 internal/workbench/embedder.go create mode 100644 internal/workbench/embedder_test.go create mode 100644 internal/workbench/retrievers.go create mode 100644 internal/workbench/retrievers_test.go create mode 100644 internal/workbench/vector_integration_test.go create mode 100644 migrations/000024_pgvector_embeddings.sql diff --git a/.env.example b/.env.example index c4feef7..0748851 100644 --- a/.env.example +++ b/.env.example @@ -80,6 +80,15 @@ S3_REGION=us-east-1 S3_USE_SSL=false S3_MAX_FILE_BYTES=134217728 +# M8 P2 向量化（本地 Ollama，默认 bge-m3）。设为 false 可整体关闭，知识库退回纯 FTS。 +# EMBEDDING_DIM 必须与迁移 000024 的 vector(1024) 列一致。 +EMBEDDINGS_ENABLED=true +OLLAMA_BASE_URL=http://ollama:11434 +EMBEDDING_MODEL=bge-m3 +EMBEDDING_DIM=1024 +EMBEDDING_BATCH_SIZE=64 +EMBEDDING_TIMEOUT=120s + # Used only by cmd/gateway-bootstrap; never commit the real value. BOOTSTRAP_ADMIN_USERNAME=admin BOOTSTRAP_ADMIN_PASSWORD= diff --git a/.gitignore b/.gitignore index 0b438b9..55aa85c 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,7 @@ deploy/production.env .vscode/ .DS_Store bin/ +backups/ coverage.out web/node_modules/ web/**/dist/ diff --git a/README.md b/README.md index 6e7021b..eaf43e5 100644 --- a/README.md +++ b/README.md @@ -30,22 +30,23 @@ AI Gateway 的全量 Go 重构工程。M0–M6 工程实现已完成，当前可 - 可扩展内容策略：Go RE2 不可变编译快照，按端点、模型/API Key 匹配，支持仅审计、阻断与提示词文本脱敏；默认保护常见 API Key、Token、密码和 secret。 - 带时间版本的模型价格与成本核算：按 Provider/模型选择价格，输入与输出 Token 分别计价，结果进入调用审计和 PostgreSQL 按日聚合。 - Prompt 分类、模板和不可变版本，支持显式变量定义、必填校验、历史版本激活与 API Key 渲染接口。 -- PostgreSQL 知识库：2 MiB 有界文本正文、段落感知重叠分块、FTS + 中文二元词片混合检索，以及可替换的 `Retriever` 接口；向量化由 M8 的 pgvector/Ollama 阶段提供。 +- PostgreSQL 知识库：2 MiB 有界文本正文、段落感知重叠分块、FTS + 中文二元词片检索；M8 起支持 pgvector + 本地 Ollama(bge-m3)的 vector/hybrid 三态语义检索，文档导入即同步向量化、Ollama 异常自动降级纯 FTS。 - 声明式 HTTP 工具：JSON Schema 基础校验、KEK 加密请求头、注册和拨号双层 SSRF 防护、禁止重定向、1 MiB 响应限制与调用记录。 - AI 应用草稿和不可变发布版本，将模型、Prompt、知识库、工具组合为 `/v1/applications/{code}/chat/completions`；所有模型轮次继续经过鉴权、配额、内容策略、路由、成本和审计。 - 独立通知 Worker 消费可靠 outbox，按精确事件或末尾 `*` 模式投递 HMAC-SHA256 Webhook；内容策略命中由审计批处理异步产生脱敏事件，失败投递可在 Art 管理端重试。 - 门户自助工作台：部门范围资产目录、Prompt 搜索/收藏、个人审计/用量/成本、模型访问申请与管理员审批。 - M8 对象存储：自托管 MinIO，上传/下载全部经网关代理（不暴露主机端口），管理端文件管理与门户个人文件仓库，`sha256` 完整性校验与严格归属隔离。 +- M8 向量化：pgvector + 本地 Ollama(bge-m3)为知识库提供 vector/hybrid 三态语义检索；HNSW 余弦索引，`EMBEDDINGS_ENABLED=false` 或 Ollama 异常时自动降级纯 FTS，不阻断文档入库。 - 门户应用托管会话：服务端加密运行凭证、单会话租约、不可变消息序列和 SHA-256 哈希链，不向浏览器暴露应用 API Key。 - 独立事实核验配置、作用域策略与事件契约，复用 Provider 加密凭据和知识库引用，为同步/异步执行器保留清晰模块边界。 - 旧 Python 源码 201 条路由全部有覆盖、替代或退役决策，未决契约缺口为 0；OpenAPI 0.10.0 覆盖全部 Go 字面量路由。 -M8 起 MinIO（对象存储）纳入基线部署并由 compose 提供，但不作为启动依赖：网关启动时 MinIO 未就绪只告警、上传请求得到明确报错，服务不会因对象存储缺失而崩溃。ClickHouse 不属于基线，审计与统计继续存放在 PostgreSQL。 +M8 起 MinIO（对象存储）与本地 Ollama（向量化）纳入基线部署并由 compose 提供，但均不作为启动依赖：MinIO 未就绪时上传得到明确报错，Ollama 未就绪时知识库入库降级（embedding 置 NULL），服务不会因二者缺失而崩溃。PostgreSQL 已切换为 `pgvector/pgvector:pg17` 镜像（数据卷兼容），`CREATE EXTENSION vector` 由迁移 `000024` 执行。ClickHouse 不属于基线，审计与统计继续存放在 PostgreSQL。 ## 本地启动 1. 复制 `.env.example` 为 `.env` 并修改密钥。 -2. 启动 PostgreSQL 和两套 Redis：`docker compose -f deploy/docker-compose.yml up -d postgres redis-critical redis-cache`。 +2. 启动 PostgreSQL(pgvector)、两套 Redis、MinIO 和 Ollama：`docker compose -f deploy/docker-compose.yml up -d postgres redis-critical redis-cache minio ollama`。 3. 执行迁移：`go run ./cmd/gateway-migrator`。 4. 设置 `BOOTSTRAP_ADMIN_PASSWORD` 后执行 `go run ./cmd/gateway-bootstrap` 创建初始管理员。 5. 启动 API：`go run ./cmd/gateway-api`。 diff --git a/cmd/gateway-api/main.go b/cmd/gateway-api/main.go index c17e0b5..e148c7c 100644 --- a/cmd/gateway-api/main.go +++ b/cmd/gateway-api/main.go @@ -202,6 +202,20 @@ func main() { pricingHandler := pricing.NewAdminHTTPHandler(pricingService, identityService) factCheckHandler := factcheck.NewAdminHTTPHandler(factcheck.NewService(db), identityService) workbenchService := workbench.NewService(db) + // M8 P2:本地 Ollama 向量化。EMBEDDINGS_ENABLED=false 时不构造 embedder, + // 知识库检索自动回退纯 FTS;Ollama 挂时入库降级(embedding 置 NULL)。 + if cfg.Embeddings.Enabled { + workbenchService.SetEmbedder(workbench.NewOllamaEmbedder(workbench.OllamaEmbedderConfig{ + BaseURL: cfg.Embeddings.BaseURL, + Model: cfg.Embeddings.Model, + Dim: cfg.Embeddings.Dim, + BatchSize: cfg.Embeddings.BatchSize, + Timeout: cfg.Embeddings.Timeout, + })) + logger.Info("knowledge embeddings enabled", "model", cfg.Embeddings.Model, "base_url", cfg.Embeddings.BaseURL) + } else { + logger.Info("knowledge embeddings disabled, knowledge retrieval uses postgres_fts only") + } toolCipher, err := cryptox.NewKeyring( cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "tool-request-headers", ) @@ -263,7 +277,7 @@ func main() { } shadowMiddleware := shadow.New(cfg.Shadow, logger) governedGateway := shadowMiddleware.Wrap(proxy) - workbenchRuntime := workbench.NewRuntimeHTTPHandler(workbenchService, toolService, workbench.NewPostgreSQLRetriever(workbenchService), apiKeyAuthenticator, governedGateway, workbench.MarketplaceDeps{ + workbenchRuntime := workbench.NewRuntimeHTTPHandler(workbenchService, toolService, workbench.NewRetriever(workbenchService, workbenchService.Embedder()), apiKeyAuthenticator, governedGateway, workbench.MarketplaceDeps{ MCPServers: mcpServerService, Skills: skillService, Employees: digitalEmployeeService, @@ -273,7 +287,7 @@ func main() { workbenchRuntime.SetLogger(logger) // Wire the fact-check engine: the admin fact-check settings/policies UI now // actually governs application answers instead of being inert configuration. - factCheckEngine := factcheck.NewEngine(db, workbench.NewFactCheckRetriever(workbench.NewPostgreSQLRetriever(workbenchService)), logger) + factCheckEngine := factcheck.NewEngine(db, workbench.NewFactCheckRetriever(workbench.NewRetriever(workbenchService, workbenchService.Embedder())), logger) workbenchRuntime.SetFactCheckEngine(factCheckEngine) portalService := portal.NewService(db, workbenchService, toolService, identityService) portalService.SetApplicationRuntime(portal.NewRuntimeCredentials(db, apiKeyRepository, applicationKeyCipher), workbenchRuntime) diff --git a/deploy/PRODUCTION.md b/deploy/PRODUCTION.md index 1930e7a..7d2b97d 100644 --- a/deploy/PRODUCTION.md +++ b/deploy/PRODUCTION.md @@ -1,9 +1,12 @@ # Production deployment This bundle builds the Go services and both Art Design Pro applications from -source. PostgreSQL, two Redis roles and MinIO (object storage, M8) are included; -ClickHouse is not required. MinIO is not a startup dependency: the gateway only -warns and refuses file uploads until the bucket is reachable. +source. PostgreSQL (with pgvector), two Redis roles, MinIO (object storage, M8) +and a local Ollama container (vectorization, M8) are included; ClickHouse is not +required. Neither MinIO nor Ollama is a startup dependency: the gateway warns +and refuses file uploads until the bucket is reachable, and knowledge-base +documents are still stored (with embedding set to NULL) when Ollama is down, +with retrieval falling back to full-text search. ## Prerequisites @@ -73,3 +76,15 @@ it removes persistent data. The bundled database URLs use `sslmode=disable` only for the private Compose network. When using an external PostgreSQL or Redis service, require TLS and use `sslmode=verify-full` / `rediss://` as supported by that service. + +### Vectorization and object storage + +- The PostgreSQL image is `pgvector/pgvector:pg17` (data-volume compatible with + `postgres:17-alpine`); migration `000024` creates the `vector` extension and + adds the HNSW embedding column. `EMBEDDING_DIM` must stay at `1024` to match + the `vector(1024)` column. +- Ollama runs locally and lazily pulls `bge-m3` (~1.2 GiB) on first embedding + request. Set `EMBEDDINGS_ENABLED=false` to disable vectorization entirely. +- Back up the `minio-data` and `ollama-models` volumes alongside PostgreSQL. +- If you previously deployed with `postgres:17-alpine`, back up the PostgreSQL + volume before switching images. diff --git a/deploy/docker-compose.production.yml b/deploy/docker-compose.production.yml index 472662c..cdaf46e 100644 --- a/deploy/docker-compose.production.yml +++ b/deploy/docker-compose.production.yml @@ -33,6 +33,13 @@ x-gateway-environment: &gateway-environment S3_REGION: ${S3_REGION:-us-east-1} S3_USE_SSL: ${S3_USE_SSL:-false} S3_MAX_FILE_BYTES: ${S3_MAX_FILE_BYTES:-134217728} + # M8 P2:本地 Ollama 向量化。生产可整体关闭(EMBEDDINGS_ENABLED=false)走纯 FTS。 + EMBEDDINGS_ENABLED: ${EMBEDDINGS_ENABLED:-true} + OLLAMA_BASE_URL: ${OLLAMA_BASE_URL:-http://ollama:11434} + EMBEDDING_MODEL: ${EMBEDDING_MODEL:-bge-m3} + EMBEDDING_DIM: ${EMBEDDING_DIM:-1024} + EMBEDDING_BATCH_SIZE: ${EMBEDDING_BATCH_SIZE:-64} + EMBEDDING_TIMEOUT: ${EMBEDDING_TIMEOUT:-120s} x-backend-service: &backend-service image: ai-gateway-go:${GATEWAY_VERSION:-0.10.0} @@ -47,8 +54,9 @@ x-backend-service: &backend-service - no-new-privileges:true services: + # M8 P2:pgvector 镜像提供 vector 扩展;数据卷与 postgres:17-alpine 兼容。 postgres: - image: postgres:17-alpine + image: pgvector/pgvector:pg17 environment: POSTGRES_DB: ${POSTGRES_DB:-gateway} POSTGRES_USER: ${POSTGRES_USER:-gateway} @@ -116,6 +124,9 @@ services: condition: service_healthy minio: condition: service_started + # 向量化不阻断 API 启动:ollama 未就绪时知识库入库降级。 + ollama: + condition: service_started healthcheck: test: ["CMD-SHELL", "wget -q -O /dev/null http://127.0.0.1:8080/readyz"] interval: 10s @@ -192,6 +203,14 @@ services: - minio-data:/data restart: unless-stopped + # M8 P2:本地向量化(默认 bge-m3)。stateful,首次 embed 惰性 pull 模型; + # 生产不需要时可移除并设 EMBEDDINGS_ENABLED=false 走纯 FTS。 + ollama: + image: ollama/ollama:latest + volumes: + - ollama-models:/root/.ollama + restart: unless-stopped + bootstrap-admin: <<: *backend-service profiles: ["tools"] @@ -209,3 +228,4 @@ volumes: postgres-data: redis-critical-data: minio-data: + ollama-models: diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index 170e667..9525ada 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -1,6 +1,8 @@ services: + # M8 P2:pgvector 镜像与 postgres:17-alpine 数据卷兼容;CREATE EXTENSION vector + # 由迁移 000024 执行(容器内 gateway 用户即 superuser)。 postgres: - image: postgres:17-alpine + image: pgvector/pgvector:pg17 environment: POSTGRES_DB: gateway POSTGRES_USER: gateway @@ -65,6 +67,13 @@ services: S3_REGION: ${S3_REGION:-us-east-1} S3_USE_SSL: ${S3_USE_SSL:-false} S3_MAX_FILE_BYTES: ${S3_MAX_FILE_BYTES:-134217728} + # M8 P2:本地 Ollama 向量化(bge-m3)。EMBEDDINGS_ENABLED=false 可整体关闭。 + EMBEDDINGS_ENABLED: ${EMBEDDINGS_ENABLED:-true} + OLLAMA_BASE_URL: ${OLLAMA_BASE_URL:-http://ollama:11434} + EMBEDDING_MODEL: ${EMBEDDING_MODEL:-bge-m3} + EMBEDDING_DIM: ${EMBEDDING_DIM:-1024} + EMBEDDING_BATCH_SIZE: ${EMBEDDING_BATCH_SIZE:-64} + EMBEDDING_TIMEOUT: ${EMBEDDING_TIMEOUT:-120s} depends_on: postgres: condition: service_healthy @@ -87,6 +96,9 @@ services: condition: service_healthy minio: condition: service_started + # 向量化不阻断 API 启动:ollama 未就绪时 knowledge 入库降级(embedding 置 NULL)。 + ollama: + condition: service_started restart: unless-stopped # M8: 对象存储。MinIO 不暴露主机端口,凭据只留在 API 容器内; @@ -101,6 +113,14 @@ services: - minio-data:/data restart: unless-stopped + # M8 P2:本地向量化。模型默认 bge-m3(~1.2GB),首次 embed 惰性 pull; + # 无需对外暴露端口,仅网关容器访问。restart=unless-stopped 以便重启后自愈。 + ollama: + image: ollama/ollama:latest + volumes: + - ollama-models:/root/.ollama + restart: unless-stopped + admin-web: build: context: .. @@ -174,3 +194,4 @@ volumes: postgres-data: redis-critical-data: minio-data: + ollama-models: diff --git a/deploy/production.env.example b/deploy/production.env.example index 6ac92fa..ae4efb2 100644 --- a/deploy/production.env.example +++ b/deploy/production.env.example @@ -48,6 +48,16 @@ S3_REGION=us-east-1 S3_USE_SSL=false S3_MAX_FILE_BYTES=134217728 +# M8 P2 向量化（本地 Ollama，默认 bge-m3）。生产可整体关闭走纯 FTS： +# 设 EMBEDDINGS_ENABLED=false 并移除 compose 的 ollama 服务。 +# EMBEDDING_DIM 必须与迁移 000024 的 vector(1024) 列一致。 +EMBEDDINGS_ENABLED=true +OLLAMA_BASE_URL=http://ollama:11434 +EMBEDDING_MODEL=bge-m3 +EMBEDDING_DIM=1024 +EMBEDDING_BATCH_SIZE=64 +EMBEDDING_TIMEOUT=120s + # Used only for the one-time bootstrap-admin command; remove after use. BOOTSTRAP_ADMIN_USERNAME=admin BOOTSTRAP_ADMIN_PASSWORD=CHANGE_ME_AT_LEAST_12_CHARACTERS diff --git a/docs/adr/0001-baseline-storage.md b/docs/adr/0001-baseline-storage.md index 95e46e6..bf152b9 100644 --- a/docs/adr/0001-baseline-storage.md +++ b/docs/adr/0001-baseline-storage.md @@ -1,7 +1,7 @@ # ADR-0001：基线存储只使用 PostgreSQL 与 Redis -状态：已接受 -日期：2026-08-10 +状态：已被 [ADR-0002](./0002-vector-storage-and-object-store.md) 取代 +日期：2026-08-10（2026-08-12 修订：对象存储与向量化基础设施入基线，见 ADR-0002） ## 决策 diff --git a/docs/adr/0002-vector-storage-and-object-store.md b/docs/adr/0002-vector-storage-and-object-store.md new file mode 100644 index 0000000..fc1e371 --- /dev/null +++ b/docs/adr/0002-vector-storage-and-object-store.md @@ -0,0 +1,27 @@ +# ADR-0002：对象存储与向量化基础设施入基线 + +状态：已接受 +日期：2026-08-12 +取代/修订：[ADR-0001](./0001-baseline-storage.md) + +## 决策 + +M8 起把 MinIO（对象存储）与 pgvector + 本地 Ollama（向量化）纳入基线，但它们都**不作为启动依赖**。 + +- **MinIO 自托管**：上传/下载全部经 `gateway-api` 代理，MinIO 不暴露主机端口；`gateway.file_objects` 元数据行以 PostgreSQL 为权威，对象未就绪时上传得到明确报错，服务不崩溃。 +- **pgvector + Ollama**：PostgreSQL 镜像切换为 `pgvector/pgvector:pg17`（数据卷兼容），`knowledge_chunks.embedding vector(1024)` 由本地 Ollama 容器（默认 `bge-m3`）产生；Ollama 未就绪时知识库入库降级（embedding 置 NULL + `knowledge_document.embedding_failed` 事件），检索自动回退纯 FTS。 +- **不引入 pgvector-go / S3 客户端到领域层**：向量以 `$1::vector` 字面量直传；对象存储经 `internal/platform/storage` 适配层隔离，领域层只依赖接口。 +- **ClickHouse 仍不属于基线**：审计与统计继续存放 PostgreSQL。 + +## 动机 + +- 知识库需要向量化语义召回（旗舰版"外部文档导入、自动切分、向量化、语义匹配召回"），并作为 M10 记忆管理、M12 语义输出的依赖基座。 +- 门户个人文件仓库与管理端文件管理需要对象存储；自托管 MinIO 保证完全内网离线。 +- 本地 Ollama 生成嵌入避免外发文本到公网 embedding API，多语言模型 bge-m3 与中文检索目标匹配。 + +## 影响 + +- 部署编排新增 `minio`、`ollama` 两个 stateful 服务与 `minio-data`、`ollama-models` 卷；nginx `client_max_body_size` 提到 256 MiB 盖过 128 MiB 上传上限。 +- 运维需备份 `minio-data`、`ollama-models` 卷；Ollama 首次拉取 bge-m3 约 1.2 GiB，`EMBEDDINGS_ENABLED=false` 可关闭。 +- `EMBEDDING_DIM` 必须与 `vector(1024)` 一致，配置校验拦截不符。 +- ADR-0001 第 4 条（MinIO/S3 不在基线）与"知识库只保存纯文本"不再成立，本 ADR 取代。 diff --git a/docs/rewrite-progress.md b/docs/rewrite-progress.md index d2bf5c1..865ecb2 100644 --- a/docs/rewrite-progress.md +++ b/docs/rewrite-progress.md @@ -122,3 +122,14 @@ - Compose（dev + production）新增 `minio` 服务与 `minio-data` 卷；nginx `client_max_body_size` 32m→256m 盖过 128 MiB 上传上限；`.env.example` / `production.env.example` 补 `S3_*`。 - 端到端验证：上传→列表→下载往返一致→删除后桶无孤儿；admin 读 portal 文件 404；`TestFileObjectLifecycle` 集成测试连真实 MinIO+PostgreSQL 通过；`go build ./...`、`go vet ./...`、全量单测通过。 - 管理端"文件管理"与门户端"文件仓库"菜单由服务端动态菜单下发。 + +## 已完成：M8 基础设施层（P2 向量化与语义检索） + +- PostgreSQL 换 `pgvector/pgvector:pg17` 镜像（数据卷兼容，先备份再切换）；迁移 `000024` `CREATE EXTENSION vector`、`knowledge_chunks.embedding vector(1024)` 列 + HNSW 余弦索引，并把 `retrieval_mode` CHECK 放宽为 `('postgres_fts','vector','hybrid')` 三态。 +- 本地 Ollama（`bge-m3`，1024 维）生成嵌入，Compose 新增 `ollama` 服务与 `ollama-models` 卷；`OllamaEmbedder` 批量调 `/api/embed`，首见 404 惰性 `/api/pull` 重试一次，按批校验维度，超时 `EMBEDDING_TIMEOUT`；`EMBEDDINGS_ENABLED=false` 可整体关闭。 +- `SemanticRetriever`（`embedding <=> $1::vector` 余弦距离 + `embedding IS NOT NULL` 过滤）与 `HybridRetriever`（FTS + 语义按 chunk 去重融合）实现；`NewRetriever(service, embedder)` 按知识库 `retrieval_mode` 分发，embedder 为 nil 时自动回退纯 FTS（不因缺少 Ollama 而报错）。 +- 文档入库/重新分块时在 vector/hybrid 模式下同步批量计算向量（每文档一次 `/api/embed` 收数组）；Ollama 故障时优雅降级：文档照常入库、embedding 置 NULL，并在事务内发 `knowledge_document.embedding_failed` 事件。 +- 踩坑并修复：pgx v5 `CopyFrom` 对未知 OID（vector）列走二进制编码，字面量随 COPY 上传会被 `vector_recv` 误读为维度数而报 `vector cannot have more than 16000 dimensions`；改为 COPY 仅基础列，随后在同一事务内用 `UPDATE ... FROM unnest($1::uuid[], $2::text[])` 批量回填向量，避免 2 MiB 文档上千条逐条 INSERT。 +- 知识库列表新增 `vectorized_chunk_count`（`count(c.embedding)`），管理端展示"已向量化切片/总切片"覆盖率；创建/编辑知识库可三态选择检索模式，未向量化分块需重新处理才被语义召回。 +- `EMBEDDINGS_*` 配置（`OLLAMA_BASE_URL` 默认 `http://ollama:11434`、`EMBEDDING_MODEL` bge-m3、`EMBEDDING_DIM` 须与 `vector(1024)` 一致、批大小、超时）写入两个 `.env.example` 与 compose anchor。 +- 单测 `embedder_test.go`/`retrievers_test.go`（httptest 假 Ollama：批量切分、维度不符、404→pull→重试、降级分发）与集成 `TestKnowledgeVectorLifecycle`（真 pgvector+Ollama：导入即向量化、语义命中、embedder 失败时入库 + `embedding_failed` 事件）全部通过；`go build ./...`、`go vet ./...` 通过。 diff --git a/docs/旗舰版需求规划与完成情况.md b/docs/旗舰版需求规划与完成情况.md index e983430..07eb65a 100644 --- a/docs/旗舰版需求规划与完成情况.md +++ b/docs/旗舰版需求规划与完成情况.md @@ -3,7 +3,7 @@ - **报告日期**: 2026-08-12 - **工程**: AI Gateway 全量 Go 重构(替代原 Python/FastAPI 网关) - **活跃工作树**: `/home/ben/ai-gateway-src/ai-gateway-go-deploy-0.10.0` -- **当前版本**: 0.10.0(Go 1.26,PostgreSQL 17 + 双 Redis,22 个迁移) +- **当前版本**: 0.10.0(Go 1.26,PostgreSQL 17+pgvector + 双 Redis + MinIO + Ollama,24 个迁移) --- @@ -12,11 +12,11 @@ | 项 | 状态 | |---|---| | 部署形态 | Docker Compose(项目名 `deploy`),gateway-api `:8080` / admin-web `:8081` / portal-web `:8082` | -| 数据层 | PostgreSQL 17(权威配置 + 审计分区)+ critical/cache 双 Redis;22 个迁移已应用 | -| 迁移 | `000001`…`000022`(含旗舰版资源市场 `000022_resource_marketplace.sql`) | -| 验证 | `go build ./...`、`go vet`、全套单测、资源市场集成测试连真实库 **全部通过** | +| 数据层 | PostgreSQL 17 + pgvector(权威配置 + 审计分区 + 向量列)+ critical/cache 双 Redis + MinIO + Ollama;24 个迁移已应用 | +| 迁移 | `000001`…`000024`(含资源市场 `000022`、MinIO `000023`、pgvector `000024`) | +| 验证 | `go build ./...`、`go vet`、全套单测、资源市场/文件管理集成测试连真实库 **全部通过** | | 前端 | Art Design Pro 管理端 + 门户端,已构建进镜像并运行 | -| 版本控制 | 无 git(目录无 `.git`),靠目录快照;建议尽快接入版本控制 | +| 版本控制 | git(remote `origin`=Gitea `superidou/ai-gateway-go`),每次改动 commit + push | --- @@ -77,7 +77,7 @@ Prompt 分类/模板/不可变版本/必填校验;知识库(2 MiB 有界正文 | 模型管理 | 对接本地模型(Ollama/vLLM) | 不支持 | ✓ | ✓ | ✅ 走 OpenAI 兼容通用 provider | | 模型管理 | 不同大模型 Token 配额、使用量统计 | 不支持 | ✓ | ✓ | ✅ quota + usage 聚合 | | 模型管理 | 大模型使用权限分级管控(用户/角色) | ✓ | ✓ | ✓ | ✅ RBAC scope | -| 知识库 | 外部文档导入、自动切分、向量化、语义匹配召回 | — | ✓ | ✓ | ⚠️ 导入/切分/检索 ✅;向量化+语义匹配 ❌(现为 FTS+中文分词) | +| 知识库 | 外部文档导入、自动切分、向量化、语义匹配召回 | — | ✓ | ✓ | ✅ M8 P2:pgvector + 本地 Ollama bge-m3,vector/hybrid 三态,导入即同步向量化(失败自动降级 FTS) | | 记忆管理 | 记忆集合:个人/部门/全局多层记忆,提炼与语义匹配召回 | — | ✓ | ✓ | ❌ | | 记忆管理 | 根据调用自动裁剪衰减片段 | — | ✓ | ✓ | ❌ | | 记忆管理 | 记忆授权:提炼内容/沉淀经验授权给其他用户 | — | ✓ | ✓ | ❌ | @@ -134,8 +134,8 @@ Prompt 分类/模板/不可变版本/必填校验;知识库(2 MiB 有界正文 ## 四、差距汇总 -- **完全未实现(❌,约 18 项)**:AI 助手、收藏、数据报表/企业报表、配置管理(env 注入)、智能体管理三项(节点/LLMTrace/会话)、记忆管理三项、渠道管理全项、多租户、供应链安全扫描、站内消息、完整审批流、License、定时任务全项、个人渠道、个人安全策略、ARM64。(M8 已落地:文件管理/MinIO、个人文件仓库) -- **部分覆盖需补齐(⚠️,约 10 项)**:平台概览看板、知识库向量化语义召回、资源/渠道审批、工具输出脱敏与大模型回答拦截替换、工具命令审批与工具限流、集群部署方案、站内消息/审批待办/任务结果、数字员工会话入口、资源权限等级三档、全类型权限申请。 +- **完全未实现(❌,约 16 项)**:AI 助手、收藏、数据报表/企业报表、配置管理(env 注入)、智能体管理三项(节点/LLMTrace/会话)、记忆管理三项、渠道管理全项、多租户、供应链安全扫描、完整审批流、License、定时任务全项、个人渠道、个人安全策略、ARM64。(M8 已落地:文件管理/MinIO、个人文件仓库、知识库向量化语义召回) +- **部分覆盖需补齐(⚠️,约 9 项)**:平台概览看板、资源/渠道审批、工具输出脱敏与大模型回答拦截替换、工具命令审批与工具限流、集群部署方案、站内消息/审批待办/任务结果、数字员工会话入口、资源权限等级三档、全类型权限申请。 --- @@ -143,7 +143,7 @@ Prompt 分类/模板/不可变版本/必填校验;知识库(2 MiB 有界正文 | 里程碑 | 内容 | 依赖 | |---|---|---| -| **M8 基础设施层** | 对象存储(MinIO)、向量化(pgvector)、定时任务调度器、站内消息 | — | +| **M8 基础设施层** | 对象存储(MinIO)、向量化(pgvector)、定时任务调度器、站内消息 | —(P1 MinIO/P2 pgvector+Ollama 已完成,待 P3 调度器 + P4 站内信) | | **M9 智能体与可观测** | LLMTrace、智能体会话、智能体节点(监控/节点池/路由)、AI 助手 | M8 | | **M10 记忆管理** | 多层记忆集合、语义召回、裁剪衰减、记忆授权 | M8(pgvector) | | **M11 渠道与审批** | 渠道管理(企业微信/个人微信/钉钉/飞书)、个人渠道、完整审批流、资源权限等级 | M8 | @@ -156,10 +156,10 @@ Prompt 分类/模板/不可变版本/必填校验;知识库(2 MiB 有界正文 ## 六、验证情况(当前基线) -- `go build ./...` ✅、`go vet ./internal/workbench/` ✅ -- 全套单测(全部包)✅;资源市场单测(含 MCP 客户端)✅ -- 集成测试 `TestMarketplaceLifecycle`、`TestWorkbenchPostgreSQLLifecycle` 连真实 PostgreSQL ✅ -- 部署冒烟:healthz/readyz ✅、admin :8081 / portal :8082 302 ✅、22 迁移应用 ✅ +- `go build ./...` ✅、`go vet ./...` ✅ +- 全套单测(全部包)✅;资源市场/文件/向量化单测(embedder/retriever)✅ +- 集成测试 `TestMarketplaceLifecycle`、`TestWorkbenchPostgreSQLLifecycle`、`TestFileObjectLifecycle`、`TestKnowledgeVectorLifecycle`(pgvector+Ollama:导入即向量化、语义命中、embedder 失败降级)连真实 PostgreSQL/MinIO/Ollama ✅ +- 部署冒烟:healthz/readyz ✅、admin :8081 / portal :8082 302 ✅、24 迁移应用 ✅、vector 模式知识库导入即向量化 + 语义检索命中 ✅ ## 七、部署与已知坑 diff --git a/internal/platform/config/config.go b/internal/platform/config/config.go index fd6ffd8..0f1e974 100644 --- a/internal/platform/config/config.go +++ b/internal/platform/config/config.go @@ -22,8 +22,9 @@ type Config struct { Audit Audit Outbox Outbox RuntimeData RuntimeData - Shadow Shadow + Shadow Shadow ObjectStorage ObjectStorage + Embeddings Embeddings } type Server struct { @@ -125,6 +126,17 @@ type ObjectStorage struct { MaxFileBytes int64 } +// Embeddings 配置本地 Ollama 向量化(默认 bge-m3)。Enabled 为 false 时网关不构造 +// OllamaEmbedder,知识库退回纯 FTS 检索,AddKnowledgeDocument 不生成 embedding。 +type Embeddings struct { + Enabled bool + BaseURL string + Model string + Dim int + BatchSize int + Timeout time.Duration +} + func Load() (Config, error) { cfg := Config{ Environment: env("APP_ENV", "local"), @@ -202,6 +214,14 @@ func Load() (Config, error) { UseSSL: boolValue("S3_USE_SSL", false), MaxFileBytes: int64Value("S3_MAX_FILE_BYTES", 128<<20), }, + Embeddings: Embeddings{ + Enabled: boolValue("EMBEDDINGS_ENABLED", true), + BaseURL: strings.TrimRight(env("OLLAMA_BASE_URL", "http://ollama:11434"), "/"), + Model: env("EMBEDDING_MODEL", "bge-m3"), + Dim: intValue("EMBEDDING_DIM", 1024), + BatchSize: intValue("EMBEDDING_BATCH_SIZE", 64), + Timeout: duration("EMBEDDING_TIMEOUT", 120*time.Second), + }, } return cfg, cfg.Validate() @@ -271,6 +291,27 @@ func (c Config) Validate() error { if c.ObjectStorage.MaxFileBytes < 1<<20 || c.ObjectStorage.MaxFileBytes > 512<<20 { errs = append(errs, errors.New("S3_MAX_FILE_BYTES must be between 1 MiB and 512 MiB")) } + if c.Embeddings.Enabled { + if err := validateHTTPURL(c.Embeddings.BaseURL); err != nil { + errs = append(errs, fmt.Errorf("OLLAMA_BASE_URL: %w", err)) + } + if c.Embeddings.Model == "" { + errs = append(errs, errors.New("EMBEDDING_MODEL is required when embeddings are enabled")) + } + if c.Embeddings.Dim < 128 || c.Embeddings.Dim > 8192 { + errs = append(errs, errors.New("EMBEDDING_DIM must be between 128 and 8192")) + } + if c.Embeddings.Dim != 1024 { + // 知识库 embedding 列固定为 vector(1024);维度不符会让入库向量报错。 + errs = append(errs, errors.New("EMBEDDING_DIM must be 1024 to match the vector(1024) column")) + } + if c.Embeddings.BatchSize < 1 || c.Embeddings.BatchSize > 512 { + errs = append(errs, errors.New("EMBEDDING_BATCH_SIZE must be between 1 and 512")) + } + if c.Embeddings.Timeout < time.Second || c.Embeddings.Timeout > 30*time.Minute { + errs = append(errs, errors.New("EMBEDDING_TIMEOUT must be between 1s and 30m")) + } + } return errors.Join(errs...) } diff --git a/internal/workbench/admin_http.go b/internal/workbench/admin_http.go index fe7c780..bb17267 100644 --- a/internal/workbench/admin_http.go +++ b/internal/workbench/admin_http.go @@ -417,7 +417,7 @@ func (h *AdminHTTPHandler) searchKnowledge(w http.ResponseWriter, r *http.Reques if !decodeAsset(w, r, &p) { return } - items, err := NewPostgreSQLRetriever(h.service).Search(r.Context(), r.PathValue("id"), p.Query, p.TopK) + items, err := NewRetriever(h.service, h.service.Embedder()).Search(r.Context(), r.PathValue("id"), p.Query, p.TopK) if err != nil { assetError(w, err) return diff --git a/internal/workbench/embedder.go b/internal/workbench/embedder.go new file mode 100644 index 0000000..caad26b --- /dev/null +++ b/internal/workbench/embedder.go @@ -0,0 +1,189 @@ +package workbench + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "time" +) + +// Embedder 把文本批量转换成固定维度向量。知识库向量化/语义检索依赖该接口, +// Embed 失败时调用方应优雅降级(文档照常入库、embedding 置 NULL)。 +type Embedder interface { + Embed(ctx context.Context, texts []string) ([][]float32, error) + Dim() int +} + +// OllamaEmbedder 调用本地 Ollama 的 /api/embed(批量),默认模型 bge-m3(1024 维)。 +// 模型尚未拉取时(/api/embed 返回 404)会先 POST /api/pull 拉取一次再重试。 +type OllamaEmbedder struct { + baseURL string + model string + dim int + batchSize int + timeout time.Duration + client *http.Client +} + +// OllamaEmbedderConfig 是 NewOllamaEmbedder 的参数;BaseURL 与 Model 已去除首尾空白。 +type OllamaEmbedderConfig struct { + BaseURL string + Model string + Dim int + BatchSize int + Timeout time.Duration +} + +func NewOllamaEmbedder(cfg OllamaEmbedderConfig) *OllamaEmbedder { + client := &http.Client{Timeout: cfg.Timeout} + if cfg.Dim <= 0 { + cfg.Dim = 1024 + } + if cfg.BatchSize <= 0 { + cfg.BatchSize = 64 + } + return &OllamaEmbedder{ + baseURL: cfg.BaseURL, + model: cfg.Model, + dim: cfg.Dim, + batchSize: cfg.BatchSize, + timeout: cfg.Timeout, + client: client, + } +} + +func (o *OllamaEmbedder) Dim() int { return o.dim } + +type ollamaEmbedResponse struct { + Embeddings [][]float32 `json:"embeddings"` +} + +type ollamaErrorResponse struct { + Error string `json:"error"` +} + +// Embed 把 texts 按 batchSize 切批调用 Ollama。任何一次调用失败都会返回错误, +// 由调用方决定是否降级(知识库入库语义)。 +func (o *OllamaEmbedder) Embed(ctx context.Context, texts []string) ([][]float32, error) { + if len(texts) == 0 { + return nil, nil + } + all := make([][]float32, 0, len(texts)) + for start := 0; start < len(texts); start += o.batchSize { + end := start + o.batchSize + if end > len(texts) { + end = len(texts) + } + batch, err := o.embedBatch(ctx, texts[start:end]) + if err != nil { + return nil, fmt.Errorf("embed batch [%d:%d]: %w", start, end, err) + } + if len(batch) != end-start { + return nil, fmt.Errorf("embed batch [%d:%d] returned %d vectors for %d texts", start, end, len(batch), end-start) + } + for _, vector := range batch { + if len(vector) != o.dim { + return nil, fmt.Errorf("embedding dimension mismatch: got %d want %d", len(vector), o.dim) + } + all = append(all, vector) + } + } + return all, nil +} + +func (o *OllamaEmbedder) embedBatch(ctx context.Context, texts []string) ([][]float32, error) { + body, err := json.Marshal(map[string]any{"model": o.model, "input": texts}) + if err != nil { + return nil, err + } + resp, err := o.do(ctx, "/api/embed", body) + if err != nil { + return nil, err + } + return resp, nil +} + +func (o *OllamaEmbedder) do(ctx context.Context, path string, body []byte) ([][]float32, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.baseURL+path, bytes.NewReader(body)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + httpResp, err := o.client.Do(req) + if err != nil { + return nil, err + } + defer httpResp.Body.Close() + raw, _ := io.ReadAll(io.LimitReader(httpResp.Body, 8<<20)) + if httpResp.StatusCode == http.StatusNotFound { + // 模型未拉取:先 pull(流式关闭)再重试一次;仍失败则返回明确错误。 + if err := o.pullModel(ctx); err != nil { + return nil, fmt.Errorf("model %s not present and pull failed: %w", o.model, err) + } + return o.retryOnce(ctx, path, body) + } + if httpResp.StatusCode != http.StatusOK { + return nil, o.decodeError(httpResp.StatusCode, raw) + } + var parsed ollamaEmbedResponse + if err := json.Unmarshal(raw, &parsed); err != nil { + return nil, fmt.Errorf("decode embed response: %w", err) + } + return parsed.Embeddings, nil +} + +func (o *OllamaEmbedder) retryOnce(ctx context.Context, path string, body []byte) ([][]float32, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.baseURL+path, bytes.NewReader(body)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + httpResp, err := o.client.Do(req) + if err != nil { + return nil, err + } + defer httpResp.Body.Close() + raw, _ := io.ReadAll(io.LimitReader(httpResp.Body, 8<<20)) + if httpResp.StatusCode != http.StatusOK { + return nil, o.decodeError(httpResp.StatusCode, raw) + } + var parsed ollamaEmbedResponse + if err := json.Unmarshal(raw, &parsed); err != nil { + return nil, fmt.Errorf("decode embed response after pull: %w", err) + } + return parsed.Embeddings, nil +} + +func (o *OllamaEmbedder) pullModel(ctx context.Context) error { + body, err := json.Marshal(map[string]any{"model": o.model, "stream": false}) + if err != nil { + return err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.baseURL+"/api/pull", bytes.NewReader(body)) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + resp, err := o.client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + raw, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20)) + if resp.StatusCode != http.StatusOK { + return o.decodeError(resp.StatusCode, raw) + } + return nil +} + +func (o *OllamaEmbedder) decodeError(status int, raw []byte) error { + var parsed ollamaErrorResponse + if err := json.Unmarshal(raw, &parsed); err == nil && parsed.Error != "" { + return errors.New(parsed.Error) + } + return fmt.Errorf("ollama request failed with status %d", status) +} diff --git a/internal/workbench/embedder_test.go b/internal/workbench/embedder_test.go new file mode 100644 index 0000000..5519b5c --- /dev/null +++ b/internal/workbench/embedder_test.go @@ -0,0 +1,107 @@ +package workbench + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// fakeOllama 模拟 Ollama 的 /api/embed 与 /api/pull。dim 固定产出向量维数, +// requireModel 为 true 时首次 /api/embed 返回 404 触发 /api/pull。 +func fakeOllama(t *testing.T, dim int, requireModel bool) *httptest.Server { + t.Helper() + pulled := false + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/pull": + pulled = true + _ = json.NewEncoder(w).Encode(map[string]any{"status": "success"}) + case "/api/embed": + if requireModel && !pulled { + w.WriteHeader(http.StatusNotFound) + _ = json.NewEncoder(w).Encode(map[string]any{"error": "model 'bge-m3' not found"}) + return + } + var body struct { + Input []string `json:"input"` + } + _ = json.NewDecoder(r.Body).Decode(&body) + embeddings := make([][]float32, len(body.Input)) + for i := range body.Input { + vector := make([]float32, dim) + for j := range vector { + vector[j] = float32(i + 1) + } + embeddings[i] = vector + } + _ = json.NewEncoder(w).Encode(map[string]any{"embeddings": embeddings}) + default: + http.NotFound(w, r) + } + })) + t.Cleanup(server.Close) + return server +} + +func TestOllamaEmbedderBatching(t *testing.T) { + server := fakeOllama(t, 4, false) + embedder := NewOllamaEmbedder(OllamaEmbedderConfig{BaseURL: server.URL, Model: "bge-m3", Dim: 4, BatchSize: 2, Timeout: 5 * time.Second}) + texts := []string{"a", "b", "c", "d", "e"} + vectors, err := embedder.Embed(context.Background(), texts) + if err != nil { + t.Fatal(err) + } + if len(vectors) != len(texts) { + t.Fatalf("got %d vectors want %d", len(vectors), len(texts)) + } + for i, vector := range vectors { + if len(vector) != 4 { + t.Fatalf("vector %d has dim %d want 4", i, len(vector)) + } + } + if embedder.Dim() != 4 { + t.Fatalf("Dim()=%d want 4", embedder.Dim()) + } +} + +func TestOllamaEmbedderDimensionMismatch(t *testing.T) { + server := fakeOllama(t, 3, false) + embedder := NewOllamaEmbedder(OllamaEmbedderConfig{BaseURL: server.URL, Model: "bge-m3", Dim: 4, BatchSize: 2, Timeout: 5 * time.Second}) + _, err := embedder.Embed(context.Background(), []string{"x"}) + if err == nil || !strings.Contains(err.Error(), "dimension mismatch") { + t.Fatalf("expected dimension mismatch error, got %v", err) + } +} + +func TestOllamaEmbedderPullsModelOn404(t *testing.T) { + server := fakeOllama(t, 4, true) + embedder := NewOllamaEmbedder(OllamaEmbedderConfig{BaseURL: server.URL, Model: "bge-m3", Dim: 4, BatchSize: 8, Timeout: 5 * time.Second}) + vectors, err := embedder.Embed(context.Background(), []string{"hello", "world"}) + if err != nil { + t.Fatalf("expected pull-then-retry to succeed, got %v", err) + } + if len(vectors) != 2 { + t.Fatalf("got %d vectors want 2", len(vectors)) + } +} + +func TestOllamaEmbedderEmptyInput(t *testing.T) { + embedder := NewOllamaEmbedder(OllamaEmbedderConfig{BaseURL: "http://127.0.0.1:1", Model: "bge-m3", Dim: 4, BatchSize: 2, Timeout: time.Second}) + vectors, err := embedder.Embed(context.Background(), nil) + if err != nil || len(vectors) != 0 { + t.Fatalf("empty input should return nil without HTTP: vectors=%d err=%v", len(vectors), err) + } +} + +func TestFormatVector(t *testing.T) { + if got := formatVector([]float32{1, 2.5, -0.25}); got != "[1,2.5,-0.25]" { + t.Fatalf("formatVector got %q", got) + } + if got := formatVector([]float32{}); got != "[]" { + t.Fatalf("formatVector empty got %q", got) + } +} diff --git a/internal/workbench/knowledge.go b/internal/workbench/knowledge.go index e566e32..a0787c7 100644 --- a/internal/workbench/knowledge.go +++ b/internal/workbench/knowledge.go @@ -5,6 +5,7 @@ import ( "crypto/sha256" "encoding/hex" "errors" + "log/slog" "regexp" "sort" "strings" @@ -24,23 +25,11 @@ type PostgreSQLRetriever struct { } } -func NewPostgreSQLRetriever(service *Service) *PostgreSQLRetriever { - return &PostgreSQLRetriever{pool: service.pool} -} - func (r *PostgreSQLRetriever) Search(ctx context.Context, knowledgeBaseID, query string, topK int) ([]SearchHit, error) { - query = strings.TrimSpace(query) - if query == "" { - return nil, errors.New("检索词不能为空") - } - if len(query) > 6000 { - return nil, errors.New("检索词过长") - } - if topK < 1 { - topK = 4 - } - if topK > 20 { - topK = 20 + var err error + query, topK, err = validateSearch(query, topK) + if err != nil { + return nil, err } tokens := searchTokens(query) rows, err := r.pool.Query(ctx, `WITH q AS (SELECT plainto_tsquery('simple',$2) AS tsq,lower($2) AS raw), tokens AS (SELECT unnest($4::text[]) AS token) @@ -114,8 +103,11 @@ func validateKnowledgeBase(k *KnowledgeBase) error { if k.RetrievalMode == "" { k.RetrievalMode = "postgres_fts" } - if k.RetrievalMode != "postgres_fts" { - return errors.New("基线仅支持 postgres_fts 检索器") + switch k.RetrievalMode { + case "postgres_fts", "vector", "hybrid": + // M8 P2:三态检索模式。vector/hybrid 需要向量化器,未启用时检索自动回退 FTS。 + default: + return errors.New("retrieval_mode 仅支持 postgres_fts / vector / hybrid") } if k.ChunkSize == 0 { k.ChunkSize = 800 @@ -132,12 +124,12 @@ func validateKnowledgeBase(k *KnowledgeBase) error { } const knowledgeBaseSelect = `SELECT k.id::text,k.name,k.description,k.retrieval_mode,k.chunk_size,k.chunk_overlap,k.department_ids::text[],k.enabled,k.revision, -count(DISTINCT d.id)::int,count(c.id)::int,k.created_at,k.updated_at +count(DISTINCT d.id)::int,count(c.id)::int,count(c.embedding)::int,k.created_at,k.updated_at FROM gateway.knowledge_bases k LEFT JOIN gateway.knowledge_documents d ON d.knowledge_base_id=k.id LEFT JOIN gateway.knowledge_chunks c ON c.document_id=d.id` func scanKnowledgeBase(row pgx.Row) (KnowledgeBase, error) { var k KnowledgeBase - err := row.Scan(&k.ID, &k.Name, &k.Description, &k.RetrievalMode, &k.ChunkSize, &k.ChunkOverlap, &k.DepartmentIDs, &k.Enabled, &k.Revision, &k.DocumentCount, &k.ChunkCount, &k.CreatedAt, &k.UpdatedAt) + err := row.Scan(&k.ID, &k.Name, &k.Description, &k.RetrievalMode, &k.ChunkSize, &k.ChunkOverlap, &k.DepartmentIDs, &k.Enabled, &k.Revision, &k.DocumentCount, &k.ChunkCount, &k.VectorizedChunkCount, &k.CreatedAt, &k.UpdatedAt) return k, mapNotFound(err) } @@ -295,6 +287,31 @@ func ChunkText(text string, size, overlap int) []string { return chunks } +// embedResult 描述文档入库时的向量化结果。 +type embedResult struct { + embeddings [][]float32 // 与 chunks 一一对应;needed=false 或 degraded=true 时为 nil + needed bool // 该知识库模式需要 embedding + degraded bool // 需要但计算失败,文档降级为纯 FTS 入库 +} + +// computeEmbeddings 在 KB 为 vector/hybrid 且已注入 embedder 时同步批量计算分块向量。 +// 失败时不阻断入库:返回 degraded=true,由调用方在事务内补发 embedding_failed 事件。 +func (s *Service) computeEmbeddings(ctx context.Context, kb KnowledgeBase, chunks []string) embedResult { + if s.embedder == nil || (kb.RetrievalMode != "vector" && kb.RetrievalMode != "hybrid") { + return embedResult{} + } + vectors, err := s.embedder.Embed(ctx, chunks) + if err != nil { + slog.Warn("knowledge embedding failed, storing document without vectors", "knowledge_base_id", kb.ID, "error", err) + return embedResult{needed: true, degraded: true} + } + if len(vectors) != len(chunks) { + slog.Warn("knowledge embedding count mismatch, storing document without vectors", "knowledge_base_id", kb.ID, "got", len(vectors), "want", len(chunks)) + return embedResult{needed: true, degraded: true} + } + return embedResult{embeddings: vectors, needed: true} +} + func (s *Service) AddKnowledgeDocument(ctx context.Context, kbID, title, sourceType, sourceURI, content, actorID string) (KnowledgeDocument, error) { title = strings.TrimSpace(title) sourceType = strings.TrimSpace(sourceType) @@ -336,22 +353,45 @@ func (s *Service) AddKnowledgeDocument(ctx context.Context, kbID, title, sourceT if err != nil { return KnowledgeDocument{}, err } + embed := s.computeEmbeddings(ctx, kb, chunks) chunkRows := make([][]any, 0, len(chunks)) + chunkIDs := make([]string, 0, len(chunks)) + embeddings := make([]string, 0, len(chunks)) for index, chunk := range chunks { chunkID, idErr := newUUID() if idErr != nil { return KnowledgeDocument{}, idErr } chunkRows = append(chunkRows, []any{chunkID, kbID, docID, index, chunk}) + // embed.embeddings 仅在成功算出向量时非 nil;降级(失败)时跳过回填。 + if embed.embeddings != nil { + chunkIDs = append(chunkIDs, chunkID) + embeddings = append(embeddings, formatVector(embed.embeddings[index])) + } } // Bulk-copy all chunks in one statement instead of one INSERT per chunk; // a 2 MiB document can split into thousands of chunks. if _, err = tx.CopyFrom(ctx, pgx.Identifier{"gateway", "knowledge_chunks"}, []string{"id", "knowledge_base_id", "document_id", "chunk_index", "content"}, pgx.CopyFromRows(chunkRows)); err != nil { return KnowledgeDocument{}, err } + // pgx CopyFrom 对未知 OID(vector)列走二进制编码,直接随 COPY 传字面量会报 + // "vector cannot have more than 16000 dimensions";因此向量在 COPY 之后用 + // unnest 批量回填(同一事务内,单条 UPDATE)。 + if embed.embeddings != nil { + // 降级时 embeddings 为空,跳过向量回填;否则空数组 UPDATE 也是无害空操作。 + if _, err = tx.Exec(ctx, `UPDATE gateway.knowledge_chunks c SET embedding = v.vec::vector + FROM unnest($1::uuid[], $2::text[]) AS v(id, vec) WHERE c.id = v.id`, chunkIDs, embeddings); err != nil { + return KnowledgeDocument{}, err + } + } if err = emit(ctx, tx, "knowledge_document.ready", "knowledge_document", docID, actorID, map[string]any{"knowledge_base_id": kbID, "chunk_count": len(chunks)}); err != nil { return KnowledgeDocument{}, err } + if embed.degraded { + if err = emit(ctx, tx, "knowledge_document.embedding_failed", "knowledge_document", docID, actorID, map[string]any{"knowledge_base_id": kbID, "chunk_count": len(chunks)}); err != nil { + return KnowledgeDocument{}, err + } + } if err = tx.Commit(ctx); err != nil { return KnowledgeDocument{}, err } @@ -420,13 +460,20 @@ func (s *Service) ReprocessKnowledgeDocument(ctx context.Context, kbID, id, acto if _, err = tx.Exec(ctx, `DELETE FROM gateway.knowledge_chunks WHERE document_id=$1`, id); err != nil { return KnowledgeDocument{}, err } + embed := s.computeEmbeddings(ctx, kb, chunks) for index, chunk := range chunks { chunkID, idErr := newUUID() if idErr != nil { return KnowledgeDocument{}, idErr } - if _, err = tx.Exec(ctx, `INSERT INTO gateway.knowledge_chunks(id,knowledge_base_id,document_id,chunk_index,content) VALUES($1,$2,$3,$4,$5)`, chunkID, kbID, id, index, chunk); err != nil { - return KnowledgeDocument{}, err + if embed.embeddings != nil { + if _, err = tx.Exec(ctx, `INSERT INTO gateway.knowledge_chunks(id,knowledge_base_id,document_id,chunk_index,content,embedding) VALUES($1,$2,$3,$4,$5,$6::vector)`, chunkID, kbID, id, index, chunk, formatVector(embed.embeddings[index])); err != nil { + return KnowledgeDocument{}, err + } + } else { + if _, err = tx.Exec(ctx, `INSERT INTO gateway.knowledge_chunks(id,knowledge_base_id,document_id,chunk_index,content) VALUES($1,$2,$3,$4,$5)`, chunkID, kbID, id, index, chunk); err != nil { + return KnowledgeDocument{}, err + } } } err = tx.QueryRow(ctx, `UPDATE gateway.knowledge_documents SET status='ready',status_message='',chunk_count=$3,updated_at=clock_timestamp() WHERE id=$1 AND knowledge_base_id=$2 RETURNING updated_at`, id, kbID, len(chunks)).Scan(&doc.UpdatedAt) @@ -437,6 +484,11 @@ func (s *Service) ReprocessKnowledgeDocument(ctx context.Context, kbID, id, acto if err = emit(ctx, tx, "knowledge_document.reprocessed", "knowledge_document", id, actorID, map[string]any{"knowledge_base_id": kbID, "chunk_count": len(chunks)}); err != nil { return KnowledgeDocument{}, err } + if embed.degraded { + if err = emit(ctx, tx, "knowledge_document.embedding_failed", "knowledge_document", id, actorID, map[string]any{"knowledge_base_id": kbID, "chunk_count": len(chunks)}); err != nil { + return KnowledgeDocument{}, err + } + } if err = tx.Commit(ctx); err != nil { return KnowledgeDocument{}, err } diff --git a/internal/workbench/retrievers.go b/internal/workbench/retrievers.go new file mode 100644 index 0000000..254902d --- /dev/null +++ b/internal/workbench/retrievers.go @@ -0,0 +1,159 @@ +package workbench + +import ( + "context" + "errors" + "sort" + "strconv" + "strings" + + "github.com/jackc/pgx/v5" +) + +// NewRetriever 按知识库的 retrieval_mode 分发检索器: +// - embedder 为 nil 时恒为 postgres_fts(向量化未启用) +// - 否则按 knowledge_bases.retrieval_mode 每次 Search 动态分发 +// (postgres_fts / vector / hybrid),管理员修改模式后无需重启即生效。 +func NewRetriever(service *Service, embedder Embedder) Retriever { + return &modeRetriever{pool: service.pool, embedder: embedder} +} + +type modeRetriever struct { + pool interface { + Query(context.Context, string, ...any) (pgx.Rows, error) + QueryRow(context.Context, string, ...any) pgx.Row + } + embedder Embedder +} + +func (r *modeRetriever) Search(ctx context.Context, knowledgeBaseID, query string, topK int) ([]SearchHit, error) { + if r.embedder == nil { + return (&PostgreSQLRetriever{pool: r.pool}).Search(ctx, knowledgeBaseID, query, topK) + } + var mode string + err := r.pool.QueryRow(ctx, `SELECT retrieval_mode FROM gateway.knowledge_bases WHERE id=$1`, knowledgeBaseID).Scan(&mode) + if err != nil { + return nil, mapNotFound(err) + } + switch mode { + case "vector": + return (&SemanticRetriever{pool: r.pool, embedder: r.embedder}).Search(ctx, knowledgeBaseID, query, topK) + case "hybrid": + return (&HybridRetriever{pool: r.pool, embedder: r.embedder}).Search(ctx, knowledgeBaseID, query, topK) + default: + return (&PostgreSQLRetriever{pool: r.pool}).Search(ctx, knowledgeBaseID, query, topK) + } +} + +// SemanticRetriever 用 pgvector 余弦距离召回最近分块。 +type SemanticRetriever struct { + pool interface { + Query(context.Context, string, ...any) (pgx.Rows, error) + } + embedder Embedder +} + +func (r *SemanticRetriever) Search(ctx context.Context, knowledgeBaseID, query string, topK int) ([]SearchHit, error) { + query, topK, err := validateSearch(query, topK) + if err != nil { + return nil, err + } + vectors, err := r.embedder.Embed(ctx, []string{query}) + if err != nil { + return nil, err + } + if len(vectors) == 0 { + return nil, errors.New("查询向量为空") + } + rows, err := r.pool.Query(ctx, `SELECT c.id::text,c.document_id::text,d.title,c.chunk_index,c.content, + (1-(c.embedding <=> $2::vector))::float8 AS score + FROM gateway.knowledge_chunks c JOIN gateway.knowledge_documents d ON d.id=c.document_id JOIN gateway.knowledge_bases k ON k.id=c.knowledge_base_id + WHERE c.knowledge_base_id=$1 AND k.enabled AND d.status='ready' AND c.embedding IS NOT NULL + ORDER BY c.embedding <=> $2::vector LIMIT $3`, knowledgeBaseID, formatVector(vectors[0]), topK) + if err != nil { + return nil, err + } + defer rows.Close() + hits := []SearchHit{} + for rows.Next() { + var h SearchHit + if err = rows.Scan(&h.ChunkID, &h.DocumentID, &h.DocumentTitle, &h.ChunkIndex, &h.Content, &h.Score); err != nil { + return nil, err + } + hits = append(hits, h) + } + return hits, rows.Err() +} + +// HybridRetriever 融合 FTS 与语义召回:分别取 topK*2,按 chunk_id 去重取高分, +// 再截断到 topK。任一子检索失败时回退到另一路,保证向量化降级仍有结果。 +type HybridRetriever struct { + pool interface { + Query(context.Context, string, ...any) (pgx.Rows, error) + } + embedder Embedder +} + +func (r *HybridRetriever) Search(ctx context.Context, knowledgeBaseID, query string, topK int) ([]SearchHit, error) { + _, topK, err := validateSearch(query, topK) + if err != nil { + return nil, err + } + fts, ftsErr := (&PostgreSQLRetriever{pool: r.pool}).Search(ctx, knowledgeBaseID, query, topK*2) + sem, semErr := (&SemanticRetriever{pool: r.pool, embedder: r.embedder}).Search(ctx, knowledgeBaseID, query, topK*2) + if ftsErr != nil && semErr != nil { + return nil, ftsErr + } + best := make(map[string]SearchHit) + for _, h := range fts { + if existing, ok := best[h.ChunkID]; !ok || h.Score > existing.Score { + best[h.ChunkID] = h + } + } + for _, h := range sem { + if existing, ok := best[h.ChunkID]; !ok || h.Score > existing.Score { + best[h.ChunkID] = h + } + } + merged := make([]SearchHit, 0, len(best)) + for _, h := range best { + merged = append(merged, h) + } + sort.Slice(merged, func(i, j int) bool { return merged[i].Score > merged[j].Score }) + if len(merged) > topK { + merged = merged[:topK] + } + return merged, nil +} + +func validateSearch(query string, topK int) (string, int, error) { + query = strings.TrimSpace(query) + if query == "" { + return "", 0, errors.New("检索词不能为空") + } + if len(query) > 6000 { + return "", 0, errors.New("检索词过长") + } + if topK < 1 { + topK = 4 + } + if topK > 20 { + topK = 20 + } + return query, topK, nil +} + +// formatVector 把 []float32 转成 pgvector 字面量字符串 "[0.1,0.2,...]", +// 直接以 $n::vector 参数传入,避免引入 pgvector-go 依赖。 +func formatVector(vector []float32) string { + var builder strings.Builder + builder.WriteByte('[') + for index, value := range vector { + if index > 0 { + builder.WriteByte(',') + } + builder.WriteString(strconv.FormatFloat(float64(value), 'g', -1, 32)) + } + builder.WriteByte(']') + return builder.String() +} diff --git a/internal/workbench/retrievers_test.go b/internal/workbench/retrievers_test.go new file mode 100644 index 0000000..65c437d --- /dev/null +++ b/internal/workbench/retrievers_test.go @@ -0,0 +1,108 @@ +package workbench + +import ( + "context" + "errors" + "testing" +) + +func TestValidateKnowledgeBaseRetrievalModes(t *testing.T) { + for _, mode := range []string{"postgres_fts", "vector", "hybrid"} { + kb := KnowledgeBase{Name: "kb", RetrievalMode: mode, ChunkSize: 800, ChunkOverlap: 100} + if err := validateKnowledgeBase(&kb); err != nil { + t.Fatalf("mode %s should be accepted, got %v", mode, err) + } + } + kb := KnowledgeBase{Name: "kb", RetrievalMode: "bm25", ChunkSize: 800, ChunkOverlap: 100} + if err := validateKnowledgeBase(&kb); err == nil { + t.Fatal("unknown retrieval mode should be rejected") + } + // 空模式回填默认值 + empty := KnowledgeBase{Name: "kb", ChunkSize: 800, ChunkOverlap: 100} + if err := validateKnowledgeBase(&empty); err != nil || empty.RetrievalMode != "postgres_fts" { + t.Fatalf("empty mode should default to postgres_fts, got %q err=%v", empty.RetrievalMode, err) + } +} + +func TestValidateSearchClamps(t *testing.T) { + if _, topK, err := validateSearch(" 关键词 ", 0); err != nil || topK != 4 { + t.Fatalf("empty topK should clamp to 4, got %d err=%v", topK, err) + } + if _, topK, err := validateSearch("关键词", 99); err != nil || topK != 20 { + t.Fatalf("topK should clamp to 20, got %d err=%v", topK, err) + } + if _, _, err := validateSearch(" ", 4); err == nil { + t.Fatal("blank query should be rejected") + } + long := make([]byte, 6001) + for i := range long { + long[i] = 'a' + } + if _, _, err := validateSearch(string(long), 4); err == nil { + t.Fatal("over-long query should be rejected") + } +} + +// failingEmbedder 总是报错,用于验证 computeEmbeddings 优雅降级。 +type failingEmbedder struct{} + +func (failingEmbedder) Embed(context.Context, []string) ([][]float32, error) { return nil, errors.New("embedding service down") } +func (failingEmbedder) Dim() int { return 1024 } + +func TestComputeEmbeddingsDispatch(t *testing.T) { + service := NewService(nil) + chunks := []string{"a", "b"} + + // 未注入 embedder:任何模式都不算向量。 + if result := service.computeEmbeddings(context.Background(), KnowledgeBase{ID: "kb", RetrievalMode: "vector"}, chunks); result.needed { + t.Fatal("no embedder should not mark embedding needed") + } + + // 已注入 embedder 但模式为 postgres_fts:不算向量。 + service.SetEmbedder(fakeGoodEmbedder{}) + if result := service.computeEmbeddings(context.Background(), KnowledgeBase{ID: "kb", RetrievalMode: "postgres_fts"}, chunks); result.needed { + t.Fatal("postgres_fts mode should not embed") + } + + // vector 模式 + 失败 embedder:needed + degraded,embedding 为 nil。 + service.SetEmbedder(failingEmbedder{}) + result := service.computeEmbeddings(context.Background(), KnowledgeBase{ID: "kb", RetrievalMode: "hybrid"}, chunks) + if !result.needed || !result.degraded || result.embeddings != nil { + t.Fatalf("expected needed+degraded with nil embeddings, got %+v", result) + } + + // vector 模式 + 正常 embedder:返回与 chunks 等长的向量。 + service.SetEmbedder(fakeGoodEmbedder{}) + result = service.computeEmbeddings(context.Background(), KnowledgeBase{ID: "kb", RetrievalMode: "vector"}, chunks) + if !result.needed || result.degraded || len(result.embeddings) != len(chunks) { + t.Fatalf("expected aligned embeddings, got %+v", result) + } +} + +type fakeGoodEmbedder struct{} + +func (fakeGoodEmbedder) Embed(_ context.Context, texts []string) ([][]float32, error) { + vectors := make([][]float32, len(texts)) + for i := range texts { + vectors[i] = make([]float32, 4) + for j := range vectors[i] { + vectors[i][j] = float32(i + 1) + } + } + return vectors, nil +} + +func (fakeGoodEmbedder) Dim() int { return 4 } + +func TestNewRetrieverNilEmbedder(t *testing.T) { + // embedder 为 nil 时 NewRetriever 返回的检索器直接走 FTS 分支,不会查询检索模式。 + service := NewService(nil) + retriever := NewRetriever(service, nil) + if retriever == nil { + t.Fatal("NewRetriever with nil embedder should not return nil") + } + // FTS 分支在无连接池时只在 query 校验处返回,不 panic。 + if _, err := retriever.Search(context.Background(), "kb", " ", 4); err == nil { + t.Fatal("blank query should fail at validation before touching the pool") + } +} diff --git a/internal/workbench/service.go b/internal/workbench/service.go index f3606e3..9ede324 100644 --- a/internal/workbench/service.go +++ b/internal/workbench/service.go @@ -12,10 +12,19 @@ import ( "github.com/jackc/pgx/v5/pgxpool" ) -type Service struct{ pool *pgxpool.Pool } +type Service struct { + pool *pgxpool.Pool + embedder Embedder +} func NewService(pool *pgxpool.Pool) *Service { return &Service{pool: pool} } +// SetEmbedder 注入向量化器(M8 P2)。为 nil 时知识库退回纯 FTS,不生成 embedding。 +func (s *Service) SetEmbedder(embedder Embedder) { s.embedder = embedder } + +// Embedder 返回当前向量化器(可能为 nil)。 +func (s *Service) Embedder() Embedder { return s.embedder } + func newUUID() (string, error) { return platformid.NewUUID() } func emit(ctx context.Context, tx pgx.Tx, eventType, aggregateType, aggregateID, actorID string, values map[string]any) error { diff --git a/internal/workbench/types.go b/internal/workbench/types.go index e2920bb..820f479 100644 --- a/internal/workbench/types.go +++ b/internal/workbench/types.go @@ -82,11 +82,12 @@ type KnowledgeBase struct { ChunkOverlap int `json:"chunk_overlap"` DepartmentIDs []string `json:"department_ids"` Enabled bool `json:"enabled"` - Revision int64 `json:"revision"` - DocumentCount int `json:"document_count"` - ChunkCount int `json:"chunk_count"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` + Revision int64 `json:"revision"` + DocumentCount int `json:"document_count"` + ChunkCount int `json:"chunk_count"` + VectorizedChunkCount int `json:"vectorized_chunk_count"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` } type KnowledgeDocument struct { diff --git a/internal/workbench/vector_integration_test.go b/internal/workbench/vector_integration_test.go new file mode 100644 index 0000000..68c1ea4 --- /dev/null +++ b/internal/workbench/vector_integration_test.go @@ -0,0 +1,90 @@ +package workbench + +import ( + "context" + "os" + "strings" + "testing" + "time" + + "aigateway.local/core/internal/platform/config" + "aigateway.local/core/internal/platform/database" +) + +// TestKnowledgeVectorLifecycle 跑真实 PostgreSQL(pgvector)+ Ollama: +// vector 模式知识库导入文档 → embedding 非空 → SemanticRetriever 语义命中; +// 并验证 embedder 失败时文档照常入库且发出 embedding_failed 事件。 +// 需要 WORKBENCH_TEST_DATABASE_URL 与 WORKBENCH_TEST_OLLAMA_URL。 +func TestKnowledgeVectorLifecycle(t *testing.T) { + databaseURL := os.Getenv("WORKBENCH_TEST_DATABASE_URL") + ollamaURL := os.Getenv("WORKBENCH_TEST_OLLAMA_URL") + if databaseURL == "" || ollamaURL == "" { + t.Skip("WORKBENCH_TEST_DATABASE_URL and WORKBENCH_TEST_OLLAMA_URL are not set") + } + ctx := context.Background() + pool, err := database.Open(ctx, config.Database{URL: databaseURL, MaxConns: 8, MinConns: 0}) + if err != nil { + t.Fatal(err) + } + defer pool.Close() + actorID := "55555555-5555-4555-8555-555555555555" + _, err = pool.Exec(ctx, `INSERT INTO gateway.admin_accounts(id,username,password_hash,role) VALUES($1,'m8-vector-test','test','superadmin') ON CONFLICT(id) DO NOTHING`, actorID) + if err != nil { + t.Fatal(err) + } + name := "m8-vector-integration-kb" + cleanup := func() { + _, _ = pool.Exec(ctx, `DELETE FROM gateway.knowledge_bases WHERE name=$1 OR name=$2`, name, name+"-fts") + } + cleanup() + defer cleanup() + + assets := NewService(pool) + assets.SetEmbedder(NewOllamaEmbedder(OllamaEmbedderConfig{BaseURL: ollamaURL, Model: "bge-m3", Dim: 1024, BatchSize: 8, Timeout: 60 * time.Second})) + + kb, err := assets.SaveKnowledgeBase(ctx, KnowledgeBase{Name: name, Description: "vector itest", RetrievalMode: "vector", ChunkSize: 300, ChunkOverlap: 40, DepartmentIDs: []string{}, Enabled: true}, actorID, true) + if err != nil { + t.Fatal(err) + } + doc, err := assets.AddKnowledgeDocument(ctx, kb.ID, "向量检索测试", "text", "", "pgvector 语义检索依赖 bge-m3 向量。Ollama 本地生成嵌入。", actorID) + if err != nil { + t.Fatal(err) + } + var embedded int + if err = pool.QueryRow(ctx, `SELECT count(*) FROM gateway.knowledge_chunks WHERE document_id=$1 AND embedding IS NOT NULL`, doc.ID).Scan(&embedded); err != nil { + t.Fatal(err) + } + if embedded == 0 || embedded != doc.ChunkCount { + t.Fatalf("expected all chunks embedded, got %d/%d", embedded, doc.ChunkCount) + } + + // 语义检索:查询词与正文无字面重合也应命中(余弦相似度)。 + hitChunks, err := (&SemanticRetriever{pool: pool, embedder: assets.Embedder()}).Search(ctx, kb.ID, "语义相似度匹配", 4) + if err != nil { + t.Fatalf("semantic search: %v", err) + } + if len(hitChunks) == 0 { + t.Fatal("expected semantic hit") + } + if !strings.Contains(hitChunks[0].Content, "pgvector") { + t.Fatalf("expected pgvector content in top hit, got %q", hitChunks[0].Content) + } + + // 降级路径:embedder 失败 → 文档照常入库 + embedding_failed 事件。 + assets.SetEmbedder(failingEmbedder{}) + kbFts, err := assets.SaveKnowledgeBase(ctx, KnowledgeBase{Name: name + "-fts", Description: "degraded", RetrievalMode: "vector", ChunkSize: 300, ChunkOverlap: 40, DepartmentIDs: []string{}, Enabled: true}, actorID, true) + if err != nil { + t.Fatal(err) + } + degradedDoc, err := assets.AddKnowledgeDocument(ctx, kbFts.ID, "降级测试", "text", "", "没有向量也能入库。", actorID) + if err != nil { + t.Fatalf("degraded add should succeed, got %v", err) + } + var failed bool + if err = pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.outbox_events WHERE event_type='knowledge_document.embedding_failed' AND aggregate_id=$1)`, degradedDoc.ID).Scan(&failed); err != nil { + t.Fatal(err) + } + if !failed { + t.Fatal("expected embedding_failed event when embedder is down") + } +} diff --git a/internal/workbench/workbench_integration_test.go b/internal/workbench/workbench_integration_test.go index 48769ed..52cfc6b 100644 --- a/internal/workbench/workbench_integration_test.go +++ b/internal/workbench/workbench_integration_test.go @@ -66,7 +66,7 @@ func TestWorkbenchPostgreSQLLifecycle(t *testing.T) { if doc.ChunkCount == 0 { t.Fatal("expected chunks") } - hits, err := NewPostgreSQLRetriever(assets).Search(ctx, kb.ID, "不可变运行时快照", 4) + hits, err := NewRetriever(assets, nil).Search(ctx, kb.ID, "不可变运行时快照", 4) if err != nil || len(hits) == 0 { t.Fatalf("hits=%#v err=%v", hits, err) } @@ -116,7 +116,7 @@ func TestWorkbenchPostgreSQLLifecycle(t *testing.T) { } _ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{"message": map[string]any{"role": "assistant", "content": "完成"}}}}) }) - runtime := NewRuntimeHTTPHandler(assets, tools, NewPostgreSQLRetriever(assets), staticPrincipalAuthenticator{}, fakeGateway, MarketplaceDeps{}) + runtime := NewRuntimeHTTPHandler(assets, tools, NewRetriever(assets, nil), staticPrincipalAuthenticator{}, fakeGateway, MarketplaceDeps{}) runtimeRequest := httptest.NewRequest(http.MethodPost, "/v1/applications/m4_app/chat/completions", bytes.NewBufferString(`{"messages":[{"role":"user","content":"不可变运行时快照是什么？"}],"variables":{"question":"架构"}}`)) runtimeRequest.Header.Set("Authorization", "Bearer test") runtimeRequest = runtimeRequest.WithContext(gateway.WithRequestID(runtimeRequest.Context(), "m4-runtime")) diff --git a/migrations/000024_pgvector_embeddings.sql b/migrations/000024_pgvector_embeddings.sql new file mode 100644 index 0000000..45b1bb2 --- /dev/null +++ b/migrations/000024_pgvector_embeddings.sql @@ -0,0 +1,21 @@ +-- 000024_pgvector_embeddings.sql +-- M8 P2:知识库向量化。在 knowledge_chunks 上增加 pgvector embedding 列,HNSW 余弦索引, +-- 并把 knowledge_bases.retrieval_mode 放宽为三态(postgres_fts / vector / hybrid)。 +-- +-- 依赖 pgvector 扩展:部署镜像需从 postgres:17-alpine 换成 pgvector/pgvector:pg17 +-- (数据卷兼容)。gateway 用户即 POSTGRES_USER(容器内 superuser),可创建扩展。 + +CREATE EXTENSION IF NOT EXISTS vector; + +ALTER TABLE gateway.knowledge_chunks ADD COLUMN IF NOT EXISTS embedding vector(1024); + +-- HNSW 余弦索引;embedding 为 NULL 的行不进入索引(旧文档未向量化)。 +CREATE INDEX IF NOT EXISTS knowledge_chunks_embedding_idx + ON gateway.knowledge_chunks USING hnsw (embedding vector_cosine_ops); + +ALTER TABLE gateway.knowledge_bases DROP CONSTRAINT IF EXISTS knowledge_bases_retrieval_mode_check; +ALTER TABLE gateway.knowledge_bases ADD CONSTRAINT knowledge_bases_retrieval_mode_check + CHECK (retrieval_mode IN ('postgres_fts', 'vector', 'hybrid')); + +COMMENT ON COLUMN gateway.knowledge_chunks.embedding IS '语义向量,由本地 Ollama bge-m3 生成;为 NULL 表示该分块尚未向量化'; +COMMENT ON CONSTRAINT knowledge_bases_retrieval_mode_check ON gateway.knowledge_bases IS '检索模式:postgres_fts 全文检索 / vector 语义向量 / hybrid 混合'; diff --git a/web/apps/admin/src/api/workbench.ts b/web/apps/admin/src/api/workbench.ts index 5f7be68..d81fb4f 100644 --- a/web/apps/admin/src/api/workbench.ts +++ b/web/apps/admin/src/api/workbench.ts @@ -16,7 +16,7 @@ export const fetchPromptVersions=(id:string)=>request.get({url: export const createPromptVersion=(id:string,params:{content:string;variables:PromptVariable[];change_note:string;activate:boolean})=>request.post({url:`/api/v1/admin/prompts/${id}/versions`,params}) export const activatePromptVersion=(id:string,version:number)=>request.post({url:`/api/v1/admin/prompts/${id}/versions/${version}/activate`}) -export interface KnowledgeBase { id:string;name:string;description:string;retrieval_mode:string;chunk_size:number;chunk_overlap:number;department_ids:string[];enabled:boolean;revision:number;document_count:number;chunk_count:number } +export interface KnowledgeBase { id:string;name:string;description:string;retrieval_mode:string;chunk_size:number;chunk_overlap:number;department_ids:string[];enabled:boolean;revision:number;document_count:number;chunk_count:number;vectorized_chunk_count:number } export interface KnowledgeDocument {id:string;knowledge_base_id:string;title:string;source_type:string;source_uri:string;content_sha256:string;status:string;status_message:string;char_count:number;chunk_count:number;created_at:string} export interface SearchHit {chunk_id:string;document_id:string;document_title:string;chunk_index:number;content:string;score:number} export type KnowledgeBaseInput=Pick diff --git a/web/apps/admin/src/views/gateway/knowledge/index.vue b/web/apps/admin/src/views/gateway/knowledge/index.vue index 0c82d8d..4dd34c9 100644 --- a/web/apps/admin/src/views/gateway/knowledge/index.vue +++ b/web/apps/admin/src/views/gateway/knowledge/index.vue @@ -1,5 +1,5 @@ -