706 lines
24 KiB
Python
706 lines
24 KiB
Python
"""U8 — TaskIQ 异步任务集成 — 文档向量化与批量索引。
|
||
|
||
包装 TaskIQ broker，提供：
|
||
- 任务参数 schema 校验（Pydantic 模型）
|
||
- per-user 并发上限
|
||
- 任务状态机（PENDING → RUNNING → COMPLETED | FAILED | CANCELLED）
|
||
- Worker liveness & sweeper（worker_heartbeat_at 超时检测）
|
||
- 错误消息净化（剥离内部路径/连接串）
|
||
- 降级模式：broker 未配置时同步执行
|
||
|
||
Redis 隔离：使用独立 db=1（与 bus/cache 的 db=0 分离），或 key 前缀 ``taskiq:``。
|
||
任务参数在 Redis 中带 TTL（默认 1 小时）。
|
||
|
||
ponytail: TaskIQ broker 为可选依赖 — 通过 TYPE_CHECKING 导入，
|
||
未安装时模块仍可加载，TaskManager 退化为同步执行。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import inspect
|
||
import logging
|
||
import re
|
||
import uuid
|
||
from datetime import datetime, timedelta, timezone
|
||
from typing import TYPE_CHECKING, Protocol
|
||
|
||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||
|
||
from agentkit.core.protocol import TaskStatus
|
||
from agentkit.rag_platform.document_processor import (
|
||
DEFAULT_CHUNK_OVERLAP,
|
||
DEFAULT_CHUNK_SIZE,
|
||
DocumentProcessor,
|
||
)
|
||
from agentkit.server.task_store import InMemoryTaskStore, TaskRecord
|
||
|
||
if TYPE_CHECKING:
|
||
from llama_index.core.embeddings import BaseEmbedding
|
||
from llama_index.vector_stores.postgres import PGVectorStore
|
||
|
||
from agentkit.rag_platform.store import KBStore
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 默认配置
|
||
DEFAULT_MAX_CONCURRENT_PER_USER = 3
|
||
DEFAULT_WORKER_TTL_SECONDS = 300 # worker 心跳超时阈值
|
||
DEFAULT_ERROR_MESSAGE_MAX_LENGTH = 500
|
||
DEFAULT_TASK_TTL_SECONDS = 3600 # 任务参数 TTL
|
||
|
||
# Redis 隔离 — 独立 db=1，与 bus/cache 的 db=0 分离
|
||
TASKIQ_REDIS_DB = 1
|
||
TASKIQ_KEY_PREFIX = "taskiq:"
|
||
|
||
|
||
async def _maybe_await(result: object) -> object:
|
||
"""统一处理 sync/async 调用结果 — InMemoryTaskStore 方法为 sync，RedisTaskStore 为 async。"""
|
||
if inspect.isawaitable(result):
|
||
return await result
|
||
return result
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 任务参数模型 — schema 校验
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class VectorizeTaskParams(BaseModel):
|
||
"""向量化任务参数 — 单文档 parse → segment → vectorize → index。"""
|
||
|
||
model_config = ConfigDict()
|
||
|
||
document_id: str = Field(min_length=1, max_length=128)
|
||
kb_id: str = Field(min_length=1, max_length=128)
|
||
file_path: str = Field(min_length=1, max_length=1024)
|
||
file_type: str = Field(min_length=1, max_length=64)
|
||
chunk_size: int = Field(default=DEFAULT_CHUNK_SIZE, ge=50, le=8192)
|
||
chunk_overlap: int = Field(default=DEFAULT_CHUNK_OVERLAP, ge=0, le=1024)
|
||
user_id: str = Field(min_length=1, max_length=128)
|
||
|
||
@field_validator("chunk_overlap")
|
||
@classmethod
|
||
def _overlap_less_than_size(cls, v: int, info: object) -> int:
|
||
"""chunk_overlap 必须小于 chunk_size。"""
|
||
size = info.data.get("chunk_size", DEFAULT_CHUNK_SIZE)
|
||
if v >= size:
|
||
raise ValueError(f"chunk_overlap ({v}) must be < chunk_size ({size})")
|
||
return v
|
||
|
||
|
||
class BatchIndexTaskParams(BaseModel):
|
||
"""批量索引任务参数 — 多文档并行向量化。"""
|
||
|
||
model_config = ConfigDict()
|
||
|
||
kb_id: str = Field(min_length=1, max_length=128)
|
||
document_ids: list[str] = Field(min_length=1, max_length=100)
|
||
user_id: str = Field(min_length=1, max_length=128)
|
||
file_paths: dict[str, str] = Field(default_factory=dict) # document_id -> file_path
|
||
file_types: dict[str, str] = Field(default_factory=dict) # document_id -> file_type
|
||
chunk_size: int = Field(default=DEFAULT_CHUNK_SIZE, ge=50, le=8192)
|
||
chunk_overlap: int = Field(default=DEFAULT_CHUNK_OVERLAP, ge=0, le=1024)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 错误消息净化
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# 匹配 Unix/Windows 文件路径
|
||
_PATH_RE = re.compile(
|
||
r"(?:/[\w.\-]+)+|" # Unix 路径 /usr/local/...
|
||
r"(?:[A-Za-z]:\\[\w.\-]+(?:\\[\w.\-]+)*)" # Windows 路径 C:\Users\...
|
||
)
|
||
# 匹配 Redis/PostgreSQL 连接串
|
||
_CONN_STR_RE = re.compile(
|
||
r"(redis://|postgresql://|postgres://|mysql://)"
|
||
r"[^\s\"']+",
|
||
re.IGNORECASE,
|
||
)
|
||
# 匹配环境变量风格的密钥
|
||
_SECRET_RE = re.compile(
|
||
r"(api[_-]?key|token|password|secret)[\s:=]+[\w\-]+",
|
||
re.IGNORECASE,
|
||
)
|
||
|
||
|
||
def sanitize_error_message(msg: str, max_length: int = DEFAULT_ERROR_MESSAGE_MAX_LENGTH) -> str:
|
||
"""净化错误消息 — 剥离内部路径、连接串、密钥。
|
||
|
||
Args:
|
||
msg: 原始错误消息
|
||
max_length: 最大长度（截断）
|
||
|
||
Returns:
|
||
净化后的错误消息
|
||
"""
|
||
if not msg:
|
||
return ""
|
||
# 连接串优先于路径 — 避免 /localhost:6379/0 被路径正则部分匹配
|
||
cleaned = _CONN_STR_RE.sub("<connection-string>", msg)
|
||
cleaned = _PATH_RE.sub("<path>", cleaned)
|
||
cleaned = _SECRET_RE.sub(r"\1=<redacted>", cleaned)
|
||
if len(cleaned) > max_length:
|
||
cleaned = cleaned[: max_length - 3] + "..."
|
||
return cleaned
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Worker liveness & sweeper
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TaskStoreProtocol(Protocol):
|
||
"""TaskStore 协议 — InMemoryTaskStore / RedisTaskStore 均满足。"""
|
||
|
||
async def create(
|
||
self, task_id: str, agent_name: str, input_data: dict, skill_name: str | None = None
|
||
) -> TaskRecord: ...
|
||
|
||
def get(self, task_id: str) -> TaskRecord | None: ...
|
||
|
||
async def update_status(
|
||
self, task_id: str, status: TaskStatus, **kwargs: object
|
||
) -> TaskRecord: ...
|
||
|
||
def list_tasks(
|
||
self, status: TaskStatus | None = None, limit: int = 100
|
||
) -> list[TaskRecord]: ...
|
||
|
||
|
||
class WorkerSweeper:
|
||
"""检测超时任务并标记为 failed。
|
||
|
||
扫描 status=RUNNING 且 worker_heartbeat_at 超过 TTL 的任务，
|
||
将其状态置为 failed，error_message="worker_timeout"。
|
||
|
||
ponytail: 升级路径 — 当前依赖 list_tasks 全量扫描（O(n)），
|
||
生产环境可改为 Redis ZSET 按心跳时间戳排序（O(log n)）。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
task_store: TaskStoreProtocol,
|
||
ttl_seconds: int = DEFAULT_WORKER_TTL_SECONDS,
|
||
) -> None:
|
||
self._store = task_store
|
||
self._ttl = ttl_seconds
|
||
|
||
async def sweep(self) -> int:
|
||
"""扫描超时任务，返回清理数量。"""
|
||
now = datetime.now(timezone.utc)
|
||
cutoff = now - timedelta(seconds=self._ttl)
|
||
cleaned = 0
|
||
|
||
running_tasks = await _maybe_await(
|
||
self._store.list_tasks(status=TaskStatus.RUNNING, limit=1000)
|
||
)
|
||
|
||
for record in running_tasks:
|
||
heartbeat = record.metadata.get("worker_heartbeat_at")
|
||
if heartbeat is None:
|
||
# 无心跳记录 — 使用 started_at 作为回退
|
||
heartbeat = record.started_at.isoformat() if record.started_at else None
|
||
if heartbeat is None:
|
||
continue
|
||
|
||
try:
|
||
hb_time = datetime.fromisoformat(heartbeat)
|
||
except (ValueError, TypeError):
|
||
continue
|
||
|
||
# 确保 timezone-aware
|
||
if hb_time.tzinfo is None:
|
||
hb_time = hb_time.replace(tzinfo=timezone.utc)
|
||
|
||
if hb_time < cutoff:
|
||
try:
|
||
await _maybe_await(
|
||
self._store.update_status(
|
||
record.task_id,
|
||
TaskStatus.FAILED,
|
||
error_message="worker_timeout",
|
||
completed_at=now,
|
||
metadata={**record.metadata, "swept_at": now.isoformat()},
|
||
)
|
||
)
|
||
cleaned += 1
|
||
logger.warning(
|
||
"Swept timed-out task %s (heartbeat=%s, cutoff=%s)",
|
||
record.task_id,
|
||
heartbeat,
|
||
cutoff.isoformat(),
|
||
)
|
||
except KeyError:
|
||
# 任务已被并发删除
|
||
continue
|
||
return cleaned
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# TaskManager — 任务提交、查询、取消
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class ConcurrencyLimitExceeded(Exception):
|
||
"""per-user 并发上限超出。"""
|
||
|
||
|
||
class TaskNotFoundError(Exception):
|
||
"""任务不存在。"""
|
||
|
||
|
||
class TaskManager:
|
||
"""异步任务管理器 — 包装 TaskIQ broker。
|
||
|
||
如果 broker 未配置或 TaskIQ 不可用，任务同步执行（降级模式）。
|
||
|
||
Args:
|
||
broker: TaskIQ broker 实例（可选）
|
||
task_store: 任务状态存储（默认 InMemoryTaskStore）
|
||
max_concurrent_per_user: per-user 并发上限
|
||
task_ttl_seconds: 任务参数 TTL（Redis 模式）
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
broker: object | None = None,
|
||
task_store: TaskStoreProtocol | None = None,
|
||
max_concurrent_per_user: int = DEFAULT_MAX_CONCURRENT_PER_USER,
|
||
task_ttl_seconds: int = DEFAULT_TASK_TTL_SECONDS,
|
||
) -> None:
|
||
self._broker = broker
|
||
self._store: TaskStoreProtocol = task_store or InMemoryTaskStore()
|
||
self._max_concurrent = max_concurrent_per_user
|
||
self._task_ttl = task_ttl_seconds
|
||
# user_id -> 当前 RUNNING 任务数
|
||
self._running_counts: dict[str, int] = {}
|
||
self._lock = asyncio.Lock()
|
||
|
||
@property
|
||
def store(self) -> TaskStoreProtocol:
|
||
"""暴露底层 task store（供 sweeper 使用）。"""
|
||
return self._store
|
||
|
||
@property
|
||
def broker(self) -> object:
|
||
"""暴露底层 broker（供 startup/shutdown 集成）。"""
|
||
return self._broker
|
||
|
||
# ── 任务提交 ────────────────────────────────────────────────
|
||
|
||
async def submit_vectorize(
|
||
self,
|
||
params: VectorizeTaskParams,
|
||
dependencies: dict[str, object] | None = None,
|
||
) -> str:
|
||
"""提交向量化任务，返回 task_id。
|
||
|
||
Args:
|
||
params: 向量化任务参数（已通过 schema 校验）
|
||
dependencies: 运行时依赖（store/vector_store/embed_model），
|
||
降级模式下直接使用；broker 模式下由 worker 注入
|
||
|
||
Raises:
|
||
ConcurrencyLimitExceeded: per-user 并发上限超出
|
||
"""
|
||
await self._check_concurrency(params.user_id)
|
||
|
||
task_id = f"vec-{uuid.uuid4()}"
|
||
input_data = params.model_dump()
|
||
# 依赖对象不序列化 — 仅在降级模式内存中传递
|
||
deps = dependencies or {}
|
||
|
||
await _maybe_await(
|
||
self._store.create(
|
||
task_id=task_id,
|
||
agent_name="rag_vectorize",
|
||
input_data=input_data,
|
||
skill_name="vectorize",
|
||
)
|
||
)
|
||
# 记录 user_id 到 metadata 便于按用户查询
|
||
record = await _maybe_await(self._store.get(task_id))
|
||
if record is not None:
|
||
record.metadata["user_id"] = params.user_id
|
||
|
||
logger.info(
|
||
"Submitted vectorize task %s for document=%s user=%s",
|
||
task_id,
|
||
params.document_id,
|
||
params.user_id,
|
||
)
|
||
|
||
if self._broker is not None:
|
||
# broker 模式 — 异步派发
|
||
await self._dispatch_to_broker("vectorize", input_data, task_id)
|
||
else:
|
||
# 降级模式 — 同步执行（不阻塞调用方，使用 asyncio.create_task）
|
||
asyncio.create_task(self._run_vectorize_degraded(task_id, params, deps))
|
||
|
||
return task_id
|
||
|
||
async def submit_batch_index(
|
||
self,
|
||
params: BatchIndexTaskParams,
|
||
dependencies: dict[str, object] | None = None,
|
||
) -> str:
|
||
"""提交批量索引任务，返回 task_id。"""
|
||
await self._check_concurrency(params.user_id)
|
||
|
||
task_id = f"batch-{uuid.uuid4()}"
|
||
input_data = params.model_dump()
|
||
deps = dependencies or {}
|
||
|
||
await _maybe_await(
|
||
self._store.create(
|
||
task_id=task_id,
|
||
agent_name="rag_batch_index",
|
||
input_data=input_data,
|
||
skill_name="batch_index",
|
||
)
|
||
)
|
||
record = await _maybe_await(self._store.get(task_id))
|
||
if record is not None:
|
||
record.metadata["user_id"] = params.user_id
|
||
|
||
logger.info(
|
||
"Submitted batch index task %s for %d documents user=%s",
|
||
task_id,
|
||
len(params.document_ids),
|
||
params.user_id,
|
||
)
|
||
|
||
if self._broker is not None:
|
||
await self._dispatch_to_broker("batch_index", input_data, task_id)
|
||
else:
|
||
asyncio.create_task(self._run_batch_index_degraded(task_id, params, deps))
|
||
|
||
return task_id
|
||
|
||
# ── 任务查询 ────────────────────────────────────────────────
|
||
|
||
async def get_task_status(self, task_id: str) -> TaskRecord | None:
|
||
"""查询任务状态。"""
|
||
return await _maybe_await(self._store.get(task_id))
|
||
|
||
async def list_tasks(
|
||
self,
|
||
user_id: str | None = None,
|
||
status: TaskStatus | None = None,
|
||
limit: int = 100,
|
||
) -> list[TaskRecord]:
|
||
"""查询任务历史 — 可按 user_id 和 status 过滤。"""
|
||
tasks = await _maybe_await(self._store.list_tasks(status=status, limit=limit))
|
||
if user_id is not None:
|
||
tasks = [t for t in tasks if t.metadata.get("user_id") == user_id]
|
||
return tasks
|
||
|
||
async def cancel_task(self, task_id: str) -> bool:
|
||
"""取消任务 — 仅 PENDING/RUNNING 可取消。
|
||
|
||
Returns:
|
||
True 如果取消成功
|
||
Raises:
|
||
TaskNotFoundError: 任务不存在
|
||
"""
|
||
record = await _maybe_await(self._store.get(task_id))
|
||
if record is None:
|
||
raise TaskNotFoundError(f"Task {task_id} not found")
|
||
|
||
if record.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED):
|
||
return False
|
||
|
||
await _maybe_await(
|
||
self._store.update_status(
|
||
task_id,
|
||
TaskStatus.CANCELLED,
|
||
completed_at=datetime.now(timezone.utc),
|
||
)
|
||
)
|
||
# 释放并发计数
|
||
user_id = record.metadata.get("user_id")
|
||
if user_id and self._running_counts.get(user_id, 0) > 0:
|
||
self._running_counts[user_id] -= 1
|
||
return True
|
||
|
||
# ── 心跳更新 ────────────────────────────────────────────────
|
||
|
||
async def update_heartbeat(self, task_id: str) -> None:
|
||
"""更新 worker 心跳时间戳 — 由 worker 周期性调用。"""
|
||
record = await _maybe_await(self._store.get(task_id))
|
||
if record is None:
|
||
return
|
||
now = datetime.now(timezone.utc).isoformat()
|
||
record.metadata["worker_heartbeat_at"] = now
|
||
|
||
# ── 内部方法 ────────────────────────────────────────────────
|
||
|
||
async def _check_concurrency(self, user_id: str) -> None:
|
||
"""检查 per-user 并发上限。"""
|
||
async with self._lock:
|
||
current = self._running_counts.get(user_id, 0)
|
||
if current >= self._max_concurrent:
|
||
raise ConcurrencyLimitExceeded(
|
||
f"User {user_id} has {current} running tasks (max={self._max_concurrent})"
|
||
)
|
||
self._running_counts[user_id] = current + 1
|
||
|
||
async def _release_concurrency(self, user_id: str) -> None:
|
||
"""释放并发计数。"""
|
||
async with self._lock:
|
||
if self._running_counts.get(user_id, 0) > 0:
|
||
self._running_counts[user_id] -= 1
|
||
|
||
async def _dispatch_to_broker(self, task_type: str, params: dict, task_id: str) -> None:
|
||
"""派发任务到 TaskIQ broker。"""
|
||
# broker.kiq 返回 TaskiqTask — 实际执行由 worker 进程完成
|
||
try:
|
||
await self._broker.kiq(task_type=task_type, params=params, task_id=task_id)
|
||
except Exception as e:
|
||
logger.error("Failed to dispatch task %s to broker: %s", task_id, e)
|
||
raise
|
||
|
||
async def _run_vectorize_degraded(
|
||
self,
|
||
task_id: str,
|
||
params: VectorizeTaskParams,
|
||
deps: dict[str, object],
|
||
) -> None:
|
||
"""降级模式 — 同步执行向量化任务（在 asyncio 任务中）。"""
|
||
now = datetime.now(timezone.utc)
|
||
try:
|
||
await _maybe_await(
|
||
self._store.update_status(task_id, TaskStatus.RUNNING, started_at=now)
|
||
)
|
||
await self.update_heartbeat(task_id)
|
||
|
||
await run_vectorize_task(params, **deps)
|
||
|
||
await _maybe_await(
|
||
self._store.update_status(
|
||
task_id,
|
||
TaskStatus.COMPLETED,
|
||
completed_at=datetime.now(timezone.utc),
|
||
progress=1.0,
|
||
)
|
||
)
|
||
except Exception as e:
|
||
sanitized = sanitize_error_message(str(e))
|
||
await _maybe_await(
|
||
self._store.update_status(
|
||
task_id,
|
||
TaskStatus.FAILED,
|
||
error_message=sanitized,
|
||
completed_at=datetime.now(timezone.utc),
|
||
)
|
||
)
|
||
logger.error("Vectorize task %s failed: %s", task_id, sanitized)
|
||
finally:
|
||
await self._release_concurrency(params.user_id)
|
||
|
||
async def _run_batch_index_degraded(
|
||
self,
|
||
task_id: str,
|
||
params: BatchIndexTaskParams,
|
||
deps: dict[str, object],
|
||
) -> None:
|
||
"""降级模式 — 同步执行批量索引任务。"""
|
||
now = datetime.now(timezone.utc)
|
||
total = len(params.document_ids)
|
||
succeeded: list[str] = []
|
||
failed: list[str] = []
|
||
|
||
try:
|
||
await _maybe_await(
|
||
self._store.update_status(task_id, TaskStatus.RUNNING, started_at=now)
|
||
)
|
||
await self.update_heartbeat(task_id)
|
||
|
||
for i, doc_id in enumerate(params.document_ids):
|
||
await self.update_heartbeat(task_id)
|
||
try:
|
||
single_params = VectorizeTaskParams(
|
||
document_id=doc_id,
|
||
kb_id=params.kb_id,
|
||
file_path=params.file_paths.get(doc_id, ""),
|
||
file_type=params.file_types.get(doc_id, "txt"),
|
||
chunk_size=params.chunk_size,
|
||
chunk_overlap=params.chunk_overlap,
|
||
user_id=params.user_id,
|
||
)
|
||
await run_vectorize_task(single_params, **deps)
|
||
succeeded.append(doc_id)
|
||
except Exception as e:
|
||
logger.warning("Batch item %s failed: %s", doc_id, e)
|
||
failed.append(doc_id)
|
||
|
||
# 更新进度
|
||
progress = (i + 1) / total
|
||
await _maybe_await(
|
||
self._store.update_status(
|
||
task_id,
|
||
TaskStatus.RUNNING,
|
||
progress=progress,
|
||
progress_message=f"Processed {i + 1}/{total}",
|
||
)
|
||
)
|
||
|
||
final_status = TaskStatus.COMPLETED if not failed else TaskStatus.PARTIALLY_COMPLETED
|
||
await _maybe_await(
|
||
self._store.update_status(
|
||
task_id,
|
||
final_status,
|
||
completed_at=datetime.now(timezone.utc),
|
||
progress=1.0,
|
||
output_data={
|
||
"succeeded": succeeded,
|
||
"failed": failed,
|
||
"total": total,
|
||
},
|
||
)
|
||
)
|
||
except Exception as e:
|
||
sanitized = sanitize_error_message(str(e))
|
||
await _maybe_await(
|
||
self._store.update_status(
|
||
task_id,
|
||
TaskStatus.FAILED,
|
||
error_message=sanitized,
|
||
completed_at=datetime.now(timezone.utc),
|
||
)
|
||
)
|
||
logger.error("Batch index task %s failed: %s", task_id, sanitized)
|
||
finally:
|
||
await self._release_concurrency(params.user_id)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 任务执行函数 — 由 worker 调用
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def run_vectorize_task(
|
||
params: VectorizeTaskParams,
|
||
store: "KBStore",
|
||
vector_store: "PGVectorStore",
|
||
embed_model: "BaseEmbedding",
|
||
) -> None:
|
||
"""向量化任务 — 由 worker 执行。
|
||
|
||
封装 DocumentProcessor.process() — parse → segment → vectorize → index。
|
||
状态转换由 DocumentProcessor 内部管理（pending → parsing → ... → indexed | failed）。
|
||
|
||
Args:
|
||
params: 向量化任务参数
|
||
store: KBStore 实例（用于文档状态更新）
|
||
vector_store: LlamaIndex PGVectorStore 实例
|
||
embed_model: LlamaIndex embedding 模型
|
||
|
||
Raises:
|
||
Exception: 管道任一阶段失败
|
||
"""
|
||
processor = DocumentProcessor(
|
||
chunk_size=params.chunk_size,
|
||
chunk_overlap=params.chunk_overlap,
|
||
)
|
||
await processor.process(
|
||
params.file_path,
|
||
params.file_type,
|
||
params.kb_id,
|
||
params.document_id,
|
||
vector_store,
|
||
embed_model,
|
||
store,
|
||
)
|
||
|
||
|
||
async def run_batch_index_task(
|
||
params: BatchIndexTaskParams,
|
||
store: "KBStore",
|
||
vector_store: "PGVectorStore",
|
||
embed_model: "BaseEmbedding",
|
||
) -> dict[str, list[str]]:
|
||
"""批量索引任务 — 顺序处理多个文档。
|
||
|
||
Returns:
|
||
{"succeeded": [...], "failed": [...]}
|
||
"""
|
||
succeeded: list[str] = []
|
||
failed: list[str] = []
|
||
|
||
for doc_id in params.document_ids:
|
||
try:
|
||
single_params = VectorizeTaskParams(
|
||
document_id=doc_id,
|
||
kb_id=params.kb_id,
|
||
file_path=params.file_paths.get(doc_id, ""),
|
||
file_type=params.file_types.get(doc_id, "txt"),
|
||
chunk_size=params.chunk_size,
|
||
chunk_overlap=params.chunk_overlap,
|
||
user_id=params.user_id,
|
||
)
|
||
await run_vectorize_task(single_params, store, vector_store, embed_model)
|
||
succeeded.append(doc_id)
|
||
except Exception as e:
|
||
logger.warning("Batch item %s failed: %s", doc_id, e)
|
||
failed.append(doc_id)
|
||
|
||
return {"succeeded": succeeded, "failed": failed}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# TaskIQ broker 工厂（可选 — 仅 Redis 可用时使用）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def create_broker(redis_url: str) -> object:
|
||
"""创建 TaskIQ Redis broker — 独立 db=1 隔离。
|
||
|
||
Args:
|
||
redis_url: Redis 连接 URL（db=0 用于 bus/cache，此处强制改为 db=1）
|
||
|
||
Returns:
|
||
配置好的 TaskiqBroker 实例
|
||
|
||
Raises:
|
||
ImportError: taskiq 未安装
|
||
"""
|
||
from urllib.parse import urlparse, urlunparse
|
||
|
||
from taskiq_redis import ListQueueBroker, RedisAsyncResultBackend
|
||
|
||
from taskiq import TaskiqBroker
|
||
|
||
# 强制使用 db=1 — 与 bus/cache 的 db=0 隔离
|
||
parsed = urlparse(redis_url)
|
||
# 替换 path 为 /1
|
||
new_path = "/1"
|
||
isolated_url = urlunparse(parsed._replace(path=new_path))
|
||
|
||
broker = TaskiqBroker(ListQueueBroker(url=isolated_url))
|
||
broker.with_result_backend(
|
||
RedisAsyncResultBackend(redis_url=isolated_url, result_ex_time=DEFAULT_TASK_TTL_SECONDS)
|
||
)
|
||
|
||
logger.info("TaskIQ broker created with Redis db=1 isolation")
|
||
return broker
|
||
|
||
|
||
__all__ = [
|
||
"BatchIndexTaskParams",
|
||
"ConcurrencyLimitExceeded",
|
||
"DEFAULT_MAX_CONCURRENT_PER_USER",
|
||
"DEFAULT_WORKER_TTL_SECONDS",
|
||
"TaskManager",
|
||
"TaskNotFoundError",
|
||
"TaskStoreProtocol",
|
||
"VectorizeTaskParams",
|
||
"WorkerSweeper",
|
||
"create_broker",
|
||
"run_batch_index_task",
|
||
"run_vectorize_task",
|
||
"sanitize_error_message",
|
||
]
|