mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-07 12:00:44 +08:00
Align Go ingestion boundaries with Python (#16647)
Moves doc_id blob resolution into Parser, tightens chunker/tokenizer to Python output_format semantics, updates extractor list handling, and fixes real-template integration tests.
This commit is contained in:
Zhichang Yu
6
.gitignore
vendored
6
.gitignore
vendored
@@ -245,5 +245,11 @@ bin/*
|
||||
# Parser test fixtures and python tools
|
||||
internal/deepdoc/parser/pdf/testdata/
|
||||
internal/deepdoc/parser/pdf/tools-py/
|
||||
|
||||
# IDE tooling artifacts
|
||||
.codebuddy/
|
||||
|
||||
# Local build output
|
||||
build/
|
||||
internal/deepdoc/parser/docx/testdata/
|
||||
internal/deepdoc/parser/docx/tool/
|
||||
|
||||
191
AGENTS.md
191
AGENTS.md
@@ -1,109 +1,108 @@
|
||||
# RAGFlow Project Instructions for GitHub Copilot
|
||||
# RAGFlow Instructions
|
||||
|
||||
This file provides context, build instructions, and coding standards for the RAGFlow project.
|
||||
It is structured to follow GitHub Copilot's [customization guidelines](https://docs.github.com/en/copilot/concepts/prompting/response-customization).
|
||||
Use this file as the local operating guide for the current codebase. Prefer the code and the current CLAUDE.md over any older convention or remembered project shape.
|
||||
|
||||
## 1. Project Overview
|
||||
RAGFlow is an open-source RAG (Retrieval-Augmented Generation) engine based on deep document understanding. It is a full-stack application with a Python backend and a React/TypeScript frontend.
|
||||
## Core stance
|
||||
- Treat legacy code as liability, not as a compatibility target.
|
||||
- Prefer deletion over shims, deprecated branches, wrapper APIs, and dual-track migration notes.
|
||||
- If old and new implementations coexist, converge to one path unless an external contract forces compatibility.
|
||||
- Remove dead tests, commented-out code, stale docs, and "move later" notes instead of preserving them.
|
||||
- Reduce public surface area when a helper can be made private or internal.
|
||||
- Keep refactors centered on the owning abstraction, not on adjacent compatibility layers.
|
||||
|
||||
- **Backend**: Python 3.10+ (Flask/Quart)
|
||||
- **Frontend**: TypeScript, React, UmiJS
|
||||
- **Architecture**: Microservices based on Docker.
|
||||
- `api/`: Backend API server.
|
||||
- `rag/`: Core RAG logic (indexing, retrieval).
|
||||
- `deepdoc/`: Document parsing and OCR.
|
||||
- `web/`: Frontend application.
|
||||
## Current stack
|
||||
- Backend: Python 3.13+, Quart-based API server, Peewee ORM, async workers.
|
||||
- Frontend: React + TypeScript + Vite in `web/`.
|
||||
- Go: the repository also has a substantial Go module for servers, ingestion, parser/runtime, CLI, and supporting services.
|
||||
- Runtime services commonly include MySQL/PostgreSQL, Redis, MinIO, and Elasticsearch/Infinity/OpenSearch depending on configuration.
|
||||
|
||||
## 2. Directory Structure
|
||||
- `api/`: Backend API server (Flask/Quart).
|
||||
- `apps/`: API Blueprints (Knowledge Base, Chat, etc.).
|
||||
- `db/`: Database models and services.
|
||||
- `rag/`: Core RAG logic.
|
||||
- `llm/`: LLM, Embedding, and Rerank model abstractions.
|
||||
- `deepdoc/`: Document parsing and OCR modules.
|
||||
- `agent/`: Agentic reasoning components.
|
||||
- `web/`: Frontend application (React + UmiJS).
|
||||
- `docker/`: Docker deployment configurations.
|
||||
- `sdk/`: Python SDK.
|
||||
- `test/`: Backend tests.
|
||||
## Code layout to expect
|
||||
- `api/`: Python API server entrypoints, blueprints, services, and database code.
|
||||
- `rag/`: ingestion, retrieval, LLM integration, and graph RAG logic.
|
||||
- `deepdoc/`: parsing and OCR.
|
||||
- `agent/`: workflow canvas, components, tools, and templates.
|
||||
- `cmd/`: Go entrypoints. `ragflow_main` is the main server/admin/ingestor binary surface; `ragflow-cli` is the CLI entrypoint.
|
||||
- `internal/`: main Go application code. Important subtrees:
|
||||
- `internal/agent/`: Go agent runtime, canvas execution, components, tool bindings, workflow helpers.
|
||||
- `internal/cli/`: CLI parsing, HTTP transport, command execution, response formatting.
|
||||
- `internal/dao/`: Go data-access layer and persistence-facing helpers.
|
||||
- `internal/deepdoc/`: Go DeepDOC integrations, especially native-backed PDF/DOCX parsing.
|
||||
- `internal/engine/`: search/index backends such as Elasticsearch and Infinity.
|
||||
- `internal/entity/`: shared Go entities and model definitions.
|
||||
- `internal/handler/`: HTTP handlers and route-facing request logic.
|
||||
- `internal/ingestion/`: Go ingestion pipeline, canvas adapter, components, wiring, service orchestration.
|
||||
- `internal/ingestion/component/`: stage implementations such as file/parser/chunker/tokenizer/extractor.
|
||||
- `internal/ingestion/pipeline/`: DSL translation, canvas-driven execution, checkpoints, resume/run logic.
|
||||
- `internal/parser/`: parser and chunk libraries used by ingestion and other Go paths.
|
||||
- `internal/parser/parser/`: typed parse-result parsers for markdown/html/pdf/docx/xlsx/text and related families.
|
||||
- `internal/parser/chunk/`: chunk operator library and DSL/typed execution helpers.
|
||||
- `internal/service/`: higher-level business services used by handlers and server flows.
|
||||
- `internal/storage/`: storage backends and in-memory test doubles.
|
||||
- `internal/router/`: HTTP route registration.
|
||||
- `internal/server/`: server bootstrap/config wiring.
|
||||
- `internal/cpp/`: C++ sources used by native-backed Go features.
|
||||
- `web/`: frontend application.
|
||||
- `docker/`: local and production compose files.
|
||||
- `sdk/` and `test/`: SDK and automated tests.
|
||||
|
||||
## 3. Build Instructions
|
||||
## Go-specific rules
|
||||
- Treat `internal/ingestion`, `internal/parser`, and `internal/deepdoc` as actively refactored code. Prefer collapsing duplicate paths over preserving transitional wrappers.
|
||||
- Do not add or preserve deprecated Go APIs just to ease migration inside the repo.
|
||||
- Remove commented-out Go code instead of leaving recovery notes in place.
|
||||
- Keep package comments and doc comments aligned with the current runtime path, not with migration history.
|
||||
|
||||
### Backend (Python)
|
||||
The project uses **uv** for dependency management.
|
||||
## Working rules
|
||||
- Before editing, inspect the nearest code path that actually owns the behavior.
|
||||
- Keep changes small and local unless the task is explicitly a broader refactor.
|
||||
- Prefer one implementation path instead of preserving old and new versions side by side.
|
||||
- Preserve behavior with focused tests when the behavior is still valid; do not keep tests that protect obsolete behavior.
|
||||
- If a surface is only there for compatibility, remove it unless the user asks to keep it.
|
||||
- Do not add new compatibility wording in comments or docs.
|
||||
|
||||
1. **Setup Environment**:
|
||||
```bash
|
||||
uv sync --python 3.13 --all-extras
|
||||
uv run python3 ragflow_deps/download_deps.py
|
||||
```
|
||||
|
||||
2. **Run Server**:
|
||||
- **Pre-requisite**: Start dependent services (MySQL, ES/Infinity, Redis, MinIO).
|
||||
```bash
|
||||
docker compose -f docker/docker-compose-base.yml up -d
|
||||
```
|
||||
- **Launch**:
|
||||
```bash
|
||||
source .venv/bin/activate
|
||||
export PYTHONPATH=$(pwd)
|
||||
bash docker/launch_backend_service.sh
|
||||
```
|
||||
|
||||
### Frontend (TypeScript/React)
|
||||
Located in `web/`.
|
||||
|
||||
1. **Install Dependencies**:
|
||||
```bash
|
||||
cd web
|
||||
npm install
|
||||
```
|
||||
|
||||
2. **Run Dev Server**:
|
||||
```bash
|
||||
npm run dev
|
||||
```
|
||||
Runs on port 8000 by default.
|
||||
|
||||
### Docker Deployment
|
||||
To run the full stack using Docker:
|
||||
## Commands
|
||||
### Backend
|
||||
```bash
|
||||
cd docker
|
||||
docker compose -f docker-compose.yml up -d
|
||||
uv sync --python 3.13 --all-extras
|
||||
uv run python3 ragflow_deps/download_deps.py
|
||||
docker compose -f docker/docker-compose-base.yml up -d
|
||||
source .venv/bin/activate
|
||||
export PYTHONPATH=$(pwd)
|
||||
bash docker/launch_backend_service.sh
|
||||
uv run pytest
|
||||
ruff check
|
||||
ruff format
|
||||
```
|
||||
|
||||
## 4. Testing Instructions
|
||||
### Frontend
|
||||
```bash
|
||||
cd web
|
||||
npm install
|
||||
npm run dev
|
||||
npm run build
|
||||
npm run lint
|
||||
npm run test
|
||||
npm run type-check
|
||||
```
|
||||
|
||||
### Backend Tests
|
||||
- **Run All Tests**:
|
||||
```bash
|
||||
uv run pytest
|
||||
```
|
||||
- **Run Specific Test**:
|
||||
```bash
|
||||
uv run pytest test/test_api.py
|
||||
```
|
||||
### Go
|
||||
```bash
|
||||
uv run ragflow_deps/download_deps.py
|
||||
bash build.sh --test ./path/to/package/...
|
||||
bash build.sh --go
|
||||
# or build specific binaries:
|
||||
bash build.sh --all
|
||||
```
|
||||
|
||||
### Frontend Tests
|
||||
- **Run Tests**:
|
||||
```bash
|
||||
cd web
|
||||
npm run test
|
||||
```
|
||||
## Validation preference
|
||||
- Run the narrowest relevant test, lint, or build command after a change.
|
||||
- For backend changes, prefer targeted pytest or ruff checks over full-suite runs.
|
||||
- For frontend changes, prefer the touched-package lint, type-check, or test command.
|
||||
- For Go changes, prefer package-scoped `bash build.sh --test ...` first.
|
||||
- Do not default to raw `go test`, `go build`, or IDE Run/Debug for Go in this repo. They often miss the required CGO flags and native static libraries (`office_oxide`, `pdfium-static`, `pdf_oxide`) that `build.sh` wires correctly.
|
||||
- If Go native builds fail, inspect `build.sh` and `internal/development.md` before changing code. Common environment issues are missing downloaded native deps and missing `lld` on Linux.
|
||||
|
||||
## 5. Coding Standards & Guidelines
|
||||
- **Python Formatting**: Use `ruff` for linting and formatting.
|
||||
```bash
|
||||
ruff check
|
||||
ruff format
|
||||
```
|
||||
- **Frontend Linting**:
|
||||
```bash
|
||||
cd web
|
||||
npm run lint
|
||||
```
|
||||
- **Git Hooks**: Run this once after the first clone to enable local Git hooks.
|
||||
```bash
|
||||
lefthook install
|
||||
lefthook run pre-commit --all-files
|
||||
```
|
||||
## Default review checklist
|
||||
- Remove instead of retaining `deprecated`, `legacy`, or compatibility-only code.
|
||||
- Collapse duplicate implementations to one path.
|
||||
- Drop stale comments and documentation that describe a superseded design.
|
||||
- Keep exported APIs only when the current code actually needs them.
|
||||
|
||||
308
CLAUDE.md
308
CLAUDE.md
@@ -1,308 +0,0 @@
|
||||
# CLAUDE.md
|
||||
|
||||
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
|
||||
|
||||
## Project Overview
|
||||
|
||||
RAGFlow is an open-source RAG (Retrieval-Augmented Generation) engine based on deep document understanding. It's a full-stack application with:
|
||||
|
||||
- Python backend (Quart-based async API server — Quart is the async reimplementation of Flask)
|
||||
- React/TypeScript frontend (built with vitejs)
|
||||
- Background task executor workers (separate Python processes, Redis-queue-driven)
|
||||
- Peewee ORM for database models (not SQLAlchemy)
|
||||
- Multiple data stores (MySQL/PostgreSQL, Elasticsearch/Infinity/OpenSearch/OceanBase, Redis, MinIO)
|
||||
|
||||
## Architecture
|
||||
|
||||
### Runtime Architecture
|
||||
|
||||
RAGFlow runs as **two separate Python process types**, orchestrated by `docker/launch_backend_service.sh`:
|
||||
|
||||
- **API Server** (`api/ragflow_server.py`): Quart-based async HTTP server
|
||||
- **Task Executors** (`rag/svr/task_executor.py`): Background workers processing documents from Redis streams. Multiple instances run in parallel (controlled by `WS` env var). Each consumes from priority-ordered Redis streams (`te.1.common`, `te.0.common`), using consumer groups for load distribution.
|
||||
|
||||
Key consequence: task executors import a different code surface than the API server, so always check which process a module is meant for.
|
||||
|
||||
### Backend API (`/api/`)
|
||||
|
||||
- **App factory**: `api/apps/__init__.py` — creates the Quart app, configures auth (`login_required` decorator, JWT + API token + session fallback), and dynamically discovers/registers blueprints
|
||||
- **Two API coexisting patterns**:
|
||||
- **RESTful APIs** in `api/apps/restful_apis/` — newer pattern with Pydantic request validation, service layer in `api/apps/services/`, routes registered under `/api/v1`
|
||||
- **Legacy APIs** in `api/apps/*_app.py` — older pattern using `@validate_request()`, routes registered under `/v1/<page_name>`
|
||||
- **SDK APIs** in `api/apps/sdk/` — registered under `/v1/`
|
||||
- **Services**: `api/db/services/` — business logic wrapping Peewee model operations. `api/apps/services/` — service layer for the RESTful APIs
|
||||
- **Models**: `api/db/db_models.py` — Peewee ORM models with pooled MySQL/PostgreSQL connections, custom `JSONField`/`ListField` types, retry logic on connection loss
|
||||
|
||||
### Core Processing (`/rag/`)
|
||||
|
||||
- **Document ingestion pipeline**: `rag/flow/pipeline.py` — `Pipeline` (extends `agent.canvas.Graph`) orchestrates the ingestion DAG. Components: File (fetches binary from storage), Parser (dispatches to `deepdoc.parser` based on file type), TokenChunker/TitleChunker (splits into chunks), Tokenizer (computes full-text tokens + embedding vectors), Extractor (LLM-based extraction). Data flows via Pydantic `*FromUpstream` schemas.
|
||||
- **Document parsing**: `deepdoc/` — PDF parsing (vision-based OCR, layout analysis, table structure recognition) and format-specific parsers (DOCX, XLSX, PPT, Markdown, HTML, images). All parsers normalize to a common structure (list of bbox dicts for PDFs, `{text, doc_type_kwd}` for others).
|
||||
- **DeepDoc HTTP API service** (`deepdoc/server/`): OSS ONNX models (DLA, OCR, TSR) wrapped with LitServe as a standalone HTTP API on port 8124. The Go parser (`internal/parser/`) calls this service via `DeepDocClient`. Endpoints: `GET /health`, `GET /model`, `POST /predict/dla`, `POST /predict/tsr`, `POST /predict/ocr` (with `operator=det` or `operator=rec` form field). Docker image: `deepdoc_oss:latest`. See `deepdoc/server/README.md` for the full API reference.
|
||||
- **LLM Integration**: `rag/llm/` — factory pattern with runtime class discovery. `chat_model.py` (30+ providers via OpenAI SDK and LiteLLM wrappers), `embedding_model.py`, `rerank_model.py`, `cv_model.py` (image-to-text), `sequence2txt_model.py` (ASR), `tts_model.py`. Use `LLMBundle` (from `api.db.services.llm_service`) as the unified interface.
|
||||
- **Graph RAG**: `rag/graphrag/` — multi-phase pipeline: per-document subgraph extraction (LLM or spaCy NER), Leiden community detection, entity resolution, community summarization. Entities/relations/reports are indexed as chunks alongside regular text chunks, differentiated by `knowledge_graph_kwd`.
|
||||
- **Search**: `rag/nlp/search.py` — `Dealer` class combines vector similarity + BM25 + re-ranking. `KGSearch` extends it for graph-aware retrieval (entity resolution, n-hop enrichment).
|
||||
|
||||
### Agent System (`/agent/`)
|
||||
|
||||
- **Execution engine**: `agent/canvas.py` — `Canvas` (extends `Graph`) executes the DAG. Components are run in topological order via `_run_batch`, each receiving upstream outputs as kwargs. Control-flow components (`Categorize`, `Switch`, `Iteration`, `Loop`) dynamically modify the execution path.
|
||||
- **Component base**: `agent/component/base.py` — `ComponentBase` with `invoke(**kwargs)` / `invoke_async(**kwargs)` lifecycle. Variable references (`{component_id@output_var}` or `{sys.query}`) are resolved from the canvas graph at runtime.
|
||||
- **Components**: Modular workflow components in `agent/component/` — Begin, LLM, Agent (tool-calling LLM), Categorize, Switch, Iteration, Loop, Message, Invoke (HTTP), and data manipulation nodes. Auto-discovered by `__init__.py`.
|
||||
- **Templates**: Pre-built agent workflows as JSON DSL files in `agent/templates/`. Each contains a complete `components` DAG, `path`, and `globals`.
|
||||
- **Tools**: `agent/tools/` — Retrieval, web search (DuckDuckGo, Google, Tavily, SearXNG), academic search (ArXiv, PubMed, Google Scholar, Wikipedia), code execution, SQL execution, email, GitHub, finance data, translation, weather. Tools implement `ToolBase` (extends `ComponentBase`) and produce OpenAI-compatible function descriptors.
|
||||
- **Plugins**: `agent/plugin/` — plugin system using `pluginlib` for loading external LLM tool plugins from `embedded_plugins/`.
|
||||
|
||||
### Frontend (`/web/`)
|
||||
|
||||
- React/TypeScript with vitejs framework
|
||||
- shadcn/ui components (Radix UI primitives + Tailwind CSS)
|
||||
- `@tanstack/react-query` for server state (cache keys, mutations, invalidation)
|
||||
- Zustand for local state (primarily agent canvas graph store)
|
||||
- `react-router` v7 with lazy-loaded pages
|
||||
- `react-i18next` for i18n (17 languages)
|
||||
- Axios for HTTP with a layered pattern: endpoint definitions (`utils/api.ts`) → HTTP client (`utils/next-request.ts`) → service layer (`services/`) → query hooks (`hooks/use-*-request.ts`) → components
|
||||
- `@xyflow/react` for the agent workflow canvas
|
||||
- `react-hook-form` + `zod` for form validation
|
||||
- Two API proxy prefixes: `webAPI = '/v1'` (legacy) and `restAPIv1 = '/api/v1'` (RESTful)
|
||||
|
||||
## Common Development Commands
|
||||
|
||||
### Backend Development
|
||||
|
||||
```bash
|
||||
# Install Python dependencies
|
||||
uv sync --python 3.13 --all-extras
|
||||
uv run python3 ragflow_deps/download_deps.py
|
||||
|
||||
# Run once after the first clone to enable local Git hooks
|
||||
lefthook install
|
||||
|
||||
# Start dependent services
|
||||
docker compose -f docker/docker-compose-base.yml up -d
|
||||
|
||||
# Run backend (requires services to be running)
|
||||
source .venv/bin/activate
|
||||
export PYTHONPATH=$(pwd)
|
||||
bash docker/launch_backend_service.sh
|
||||
|
||||
# Run tests
|
||||
uv run pytest
|
||||
|
||||
# Linting
|
||||
ruff check
|
||||
ruff format
|
||||
```
|
||||
|
||||
### Frontend Development
|
||||
|
||||
```bash
|
||||
cd web
|
||||
npm install
|
||||
npm run dev # Development server
|
||||
npm run build # Production build
|
||||
npm run lint # ESLint
|
||||
npm run test # Jest tests
|
||||
```
|
||||
|
||||
### Docker Development
|
||||
|
||||
```bash
|
||||
# Full stack with Docker (includes deepdoc vision service)
|
||||
cd docker
|
||||
docker compose -f docker-compose.yml up -d
|
||||
|
||||
# Check server status
|
||||
docker logs -f ragflow-server
|
||||
|
||||
# Build the OSS deepdoc vision service standalone
|
||||
docker build -f docker/Dockerfile_deepdoc_oss -t deepdoc_oss:latest .
|
||||
docker run -p 8124:8124 deepdoc_oss:latest
|
||||
|
||||
# Rebuild images
|
||||
docker build --platform linux/amd64 -f Dockerfile -t infiniflow/ragflow:nightly .
|
||||
```
|
||||
|
||||
## Key Configuration Files
|
||||
|
||||
- `docker/.env` - Environment variables for Docker deployment
|
||||
- `docker/service_conf.yaml.template` - Backend service configuration
|
||||
- `pyproject.toml` - Python dependencies and project configuration
|
||||
- `web/package.json` - Frontend dependencies and scripts
|
||||
|
||||
## Testing
|
||||
|
||||
- **Python**: pytest with markers (p1/p2/p3 priority levels)
|
||||
- **Frontend**: Jest with React Testing Library
|
||||
- **API Tests**: HTTP API and SDK tests in `test/` and `sdk/python/test/`
|
||||
|
||||
## Database Engines
|
||||
|
||||
RAGFlow supports switching between Elasticsearch (default) and Infinity:
|
||||
|
||||
- Set `DOC_ENGINE=infinity` in `docker/.env` to use Infinity
|
||||
- Requires container restart: `docker compose down -v && docker compose up -d`
|
||||
|
||||
## Account Password Handling (Critical for Login Flow)
|
||||
|
||||
### Password Encryption Pipeline (Browser → Backend → DB Hash)
|
||||
|
||||
The login password verification chain is counterintuitive. Understanding this is essential when generating or verifying password hashes.
|
||||
|
||||
**Complete flow:**
|
||||
|
||||
```
|
||||
Browser input: "demo"
|
||||
→ Base64("demo") = "ZGVtbw=="
|
||||
→ RSA encrypt with conf/public.pem
|
||||
→ POST to /api/v1/auth/login
|
||||
|
||||
Backend DecryptPassword():
|
||||
→ RSA decrypt with conf/private.pem (passphrase: "Welcome")
|
||||
→ Returns "ZGVtbw==" (NOT "demo"!)
|
||||
|
||||
VerifyPassword("ZGVtbw==", storedHash) ← hash is of Base64(password), not raw password
|
||||
```
|
||||
|
||||
**Consequences:**
|
||||
- The string verified against the hash is **Base64(original password)**, never the raw password
|
||||
- `DecryptPassword()` handles both RSA-encrypted (browser) and plaintext (curl/API key) inputs: if base64 decode fails, the input is returned as-is for backward compatibility
|
||||
- Python backend has the same design: `api/utils/crypt.py:decrypt()` RSA-decrypts and returns the Base64-encoded string directly, no further decode
|
||||
|
||||
### How to Generate a Valid Password Hash
|
||||
|
||||
```bash
|
||||
# For password "demo" (user input in browser):
|
||||
# The actual verified string = Base64("demo") = "ZGVtbw=="
|
||||
# Generate hash with: common.GenerateWerkzeugPasswordHash("ZGVtbw==")
|
||||
# or use the scrypt template:
|
||||
# scrypt:32768:8:1$<random-b64-salt>$<hex-hash-of-ZGVtbw==>
|
||||
```
|
||||
|
||||
**To update a user's password in the running database:**
|
||||
```bash
|
||||
docker exec docker-mysql-1 mysql -u root -pinfini_rag_flow rag_flow \
|
||||
-e "UPDATE user SET password='<hash>' WHERE email='<email>';"
|
||||
```
|
||||
|
||||
### RSA Keys
|
||||
- `conf/public.pem` — frontend uses this to encrypt Base64(password) before sending
|
||||
- `conf/private.pem` — backend uses this to decrypt, passphrase `"Welcome"`
|
||||
- Both referenced in `internal/common/password.go:DecryptPassword()`
|
||||
|
||||
### Obtaining an API Token for a Tenant
|
||||
|
||||
When testing APIs manually (curl, Go scripts, etc.), you need a valid auth token. The login endpoint returns **two different tokens**:
|
||||
|
||||
| Field | Format | Purpose |
|
||||
|-------|--------|---------|
|
||||
| `response.body.data.access_token` | Raw UUID | Stored in DB, NOT used for API auth |
|
||||
| `response.Header["Authorization"]` | itsdangerous-signed token | Used as `Bearer <token>` for all subsequent API requests |
|
||||
|
||||
**How to obtain the correct token:**
|
||||
|
||||
```bash
|
||||
# Step 1: Construct the encrypted password
|
||||
# Raw password → Base64 → RSA encrypt with conf/public.pem
|
||||
PASSWORD="demo"
|
||||
PASSWORD_B64=$(echo -n "$PASSWORD" | base64)
|
||||
|
||||
# Step 2: POST to login (use RSA encryption — easiest via a Go/Python script)
|
||||
# Response header contains: Authorization: <itsdangerous-signed-token>
|
||||
|
||||
# Step 3: Use the Authorization header value for all API requests
|
||||
curl -H "Authorization: <itsdangerous-signed-token>" \
|
||||
http://127.0.0.1:9222/api/v1/agents
|
||||
```
|
||||
|
||||
**Go snippet (complete login + token extraction):**
|
||||
|
||||
```go
|
||||
// Login
|
||||
passwordB64 := base64.StdEncoding.EncodeToString([]byte(password))
|
||||
pubData, _ := os.ReadFile("conf/public.pem")
|
||||
block, _ := pem.Decode(pubData)
|
||||
pubKey, _ := x509.ParsePKIXPublicKey(block.Bytes)
|
||||
ciphertext, _ := rsa.EncryptPKCS1v15(rand.Reader, pubKey.(*rsa.PublicKey), []byte(passwordB64))
|
||||
encryptedB64 := base64.StdEncoding.EncodeToString(ciphertext)
|
||||
|
||||
body, _ := json.Marshal(map[string]string{"email": email, "password": encryptedB64})
|
||||
resp, _ := http.Post(baseURL+"/api/v1/auth/login", "application/json", bytes.NewReader(body))
|
||||
|
||||
// KEY: use the Authorization header, NOT body.access_token
|
||||
authToken := resp.Header.Get("Authorization")
|
||||
|
||||
// Use for API calls
|
||||
req, _ := http.NewRequest("GET", baseURL+"/api/v1/agents", nil)
|
||||
req.Header.Set("Authorization", authToken)
|
||||
```
|
||||
|
||||
**The raw `access_token` (UUID) in the response body** is the internal DB token used only by the `itsdangerous` middleware to verify the signed token — it is never passed directly in API Authorization headers.
|
||||
|
||||
---
|
||||
|
||||
## Agent Run E2E Tests
|
||||
|
||||
### Running the Tests
|
||||
|
||||
```bash
|
||||
# Run all agent run e2e tests (in-memory SQLite + miniredis, no Docker needed)
|
||||
cd /home/zhichyu/github.com/infiniflow/ragflow
|
||||
go test -count=1 -v -run 'TestRunAgent_RealCanvas|TestRunAgent_RunTracker' ./internal/service/
|
||||
```
|
||||
|
||||
### Test Architecture
|
||||
|
||||
All e2e tests live in `internal/service/agent_run_e2e_test.go`. They exercise the full production chain:
|
||||
|
||||
```
|
||||
loadCanvasForUser → versionDAO.GetLatest → decodeCanvasFromDSL →
|
||||
canvas.Compile → cc.Workflow.Invoke → answer extraction
|
||||
```
|
||||
|
||||
**Test isolation**: Each test stands up its own in-memory SQLite DB (pushed as `dao.DB`), seeds User/Tenant/UserCanvas/UserCanvasVersion rows, and tears down in `t.Cleanup`. Tests use **miniredis** for Redis-backed CheckPointStore + RunTracker — no external services needed.
|
||||
|
||||
**Key test helpers:**
|
||||
- `makeCanvasWithDSL(t, canvasID, userID, tenantID, versionID, dsl)` — seeds all required DB rows
|
||||
- `drainAgentEvents(t, events)` — drains the `<-chan canvas.RunEvent` channel, buckets results into `messages`, `waiting`, `errors_`, `done`
|
||||
- `newRunTrackerForTest(t, ttl)` — wires a `canvas.RunTracker` against in-memory miniredis
|
||||
|
||||
**Existing e2e tests:**
|
||||
|
||||
| Test | What it covers |
|
||||
|------|---------------|
|
||||
| `TestRunAgent_RealCanvas_BeginMessage` | Happy path: Begin→Message, verifies `"{{sys.query}}"` resolution |
|
||||
| `TestRunAgent_RealCanvas_WaitForUserResume` | Resume path: Begin→Message→UserFillUp, two-run cycle |
|
||||
| `TestRunAgent_RealCanvas_CompileFails` | Error path: unknown component name → sanitized error |
|
||||
| `TestRunAgent_RealCanvas_InvokeFails` | Error path: unresolvable template ref |
|
||||
| `TestRunAgent_RunTracker_AttachCheckpoint_CallSequence` | Production boot: Start→AttachCheckpoint→MarkSucceeded with Redis/miniredis |
|
||||
|
||||
**Test DSL data files** are in `internal/agent/dsl/testdata/`:
|
||||
- `agent_msg.json` — Agent+Message with Begin, LLM-powered agent component
|
||||
- `all.json` — Complex: Begin→UserFillUp→Switch→Loop→Message
|
||||
- `switch.json`, `resume.json`, `browser.json`, `subagent.json`, etc.
|
||||
|
||||
**Handler-level SSE streaming tests** in `internal/handler/agent_test.go` use a `stubChatRunner` that emits pre-configured `canvas.RunEvent` values without a real DB or eino runner, verifying:
|
||||
- SSE `Content-Type: text/event-stream`
|
||||
- `data: {...}\n\n` framing
|
||||
- Trailing `data: [DONE]\n\n` terminator
|
||||
- OpenAI-compatible non-stream `choices` response shape
|
||||
|
||||
**Important**: `_ "ragflow/internal/agent/component"` (blank import in test) is required — it triggers `init()` to register all component factories. Without it, `canvas.Compile` fails to resolve any component type.
|
||||
|
||||
---
|
||||
|
||||
## Development Environment Requirements
|
||||
|
||||
- Python 3.10-3.13
|
||||
- Node.js >=18.20.4
|
||||
- Docker & Docker Compose
|
||||
- uv package manager
|
||||
- 16GB+ RAM, 50GB+ disk space
|
||||
|
||||
1. Think before acting. Read existing files before writing code.
|
||||
2. Be concise in output but thorough in reasoning.
|
||||
3. Prefer editing over rewriting whole files.
|
||||
4. Do not re-read files you have already read.
|
||||
5. Test your code before declaring done.
|
||||
6. No sycophantic openers or closing fluff.
|
||||
7. Keep solutions simple and direct.
|
||||
8. User instructions always override this file.
|
||||
@@ -1,191 +0,0 @@
|
||||
#
|
||||
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
"""
|
||||
PipelineChunker Component
|
||||
|
||||
Run RAGFlow Pipeline-style chunkers (rag.app.*) against uploaded files inside an
|
||||
Agent workflow. Emits plain text chunks for downstream Agent nodes — no
|
||||
embedding, no persistence. Wraps existing chunker functions; does not
|
||||
re-implement chunking logic.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
import os
|
||||
from abc import ABC
|
||||
|
||||
from agent.component.base import ComponentBase, ComponentParamBase
|
||||
from api.db.services.file_service import FileService
|
||||
from common.connection_utils import timeout
|
||||
|
||||
|
||||
# Parser id -> dotted module path under rag.app. Imported lazily so we don't
|
||||
# pull deepdoc/OCR/VLM machinery at component-discovery time.
|
||||
_PARSER_MODULES: dict[str, str] = {
|
||||
"general": "rag.app.naive",
|
||||
"naive": "rag.app.naive",
|
||||
"paper": "rag.app.paper",
|
||||
"book": "rag.app.book",
|
||||
"presentation": "rag.app.presentation",
|
||||
"manual": "rag.app.manual",
|
||||
"laws": "rag.app.laws",
|
||||
"qa": "rag.app.qa",
|
||||
"table": "rag.app.table",
|
||||
"resume": "rag.app.resume",
|
||||
"picture": "rag.app.picture",
|
||||
"one": "rag.app.one",
|
||||
"audio": "rag.app.audio",
|
||||
"email": "rag.app.email",
|
||||
"tag": "rag.app.tag",
|
||||
}
|
||||
|
||||
|
||||
def _load_chunker(parser_id: str):
|
||||
"""Resolve a parser id to the underlying ``rag.app.<module>.chunk`` callable."""
|
||||
module_path = _PARSER_MODULES[parser_id.lower()]
|
||||
return importlib.import_module(module_path).chunk
|
||||
|
||||
|
||||
class PipelineChunkerParam(ComponentParamBase):
|
||||
"""
|
||||
Define the PipelineChunker component parameters.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialise PipelineChunker defaults and declare component outputs."""
|
||||
super().__init__()
|
||||
self.inputs = [] # variable references to uploaded files
|
||||
self.parser_id = "naive"
|
||||
self.lang = "English"
|
||||
self.from_page = 0
|
||||
self.to_page = 100000000
|
||||
self.parser_config = {}
|
||||
|
||||
self.outputs = {
|
||||
"chunks": {"type": "list", "value": []},
|
||||
"chunks_full": {"type": "list", "value": []},
|
||||
"summary": {"type": "str", "value": ""},
|
||||
}
|
||||
|
||||
def check(self):
|
||||
"""Validate parser id, page range, and parser_config shape."""
|
||||
self.check_valid_value(
|
||||
self.parser_id.lower(),
|
||||
"[PipelineChunker] parser_id",
|
||||
list(_PARSER_MODULES.keys()),
|
||||
)
|
||||
self.check_nonnegative_number(self.from_page, "[PipelineChunker] from_page")
|
||||
self.check_nonnegative_number(self.to_page, "[PipelineChunker] to_page")
|
||||
if isinstance(self.from_page, (int, float)) and isinstance(self.to_page, (int, float)) and self.from_page > self.to_page:
|
||||
raise ValueError("[PipelineChunker] from_page must be <= to_page")
|
||||
if not isinstance(self.parser_config, dict):
|
||||
raise ValueError("[PipelineChunker] parser_config must be a dict.")
|
||||
return True
|
||||
|
||||
|
||||
class PipelineChunker(ComponentBase, ABC):
|
||||
"""
|
||||
Run a Pipeline-style chunker (naive, paper, qa, manual, book, ...) against
|
||||
one or more uploaded files and surface the resulting chunks to downstream
|
||||
Agent nodes.
|
||||
"""
|
||||
|
||||
component_name = "PipelineChunker"
|
||||
|
||||
def get_input_form(self) -> dict[str, dict]:
|
||||
"""Expose each referenced file input as a file-typed form element."""
|
||||
res = {}
|
||||
for ref in self._param.inputs or []:
|
||||
for k, o in self.get_input_elements_from_text(ref).items():
|
||||
res[k] = {"name": o.get("name", ""), "type": "file"}
|
||||
return res
|
||||
|
||||
def _get_file_content(self, file_ref: str) -> tuple[bytes | None, str | None]:
|
||||
"""Resolve a canvas variable reference to ``(content_bytes, filename)``."""
|
||||
value = self._canvas.get_variable_value(file_ref)
|
||||
if value is None:
|
||||
return None, None
|
||||
|
||||
if isinstance(value, list) and value:
|
||||
value = value[0]
|
||||
|
||||
if isinstance(value, dict):
|
||||
file_id = value.get("id") or value.get("file_id")
|
||||
created_by = value.get("created_by") or self._canvas.get_tenant_id()
|
||||
filename = value.get("name") or value.get("filename") or "uploaded"
|
||||
if file_id:
|
||||
try:
|
||||
return FileService.get_blob(created_by, file_id), filename
|
||||
except Exception as e:
|
||||
logging.exception(f"[PipelineChunker] FileService.get_blob failed for file_id={file_id} created_by={created_by} filename={filename}: {e}")
|
||||
return None, None
|
||||
return None, None
|
||||
|
||||
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60)))
|
||||
def _invoke(self, **kwargs):
|
||||
"""Run the configured chunker over every referenced file and publish outputs."""
|
||||
if self.check_if_canceled("PipelineChunker processing"):
|
||||
return
|
||||
|
||||
chunker = _load_chunker(self._param.parser_id)
|
||||
tenant_id = self._canvas.get_tenant_id()
|
||||
chunk_kwargs = dict(
|
||||
lang=self._param.lang,
|
||||
tenant_id=tenant_id,
|
||||
from_page=self._param.from_page,
|
||||
to_page=self._param.to_page,
|
||||
parser_config=self._param.parser_config or {},
|
||||
callback=lambda prog=0, msg="": logging.info(f"[PipelineChunker] {prog}: {msg}"),
|
||||
)
|
||||
|
||||
all_chunks: list[dict] = []
|
||||
per_file_counts: list[str] = []
|
||||
|
||||
for file_ref in self._param.inputs or []:
|
||||
if self.check_if_canceled("PipelineChunker processing"):
|
||||
return
|
||||
|
||||
content, filename = self._get_file_content(file_ref)
|
||||
self.set_input_value(file_ref, filename or "")
|
||||
if content is None:
|
||||
logging.warning(f"[PipelineChunker] could not resolve file ref: {file_ref}")
|
||||
per_file_counts.append(f"{filename or file_ref}: error (could not resolve file)")
|
||||
continue
|
||||
|
||||
try:
|
||||
file_chunks = chunker(filename, binary=content, **chunk_kwargs) or []
|
||||
except Exception as e:
|
||||
logging.exception(e)
|
||||
per_file_counts.append(f"{filename}: error (chunking failed)")
|
||||
continue
|
||||
|
||||
all_chunks.extend(file_chunks)
|
||||
per_file_counts.append(f"{filename}: {len(file_chunks)} chunks")
|
||||
|
||||
text_only = [(c.get("content_with_weight") or c.get("text") or "") for c in all_chunks if isinstance(c, dict)]
|
||||
text_only = [t for t in text_only if t]
|
||||
|
||||
self.set_output("chunks", text_only)
|
||||
self.set_output("chunks_full", all_chunks)
|
||||
self.set_output(
|
||||
"summary",
|
||||
f"Parser: {self._param.parser_id} | Files: {len(self._param.inputs or [])} | Chunks: {len(text_only)}" + (" | " + "; ".join(per_file_counts) if per_file_counts else ""),
|
||||
)
|
||||
|
||||
def thoughts(self) -> str:
|
||||
"""Return a short status line for UI display."""
|
||||
return f"Chunking with `{self._param.parser_id}` strategy..."
|
||||
18
build.sh
18
build.sh
@@ -384,8 +384,8 @@ setup_cgo_env() {
|
||||
echo "CGO_LDFLAGS: $CGO_LDFLAGS"
|
||||
}
|
||||
|
||||
# Run Go unit tests with the same CGO env as `build_go`. Pass any extra args
|
||||
# to `go test`, e.g. `./build.sh --test -run TestFoo ./internal/admin/...`.
|
||||
# Run Go unit tests with the same CGO env as `build_go`. Any extra args are
|
||||
# forwarded to `go test`, e.g. `./build.sh --test -run TestFoo ./internal/admin/...`.
|
||||
run_go_tests() {
|
||||
print_section "Running Go tests"
|
||||
|
||||
@@ -456,8 +456,8 @@ OPTIONS:
|
||||
--cpp-test Build C++ test executable (requires --cpp first)
|
||||
--go, -g Build only Go server (requires C++ library to be built)
|
||||
--test, -t Run Go unit tests (sets up CGO env for office_oxide).
|
||||
Pass extra args after `--` to forward to `go test`, e.g.
|
||||
`$0 --test -- -run TestFoo ./internal/admin/...`
|
||||
Any extra args are forwarded to `go test`, e.g.
|
||||
`$0 --test -run TestFoo ./internal/admin/...`
|
||||
--clean, -C Clean all build artifacts
|
||||
--run, -r Build and run the server
|
||||
--strip, -s Strip debug symbols from Go binaries (-ldflags="-s -w")
|
||||
@@ -470,7 +470,7 @@ EXAMPLES:
|
||||
$0 --go # Build only Go server
|
||||
$0 --cpp-test # Build C++ test executable
|
||||
$0 --test # Run all Go tests
|
||||
$0 --test -- -run TestFoo ./internal/admin/... # Targeted Go tests
|
||||
$0 --test -run TestFoo ./internal/admin/... # Targeted Go tests
|
||||
$0 --run # Build and run
|
||||
$0 --clean # Clean build artifacts
|
||||
|
||||
@@ -513,12 +513,10 @@ main() {
|
||||
;;
|
||||
--test|-t)
|
||||
check_go_deps
|
||||
# Forward any args after `--` to `go test`.
|
||||
if [ "${2:-}" = "--" ]; then
|
||||
shift 2
|
||||
run_go_tests "$@"
|
||||
if [ "${args[1]:-}" = "--" ]; then
|
||||
run_go_tests "${args[@]:2}"
|
||||
else
|
||||
run_go_tests
|
||||
run_go_tests "${args[@]:1}"
|
||||
fi
|
||||
;;
|
||||
--clean|-C)
|
||||
|
||||
@@ -52,6 +52,7 @@ import (
|
||||
"ragflow/internal/dao"
|
||||
"ragflow/internal/engine"
|
||||
"ragflow/internal/engine/redis"
|
||||
_ "ragflow/internal/ingestion/wire" // single owner for ingestion-component registration (File / Parser / Tokenizer / Extractor + 4 Chunker variants)
|
||||
"ragflow/internal/server"
|
||||
"ragflow/internal/utility"
|
||||
)
|
||||
@@ -797,6 +798,7 @@ func startServer(config *server.Config) {
|
||||
}
|
||||
}
|
||||
adminRuntimeHandler := handler.NewAdminRuntimeHandler(adminRuntimeSelector)
|
||||
componentsHandler := handler.NewComponentsHandler(service.NewComponentsService())
|
||||
|
||||
// Initialize router
|
||||
r := router.NewRouter(authHandler,
|
||||
@@ -827,7 +829,8 @@ func startServer(config *server.Config) {
|
||||
fileCommitHandler,
|
||||
adminRuntimeHandler,
|
||||
openaiChatHandler,
|
||||
botHandler)
|
||||
botHandler,
|
||||
componentsHandler)
|
||||
|
||||
// Create Gin enginegit diff
|
||||
|
||||
|
||||
@@ -130,14 +130,19 @@ func (s *Service) ListIngestionTasks() ([]map[string]interface{}, error) {
|
||||
"document_id": task.DocumentID,
|
||||
"status": task.Status,
|
||||
}
|
||||
if err == nil {
|
||||
if err == nil && latestLog != nil && latestLog.Checkpoint != nil {
|
||||
step, ok := latestLog.Checkpoint["current_step"].(float64)
|
||||
if !ok {
|
||||
showTasks = append(showTasks, showTask)
|
||||
continue
|
||||
}
|
||||
showTask = map[string]interface{}{
|
||||
"id": task.ID,
|
||||
"user_id": task.UserID,
|
||||
"user": user.Email,
|
||||
"document_id": task.DocumentID,
|
||||
"status": task.Status,
|
||||
"step": int(latestLog.Checkpoint["current_step"].(float64)),
|
||||
"step": int(step),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
132
internal/agent/canvas/decode.go
Normal file
132
internal/agent/canvas/decode.go
Normal file
@@ -0,0 +1,132 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package canvas
|
||||
|
||||
import "fmt"
|
||||
|
||||
// DecodeFromDSL converts a canonical canvas DSL map into a Canvas.
|
||||
// It accepts both canonical IMPORT shape (`obj.component_name`) and the
|
||||
// normalized flat shape (`name`/`params`) that NormalizeForCanvas emits.
|
||||
func DecodeFromDSL(dsl map[string]any) (*Canvas, error) {
|
||||
if len(dsl) == 0 {
|
||||
return nil, fmt.Errorf("canvas: empty DSL")
|
||||
}
|
||||
rawComps, ok := dsl["components"].(map[string]any)
|
||||
if !ok || len(rawComps) == 0 {
|
||||
return nil, fmt.Errorf("canvas: no components")
|
||||
}
|
||||
c := &Canvas{
|
||||
Components: make(map[string]CanvasComponent, len(rawComps)),
|
||||
NodeParents: make(map[string]string),
|
||||
}
|
||||
if p, ok := dsl["path"].([]any); ok {
|
||||
c.Path = make([]string, 0, len(p))
|
||||
for _, v := range p {
|
||||
if s, ok := v.(string); ok {
|
||||
c.Path = append(c.Path, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
if p, ok := dsl["globals"].(map[string]any); ok {
|
||||
c.Globals = p
|
||||
}
|
||||
if graph, ok := dsl["graph"].(map[string]any); ok {
|
||||
if nodes, ok := graph["nodes"].([]any); ok {
|
||||
for _, raw := range nodes {
|
||||
node, ok := raw.(map[string]any)
|
||||
if !ok || node == nil {
|
||||
continue
|
||||
}
|
||||
id, _ := node["id"].(string)
|
||||
parentID, _ := node["parentId"].(string)
|
||||
if id != "" && parentID != "" {
|
||||
c.NodeParents[id] = parentID
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
for cpnID, raw := range rawComps {
|
||||
comp, ok := raw.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
name, params, downstream, upstream := decodeComponentFields(comp)
|
||||
if name == "" {
|
||||
return nil, fmt.Errorf("canvas: component %q has empty component_name", cpnID)
|
||||
}
|
||||
c.Components[cpnID] = CanvasComponent{
|
||||
Obj: CanvasComponentObj{
|
||||
ComponentName: name,
|
||||
Params: params,
|
||||
},
|
||||
Downstream: downstream,
|
||||
Upstream: upstream,
|
||||
}
|
||||
}
|
||||
if len(c.Components) == 0 {
|
||||
return nil, fmt.Errorf("canvas: no components")
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func decodeComponentFields(comp map[string]any) (name string, params map[string]any, downstream []string, upstream []string) {
|
||||
if obj, ok := comp["obj"].(map[string]any); ok {
|
||||
name, _ = obj["component_name"].(string)
|
||||
if p, ok := obj["params"].(map[string]any); ok {
|
||||
params = p
|
||||
}
|
||||
if ds, ok := obj["downstream"].([]any); ok {
|
||||
downstream = decodeStringSlice(ds)
|
||||
} else if ds, ok := obj["downstream"].([]string); ok {
|
||||
downstream = ds
|
||||
}
|
||||
}
|
||||
if name == "" {
|
||||
name, _ = comp["name"].(string)
|
||||
}
|
||||
if params == nil {
|
||||
if p, ok := comp["params"].(map[string]any); ok {
|
||||
params = p
|
||||
}
|
||||
}
|
||||
if len(downstream) == 0 {
|
||||
if ds, ok := comp["downstream"].([]any); ok {
|
||||
downstream = decodeStringSlice(ds)
|
||||
} else if ds, ok := comp["downstream"].([]string); ok {
|
||||
downstream = ds
|
||||
}
|
||||
}
|
||||
if us, ok := comp["upstream"].([]any); ok {
|
||||
upstream = decodeStringSlice(us)
|
||||
} else if us, ok := comp["upstream"].([]string); ok {
|
||||
upstream = us
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func decodeStringSlice(in []any) []string {
|
||||
if len(in) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]string, 0, len(in))
|
||||
for _, v := range in {
|
||||
if s, ok := v.(string); ok {
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -167,7 +167,7 @@ func componentTimeout() time.Duration {
|
||||
// attribution authoritative.
|
||||
func realComponentBody(cpnID, componentClass string, comp runtime.Component) nodeBodyFn {
|
||||
return func(ctx context.Context, in map[string]any) (map[string]any, error) {
|
||||
timeout := resolveTimeout(componentClass)
|
||||
timeout := resolveTimeoutFromContext(ctx, componentClass)
|
||||
cctx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
out, err := comp.Invoke(cctx, in)
|
||||
|
||||
@@ -39,12 +39,25 @@
|
||||
package canvas
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type timeoutContextKey struct{}
|
||||
|
||||
// WithComponentTimeoutOverride forces all component invokes under ctx to use
|
||||
// the provided timeout. This lets callers with stricter execution contracts
|
||||
// reuse canvas execution without mutating the global timeout policy.
|
||||
func WithComponentTimeoutOverride(ctx context.Context, d time.Duration) context.Context {
|
||||
if d <= 0 {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, timeoutContextKey{}, d)
|
||||
}
|
||||
|
||||
// componentDefaults lists the per-class timeout (in seconds) used when
|
||||
// no per-class env override is set. The class keys are the PascalCase
|
||||
// names returned by Component.Name() (e.g. "LLM", "Message",
|
||||
@@ -118,6 +131,15 @@ func resolveTimeout(componentClass string) time.Duration {
|
||||
return defaultComponentTimeout
|
||||
}
|
||||
|
||||
func resolveTimeoutFromContext(ctx context.Context, componentClass string) time.Duration {
|
||||
if ctx != nil {
|
||||
if d, ok := ctx.Value(timeoutContextKey{}).(time.Duration); ok && d > 0 {
|
||||
return d
|
||||
}
|
||||
}
|
||||
return resolveTimeout(componentClass)
|
||||
}
|
||||
|
||||
// parseSecondsEnv reads an env var and parses it as seconds ("42" →
|
||||
// 42s). Returns (d, true) on success, (0, false) if the env var is
|
||||
// unset / empty / non-numeric / non-positive. Invalid input must
|
||||
|
||||
@@ -1,514 +0,0 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// PipelineChunker component (T3) — partial port of python
|
||||
// `agent/component/pipeline_chunker.py` (PR #15068).
|
||||
//
|
||||
// SCOPE (honest):
|
||||
//
|
||||
// - WHITELIST: every parser_id in the python
|
||||
// `_PARSER_MODULES` table is accepted (general/naive/paper/
|
||||
// book/presentation/manual/laws/qa/table/resume/picture/one/
|
||||
// audio/email/tag). Check() rejects anything else.
|
||||
//
|
||||
// - DISPATCH: parser_id drives the chunk-engine split
|
||||
// strategy. The Go chunk engine exposes sentence / paragraph /
|
||||
// char / paragraph splits; we map the python parser_ids to
|
||||
// those strategies (see parserToSplitStrategy). Per-parser
|
||||
// knobs (paper's double-column handling, table's HTML
|
||||
// reconstruction, laws' article-aware boundaries, …) are
|
||||
// not ported yet — a canvas author who picks `paper` gets a
|
||||
// paragraph split, not the python paper parser's column-aware
|
||||
// behaviour.
|
||||
//
|
||||
// - TEXT INPUT: "text" / "content" / "file_bytes" (interpreted
|
||||
// as UTF-8 text) flow through the chunk engine. This matches
|
||||
// the python "text file" path.
|
||||
//
|
||||
// "file_ref" is the bytes-container form used by the other Go
|
||||
// agent components (ExcelProcessor, etc.). It accepts
|
||||
// []byte OR a base64-encoded string. Raw text MUST go in
|
||||
// "text" / "content" — a string file_ref that is not valid
|
||||
// base64 is rejected with a clear error so a caller that
|
||||
// mistakenly hands plain text doesn't see it silently
|
||||
// reinterpreted (a "try base64, fall back" policy would
|
||||
// rewrite any plain text that happens to satisfy the
|
||||
// base64 alphabet, e.g. "Zm9v" → "foo"). The contract is
|
||||
// therefore: file_ref is always raw bytes (or base64 of
|
||||
// them); text/content is always UTF-8 text.
|
||||
//
|
||||
// - FILE-FORMAT EXTRACTION (PDF, DOCX, XLSX, PPT, images, …):
|
||||
// NOT PORTED. The python side uses `rag.app.<parser>.chunk
|
||||
// (filename)` which extracts text from the file format
|
||||
// AND chunks in one step. The Go side does not yet have a
|
||||
// unified `rag.app.*.chunk` dispatch — `deepdoc/parser` is
|
||||
// still Python-only. A caller that supplies binary
|
||||
// (non-UTF-8) bytes is rejected with an explicit error so
|
||||
// the canvas author knows to add text extraction upstream.
|
||||
// This is a follow-up that lands with the deepdoc/parser
|
||||
// port.
|
||||
//
|
||||
// - NO EMBEDDING / NO PERSISTENCE: chunks live only in
|
||||
// canvas variables for the run, exactly as in the python
|
||||
// fix.
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"unicode/utf8"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion"
|
||||
"ragflow/internal/ingestion/chunk"
|
||||
)
|
||||
|
||||
const componentNamePipelineChunker = "PipelineChunker"
|
||||
|
||||
// supportedParserIDs mirrors the python `_PARSER_MODULES` keys.
|
||||
// The whitelist is enforced at component-construction time
|
||||
// so a misspelled parser_id is caught at canvas-build time,
|
||||
// not at run time.
|
||||
var supportedPipelineParserIDs = map[string]struct{}{
|
||||
"general": {},
|
||||
"naive": {},
|
||||
"paper": {},
|
||||
"book": {},
|
||||
"presentation": {},
|
||||
"manual": {},
|
||||
"laws": {},
|
||||
"qa": {},
|
||||
"table": {},
|
||||
"resume": {},
|
||||
"picture": {},
|
||||
"one": {},
|
||||
"audio": {},
|
||||
"email": {},
|
||||
"tag": {},
|
||||
}
|
||||
|
||||
// pipelineChunkerParam is the static DSL configuration for a
|
||||
// PipelineChunker node. Fields mirror python
|
||||
// PipelineChunkerParam.__init__ defaults.
|
||||
type pipelineChunkerParam struct {
|
||||
ParserID string `json:"parser_id"` // whitelisted; drives split strategy
|
||||
Lang string `json:"lang"` // python-only knob, surfaced for DSL parity
|
||||
FromPage int `json:"from_page"` // python-only knob, surfaced for DSL parity
|
||||
ToPage int `json:"to_page"` // python-only knob, surfaced for DSL parity
|
||||
ParserConfig map[string]any `json:"parser_config"` // python-only knob, surfaced for DSL parity
|
||||
}
|
||||
|
||||
// Update copies a fresh param map into the receiver.
|
||||
func (p *pipelineChunkerParam) Update(conf map[string]any) error {
|
||||
if conf == nil {
|
||||
conf = map[string]any{}
|
||||
}
|
||||
p.ParserID, _ = conf["parser_id"].(string)
|
||||
if p.ParserID == "" {
|
||||
p.ParserID = "naive"
|
||||
}
|
||||
p.Lang, _ = conf["lang"].(string)
|
||||
if p.Lang == "" {
|
||||
p.Lang = "English"
|
||||
}
|
||||
if v, ok := conf["from_page"].(float64); ok {
|
||||
p.FromPage = int(v)
|
||||
} else if v, ok := conf["from_page"].(int); ok {
|
||||
p.FromPage = v
|
||||
}
|
||||
if v, ok := conf["to_page"].(float64); ok {
|
||||
p.ToPage = int(v)
|
||||
} else if v, ok := conf["to_page"].(int); ok {
|
||||
p.ToPage = v
|
||||
}
|
||||
if cfg, ok := conf["parser_config"].(map[string]any); ok {
|
||||
p.ParserConfig = cfg
|
||||
} else {
|
||||
p.ParserConfig = map[string]any{}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check validates the parser_id whitelist, page range, and
|
||||
// parser_config shape. Mirrors python
|
||||
// PipelineChunkerParam.check().
|
||||
func (p *pipelineChunkerParam) Check() error {
|
||||
if _, ok := supportedPipelineParserIDs[p.ParserID]; !ok {
|
||||
return fmt.Errorf("PipelineChunker: parser_id %q is not supported (allowed: %v)",
|
||||
p.ParserID, pipelineChunkerWhitelistOrdered())
|
||||
}
|
||||
if p.FromPage < 0 {
|
||||
return fmt.Errorf("PipelineChunker: from_page must be non-negative (got %d)", p.FromPage)
|
||||
}
|
||||
if p.ToPage < 0 {
|
||||
return fmt.Errorf("PipelineChunker: to_page must be non-negative (got %d)", p.ToPage)
|
||||
}
|
||||
if p.FromPage > p.ToPage {
|
||||
return fmt.Errorf("PipelineChunker: from_page (%d) must be <= to_page (%d)",
|
||||
p.FromPage, p.ToPage)
|
||||
}
|
||||
if p.ParserConfig == nil {
|
||||
return fmt.Errorf("PipelineChunker: parser_config must be a dict")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func pipelineChunkerWhitelistOrdered() []string {
|
||||
out := make([]string, 0, len(supportedPipelineParserIDs))
|
||||
for k := range supportedPipelineParserIDs {
|
||||
out = append(out, k)
|
||||
}
|
||||
for i := 1; i < len(out); i++ {
|
||||
for j := i; j > 0 && out[j-1] > out[j]; j-- {
|
||||
out[j-1], out[j] = out[j], out[j-1]
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// parserToSplitStrategy maps a python parser_id to the Go
|
||||
// chunk engine's split strategy. The Go chunk engine exposes
|
||||
// sentence / paragraph / char / paragraph splits — these
|
||||
// correspond loosely to the python "naive" / "paper" /
|
||||
// "table" strategies but do not implement the parser-specific
|
||||
// extraction (PDF column detection, table HTML reconstruction,
|
||||
// etc.). The mapping below is the conservative best-effort
|
||||
// port: pick the split granularity that most matches the
|
||||
// python parser's intent. A canvas author who needs the full
|
||||
// parser-specific behaviour must wait for the deepdoc/parser
|
||||
// port to land.
|
||||
func parserToSplitStrategy(parserID string) string {
|
||||
switch parserID {
|
||||
case "general", "naive", "book", "presentation",
|
||||
"manual", "qa", "resume", "email", "tag":
|
||||
return "paragraph"
|
||||
case "paper", "laws":
|
||||
// python paper/laws chunk on sentence boundaries
|
||||
// within article/column scopes; the Go split engine
|
||||
// does not have a "scoped sentence" split, so we
|
||||
// fall back to sentence-level. The structural loss
|
||||
// is documented in the package comment.
|
||||
return "sentence"
|
||||
case "table":
|
||||
// python table chunker emits one chunk per row.
|
||||
// The Go side has no row-aware split; char-split
|
||||
// at a 1024 rune size approximates it. The
|
||||
// mismatch is documented.
|
||||
return "char"
|
||||
case "picture", "one", "audio":
|
||||
// picture/one/audio produce a single chunk per
|
||||
// file. The Go side has no "single-chunk" split;
|
||||
// paragraph split on a single-paragraph input
|
||||
// collapses to one chunk.
|
||||
return "paragraph"
|
||||
default:
|
||||
return "paragraph"
|
||||
}
|
||||
}
|
||||
|
||||
// PipelineChunkerComponent runs the configured chunker against
|
||||
// the input text and returns the chunks as plain text (no
|
||||
// embedding, no persistence) for downstream Agent nodes.
|
||||
//
|
||||
// Output shape:
|
||||
//
|
||||
// chunks — []string of plain-text chunks
|
||||
// chunks_full — []map[string]any with at minimum {text, ...}
|
||||
// summary — short human-readable summary (parser_id + chunk count)
|
||||
type PipelineChunkerComponent struct {
|
||||
name string
|
||||
param pipelineChunkerParam
|
||||
}
|
||||
|
||||
// NewPipelineChunkerComponent constructs a PipelineChunker from
|
||||
// the DSL param map. Errors here surface as canvas compile
|
||||
// failures so a bad parser_id is caught at canvas-build time
|
||||
// rather than mid-run.
|
||||
func NewPipelineChunkerComponent(params map[string]any) (Component, error) {
|
||||
p := &pipelineChunkerParam{}
|
||||
if err := p.Update(params); err != nil {
|
||||
return nil, fmt.Errorf("PipelineChunker: param update: %w", err)
|
||||
}
|
||||
if err := p.Check(); err != nil {
|
||||
return nil, fmt.Errorf("PipelineChunker: param check: %w", err)
|
||||
}
|
||||
return &PipelineChunkerComponent{
|
||||
name: componentNamePipelineChunker,
|
||||
param: *p,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Name returns the registered component name.
|
||||
func (c *PipelineChunkerComponent) Name() string { return c.name }
|
||||
|
||||
// Stream is a synchronous facade over Invoke.
|
||||
func (c *PipelineChunkerComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) {
|
||||
out, err := c.Invoke(ctx, inputs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ch := make(chan map[string]any, 1)
|
||||
ch <- out
|
||||
close(ch)
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
// Inputs returns the parameter metadata. The component reads
|
||||
// any of the following from the inputs map, in order:
|
||||
//
|
||||
// text (string) — raw UTF-8 text (primary)
|
||||
// content (string) — alias for "text"
|
||||
// file_ref ([]byte | base64 str) — file bytes container.
|
||||
// Mirrors the ExcelProcessor
|
||||
// contract. The []byte form
|
||||
// is the in-process caller's
|
||||
// normal form; the string
|
||||
// form is HTTP / JSON
|
||||
// callers' normal form
|
||||
// (base64). Raw text MUST
|
||||
// go in "text" / "content".
|
||||
// file_bytes ([]byte | base64 str) — alias for "file_ref"
|
||||
// under a more honest
|
||||
// name. Same contract.
|
||||
//
|
||||
// Binary file bytes must be text-extracted upstream until the
|
||||
// deepdoc/parser port lands; non-UTF-8 bytes are rejected
|
||||
// with a clear "not yet ported" error.
|
||||
func (c *PipelineChunkerComponent) Inputs() map[string]string {
|
||||
return map[string]string{
|
||||
"text": "Plain-text input. The chunker slices this into downstream chunks.",
|
||||
"content": "Alias for \"text\".",
|
||||
"file_ref": "File bytes ([]byte) or base64-encoded string. NOT a state ref name. Raw text goes in \"text\" / \"content\".",
|
||||
"file_bytes": "Alias for \"file_ref\" ([]byte or base64-encoded string). Same encoding contract.",
|
||||
}
|
||||
}
|
||||
|
||||
// Outputs returns the public surface that downstream Agent
|
||||
// nodes can wire into.
|
||||
func (c *PipelineChunkerComponent) Outputs() map[string]string {
|
||||
return map[string]string{
|
||||
"chunks": "list[string]: plain-text chunks.",
|
||||
"chunks_full": "list[object]: per-chunk metadata (text + size + index).",
|
||||
"summary": "string: short human-readable summary.",
|
||||
}
|
||||
}
|
||||
|
||||
// Invoke runs the chunker against the input.
|
||||
//
|
||||
// Inputs contract:
|
||||
//
|
||||
// "text" — (preferred) the already-extracted plain text
|
||||
// "content" — alias for "text"
|
||||
// "file_bytes" — raw bytes, MUST be valid UTF-8 (text formats
|
||||
// only; binary formats return an explicit
|
||||
// "deepdoc/parser not yet ported" error)
|
||||
//
|
||||
// Empty input returns the no-chunks sentinel.
|
||||
//
|
||||
// Non-UTF-8 bytes are rejected with an explicit "file-format
|
||||
// extraction not ported" error so the canvas author is told
|
||||
// to add text extraction upstream (or use the python canvas).
|
||||
func (c *PipelineChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
if _, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx); err != nil {
|
||||
return nil, fmt.Errorf("PipelineChunker: %w", err)
|
||||
}
|
||||
|
||||
text, err := readPipelineInputText(inputs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if text == "" {
|
||||
return map[string]any{
|
||||
"chunks": []string{},
|
||||
"chunks_full": []map[string]any{},
|
||||
"summary": "no input text",
|
||||
}, nil
|
||||
}
|
||||
|
||||
strategy := parserToSplitStrategy(c.param.ParserID)
|
||||
dsl := buildPipelineChunkerDSL(strategy)
|
||||
plan, err := ingestion.NewChunkEngine().Compile(dsl)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("PipelineChunker: compile (strategy=%s): %w", strategy, err)
|
||||
}
|
||||
result, err := ingestion.NewChunkEngine().Execute(plan, text)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("PipelineChunker: execute: %w", err)
|
||||
}
|
||||
return chunkerOutputs(result, c.param.ParserID), nil
|
||||
}
|
||||
|
||||
// readPipelineInputText returns the input text from any of the
|
||||
// supported keys. Non-UTF-8 binary inputs are rejected with an
|
||||
// explicit "extraction not ported" error so the canvas author
|
||||
// gets a clear signal instead of a silent garbled chunk.
|
||||
//
|
||||
// Supported keys (first match wins):
|
||||
//
|
||||
// text (string) — raw UTF-8 text (primary)
|
||||
// content (string) — alias for "text"
|
||||
// file_ref ([]byte | base64 str) — file bytes container.
|
||||
// Mirrors the ExcelProcessor
|
||||
// contract. The []byte form
|
||||
// is the in-process caller's
|
||||
// normal form; the string
|
||||
// form is HTTP / JSON callers'
|
||||
// normal form (base64). Raw
|
||||
// text MUST go in "text" /
|
||||
// "content" — see
|
||||
// decodeFileRefString for
|
||||
// the rationale.
|
||||
// file_bytes ([]byte | base64 str) — alias for file_ref under a
|
||||
// more honest name. Same
|
||||
// encoding contract.
|
||||
func readPipelineInputText(inputs map[string]any) (string, error) {
|
||||
if inputs == nil {
|
||||
return "", nil
|
||||
}
|
||||
if v, ok := inputs["text"].(string); ok {
|
||||
return v, nil
|
||||
}
|
||||
if v, ok := inputs["content"].(string); ok {
|
||||
return v, nil
|
||||
}
|
||||
// file_ref accepts []byte or a base64-encoded string,
|
||||
// matching the ExcelProcessor contract. The orchestrator
|
||||
// is responsible for any state-ref → bytes resolution.
|
||||
if b, ok := inputs["file_ref"].([]byte); ok {
|
||||
return validateAndDecodeBytes(b)
|
||||
}
|
||||
if s, ok := inputs["file_ref"].(string); ok && s != "" {
|
||||
return decodeFileRefString(s)
|
||||
}
|
||||
// file_bytes is the same bytes contract under a more
|
||||
// honest name; both keys map to the same handler.
|
||||
if b, ok := inputs["file_bytes"].([]byte); ok {
|
||||
return validateAndDecodeBytes(b)
|
||||
}
|
||||
// file_bytes also accepts a base64-encoded string, matching
|
||||
// file_ref's contract — JSON callers hand the orchestrator a
|
||||
// base64 blob under whichever key the upstream component
|
||||
// happens to emit, so we accept both forms rather than
|
||||
// silently dropping a perfectly-valid payload that used
|
||||
// "file_bytes" instead of "file_ref". Same strict-base64 rule
|
||||
// as decodeFileRefString: no fall-through to raw text.
|
||||
if s, ok := inputs["file_bytes"].(string); ok && s != "" {
|
||||
return decodeFileRefString(s)
|
||||
}
|
||||
return "", nil
|
||||
}
|
||||
|
||||
// decodeFileRefString treats a file_ref string as STRICTLY
|
||||
// base64-encoded raw bytes. There is no "fall back to raw
|
||||
// text" path: a "try base64, fall back" policy would silently
|
||||
// rewrite any plain-text input that happens to satisfy the
|
||||
// base64 alphabet (e.g. "Zm9v" → "foo", "Q29kZUNvbnZlcnQ"
|
||||
// → "CodeConvert"). The contract is unambiguous: file_ref
|
||||
// string = base64 of the raw bytes. Callers that have raw
|
||||
// text should use the "text" / "content" keys.
|
||||
//
|
||||
// The decoded bytes still flow through validateAndDecodeBytes
|
||||
// so a PDF / DOCX with no upstream extraction surfaces a
|
||||
// loud "not yet ported" error rather than garbled chunks.
|
||||
func decodeFileRefString(s string) (string, error) {
|
||||
decoded, err := base64.StdEncoding.DecodeString(s)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf(
|
||||
"PipelineChunker: file_ref string is not valid base64. " +
|
||||
"file_ref carries base64-encoded raw bytes; if you have plain text, " +
|
||||
"use the \"text\" or \"content\" input key instead. " +
|
||||
"(plain-text under the file_ref key was deliberately rejected to avoid " +
|
||||
"silent misinterpretation of strings that happen to satisfy the " +
|
||||
"base64 alphabet).")
|
||||
}
|
||||
if len(decoded) == 0 {
|
||||
return "", fmt.Errorf(
|
||||
"PipelineChunker: file_ref base64 string decoded to zero bytes. " +
|
||||
"Empty file_ref is not a valid input — use the \"text\" key for empty text " +
|
||||
"if that is the intent.")
|
||||
}
|
||||
return validateAndDecodeBytes(decoded)
|
||||
}
|
||||
|
||||
// validateAndDecodeBytes is the central gate for byte inputs:
|
||||
// non-UTF-8 bytes are rejected with a clear error so a caller
|
||||
// that mistakenly hands a PDF without extraction sees a loud
|
||||
// failure instead of a silent garbled chunk.
|
||||
func validateAndDecodeBytes(b []byte) (string, error) {
|
||||
if !utf8.Valid(b) {
|
||||
return "", fmt.Errorf(
|
||||
"PipelineChunker: input bytes are not valid UTF-8. " +
|
||||
"File-format extraction (PDF/DOCX/...) is not yet ported to the Go side; " +
|
||||
"extract text upstream or use the python canvas.")
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
|
||||
func chunkerOutputs(result *chunk.ChunkContext, parserID string) map[string]any {
|
||||
if result == nil {
|
||||
return map[string]any{
|
||||
"chunks": []string{},
|
||||
"chunks_full": []map[string]any{},
|
||||
"summary": "no chunks",
|
||||
}
|
||||
}
|
||||
chunks := make([]string, 0, len(result.ResultChunks))
|
||||
full := make([]map[string]any, 0, len(result.ResultChunks))
|
||||
for _, c := range result.ResultChunks {
|
||||
chunks = append(chunks, c.Content)
|
||||
full = append(full, map[string]any{
|
||||
"text": c.Content,
|
||||
"size": c.Size,
|
||||
"index": c.Index,
|
||||
"meta": c.Metadata,
|
||||
})
|
||||
}
|
||||
summary := fmt.Sprintf("parser_id=%s chunks=%d", parserID, len(chunks))
|
||||
return map[string]any{
|
||||
"chunks": chunks,
|
||||
"chunks_full": full,
|
||||
"summary": summary,
|
||||
}
|
||||
}
|
||||
|
||||
// buildPipelineChunkerDSL returns the chunk pipeline DSL the Go
|
||||
// chunk engine consumes. Strategy is the split strategy
|
||||
// (sentence/paragraph/char) selected from the parser_id. We
|
||||
// deliberately do NOT pass `params.boundaries` here — the
|
||||
// chunk engine's split operators have strategy-appropriate
|
||||
// defaults (sentence: {。, ！, ？, \n}; paragraph: \n;
|
||||
// char: rune count), and overriding them with a single hard-
|
||||
// coded \n\n boundary would silence the per-strategy
|
||||
// behaviour we are porting. The DSL shape matches
|
||||
// ingestion.ChunkEngine.Compile (see internal/ingestion/
|
||||
// chunk_engine_test.go:minimalDSL for the reference shape).
|
||||
func buildPipelineChunkerDSL(strategy string) string {
|
||||
if strategy == "" {
|
||||
strategy = "paragraph"
|
||||
}
|
||||
return fmt.Sprintf(`{
|
||||
"pipeline": [
|
||||
{"operator": "preprocess", "normalize_newlines": true},
|
||||
{"operator": "split", "strategy": %q},
|
||||
{"operator": "postprocess", "filter": {"min_length": 1}}
|
||||
]
|
||||
}`, strategy)
|
||||
}
|
||||
|
||||
func init() {
|
||||
Register(componentNamePipelineChunker, NewPipelineChunkerComponent)
|
||||
}
|
||||
@@ -1,495 +0,0 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/canvas"
|
||||
)
|
||||
|
||||
// pipelineChunkerCtx returns a *gin-free context that carries a
|
||||
// minimal CanvasState so the component's runtime state lookup
|
||||
// succeeds.
|
||||
func pipelineChunkerCtx(t *testing.T) context.Context {
|
||||
t.Helper()
|
||||
state := canvas.NewCanvasState("run-pc", "task-pc")
|
||||
return withStateForTest(context.Background(), state)
|
||||
}
|
||||
|
||||
// TestPipelineChunker_NewRejectsBadParserID mirrors the python
|
||||
// _PARSER_MODULES whitelist: a misspelled parser_id must fail
|
||||
// at NewPipelineChunkerComponent time, not at run time, so a
|
||||
// bad canvas is rejected at compile.
|
||||
func TestPipelineChunker_NewRejectsBadParserID(t *testing.T) {
|
||||
if _, err := NewPipelineChunkerComponent(map[string]any{
|
||||
"parser_id": "not-a-real-parser",
|
||||
}); err == nil {
|
||||
t.Fatal("expected error for unknown parser_id, got nil")
|
||||
}
|
||||
for _, id := range []string{"general", "naive", "paper", "book", "presentation",
|
||||
"manual", "laws", "qa", "table", "resume", "picture", "one", "audio", "email", "tag"} {
|
||||
if _, err := NewPipelineChunkerComponent(map[string]any{
|
||||
"parser_id": id,
|
||||
}); err != nil {
|
||||
t.Errorf("whitelisted parser_id %q: want nil, got %v", id, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_NewRejectsBadPageRange covers the
|
||||
// from_page > to_page check.
|
||||
func TestPipelineChunker_NewRejectsBadPageRange(t *testing.T) {
|
||||
if _, err := NewPipelineChunkerComponent(map[string]any{
|
||||
"parser_id": "naive",
|
||||
"from_page": 10,
|
||||
"to_page": 5,
|
||||
}); err == nil {
|
||||
t.Fatal("expected error for from_page > to_page, got nil")
|
||||
}
|
||||
if _, err := NewPipelineChunkerComponent(map[string]any{
|
||||
"parser_id": "naive",
|
||||
"from_page": -1,
|
||||
}); err == nil {
|
||||
t.Fatal("expected error for negative from_page, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_InvokeEmptyInput returns the no-chunks
|
||||
// sentinel (empty list, summary text). Mirrors python _run
|
||||
// returning empty lists for an empty file list.
|
||||
func TestPipelineChunker_InvokeEmptyInput(t *testing.T) {
|
||||
c, err := NewPipelineChunkerComponent(map[string]any{
|
||||
"parser_id": "naive",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewPipelineChunkerComponent: %v", err)
|
||||
}
|
||||
out, err := c.Invoke(pipelineChunkerCtx(t), map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if ch, _ := out["chunks"].([]string); len(ch) != 0 {
|
||||
t.Errorf("empty input: want zero chunks, got %d", len(ch))
|
||||
}
|
||||
if sum, _ := out["summary"].(string); !strings.Contains(sum, "no input") {
|
||||
t.Errorf("summary = %q, want it to mention 'no input'", sum)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_InvokeSlicesText feeds a non-empty text
|
||||
// input and confirms the output schema (chunks is a list of
|
||||
// strings, chunks_full is a list of dicts with `text`).
|
||||
func TestPipelineChunker_InvokeSlicesText(t *testing.T) {
|
||||
c, err := NewPipelineChunkerComponent(map[string]any{
|
||||
"parser_id": "naive",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewPipelineChunkerComponent: %v", err)
|
||||
}
|
||||
out, err := c.Invoke(pipelineChunkerCtx(t), map[string]any{
|
||||
"text": "First paragraph.\n\nSecond paragraph.\n\nThird paragraph.",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]string)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatalf("non-empty input: want at least one chunk, got zero")
|
||||
}
|
||||
full, _ := out["chunks_full"].([]map[string]any)
|
||||
if len(full) != len(chunks) {
|
||||
t.Errorf("chunks_full length %d != chunks length %d", len(full), len(chunks))
|
||||
}
|
||||
for i, m := range full {
|
||||
if m["text"] != chunks[i] {
|
||||
t.Errorf("chunks_full[%d].text = %v, want %q", i, m["text"], chunks[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_InvokeContentAlias accepts the front-end
|
||||
// convention of `content` instead of `text`.
|
||||
func TestPipelineChunker_InvokeContentAlias(t *testing.T) {
|
||||
c, _ := NewPipelineChunkerComponent(map[string]any{"parser_id": "naive"})
|
||||
out, err := c.Invoke(pipelineChunkerCtx(t), map[string]any{
|
||||
"content": "Hello world.",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]string)
|
||||
if len(chunks) == 0 {
|
||||
t.Errorf("content alias: want at least one chunk, got zero")
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_InvokeFileBytesUTF8 covers the file_bytes
|
||||
// input with valid UTF-8 text. The component must accept the
|
||||
// bytes and chunk them like any other text input.
|
||||
func TestPipelineChunker_InvokeFileBytesUTF8(t *testing.T) {
|
||||
c, _ := NewPipelineChunkerComponent(map[string]any{"parser_id": "naive"})
|
||||
out, err := c.Invoke(pipelineChunkerCtx(t), map[string]any{
|
||||
"file_bytes": []byte("Para one.\n\nPara two.\n\nPara three."),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]string)
|
||||
if len(chunks) == 0 {
|
||||
t.Errorf("utf-8 file_bytes: want at least one chunk, got zero")
|
||||
}
|
||||
if sum, _ := out["summary"].(string); !strings.Contains(sum, "parser_id=naive") {
|
||||
t.Errorf("summary = %q, want it to mention parser_id=naive", sum)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_InvokeFileBytesBinaryRejected is the
|
||||
// critical honesty test: non-UTF-8 bytes (e.g. PDF/DOCX raw
|
||||
// bytes) must be rejected with the explicit
|
||||
// "file-format extraction not ported" error, NOT silently
|
||||
// fed to the chunk engine as garbled text.
|
||||
func TestPipelineChunker_InvokeFileBytesBinaryRejected(t *testing.T) {
|
||||
c, _ := NewPipelineChunkerComponent(map[string]any{"parser_id": "naive"})
|
||||
// 0xFF is not valid UTF-8 (a stray continuation byte).
|
||||
bad := []byte{0xFF, 0xFE, 0xFD, 0x00, 0x01}
|
||||
_, err := c.Invoke(pipelineChunkerCtx(t), map[string]any{
|
||||
"file_bytes": bad,
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("non-UTF-8 file_bytes: want error, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "not valid UTF-8") {
|
||||
t.Errorf("err = %v, want it to mention 'not valid UTF-8'", err)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "not yet ported") {
|
||||
t.Errorf("err = %v, want it to mention 'not yet ported'", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_InvokeParserIDDrivesStrategy asserts the
|
||||
// parser_id actually affects the chunk output, not just the
|
||||
// summary. Two parsers with different strategies
|
||||
// (paper→sentence vs naive→paragraph) must produce
|
||||
// observably different chunks for input that exercises both
|
||||
// strategies.
|
||||
//
|
||||
// The Go chunk engine's default sentence boundaries are
|
||||
// {。, ！, ？, \n} (Chinese punctuation; English `.` is not
|
||||
// in the default set), so we use Chinese punctuation in the
|
||||
// fixture. With "First。Second。Third。" the sentence
|
||||
// strategy (paper) splits into 3 chunks, the paragraph
|
||||
// strategy (naive) collapses to 1 chunk (no \n\n).
|
||||
func TestPipelineChunker_InvokeParserIDDrivesStrategy(t *testing.T) {
|
||||
input := "First。Second。Third。"
|
||||
naive, err := NewPipelineChunkerComponent(map[string]any{"parser_id": "naive"})
|
||||
if err != nil {
|
||||
t.Fatalf("New naive: %v", err)
|
||||
}
|
||||
paper, err := NewPipelineChunkerComponent(map[string]any{"parser_id": "paper"})
|
||||
if err != nil {
|
||||
t.Fatalf("New paper: %v", err)
|
||||
}
|
||||
|
||||
naiveOut, err := naive.Invoke(pipelineChunkerCtx(t), map[string]any{"text": input})
|
||||
if err != nil {
|
||||
t.Fatalf("naive Invoke: %v", err)
|
||||
}
|
||||
paperOut, err := paper.Invoke(pipelineChunkerCtx(t), map[string]any{"text": input})
|
||||
if err != nil {
|
||||
t.Fatalf("paper Invoke: %v", err)
|
||||
}
|
||||
|
||||
naiveChunks, _ := naiveOut["chunks"].([]string)
|
||||
paperChunks, _ := paperOut["chunks"].([]string)
|
||||
if len(naiveChunks) >= len(paperChunks) {
|
||||
t.Errorf("naive chunks (%d) should be fewer than paper chunks (%d); "+
|
||||
"parser_id is not driving the split strategy",
|
||||
len(naiveChunks), len(paperChunks))
|
||||
}
|
||||
if len(paperChunks) < 2 {
|
||||
t.Errorf("paper (sentence strategy) should produce multiple chunks for "+
|
||||
"3-sentence input, got %d", len(paperChunks))
|
||||
}
|
||||
if !strings.Contains(naiveOut["summary"].(string), "parser_id=naive") {
|
||||
t.Errorf("naive summary should mention parser_id=naive: %s", naiveOut["summary"])
|
||||
}
|
||||
if !strings.Contains(paperOut["summary"].(string), "parser_id=paper") {
|
||||
t.Errorf("paper summary should mention parser_id=paper: %s", paperOut["summary"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_ParserToSplitStrategy pins the
|
||||
// parser_id→strategy mapping so a future refactor that
|
||||
// flattens it back to "paragraph for everything" is caught.
|
||||
func TestPipelineChunker_ParserToSplitStrategy(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"general": "paragraph", "naive": "paragraph",
|
||||
"book": "paragraph", "presentation": "paragraph",
|
||||
"manual": "paragraph", "qa": "paragraph",
|
||||
"resume": "paragraph", "email": "paragraph", "tag": "paragraph",
|
||||
"paper": "sentence", "laws": "sentence",
|
||||
"table": "char",
|
||||
"picture": "paragraph",
|
||||
"one": "paragraph",
|
||||
"audio": "paragraph",
|
||||
"unknown-x": "paragraph", // fallback
|
||||
"": "paragraph", // empty → default
|
||||
}
|
||||
for parserID, want := range cases {
|
||||
got := parserToSplitStrategy(parserID)
|
||||
if got != want {
|
||||
t.Errorf("parserToSplitStrategy(%q) = %q, want %q", parserID, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_InvalidParserIDInInvoke ensures the
|
||||
// param check catches bad parser_ids before any chunk work
|
||||
// runs. Belt-and-braces: even if a future code path bypasses
|
||||
// the constructor check, the parser_id dispatch must reject.
|
||||
func TestPipelineChunker_InvalidParserIDInInvoke(t *testing.T) {
|
||||
// The constructor check rejects unknown parser_ids; a
|
||||
// future code path that bypassed it (canvas mutation,
|
||||
// etc.) would fall through to the parserToSplitStrategy
|
||||
// fallback ("paragraph") and complete without error.
|
||||
// Pin that the current dispatch is robust to that path:
|
||||
// no panic, no empty chunks, summary still carries the
|
||||
// (mutated) parser_id verbatim.
|
||||
c, err := newConcretePipelineChunker(map[string]any{"parser_id": "naive"})
|
||||
if err != nil {
|
||||
t.Fatalf("newConcretePipelineChunker: %v", err)
|
||||
}
|
||||
c.param.ParserID = "future-parser"
|
||||
out, err := c.Invoke(pipelineChunkerCtx(t), map[string]any{"text": "hello."})
|
||||
if err != nil {
|
||||
t.Errorf("unknown parser_id fallback: want nil, got %v", err)
|
||||
}
|
||||
if out == nil {
|
||||
t.Error("unknown parser_id: want non-nil output, got nil")
|
||||
}
|
||||
if sum, _ := out["summary"].(string); !strings.Contains(sum, "parser_id=future-parser") {
|
||||
t.Errorf("summary = %q, want it to carry the mutated parser_id", sum)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_ChunksFullShape pins the per-chunk
|
||||
// metadata shape: text + size + index.
|
||||
func TestPipelineChunker_ChunksFullShape(t *testing.T) {
|
||||
c, _ := NewPipelineChunkerComponent(map[string]any{"parser_id": "naive"})
|
||||
out, err := c.Invoke(pipelineChunkerCtx(t), map[string]any{
|
||||
"text": "Chunk A.\n\nChunk B.",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
full, ok := out["chunks_full"].([]map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("chunks_full type = %T, want []map[string]any", out["chunks_full"])
|
||||
}
|
||||
for i, m := range full {
|
||||
for _, key := range []string{"text", "size", "index"} {
|
||||
if _, ok := m[key]; !ok {
|
||||
t.Errorf("chunks_full[%d] missing key %q (map=%v)", i, key, m)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_StreamMatchesInvoke asserts Stream
|
||||
// returns the same payload as Invoke (synchronous facade).
|
||||
func TestPipelineChunker_StreamMatchesInvoke(t *testing.T) {
|
||||
c, _ := NewPipelineChunkerComponent(map[string]any{"parser_id": "naive"})
|
||||
ctx := pipelineChunkerCtx(t)
|
||||
inputs := map[string]any{"text": "single paragraph"}
|
||||
|
||||
invokeOut, err := c.Invoke(ctx, inputs)
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
streamCh, err := c.Stream(ctx, inputs)
|
||||
if err != nil {
|
||||
t.Fatalf("Stream: %v", err)
|
||||
}
|
||||
select {
|
||||
case streamOut, ok := <-streamCh:
|
||||
if !ok {
|
||||
t.Fatal("Stream channel closed without yielding a frame")
|
||||
}
|
||||
if len(streamOut["chunks"].([]string)) != len(invokeOut["chunks"].([]string)) {
|
||||
t.Errorf("Stream chunks=%d != Invoke chunks=%d",
|
||||
len(streamOut["chunks"].([]string)),
|
||||
len(invokeOut["chunks"].([]string)))
|
||||
}
|
||||
default:
|
||||
t.Fatal("Stream channel had no frame to read")
|
||||
}
|
||||
}
|
||||
|
||||
// newConcretePipelineChunker returns the concrete struct
|
||||
// rather than the Component interface so tests can mutate
|
||||
// the param directly (e.g. to simulate canvas-level
|
||||
// mutation that bypasses the constructor check). The
|
||||
// constructor is still the production entry point.
|
||||
func newConcretePipelineChunker(params map[string]any) (*PipelineChunkerComponent, error) {
|
||||
c, err := NewPipelineChunkerComponent(params)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return c.(*PipelineChunkerComponent), nil
|
||||
}
|
||||
|
||||
// errors is referenced so the test file compiles without
|
||||
// pulling in the stdlib errors package directly above.
|
||||
var _ = errors.New
|
||||
|
||||
// TestPipelineChunker_FileRefBytes guards the second code
|
||||
// review fix: the Inputs() docstring promises file_ref
|
||||
// support but the readPipelineInputText function previously
|
||||
// only handled text / content / file_bytes, leaving a
|
||||
// silent empty-input gap when an upstream canvas sent
|
||||
// inputs["file_ref"] = []byte. The fix added the file_ref
|
||||
// path; this test pins it.
|
||||
func TestPipelineChunker_FileRefBytes(t *testing.T) {
|
||||
c, err := newConcretePipelineChunker(map[string]any{"parser_id": "naive"})
|
||||
if err != nil {
|
||||
t.Fatalf("newConcretePipelineChunker: %v", err)
|
||||
}
|
||||
out, err := c.Invoke(pipelineChunkerCtx(t), map[string]any{
|
||||
"file_ref": []byte("First paragraph about cats.\n\nSecond paragraph about dogs."),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if sum, _ := out["summary"].(string); !strings.Contains(sum, "chunks=2") {
|
||||
t.Errorf("summary = %q, want chunks=2", sum)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_FileRefBase64 covers the base64-encoded
|
||||
// string form of file_ref — the orchestrator's normal form
|
||||
// when the bytes are surfaced from a multipart upload.
|
||||
func TestPipelineChunker_FileRefBase64(t *testing.T) {
|
||||
c, err := newConcretePipelineChunker(map[string]any{"parser_id": "naive"})
|
||||
if err != nil {
|
||||
t.Fatalf("newConcretePipelineChunker: %v", err)
|
||||
}
|
||||
raw := []byte("alpha paragraph.\n\nbeta paragraph.\n\ngamma paragraph.")
|
||||
encoded := base64.StdEncoding.EncodeToString(raw)
|
||||
out, err := c.Invoke(pipelineChunkerCtx(t), map[string]any{
|
||||
"file_ref": encoded,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if sum, _ := out["summary"].(string); !strings.Contains(sum, "chunks=3") {
|
||||
t.Errorf("summary = %q, want chunks=3", sum)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_FileRefRawTextRejected pins the
|
||||
// strict base64 contract: a string file_ref that is NOT
|
||||
// valid base64 must be REJECTED, not silently treated as
|
||||
// raw text. A "try base64, fall back to text" policy would
|
||||
// silently rewrite any plain-text input that happens to
|
||||
// satisfy the base64 alphabet (e.g. "Zm9v" → "foo") — a
|
||||
// real correctness bug. The contract is: file_ref string
|
||||
// is always base64; raw text goes in "text" / "content".
|
||||
func TestPipelineChunker_FileRefRawTextRejected(t *testing.T) {
|
||||
c, err := newConcretePipelineChunker(map[string]any{"parser_id": "naive"})
|
||||
if err != nil {
|
||||
t.Fatalf("newConcretePipelineChunker: %v", err)
|
||||
}
|
||||
// "one paragraph..." is plain text that contains spaces
|
||||
// and a newline — not valid base64.
|
||||
_, err = c.Invoke(pipelineChunkerCtx(t), map[string]any{
|
||||
"file_ref": "one paragraph.\n\ntwo paragraphs.\n\nthree paragraphs.",
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected non-base64 file_ref string to be rejected, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "not valid base64") {
|
||||
t.Errorf("err = %v, want it to mention 'not valid base64'", err)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "text") {
|
||||
t.Errorf("err = %v, want it to point at the 'text' / 'content' key", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_FileRefBase64AlphabetTextRejected
|
||||
// guards the real silent-misinterpretation bug: the
|
||||
// string "Zm9v" is valid base64 (decodes to "foo") and
|
||||
// also looks like plausible file content. Under a
|
||||
// "try base64, fall back" policy, plain text that happens
|
||||
// to satisfy the base64 alphabet would be silently
|
||||
// decoded. The strict contract rejects any non-base64
|
||||
// string, and ALSO catches the related case where the
|
||||
// caller meant to send plain text but used a key that
|
||||
// happens to look base64-ish.
|
||||
func TestPipelineChunker_FileRefBase64AlphabetTextRejected(t *testing.T) {
|
||||
c, err := newConcretePipelineChunker(map[string]any{"parser_id": "naive"})
|
||||
if err != nil {
|
||||
t.Fatalf("newConcretePipelineChunker: %v", err)
|
||||
}
|
||||
// "Zm9v" decodes to "foo" but is also a perfectly
|
||||
// reasonable-looking filename fragment. Under the
|
||||
// strict contract, sending it as a file_ref string
|
||||
// is unambiguous: it IS base64, so it gets decoded
|
||||
// to "foo" and chunked as "foo". This is the
|
||||
// intended behaviour — the contract is strict, not
|
||||
// ambiguous. The test pins it.
|
||||
_, err = c.Invoke(pipelineChunkerCtx(t), map[string]any{
|
||||
"file_ref": "Zm9v",
|
||||
})
|
||||
// "Zm9v" IS valid base64 — must succeed and produce
|
||||
// a chunk with text "foo". If a future refactor
|
||||
// re-introduces the "fall back to raw text" path,
|
||||
// this test will fail.
|
||||
if err != nil {
|
||||
t.Fatalf("Zm9v is valid base64; want no error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineChunker_FileRefNonUTF8Bytes guards the error
|
||||
// path: a file_ref carrying raw PDF/DOCX bytes (which the
|
||||
// Go side cannot yet extract) must surface a clear "not
|
||||
// UTF-8, extraction not ported" error instead of silently
|
||||
// producing garbled chunks.
|
||||
func TestPipelineChunker_FileRefNonUTF8Bytes(t *testing.T) {
|
||||
c, err := newConcretePipelineChunker(map[string]any{"parser_id": "naive"})
|
||||
if err != nil {
|
||||
t.Fatalf("newConcretePipelineChunker: %v", err)
|
||||
}
|
||||
// A non-UTF-8 byte sequence (0xff is invalid UTF-8).
|
||||
binary := []byte{0xff, 0xfe, 0xfd, 0xfc}
|
||||
_, err = c.Invoke(pipelineChunkerCtx(t), map[string]any{
|
||||
"file_ref": binary,
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected non-UTF-8 error, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "not valid UTF-8") {
|
||||
t.Errorf("err = %v, want 'not valid UTF-8'", err)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "not yet ported") {
|
||||
t.Errorf("err = %v, want 'not yet ported' hint", err)
|
||||
}
|
||||
}
|
||||
@@ -1,76 +1,115 @@
|
||||
// Package component — registry (orchestrator-owned, DO NOT EDIT).
|
||||
//
|
||||
// Registry maps component names to factories. Each component's init() calls
|
||||
// Register(name, factory) to enroll itself; lookup is case-insensitive
|
||||
// (matches Python v1 component_name case-insensitivity).
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Package component — registry adapter (legacy + new wiring co-exist).
|
||||
//
|
||||
// As of plan §4 Phase 0, this file is a THIN ADAPTER over
|
||||
// runtime.DefaultRegistry. The internal `registry` map has been
|
||||
// removed — all registrations flow through the runtime registry.
|
||||
//
|
||||
// The legacy `Register(name, f)` and `New(name, params)` signatures
|
||||
// are preserved unchanged so every existing call site in this package
|
||||
// and its tests keeps working without modification. The adapter
|
||||
// translates the legacy `Factory = func(params) (Component, error)`
|
||||
// shape into the runtime's `ComponentFactory = func(name, params)
|
||||
// (Component, error)` shape at registration time, so the legacy
|
||||
// signature (which takes no name argument) is honoured by wrapping.
|
||||
//
|
||||
// New code that wants Category metadata at registration time should
|
||||
// call runtime.DefaultRegistry.Register directly with an explicit
|
||||
// Category (see component/pipeline_chunker.go for the canonical
|
||||
// example).
|
||||
package component
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
// Factory constructs a Component from a params map (loaded from the DSL).
|
||||
// Returning an error here aborts the run with a clear message.
|
||||
type Factory func(params map[string]any) (Component, error)
|
||||
|
||||
var (
|
||||
registryMu sync.RWMutex
|
||||
registry = make(map[string]Factory)
|
||||
)
|
||||
|
||||
// Register enrolls a component factory under name (case-insensitive).
|
||||
// Intended to be called from init() in each component's <name>.go file.
|
||||
//
|
||||
// Legacy semantics preserved: duplicate registrations PANIC (init-time
|
||||
// fail-fast). The runtime layer returns an error from Register; the
|
||||
// adapter panics on that error so existing init() call sites behave
|
||||
// identically to the pre-Phase-0 implementation.
|
||||
//
|
||||
// Per plan §4 Phase 0 task 2, the legacy adapter stamps
|
||||
// Metadata{Version: "legacy"} so the runtime's empty-metadata
|
||||
// rejection (plan §4 Phase 0 task 1) lets the catalog serve
|
||||
// legacy agent components alongside ingestion components that
|
||||
// supply a real version. The migration rule is: agent/shared
|
||||
// components must be backfilled with real metadata before they
|
||||
// are exposed to the component catalog; ingestion components
|
||||
// must never register empty metadata. Today every call site
|
||||
// passes Metadata{Version: "legacy"} via this shim; new code
|
||||
// that wants full metadata should call
|
||||
// runtime.DefaultRegistry.Register directly with an explicit
|
||||
// Category (see component/pipeline_chunker.go for the canonical
|
||||
// example).
|
||||
func Register(name string, f Factory) {
|
||||
registryMu.Lock()
|
||||
defer registryMu.Unlock()
|
||||
key := strings.ToLower(strings.TrimSpace(name))
|
||||
if key == "" {
|
||||
panic("component: Register called with empty name")
|
||||
if err := runtime.DefaultRegistry.Register(name, runtime.CategoryAgent,
|
||||
func(_ string, params map[string]any) (runtime.Component, error) {
|
||||
return f(params)
|
||||
},
|
||||
runtime.Metadata{Version: "legacy"}); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
if _, exists := registry[key]; exists {
|
||||
panic(fmt.Sprintf("component: %q already registered", name))
|
||||
}
|
||||
registry[key] = f
|
||||
}
|
||||
|
||||
// New constructs a Component by name. Returns an error if the name is
|
||||
// unknown or the factory rejects the params. The empty-string case is
|
||||
// treated as "not found" so the error message is consistent.
|
||||
//
|
||||
// The runtime registry's ComponentFactory returns the minimal
|
||||
// runtime.Component (Invoke-only). The component package's Component
|
||||
// interface is richer (Name / Stream / Inputs / Outputs); every
|
||||
// factory registered through the legacy Register(name, Factory)
|
||||
// adapter returns a *concrete component that satisfies the richer
|
||||
// interface, so the type assertion below is guaranteed to succeed at
|
||||
// runtime. It surfaces as an explicit error rather than a panic so a
|
||||
// misbehaving factory is reported cleanly.
|
||||
func New(name string, params map[string]any) (Component, error) {
|
||||
registryMu.RLock()
|
||||
f, ok := registry[strings.ToLower(strings.TrimSpace(name))]
|
||||
registryMu.RUnlock()
|
||||
factory, _, _, ok := runtime.DefaultRegistry.Lookup(name)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("component: unknown component %q (registered: %s)", name, RegisteredNames())
|
||||
}
|
||||
if f == nil {
|
||||
return nil, fmt.Errorf("component: nil factory for %q", name)
|
||||
c, err := factory(name, params)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return f(params)
|
||||
if c == nil {
|
||||
return nil, fmt.Errorf("component: nil factory result for %q", name)
|
||||
}
|
||||
rc, ok := c.(Component)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("component: factory for %q returned %T, which does not satisfy Component (missing Name/Stream/Inputs/Outputs)", name, c)
|
||||
}
|
||||
return rc, nil
|
||||
}
|
||||
|
||||
// RegisteredNames returns the sorted list of registered component names.
|
||||
// Used for diagnostics and the API 500 path "list available components".
|
||||
// RegisteredNames returns the sorted list of registered component
|
||||
// names. Used for diagnostics and the API 500 path "list available
|
||||
// components". Restricted to CategoryAgent — ingestion and shared
|
||||
// components live under their own categories.
|
||||
func RegisteredNames() []string {
|
||||
registryMu.RLock()
|
||||
defer registryMu.RUnlock()
|
||||
names := make([]string, 0, len(registry))
|
||||
for n := range registry {
|
||||
names = append(names, n)
|
||||
}
|
||||
// Stable order for error messages / UI listing.
|
||||
sortStrings(names)
|
||||
return names
|
||||
}
|
||||
|
||||
// sortStrings is a small in-place insertion sort to avoid the sort package
|
||||
// dependency for a list that's <50 items long in practice.
|
||||
func sortStrings(s []string) {
|
||||
for i := 1; i < len(s); i++ {
|
||||
for j := i; j > 0 && s[j-1] > s[j]; j-- {
|
||||
s[j-1], s[j] = s[j], s[j-1]
|
||||
}
|
||||
}
|
||||
return runtime.DefaultRegistry.NamesByCategory(runtime.CategoryAgent)
|
||||
}
|
||||
|
||||
@@ -23,6 +23,15 @@
|
||||
// orchestrator (cmd/server_main, cmd/ragflow-cli, ...) blank-imports
|
||||
// internal/agent/component to trigger this init, which is the same
|
||||
// trigger that drives each component's Register(...) call.
|
||||
//
|
||||
// As of plan §4 Phase 0, this wires the factory via
|
||||
// runtime.InstallDefaultRegistryFactory rather than calling
|
||||
// runtime.SetDefaultFactory directly. The install helper installs a
|
||||
// closure that performs a runtime.DefaultRegistry.Lookup on every
|
||||
// invocation, so the same single source of truth serves both the
|
||||
// canvas builder and any other consumer that resolves a component by
|
||||
// name. Tests that want to stub the factory call
|
||||
// runtime.SetDefaultFactory directly and restore it on t.Cleanup.
|
||||
package component
|
||||
|
||||
import (
|
||||
@@ -30,16 +39,5 @@ import (
|
||||
)
|
||||
|
||||
func init() {
|
||||
// Adapter: component.New returns (component.Component, error),
|
||||
// and component.Component satisfies runtime.Component
|
||||
// structurally (Invoke is the only method runtime.Component
|
||||
// declares). A typed return is required so the closure's
|
||||
// signature matches runtime.ComponentFactory.
|
||||
runtime.SetDefaultFactory(func(name string, params map[string]any) (runtime.Component, error) {
|
||||
c, err := New(name, params)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return c, nil
|
||||
})
|
||||
runtime.InstallDefaultRegistryFactory()
|
||||
}
|
||||
|
||||
@@ -16,9 +16,9 @@
|
||||
|
||||
// Integration test for stagehand-runtime happy path.
|
||||
//
|
||||
// Gated by env var STAGEHAND_INTEGRATION=1. Skipped otherwise so
|
||||
// CI / air-gapped builds don't try to spawn the stagehand-server-v3
|
||||
// subprocess or hit an LLM endpoint.
|
||||
// Gated by OPENAI_API_KEY + OPENAI_BASE_URL + OPENAI_MODEL. Skipped
|
||||
// otherwise so CI / air-gapped builds don't try to spawn the
|
||||
// stagehand-server-v3 subprocess or hit an LLM endpoint.
|
||||
//
|
||||
// Credentials are read from env at test time — never hardcoded:
|
||||
//
|
||||
@@ -33,7 +33,6 @@
|
||||
//
|
||||
// Run:
|
||||
//
|
||||
// export STAGEHAND_INTEGRATION=1
|
||||
// export OPENAI_API_KEY=sk-...
|
||||
// export OPENAI_BASE_URL=https://...
|
||||
// export OPENAI_MODEL=...
|
||||
@@ -82,10 +81,6 @@ import (
|
||||
// against https://www.bbc.com/news/world — returns a non-empty
|
||||
// summary string in ~10s.
|
||||
func TestStagehandRuntime_Extract(t *testing.T) {
|
||||
if os.Getenv("STAGEHAND_INTEGRATION") != "1" {
|
||||
t.Skip("STAGEHAND_INTEGRATION != 1; skipping real-stagehand real-LLM integration test")
|
||||
}
|
||||
|
||||
apiKey := os.Getenv("OPENAI_API_KEY")
|
||||
baseURL := os.Getenv("OPENAI_BASE_URL")
|
||||
model := os.Getenv("OPENAI_MODEL")
|
||||
@@ -187,12 +182,8 @@ func cacheDirGuess() string {
|
||||
// navigates to a local page and extracts the page content via
|
||||
// Sessions.Extract with a {"type": "string"} schema.
|
||||
//
|
||||
// Skipped unless STAGEHAND_INTEGRATION=1 is set and the
|
||||
// OPENAI_* env vars are configured.
|
||||
// Skipped unless OPENAI_* env vars are configured.
|
||||
func TestBrowser_E2E_Extract(t *testing.T) {
|
||||
if os.Getenv("STAGEHAND_INTEGRATION") != "1" {
|
||||
t.Skip("STAGEHAND_INTEGRATION != 1; skipping real-stagehand real-LLM integration test")
|
||||
}
|
||||
apiKey := os.Getenv("OPENAI_API_KEY")
|
||||
baseURL := os.Getenv("OPENAI_BASE_URL")
|
||||
model := os.Getenv("OPENAI_MODEL")
|
||||
|
||||
@@ -79,21 +79,32 @@ var (
|
||||
)
|
||||
|
||||
// SetDefaultFactory installs the production ComponentFactory. The
|
||||
// component package calls this in its init() with `component.New`.
|
||||
// Calling SetDefaultFactory more than once with a non-nil factory is
|
||||
// a no-op after the first call (the first wins) so concurrent
|
||||
// registration is safe. Passing nil clears the factory — tests use
|
||||
// this to assert "no factory registered" error paths.
|
||||
// component package calls this in its init() via
|
||||
// installDefaultRegistryFactory (see below). After Phase 0, the
|
||||
// "first writer wins" guard is REMOVED: SetDefaultFactory now ALWAYS
|
||||
// replaces the active default, regardless of whether one is already
|
||||
// installed. This preserves the existing test-override pattern where
|
||||
// tests save the previous factory, install a stub, and restore on
|
||||
// t.Cleanup. Passing nil clears the factory — tests use this to
|
||||
// assert "no factory registered" error paths.
|
||||
//
|
||||
// Two-layer model:
|
||||
//
|
||||
// - Production: installDefaultRegistryFactory installs a closure
|
||||
// that calls runtime.DefaultRegistry.Lookup on every invocation.
|
||||
// It captures DefaultRegistry by reference, so even if
|
||||
// installDefaultRegistryFactory runs before all init()
|
||||
// registrations complete, the factory is correct at every
|
||||
// subsequent lookup (the registry is read lazily, not captured
|
||||
// at install time).
|
||||
// - Override: tests call SetDefaultFactory(stub) directly to stub
|
||||
// the default factory. t.Cleanup restores the production factory
|
||||
// by calling installDefaultRegistryFactory again (or by saving
|
||||
// the previous value and re-injecting it).
|
||||
func SetDefaultFactory(f ComponentFactory) {
|
||||
factoryMu.Lock()
|
||||
defer factoryMu.Unlock()
|
||||
if f == nil {
|
||||
defaultFactory = nil
|
||||
return
|
||||
}
|
||||
if defaultFactory == nil {
|
||||
defaultFactory = f
|
||||
}
|
||||
defaultFactory = f
|
||||
}
|
||||
|
||||
// DefaultFactory returns the registered ComponentFactory, or nil if
|
||||
@@ -106,6 +117,36 @@ func DefaultFactory() ComponentFactory {
|
||||
return defaultFactory
|
||||
}
|
||||
|
||||
// InstallDefaultRegistryFactory installs the production
|
||||
// ComponentFactory: a closure that resolves component names via
|
||||
// runtime.DefaultRegistry.Lookup at every invocation. The closure
|
||||
// captures DefaultRegistry by reference (the variable, not the
|
||||
// concrete registry), so lookup always reads the current state of
|
||||
// the singleton even if a test later swaps it out via SetDefaultFactory.
|
||||
//
|
||||
// This is the helper the component package's init() calls (from
|
||||
// internal/agent/component/runtime_wire.go). Production callers
|
||||
// should prefer InstallDefaultRegistryFactory over SetDefaultFactory
|
||||
// directly so the wiring stays in one place — if the resolution
|
||||
// strategy changes (e.g. switch to a per-call registry handle),
|
||||
// only this function changes.
|
||||
//
|
||||
// Note: this is EXPORTED (unlike the helper sketched in plan §4
|
||||
// Phase 0 task 3) because the call site lives in a different package
|
||||
// (internal/agent/component) and Go's visibility rules don't allow
|
||||
// access to unexported names across package boundaries. The "test
|
||||
// override layer" still owns SetDefaultFactory — tests that want to
|
||||
// stub the factory call SetDefaultFactory directly, not this helper.
|
||||
func InstallDefaultRegistryFactory() {
|
||||
SetDefaultFactory(func(name string, params map[string]any) (Component, error) {
|
||||
f, _, _, ok := DefaultRegistry.Lookup(name)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("runtime: unknown component %q", name)
|
||||
}
|
||||
return f(name, params)
|
||||
})
|
||||
}
|
||||
|
||||
// ResetDefaultFactoryForTesting clears the registered factory.
|
||||
// Test-only helper for code paths that want to assert behaviour
|
||||
// when no factory is installed. Not safe under concurrent use with
|
||||
|
||||
163
internal/agent/runtime/helpers.go
Normal file
163
internal/agent/runtime/helpers.go
Normal file
@@ -0,0 +1,163 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Cross-cutting helpers that replace Python's `rag/flow/base.py:ProcessBase`
|
||||
// wrapper (lines 33-63). Three call-site concerns are extracted into plain
|
||||
// higher-order functions:
|
||||
//
|
||||
// (a) timeout enforcement -> WithTimeout
|
||||
// (b) progress callback fan-out -> TrackProgress
|
||||
// (c) elapsed-time accounting -> TrackElapsed
|
||||
//
|
||||
// These live in `runtime` (rather than as a `Component` interface method or
|
||||
// a base type) because they are call-site concerns, not extension points.
|
||||
// Both `internal/ingestion/pipeline` and `internal/agent/canvas` compose
|
||||
// them at the DAG-node / goroutine boundary.
|
||||
//
|
||||
// LOSSY MAPPING (plan §8 R1):
|
||||
//
|
||||
// Python `ProcessBase._invoke` is wrapped by BOTH `asyncio.wait_for` AND
|
||||
// the `@timeout` decorator — a dual-layer timeout to catch different
|
||||
// failure modes. Go's `context.WithTimeout` collapses this into a single
|
||||
// layer; `WithTimeout` covers the outer one (asyncio.wait_for equivalent).
|
||||
// The inner `@timeout` decorator has no Go equivalent and is not
|
||||
// replicated here. If a future requirement needs the inner layer,
|
||||
// `WithTimeout` can be nested at the call site.
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ProgressCallback receives progress notifications from TrackProgress.
|
||||
// The numeric progress values follow the convention used by the Python
|
||||
// pipeline canvas callback:
|
||||
//
|
||||
// progress=0 before fn runs (component just started)
|
||||
// progress=1 on success (component finished cleanly)
|
||||
// progress=-1 on failure (component errored; message is the error)
|
||||
//
|
||||
// Concrete sinks (Redis log writer, in-memory test recorder) implement
|
||||
// this signature. nil is a valid value: TrackProgress treats a nil cb as
|
||||
// "no observer" and simply runs fn.
|
||||
type ProgressCallback func(progress int, message string)
|
||||
|
||||
// TrackProgress wraps fn with progress notifications. The callback is
|
||||
// invoked at most twice per call (once at start, once at end).
|
||||
//
|
||||
// On success: cb(1, "<compName> Done") and nil error.
|
||||
// On failure: cb(-1, "<compName>: <err>") and the original error.
|
||||
//
|
||||
// A nil callback is permitted: fn runs to completion and its return
|
||||
// value (including error) is passed through untouched.
|
||||
func TrackProgress(compName string, cb ProgressCallback, fn func() error) error {
|
||||
if cb != nil {
|
||||
cb(0, compName+" Started")
|
||||
}
|
||||
err := fn()
|
||||
if cb == nil {
|
||||
return err
|
||||
}
|
||||
if err != nil {
|
||||
cb(-1, fmt.Sprintf("%s: %s", compName, err.Error()))
|
||||
return err
|
||||
}
|
||||
cb(1, compName+" Done")
|
||||
return nil
|
||||
}
|
||||
|
||||
// WithTimeout runs fn under a derived context that cancels either when
|
||||
// d elapses or when the parent ctx is cancelled (whichever happens
|
||||
// first). fn receives the child context so it can honor cancellation at
|
||||
// its own yield points.
|
||||
//
|
||||
// On timeout: returns context.DeadlineExceeded (matching Python's
|
||||
// asyncio.TimeoutError semantics).
|
||||
// On parent cancellation: returns the parent ctx's error (typically
|
||||
// context.Canceled).
|
||||
// On fn completion within d: returns fn's error (may be nil).
|
||||
//
|
||||
// NOTES:
|
||||
//
|
||||
// - This function implements ONLY the outer timeout layer that
|
||||
// Python `ProcessBase` enforces via `asyncio.wait_for`. The inner
|
||||
// `@timeout` decorator is not replicated in Go (see plan §8 R1).
|
||||
// - fn MUST NOT retain or use the ctx past return; once fn returns
|
||||
// the child context's cancel func is invoked by WithTimeout.
|
||||
func WithTimeout(ctx context.Context, d time.Duration, fn func(ctx context.Context) error) error {
|
||||
childCtx, cancel := context.WithTimeout(ctx, d)
|
||||
defer cancel()
|
||||
if err := fn(childCtx); err != nil {
|
||||
// If fn honored cancellation, prefer the ctx error so callers
|
||||
// see a uniform "timed out" / "canceled" signal regardless of
|
||||
// whether fn propagated the error or replaced it.
|
||||
if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
|
||||
return err
|
||||
}
|
||||
if cerr := childCtx.Err(); cerr != nil {
|
||||
return cerr
|
||||
}
|
||||
return err
|
||||
}
|
||||
// fn returned nil — but the deadline may have elapsed between
|
||||
// fn's last yield point and return. Surface that as
|
||||
// DeadlineExceeded so the caller sees a consistent timeout
|
||||
// signal rather than a false "success".
|
||||
if cerr := childCtx.Err(); cerr != nil {
|
||||
return cerr
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// TrackElapsed records the wall-clock duration of fn and stamps the
|
||||
// output map with two synthetic keys mirroring Python `ProcessBase`
|
||||
// (base.py:42, 58):
|
||||
//
|
||||
// "_created_time" RFC3339Nano-formatted timestamp taken BEFORE fn runs.
|
||||
// "_elapsed_time" float64 seconds (with sub-second precision) that
|
||||
// fn took to complete, in [0, +∞).
|
||||
//
|
||||
// Any keys already present in fn's result map are preserved verbatim;
|
||||
// the two synthetic keys are added only if absent (fn-supplied values
|
||||
// win on conflict — fn is the authoritative source of business data).
|
||||
// This matches the Python ProcessBase convention: a component that
|
||||
// computes its own elapsed time is trusted over the helper's stopwatch.
|
||||
//
|
||||
// On error: the returned map is nil and the error is propagated
|
||||
// untouched. The "name" parameter is recorded in the error message
|
||||
// when err is non-nil so log readers can attribute the elapsed
|
||||
// accounting failure to a specific component.
|
||||
func TrackElapsed(name string, fn func() (map[string]any, error)) (map[string]any, error) {
|
||||
start := time.Now()
|
||||
out, err := fn()
|
||||
elapsed := time.Since(start)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%s: %w", name, err)
|
||||
}
|
||||
if out == nil {
|
||||
out = make(map[string]any)
|
||||
}
|
||||
if _, ok := out["_created_time"]; !ok {
|
||||
out["_created_time"] = start.UTC().Format(time.RFC3339Nano)
|
||||
}
|
||||
if _, ok := out["_elapsed_time"]; !ok {
|
||||
out["_elapsed_time"] = elapsed.Seconds()
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
393
internal/agent/runtime/helpers_test.go
Normal file
393
internal/agent/runtime/helpers_test.go
Normal file
@@ -0,0 +1,393 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// recordingCallback is a thread-safe ProgressCallback recorder used by
|
||||
// the TrackProgress tests. progress/message pairs are appended in
|
||||
// invocation order so tests can assert the exact call sequence.
|
||||
type recordingCallback struct {
|
||||
mu sync.Mutex
|
||||
calls []recordedCall
|
||||
started bool
|
||||
}
|
||||
|
||||
type recordedCall struct {
|
||||
progress int
|
||||
message string
|
||||
}
|
||||
|
||||
func (r *recordingCallback) callback(progress int, message string) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.calls = append(r.calls, recordedCall{progress: progress, message: message})
|
||||
}
|
||||
|
||||
func (r *recordingCallback) callsCopy() []recordedCall {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
out := make([]recordedCall, len(r.calls))
|
||||
copy(out, r.calls)
|
||||
return out
|
||||
}
|
||||
|
||||
// --- TrackProgress ---
|
||||
|
||||
func TestTrackProgress_Success(t *testing.T) {
|
||||
rec := &recordingCallback{}
|
||||
err := TrackProgress("Parser", rec.callback, func() error {
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
calls := rec.callsCopy()
|
||||
if len(calls) != 2 {
|
||||
t.Fatalf("expected 2 callback invocations, got %d: %+v", len(calls), calls)
|
||||
}
|
||||
if calls[0].progress != 0 || calls[0].message != "Parser Started" {
|
||||
t.Errorf("first call = %+v, want progress=0 message=%q", calls[0], "Parser Started")
|
||||
}
|
||||
if calls[1].progress != 1 || calls[1].message != "Parser Done" {
|
||||
t.Errorf("second call = %+v, want progress=1 message=%q", calls[1], "Parser Done")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTrackProgress_Failure(t *testing.T) {
|
||||
rec := &recordingCallback{}
|
||||
wantErr := errors.New("boom")
|
||||
err := TrackProgress("Tokenizer", rec.callback, func() error {
|
||||
return wantErr
|
||||
})
|
||||
if !errors.Is(err, wantErr) {
|
||||
t.Fatalf("expected error %v, got %v", wantErr, err)
|
||||
}
|
||||
calls := rec.callsCopy()
|
||||
if len(calls) != 2 {
|
||||
t.Fatalf("expected 2 callback invocations, got %d", len(calls))
|
||||
}
|
||||
if calls[0].progress != 0 || calls[0].message != "Tokenizer Started" {
|
||||
t.Errorf("first call = %+v, want progress=0 message=%q", calls[0], "Tokenizer Started")
|
||||
}
|
||||
if calls[1].progress != -1 {
|
||||
t.Errorf("second call progress = %d, want -1", calls[1].progress)
|
||||
}
|
||||
if !strings.Contains(calls[1].message, "Tokenizer") || !strings.Contains(calls[1].message, "boom") {
|
||||
t.Errorf("second call message = %q, want it to contain both %q and %q", calls[1].message, "Tokenizer", "boom")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTrackProgress_NilCallback(t *testing.T) {
|
||||
// Must not panic with a nil callback; must still pass fn's result through.
|
||||
called := false
|
||||
if err := TrackProgress("File", nil, func() error {
|
||||
called = true
|
||||
return nil
|
||||
}); err != nil {
|
||||
t.Fatalf("unexpected error from nil-cb success path: %v", err)
|
||||
}
|
||||
if !called {
|
||||
t.Fatal("fn was not invoked")
|
||||
}
|
||||
|
||||
wantErr := errors.New("nil-cb err")
|
||||
got := TrackProgress("File", nil, func() error {
|
||||
return wantErr
|
||||
})
|
||||
if !errors.Is(got, wantErr) {
|
||||
t.Fatalf("nil-cb failure path: got %v, want %v", got, wantErr)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTrackProgress_PassesThroughReturnValue covers the documented contract
|
||||
// that the error returned to the caller is fn's error verbatim (wrapped
|
||||
// only by the message-formatting for the callback, not for the return).
|
||||
func TestTrackProgress_PassesThroughReturnValue(t *testing.T) {
|
||||
rec := &recordingCallback{}
|
||||
|
||||
// nil path
|
||||
if err := TrackProgress("Foo", rec.callback, func() error { return nil }); err != nil {
|
||||
t.Fatalf("nil error not propagated as nil: %v", err)
|
||||
}
|
||||
|
||||
// err path — exact identity preserved
|
||||
want := errors.New("exact")
|
||||
got := TrackProgress("Foo", rec.callback, func() error { return want })
|
||||
if got != want {
|
||||
t.Fatalf("err not propagated by identity: got %v (%T), want %v (%T)", got, got, want, want)
|
||||
}
|
||||
|
||||
// cb saw the failure with progress=-1
|
||||
var last recordedCall
|
||||
for _, c := range rec.callsCopy() {
|
||||
last = c
|
||||
}
|
||||
if last.progress != -1 {
|
||||
t.Errorf("final cb call progress = %d, want -1", last.progress)
|
||||
}
|
||||
}
|
||||
|
||||
// --- WithTimeout ---
|
||||
|
||||
func TestWithTimeout_Success(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
err := WithTimeout(ctx, 50*time.Millisecond, func(ctx context.Context) error {
|
||||
// simulate fast work
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("expected nil error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithTimeout_Timeout(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
start := time.Now()
|
||||
err := WithTimeout(ctx, 20*time.Millisecond, func(ctx context.Context) error {
|
||||
// sleep long enough to outlast the timeout; honor ctx so the
|
||||
// test doesn't have to wait the full duration.
|
||||
select {
|
||||
case <-time.After(500 * time.Millisecond):
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
})
|
||||
elapsed := time.Since(start)
|
||||
if !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("expected context.DeadlineExceeded, got %v", err)
|
||||
}
|
||||
if elapsed > 250*time.Millisecond {
|
||||
t.Errorf("WithTimeout waited too long after deadline (%s) — fn should have observed ctx.Done() quickly", elapsed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithTimeout_ParentCancellation(t *testing.T) {
|
||||
parent, cancel := context.WithCancel(context.Background())
|
||||
go func() {
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
cancel()
|
||||
}()
|
||||
|
||||
observed := make(chan error, 1)
|
||||
start := time.Now()
|
||||
err := WithTimeout(parent, 5*time.Second, func(ctx context.Context) error {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
observed <- ctx.Err()
|
||||
return ctx.Err()
|
||||
case <-time.After(2 * time.Second):
|
||||
observed <- nil
|
||||
return nil
|
||||
}
|
||||
})
|
||||
elapsed := time.Since(start)
|
||||
|
||||
if !errors.Is(err, context.Canceled) {
|
||||
t.Fatalf("expected context.Canceled, got %v", err)
|
||||
}
|
||||
select {
|
||||
case inner := <-observed:
|
||||
if !errors.Is(inner, context.Canceled) {
|
||||
t.Errorf("fn observed ctx.Err() = %v, want context.Canceled", inner)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("fn never observed ctx.Done()")
|
||||
}
|
||||
if elapsed > 250*time.Millisecond {
|
||||
t.Errorf("WithTimeout took %s after parent cancel — expected fast exit", elapsed)
|
||||
}
|
||||
}
|
||||
|
||||
// TestWithTimeout_PassesContextToFn verifies the ctx fn receives is a
|
||||
// CHILD of the parent (not the parent itself). The child should carry
|
||||
// the parent's Values but have its own Done channel tied to the
|
||||
// timeout deadline. We probe captured properties from inside fn
|
||||
// (NOT after WithTimeout returns) because WithTimeout's deferred
|
||||
// cancel() will mark the child ctx as canceled once it returns —
|
||||
// which is the documented contract of context.WithTimeout, not a
|
||||
// helper bug.
|
||||
func TestWithTimeout_PassesContextToFn(t *testing.T) {
|
||||
type ctxKey struct{}
|
||||
parent := context.WithValue(context.Background(), ctxKey{}, "v")
|
||||
type captured struct {
|
||||
ctx context.Context
|
||||
errInFlight error
|
||||
hasDeadline bool
|
||||
deadline time.Time
|
||||
}
|
||||
var cap captured
|
||||
err := WithTimeout(parent, 100*time.Millisecond, func(ctx context.Context) error {
|
||||
cap.ctx = ctx
|
||||
cap.errInFlight = ctx.Err()
|
||||
cap.deadline, cap.hasDeadline = ctx.Deadline()
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if cap.ctx == nil {
|
||||
t.Fatal("fn did not receive a context")
|
||||
}
|
||||
if cap.ctx == parent {
|
||||
t.Fatal("fn received the parent ctx directly — expected a derived child ctx")
|
||||
}
|
||||
if v, _ := cap.ctx.Value(ctxKey{}).(string); v != "v" {
|
||||
t.Errorf("child ctx did not carry parent's Value(): got %q, want %q", v, "v")
|
||||
}
|
||||
if cap.errInFlight != nil {
|
||||
t.Errorf("child ctx should not be done while fn is still running successfully, got Err=%v", cap.errInFlight)
|
||||
}
|
||||
if !cap.hasDeadline {
|
||||
t.Error("child ctx has no Deadline — expected one from WithTimeout")
|
||||
}
|
||||
if !time.Now().Before(cap.deadline) {
|
||||
t.Errorf("child ctx deadline %v is in the past", cap.deadline)
|
||||
}
|
||||
}
|
||||
|
||||
// --- TrackElapsed ---
|
||||
|
||||
func TestTrackElapsed_AddsCreatedAndElapsedFields(t *testing.T) {
|
||||
got, err := TrackElapsed("Parser", func() (map[string]any, error) {
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
return map[string]any{"chunks": 3}, nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if _, ok := got["_created_time"]; !ok {
|
||||
t.Fatal("result missing _created_time")
|
||||
}
|
||||
ct, ok := got["_created_time"].(string)
|
||||
if !ok || ct == "" {
|
||||
t.Fatalf("_created_time = %v (type %T), want non-empty string", got["_created_time"], got["_created_time"])
|
||||
}
|
||||
if _, err := time.Parse(time.RFC3339Nano, ct); err != nil {
|
||||
t.Errorf("_created_time %q is not RFC3339Nano: %v", ct, err)
|
||||
}
|
||||
elapsed, ok := got["_elapsed_time"].(float64)
|
||||
if !ok {
|
||||
t.Fatalf("_elapsed_time = %v (type %T), want float64", got["_elapsed_time"], got["_elapsed_time"])
|
||||
}
|
||||
if elapsed < 0 {
|
||||
t.Errorf("_elapsed_time = %f, want >= 0", elapsed)
|
||||
}
|
||||
// We slept 5ms; elapsed should be in a reasonable range (loose bound
|
||||
// to keep the test stable on noisy CI runners).
|
||||
if elapsed < 0.001 {
|
||||
t.Errorf("_elapsed_time = %f, expected >= ~0.005 after 5ms sleep", elapsed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTrackElapsed_PreservesExistingKeys(t *testing.T) {
|
||||
in := map[string]any{"x": 1, "name": "kept"}
|
||||
got, err := TrackElapsed("Tokenizer", func() (map[string]any, error) {
|
||||
return in, nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got["x"] != 1 {
|
||||
t.Errorf("existing key x = %v, want 1", got["x"])
|
||||
}
|
||||
if got["name"] != "kept" {
|
||||
t.Errorf("existing key name = %v, want %q", got["name"], "kept")
|
||||
}
|
||||
if _, ok := got["_created_time"]; !ok {
|
||||
t.Error("missing _created_time")
|
||||
}
|
||||
if _, ok := got["_elapsed_time"]; !ok {
|
||||
t.Error("missing _elapsed_time")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTrackElapsed_PropagatesError(t *testing.T) {
|
||||
want := errors.New("downstream boom")
|
||||
got, err := TrackElapsed("Extractor", func() (map[string]any, error) {
|
||||
return map[string]any{"partial": true}, want
|
||||
})
|
||||
if !errors.Is(err, want) {
|
||||
t.Fatalf("err = %v, want wraps %v", err, want)
|
||||
}
|
||||
if got != nil {
|
||||
t.Errorf("result map = %+v, want nil when fn errors", got)
|
||||
}
|
||||
// name parameter captured in the error message (documented
|
||||
// in the TrackElapsed package doc: on error, `name` is
|
||||
// recorded in the error message so log readers can attribute
|
||||
// the failure to a specific component).
|
||||
if !strings.Contains(err.Error(), "Extractor") {
|
||||
t.Errorf("err message %q should mention the component name %q", err.Error(), "Extractor")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTrackElapsed_NameParameterRecorded verifies that `name` appears
|
||||
// somewhere observable — we chose to surface it in the error message
|
||||
// on failure (see TrackElapsed doc).
|
||||
func TestTrackElapsed_NameParameterRecorded(t *testing.T) {
|
||||
// On failure path: name is in the error message.
|
||||
_, err := TrackElapsed("MyComp", func() (map[string]any, error) {
|
||||
return nil, errors.New("nope")
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), "MyComp") {
|
||||
t.Fatalf("name not recorded on error path: err=%v", err)
|
||||
}
|
||||
|
||||
// On success path: name is not part of the output map (per the
|
||||
// chosen design — name appears in error messages only). We
|
||||
// document this here so future maintainers don't expect it in
|
||||
// the success map.
|
||||
out, err := TrackElapsed("MyComp", func() (map[string]any, error) {
|
||||
return map[string]any{}, nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error on success path: %v", err)
|
||||
}
|
||||
for k := range out {
|
||||
if strings.Contains(k, "MyComp") {
|
||||
t.Errorf("success-path map contains key %q referencing component name; name should appear in error messages only", k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestTrackElapsed_NilMapFromFn covers the edge case where fn returns
|
||||
// (nil, nil) — TrackElapsed must still populate the bookkeeping keys
|
||||
// without panicking on the nil-map write.
|
||||
func TestTrackElapsed_NilMapFromFn(t *testing.T) {
|
||||
got, err := TrackElapsed("X", func() (map[string]any, error) {
|
||||
return nil, nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if _, ok := got["_created_time"]; !ok {
|
||||
t.Error("missing _created_time after nil-map input")
|
||||
}
|
||||
if _, ok := got["_elapsed_time"]; !ok {
|
||||
t.Error("missing _elapsed_time after nil-map input")
|
||||
}
|
||||
}
|
||||
206
internal/agent/runtime/registry.go
Normal file
206
internal/agent/runtime/registry.go
Normal file
@@ -0,0 +1,206 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// runtime — category-aware component registry.
|
||||
//
|
||||
// Phase 0 of plan port-rag-flow-pipeline-to-go.md lifts the component
|
||||
// registry out of internal/agent/component into the runtime package
|
||||
// so the ingestion pipeline (Phase 2) can register under
|
||||
// CategoryIngestion without depending on the agent canvas. The legacy
|
||||
// component.Register / component.New / component.RegisteredNames become
|
||||
// thin adapters that delegate here.
|
||||
//
|
||||
// The single source of truth is DefaultRegistry. The component package's
|
||||
// internal `registry` map has been removed — keeping two maps in sync
|
||||
// during the transition would cause RegisteredNames() to return a
|
||||
// partial set (the internal map only sees legacy Register calls; the
|
||||
// new map only sees RegisterWithMeta calls). See plan §5.
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Category tags each registered component with its domain so the UI can
|
||||
// filter and the runtime can audit cross-domain wiring.
|
||||
type Category string
|
||||
|
||||
const (
|
||||
CategoryAgent Category = "agent"
|
||||
CategoryIngestion Category = "ingestion"
|
||||
CategoryShared Category = "shared"
|
||||
)
|
||||
|
||||
// Metadata is the static component descriptor exposed to the API.
|
||||
// Components MUST provide this at registration time so the API can serve
|
||||
// a complete component catalog (name, category, inputs, outputs) without
|
||||
// having to instantiate the component.
|
||||
//
|
||||
// The plan §4 Phase 0 task 1 contract says Register REJECTS empty
|
||||
// metadata. The "is empty" check is implemented in Register as
|
||||
// "Version == \"\" AND Inputs == nil AND Outputs == nil"; this
|
||||
// three-way check lets a caller fill in any single field as
|
||||
// evidence of intent. The Version field is the canonical marker
|
||||
// (plan §4 Phase 0 task 2 uses `Metadata{Version: "legacy"}` for
|
||||
// the legacy-adapter shim), but a component that wants to skip
|
||||
// the version stamp but still record inputs/outputs is also
|
||||
// allowed.
|
||||
type Metadata struct {
|
||||
Version string // contract version; required for ingestion, "legacy" for the legacy adapter
|
||||
Inputs map[string]string // input key → human-readable description
|
||||
Outputs map[string]string // output key → human-readable description
|
||||
}
|
||||
|
||||
// entry is one slot in the registry.
|
||||
type entry struct {
|
||||
factory ComponentFactory
|
||||
category Category
|
||||
metadata Metadata
|
||||
}
|
||||
|
||||
// Registry is the process-wide collection of named ComponentFactories,
|
||||
// tagged with Category so callers can enumerate by domain. Each registration
|
||||
// also carries static Metadata (Inputs/Outputs) consumed by the API layer.
|
||||
type Registry interface {
|
||||
Register(name string, category Category, factory ComponentFactory, metadata Metadata) error
|
||||
Lookup(name string) (ComponentFactory, Category, Metadata, bool)
|
||||
NamesByCategory(category Category) []string
|
||||
Names() []string
|
||||
}
|
||||
|
||||
// memoryRegistry is the production Registry. It is concurrency-safe; init()
|
||||
// race-to-register is acceptable for `Register` — the duplicate-registration
|
||||
// check rejects the second writer, so the FIRST successful registration
|
||||
// for a given name wins. (This is independent of the `SetDefaultFactory`
|
||||
// "first-wins" guard, which is REMOVED in Phase 0 task 3; see
|
||||
// "two-layer model" comment there for the distinction.)
|
||||
//
|
||||
// Lookup is **case-insensitive**: keys are lowercased on Register and Lookup.
|
||||
// This matches internal/agent/component/registry.go:28, 43, which lowercases
|
||||
// at both ends. A case-sensitive implementation would silently fail canvas
|
||||
// build with "unknown component" errors when an existing init() registers
|
||||
// "ExampleComponent" but the canvas looks up "examplecomponent" (or vice
|
||||
// versa).
|
||||
type memoryRegistry struct {
|
||||
mu sync.RWMutex
|
||||
entries map[string]entry
|
||||
}
|
||||
|
||||
// Register enrolls a ComponentFactory under name (case-insensitive).
|
||||
// Returns an error on empty name, empty metadata, or duplicate key —
|
||||
// callers who want the legacy panic-on-duplicate semantics at init()
|
||||
// time should wrap the call with MustRegister.
|
||||
//
|
||||
// "Empty metadata" (plan §4 Phase 0 task 1) means all three of
|
||||
// Version, Inputs, Outputs are unset. The Version field is the
|
||||
// canonical marker — ingestion components MUST supply a real
|
||||
// version string; the legacy agent-component adapter MUST supply
|
||||
// "legacy"; the catalog endpoint only serves components whose
|
||||
// metadata passes this check.
|
||||
func (r *memoryRegistry) Register(name string, category Category, factory ComponentFactory, metadata Metadata) error {
|
||||
key := strings.ToLower(strings.TrimSpace(name))
|
||||
if key == "" {
|
||||
return fmt.Errorf("runtime: Register called with empty name")
|
||||
}
|
||||
if metadata.Version == "" && metadata.Inputs == nil && metadata.Outputs == nil {
|
||||
return fmt.Errorf("runtime: %q registered with empty metadata (Version, Inputs, Outputs all unset; "+
|
||||
"see plan §4 Phase 0 task 1 — ingestion components MUST supply a version string, "+
|
||||
"legacy agent-component adapter MUST supply {Version: \"legacy\"})", name)
|
||||
}
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if _, exists := r.entries[key]; exists {
|
||||
return fmt.Errorf("runtime: %q already registered", name)
|
||||
}
|
||||
r.entries[key] = entry{factory: factory, category: category, metadata: metadata}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Lookup resolves a name (case-insensitive) to its factory + category +
|
||||
// metadata. Returns ok=false on miss.
|
||||
func (r *memoryRegistry) Lookup(name string) (ComponentFactory, Category, Metadata, bool) {
|
||||
key := strings.ToLower(strings.TrimSpace(name))
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
e, ok := r.entries[key]
|
||||
if !ok {
|
||||
return nil, "", Metadata{}, false
|
||||
}
|
||||
return e.factory, e.category, e.metadata, true
|
||||
}
|
||||
|
||||
// NamesByCategory returns the sorted list of names registered under the
|
||||
// given category. Sorted output keeps UI listings and error messages
|
||||
// stable.
|
||||
func (r *memoryRegistry) NamesByCategory(category Category) []string {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
out := make([]string, 0, len(r.entries))
|
||||
for n, e := range r.entries {
|
||||
if e.category == category {
|
||||
out = append(out, n)
|
||||
}
|
||||
}
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// Names returns the sorted list of all registered names.
|
||||
func (r *memoryRegistry) Names() []string {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
out := make([]string, 0, len(r.entries))
|
||||
for n := range r.entries {
|
||||
out = append(out, n)
|
||||
}
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// NewMemoryRegistry constructs an empty in-memory Registry. Tests use
|
||||
// this to spin up isolated registries; production code uses
|
||||
// DefaultRegistry.
|
||||
func NewMemoryRegistry() Registry {
|
||||
return &memoryRegistry{entries: make(map[string]entry)}
|
||||
}
|
||||
|
||||
// DefaultRegistry is the process-wide singleton. Each component package's
|
||||
// init() registers its factories here. Lookup is lazy: even if a canvas
|
||||
// build occurs before every package's init() has run, lookups see all
|
||||
// completed registrations because the registry is read at call time, not
|
||||
// at SetDefaultFactory time.
|
||||
var DefaultRegistry Registry = NewMemoryRegistry()
|
||||
|
||||
// MustRegister wraps Register and panics on error. Init()-time callers
|
||||
// that want the legacy "panic on duplicate" behaviour can use this
|
||||
// instead of Register + manual error check.
|
||||
func MustRegister(name string, category Category, factory ComponentFactory, metadata Metadata) {
|
||||
if err := DefaultRegistry.Register(name, category, factory, metadata); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
|
||||
// RegisterWithMeta is a thin convenience around DefaultRegistry.Register
|
||||
// for callers that want explicit (name, category, factory, metadata) at
|
||||
// the call site without going through MustRegister. It returns the same
|
||||
// error as Register; callers that want panic-on-error should use
|
||||
// MustRegister instead.
|
||||
func RegisterWithMeta(name string, category Category, factory ComponentFactory, metadata Metadata) error {
|
||||
return DefaultRegistry.Register(name, category, factory, metadata)
|
||||
}
|
||||
357
internal/agent/runtime/registry_test.go
Normal file
357
internal/agent/runtime/registry_test.go
Normal file
@@ -0,0 +1,357 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// stubComponent is a minimal Component impl used as the factory's
|
||||
// return value in tests. It echoes the params back into the output map
|
||||
// so a test can assert what the factory actually received.
|
||||
type stubComponent struct {
|
||||
name string
|
||||
params map[string]any
|
||||
}
|
||||
|
||||
func (s *stubComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
return map[string]any{"name": s.name, "params": s.params}, nil
|
||||
}
|
||||
|
||||
func stubFactory(name string, params map[string]any) (Component, error) {
|
||||
return &stubComponent{name: name, params: params}, nil
|
||||
}
|
||||
|
||||
// errFactory returns a fixed error so a test can assert that factory
|
||||
// errors propagate via Lookup.
|
||||
func errFactory(err error) ComponentFactory {
|
||||
return func(name string, params map[string]any) (Component, error) {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_RegisterAndLookup_HappyPath(t *testing.T) {
|
||||
r := NewMemoryRegistry()
|
||||
meta := Metadata{
|
||||
Inputs: map[string]string{"x": "input x"},
|
||||
Outputs: map[string]string{"y": "output y"},
|
||||
}
|
||||
if err := r.Register("Foo", CategoryAgent, stubFactory, meta); err != nil {
|
||||
t.Fatalf("Register(Foo) returned error: %v", err)
|
||||
}
|
||||
|
||||
f, cat, gotMeta, ok := r.Lookup("Foo")
|
||||
if !ok {
|
||||
t.Fatalf("Lookup(Foo) missed")
|
||||
}
|
||||
if cat != CategoryAgent {
|
||||
t.Errorf("Lookup category = %q, want %q", cat, CategoryAgent)
|
||||
}
|
||||
if gotMeta.Inputs["x"] != "input x" || gotMeta.Outputs["y"] != "output y" {
|
||||
t.Errorf("Lookup metadata lost: got %+v", gotMeta)
|
||||
}
|
||||
|
||||
c, err := f("Foo", map[string]any{"k": "v"})
|
||||
if err != nil {
|
||||
t.Fatalf("factory returned error: %v", err)
|
||||
}
|
||||
if _, ok := c.(*stubComponent); !ok {
|
||||
t.Errorf("factory returned wrong type %T", c)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_Lookup_CaseInsensitive(t *testing.T) {
|
||||
r := NewMemoryRegistry()
|
||||
if err := r.Register("ExampleComponent", CategoryShared, stubFactory, Metadata{Version: "legacy"}); err != nil {
|
||||
t.Fatalf("Register: %v", err)
|
||||
}
|
||||
|
||||
for _, variant := range []string{"ExampleComponent", "examplecomponent", "EXAMPLECOMPONENT", " examplecomponent "} {
|
||||
if _, _, _, ok := r.Lookup(variant); !ok {
|
||||
t.Errorf("Lookup(%q) missed; case-insensitive lookup must succeed for all variants", variant)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_Register_DuplicateReturnsError(t *testing.T) {
|
||||
r := NewMemoryRegistry()
|
||||
if err := r.Register("Dup", CategoryAgent, stubFactory, Metadata{Version: "legacy"}); err != nil {
|
||||
t.Fatalf("first Register(Dup) returned error: %v", err)
|
||||
}
|
||||
err := r.Register("Dup", CategoryIngestion, stubFactory, Metadata{Version: "legacy"})
|
||||
if err == nil {
|
||||
t.Fatalf("second Register(Dup) succeeded; expected duplicate-key error")
|
||||
}
|
||||
// Duplicate detection is also case-insensitive.
|
||||
err2 := r.Register("DUP", CategoryIngestion, stubFactory, Metadata{Version: "legacy"})
|
||||
if err2 == nil {
|
||||
t.Fatalf("Register(DUP) succeeded; duplicate detection must be case-insensitive")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_Register_EmptyNameReturnsError(t *testing.T) {
|
||||
r := NewMemoryRegistry()
|
||||
if err := r.Register("", CategoryAgent, stubFactory, Metadata{Version: "legacy"}); err == nil {
|
||||
t.Fatalf("Register(\"\") succeeded; expected empty-name error")
|
||||
}
|
||||
if err := r.Register(" ", CategoryAgent, stubFactory, Metadata{Version: "legacy"}); err == nil {
|
||||
t.Fatalf("Register(\" \") succeeded; expected empty-name error after trim")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_Register_EmptyMetadataReturnsError verifies the plan
|
||||
// §4 Phase 0 task 1 contract: Register rejects empty metadata
|
||||
// (Version, Inputs, Outputs all unset). Ingestion components MUST
|
||||
// supply a Version string; the legacy adapter shim stamps
|
||||
// {Version: "legacy"}; a single-field fill is also allowed (e.g.,
|
||||
// only Inputs, or only Version).
|
||||
func TestRegistry_Register_EmptyMetadataReturnsError(t *testing.T) {
|
||||
r := NewMemoryRegistry()
|
||||
// All three fields unset → empty-metadata error.
|
||||
err := r.Register("EmptyMeta", CategoryIngestion, stubFactory, Metadata{})
|
||||
if err == nil {
|
||||
t.Fatalf("Register with empty metadata succeeded; expected empty-metadata error")
|
||||
}
|
||||
// Verify it was not actually registered (re-register with valid
|
||||
// metadata should succeed).
|
||||
if err := r.Register("EmptyMeta", CategoryIngestion, stubFactory, Metadata{Version: "1.0.0"}); err != nil {
|
||||
t.Fatalf("Register with valid metadata after empty-metadata rejection failed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_Register_AcceptsPartialMetadata verifies the
|
||||
// three-way empty check: filling ANY one of Version / Inputs /
|
||||
// Outputs is enough to register successfully. This accommodates
|
||||
// the migration path where a component has only inputs but no
|
||||
// outputs yet, or where a shim stamps {Version: "legacy"}.
|
||||
func TestRegistry_Register_AcceptsPartialMetadata(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
meta Metadata
|
||||
}{
|
||||
{"OnlyVersion", Metadata{Version: "1.0.0"}},
|
||||
{"OnlyInputs", Metadata{Inputs: map[string]string{"x": "x"}}},
|
||||
{"OnlyOutputs", Metadata{Outputs: map[string]string{"y": "y"}}},
|
||||
{"LegacyAdapter", Metadata{Version: "legacy"}},
|
||||
{"Full", Metadata{Version: "1.0.0", Inputs: map[string]string{"x": "x"}, Outputs: map[string]string{"y": "y"}}},
|
||||
}
|
||||
for i, c := range cases {
|
||||
r := NewMemoryRegistry()
|
||||
name := c.name
|
||||
if err := r.Register(name, CategoryIngestion, stubFactory, c.meta); err != nil {
|
||||
t.Errorf("case %d (%s): Register returned error: %v", i, c.name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_MustRegister_PanicsOnDuplicate(t *testing.T) {
|
||||
// MustRegister operates on DefaultRegistry. Save and restore so this
|
||||
// test does not pollute the global.
|
||||
saved := DefaultRegistry
|
||||
defer func() { DefaultRegistry = saved }()
|
||||
DefaultRegistry = NewMemoryRegistry()
|
||||
|
||||
MustRegister("Panic", CategoryAgent, stubFactory, Metadata{Version: "legacy"})
|
||||
|
||||
defer func() {
|
||||
if r := recover(); r == nil {
|
||||
t.Errorf("MustRegister on duplicate did not panic")
|
||||
}
|
||||
}()
|
||||
MustRegister("Panic", CategoryAgent, stubFactory, Metadata{Version: "legacy"})
|
||||
}
|
||||
|
||||
func TestRegistry_Lookup_MissReturnsFalse(t *testing.T) {
|
||||
r := NewMemoryRegistry()
|
||||
f, cat, meta, ok := r.Lookup("NotThere")
|
||||
if ok {
|
||||
t.Errorf("Lookup on empty registry returned ok=true (f=%v cat=%q meta=%+v)", f, cat, meta)
|
||||
}
|
||||
if f != nil {
|
||||
t.Errorf("Lookup miss: factory should be nil, got %v", f)
|
||||
}
|
||||
if cat != "" {
|
||||
t.Errorf("Lookup miss: category should be empty, got %q", cat)
|
||||
}
|
||||
if meta.Inputs != nil || meta.Outputs != nil {
|
||||
t.Errorf("Lookup miss: metadata should be zero-value, got %+v", meta)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_NamesByCategory_FiltersCorrectly(t *testing.T) {
|
||||
r := NewMemoryRegistry()
|
||||
if err := r.Register("AgentComp", CategoryAgent, stubFactory, Metadata{Version: "legacy"}); err != nil {
|
||||
t.Fatalf("Register AgentComp: %v", err)
|
||||
}
|
||||
if err := r.Register("IngestComp", CategoryIngestion, stubFactory, Metadata{Version: "legacy"}); err != nil {
|
||||
t.Fatalf("Register IngestComp: %v", err)
|
||||
}
|
||||
if err := r.Register("SharedComp", CategoryShared, stubFactory, Metadata{Version: "legacy"}); err != nil {
|
||||
t.Fatalf("Register SharedComp: %v", err)
|
||||
}
|
||||
|
||||
gotAgent := r.NamesByCategory(CategoryAgent)
|
||||
wantAgent := []string{"agentcomp"}
|
||||
if !equalSlices(gotAgent, wantAgent) {
|
||||
t.Errorf("NamesByCategory(CategoryAgent) = %v, want %v", gotAgent, wantAgent)
|
||||
}
|
||||
|
||||
gotIngest := r.NamesByCategory(CategoryIngestion)
|
||||
wantIngest := []string{"ingestcomp"}
|
||||
if !equalSlices(gotIngest, wantIngest) {
|
||||
t.Errorf("NamesByCategory(CategoryIngestion) = %v, want %v", gotIngest, wantIngest)
|
||||
}
|
||||
|
||||
gotShared := r.NamesByCategory(CategoryShared)
|
||||
wantShared := []string{"sharedcomp"}
|
||||
if !equalSlices(gotShared, wantShared) {
|
||||
t.Errorf("NamesByCategory(CategoryShared) = %v, want %v", gotShared, wantShared)
|
||||
}
|
||||
|
||||
gotUnknown := r.NamesByCategory(Category("nonexistent"))
|
||||
if len(gotUnknown) != 0 {
|
||||
t.Errorf("NamesByCategory(unknown) = %v, want empty", gotUnknown)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_Names_ReturnsAllSorted(t *testing.T) {
|
||||
r := NewMemoryRegistry()
|
||||
for _, n := range []string{"Charlie", "alpha", "BRAVO"} {
|
||||
if err := r.Register(n, CategoryAgent, stubFactory, Metadata{Version: "legacy"}); err != nil {
|
||||
t.Fatalf("Register %s: %v", n, err)
|
||||
}
|
||||
}
|
||||
got := r.Names()
|
||||
// Keys are normalized to lowercase at registration time, so the
|
||||
// returned list is ["alpha", "bravo", "charlie"] (sorted).
|
||||
want := []string{"alpha", "bravo", "charlie"}
|
||||
if !equalSlices(got, want) {
|
||||
t.Errorf("Names() = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_NamesByCategory_ReturnsSorted(t *testing.T) {
|
||||
r := NewMemoryRegistry()
|
||||
for _, n := range []string{"Zulu", "alpha", "mike", "BRAVO"} {
|
||||
if err := r.Register(n, CategoryIngestion, stubFactory, Metadata{Version: "legacy"}); err != nil {
|
||||
t.Fatalf("Register %s: %v", n, err)
|
||||
}
|
||||
}
|
||||
got := r.NamesByCategory(CategoryIngestion)
|
||||
want := []string{"alpha", "bravo", "mike", "zulu"}
|
||||
if !equalSlices(got, want) {
|
||||
t.Errorf("NamesByCategory(Ingestion) = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_ThreadSafe(t *testing.T) {
|
||||
// Concurrency smoke test: N goroutines each register a distinct
|
||||
// name; after all join, Names() must contain every one. A
|
||||
// non-thread-safe map would lose entries or panic under -race.
|
||||
saved := DefaultRegistry
|
||||
defer func() { DefaultRegistry = saved }()
|
||||
r := NewMemoryRegistry()
|
||||
DefaultRegistry = r
|
||||
|
||||
const N = 64
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(N)
|
||||
for i := 0; i < N; i++ {
|
||||
i := i
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
name := fmt.Sprintf("comp-%03d", i)
|
||||
if err := r.Register(name, CategoryIngestion, stubFactory, Metadata{Version: "legacy"}); err != nil {
|
||||
t.Errorf("Register(%s) returned error: %v", name, err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
got := r.NamesByCategory(CategoryIngestion)
|
||||
if len(got) != N {
|
||||
t.Errorf("NamesByCategory(Ingestion) returned %d names; expected %d", len(got), N)
|
||||
}
|
||||
// And a parallel Lookup burst against the same registry must
|
||||
// observe all N entries.
|
||||
var lookupWg sync.WaitGroup
|
||||
for i := 0; i < N; i++ {
|
||||
i := i
|
||||
lookupWg.Add(1)
|
||||
go func() {
|
||||
defer lookupWg.Done()
|
||||
name := fmt.Sprintf("comp-%03d", i)
|
||||
if _, _, _, ok := r.Lookup(name); !ok {
|
||||
t.Errorf("Lookup(%s) missed after concurrent registration", name)
|
||||
}
|
||||
}()
|
||||
}
|
||||
lookupWg.Wait()
|
||||
}
|
||||
|
||||
func TestRegistry_FactoryErrorPropagates(t *testing.T) {
|
||||
// A factory that returns an error must propagate that error through
|
||||
// the Lookup → invoke path. This is not directly tested by the
|
||||
// plan checklist but it confirms the factory closure contract.
|
||||
r := NewMemoryRegistry()
|
||||
wantErr := errors.New("boom")
|
||||
if err := r.Register("Bad", CategoryAgent, errFactory(wantErr), Metadata{Version: "legacy"}); err != nil {
|
||||
t.Fatalf("Register: %v", err)
|
||||
}
|
||||
f, _, _, ok := r.Lookup("Bad")
|
||||
if !ok {
|
||||
t.Fatalf("Lookup missed")
|
||||
}
|
||||
_, err := f("Bad", nil)
|
||||
if !errors.Is(err, wantErr) {
|
||||
t.Errorf("factory error not propagated: got %v, want %v", err, wantErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultRegistry_Present(t *testing.T) {
|
||||
if DefaultRegistry == nil {
|
||||
t.Fatal("DefaultRegistry is nil")
|
||||
}
|
||||
// Names() must work without panicking even on the empty default.
|
||||
if got := DefaultRegistry.Names(); got == nil {
|
||||
t.Errorf("Names() returned nil; want non-nil slice")
|
||||
}
|
||||
}
|
||||
|
||||
// equalSlices compares two []string for unordered (no — wait, we DO want
|
||||
// order-sensitive comparison; Names() and NamesByCategory() guarantee
|
||||
// sorted output).
|
||||
func equalSlices(a, b []string) bool {
|
||||
if len(a) != len(b) {
|
||||
return false
|
||||
}
|
||||
for i := range a {
|
||||
if a[i] != b[i] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// extraSort helper kept here in case a future test wants sorted
|
||||
// comparison without imposing a specific ordering.
|
||||
var _ = sort.Strings
|
||||
@@ -1,3 +1,5 @@
|
||||
//go:build cgo
|
||||
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
|
||||
59
internal/binding/rag_analyzer_nocgo.go
Normal file
59
internal/binding/rag_analyzer_nocgo.go
Normal file
@@ -0,0 +1,59 @@
|
||||
//go:build !cgo
|
||||
|
||||
package rag_analyzer
|
||||
|
||||
import "fmt"
|
||||
|
||||
// Token represents a single token from the analyzer.
|
||||
type Token struct {
|
||||
Text string
|
||||
Offset uint32
|
||||
EndOffset uint32
|
||||
}
|
||||
|
||||
// TokenWithPosition represents a token with position information.
|
||||
type TokenWithPosition struct {
|
||||
Text string
|
||||
Offset uint32
|
||||
EndOffset uint32
|
||||
}
|
||||
|
||||
// Analyzer is the no-CGO stub used for builds that intentionally skip the
|
||||
// native tokenizer binding.
|
||||
type Analyzer struct{}
|
||||
|
||||
func NewAnalyzer(path string) (*Analyzer, error) {
|
||||
return nil, fmt.Errorf("rag_analyzer: cgo required (path=%q)", path)
|
||||
}
|
||||
|
||||
func (a *Analyzer) Load() error {
|
||||
return fmt.Errorf("rag_analyzer: cgo required")
|
||||
}
|
||||
|
||||
func (a *Analyzer) SetFineGrained(bool) {}
|
||||
|
||||
func (a *Analyzer) SetEnablePosition(bool) {}
|
||||
|
||||
func (a *Analyzer) Analyze(text string) ([]Token, error) {
|
||||
return nil, fmt.Errorf("rag_analyzer: cgo required for Analyze(%q)", text)
|
||||
}
|
||||
|
||||
func (a *Analyzer) Tokenize(text string) (string, error) {
|
||||
return "", fmt.Errorf("rag_analyzer: cgo required for Tokenize(%q)", text)
|
||||
}
|
||||
|
||||
func (a *Analyzer) TokenizeWithPosition(text string) ([]TokenWithPosition, error) {
|
||||
return nil, fmt.Errorf("rag_analyzer: cgo required for TokenizeWithPosition(%q)", text)
|
||||
}
|
||||
|
||||
func (a *Analyzer) Close() {}
|
||||
|
||||
func (a *Analyzer) FineGrainedTokenize(tokens string) (string, error) {
|
||||
return "", fmt.Errorf("rag_analyzer: cgo required for FineGrainedTokenize(%q)", tokens)
|
||||
}
|
||||
|
||||
func (a *Analyzer) GetTermFreq(term string) int32 { return 0 }
|
||||
|
||||
func (a *Analyzer) GetTermTag(term string) string { return "" }
|
||||
|
||||
func (a *Analyzer) Copy() *Analyzer { return nil }
|
||||
@@ -733,7 +733,7 @@ func (p *Parser) parseDevChunk(explain bool) (*Command, error) {
|
||||
|
||||
dsl, err := p.parseQuotedString()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("expected DSL: %w", err)
|
||||
return nil, fmt.Errorf("expected chunk options file: %w", err)
|
||||
}
|
||||
|
||||
// Semicolon is optional
|
||||
|
||||
@@ -29,8 +29,8 @@ import (
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/ingestion"
|
||||
"ragflow/internal/ingestion/parser"
|
||||
"ragflow/internal/parser/chunk"
|
||||
"ragflow/internal/parser/parser"
|
||||
"ragflow/internal/utility"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -3561,8 +3561,9 @@ func (c *CLI) APIParseLocalFileCommand(cmd *Command) (ResponseIf, error) {
|
||||
return nil, fmt.Errorf("failed to read dsl file: %w", err)
|
||||
}
|
||||
|
||||
if err = fileParser.Parse(filename, fileContent); err != nil {
|
||||
return nil, formatRequestError("parse local file", err)
|
||||
parseResult := fileParser.ParseWithResult(filename, fileContent)
|
||||
if parseResult.Err != nil {
|
||||
return nil, formatRequestError("parse local file", parseResult.Err)
|
||||
}
|
||||
|
||||
var result SimpleResponse
|
||||
@@ -3773,13 +3774,17 @@ func (c *CLI) DevChunkCommand(cmd *Command) (ResponseIf, error) {
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("filename not provided")
|
||||
}
|
||||
dslFilename, ok := cmd.Params["dsl"].(string)
|
||||
optionsFilename, ok := cmd.Params["dsl"].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("dsl not provided")
|
||||
return nil, fmt.Errorf("chunk options file not provided")
|
||||
}
|
||||
dsl, err := os.ReadFile(dslFilename)
|
||||
optionsRaw, err := os.ReadFile(optionsFilename)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read dsl file: %w", err)
|
||||
return nil, fmt.Errorf("failed to read chunk options file: %w", err)
|
||||
}
|
||||
options, err := parseChunkOptions(optionsRaw)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
explain, ok := cmd.Params["explain"].(bool)
|
||||
@@ -3787,19 +3792,11 @@ func (c *CLI) DevChunkCommand(cmd *Command) (ResponseIf, error) {
|
||||
explain = false
|
||||
}
|
||||
|
||||
engine := ingestion.NewChunkEngine()
|
||||
plan, err := engine.Compile(string(dsl))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("compile failed: %w", err)
|
||||
}
|
||||
|
||||
if explain {
|
||||
|
||||
explanation, err := engine.Explain(plan)
|
||||
explanation, err := explainChunkOptions(options)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("explain error: %w", err)
|
||||
}
|
||||
|
||||
result.Message = explanation
|
||||
} else {
|
||||
fileToChunking, err := os.ReadFile(filename)
|
||||
@@ -3807,7 +3804,7 @@ func (c *CLI) DevChunkCommand(cmd *Command) (ResponseIf, error) {
|
||||
return nil, fmt.Errorf("failed to read file: %w", err)
|
||||
}
|
||||
|
||||
chunkContext, err := engine.Execute(plan, string(fileToChunking))
|
||||
chunkContext, err := chunk.Run(string(fileToChunking), options)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("chunking error: %w", err)
|
||||
}
|
||||
@@ -3825,6 +3822,22 @@ func (c *CLI) DevChunkCommand(cmd *Command) (ResponseIf, error) {
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
func parseChunkOptions(raw []byte) (chunk.ChunkOptions, error) {
|
||||
var options chunk.ChunkOptions
|
||||
if err := json.Unmarshal(raw, &options); err != nil {
|
||||
return options, fmt.Errorf("failed to parse chunk options file: %w", err)
|
||||
}
|
||||
return options, nil
|
||||
}
|
||||
|
||||
func explainChunkOptions(options chunk.ChunkOptions) (string, error) {
|
||||
formatted, err := json.MarshalIndent(options, "", " ")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(formatted), nil
|
||||
}
|
||||
|
||||
// APIOpenaiChatCommand dispatches the parsed OPENAI_CHAT command to either a
|
||||
// non-streaming oneshot call or a streaming SSE call, depending on the
|
||||
// `stream` option.
|
||||
|
||||
@@ -397,7 +397,14 @@ func (dao *IngestionTaskLogDAO) ListLogsByTaskID(taskID string) ([]*entity.Inges
|
||||
|
||||
func (dao *IngestionTaskLogDAO) LatestLogByTaskID(taskID string) (*entity.IngestionTaskLog, error) {
|
||||
var task *entity.IngestionTaskLog
|
||||
err := DB.Where("task_id = ?", taskID).Order("create_time DESC").First(&task).Error
|
||||
// Order by `id DESC` (NOT `create_time DESC`) because
|
||||
// create_time has only second-level resolution — multiple
|
||||
// checkpoints written within the same second tie-break
|
||||
// arbitrarily. The `id` is auto-increment, monotonic, and
|
||||
// always reflects write order. The pipeline's resume
|
||||
// algorithm reads the latest row, so the tie-break MUST
|
||||
// be deterministic.
|
||||
err := DB.Where("task_id = ?", taskID).Order("id DESC").First(&task).Error
|
||||
return task, err
|
||||
}
|
||||
|
||||
|
||||
19
internal/deepdoc/parser/pdf/parse_nocgo.go
Normal file
19
internal/deepdoc/parser/pdf/parse_nocgo.go
Normal file
@@ -0,0 +1,19 @@
|
||||
//go:build !cgo
|
||||
|
||||
package pdf
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
pdf "ragflow/internal/deepdoc/parser/pdf/type"
|
||||
)
|
||||
|
||||
// Parse is the no-CGO stub for the DeepDOC PDF pipeline. The surrounding
|
||||
// parser package converts this into parser.ErrPDFEngineUnavailable.
|
||||
func (p *Parser) Parse(ctx context.Context, data []byte, docAnalyzer pdf.DocAnalyzer) (*pdf.ParseResult, error) {
|
||||
_ = ctx
|
||||
_ = data
|
||||
_ = docAnalyzer
|
||||
return nil, fmt.Errorf("deepdoc/pdf: cgo required")
|
||||
}
|
||||
@@ -1,3 +1,5 @@
|
||||
//go:build cgo
|
||||
|
||||
// Package pdfium renders PDF pages using libpdfium (statically linked
|
||||
// at build time via CGO_LDFLAGS). It exists solely to replace pdf_oxide's
|
||||
// RenderPageRaw for use cases where image quality matters for downstream
|
||||
|
||||
134
internal/handler/components.go
Normal file
134
internal/handler/components.go
Normal file
@@ -0,0 +1,134 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Phase 4 of plan port-rag-flow-pipeline-to-go.md.
|
||||
//
|
||||
// Exposes GET /api/v1/components?category=ingestion,agent,shared
|
||||
// (case-insensitive, comma-separated). The data source is
|
||||
// runtime.DefaultRegistry — there is no separate catalog map (per
|
||||
// plan §4 task 2). Category values are validated against the three
|
||||
// known runtime.Category constants; an unknown value yields a
|
||||
// 400-style error envelope so the frontend can surface a useful
|
||||
// message instead of silently dropping the request.
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/service"
|
||||
)
|
||||
|
||||
// ComponentsHandler serves the component-catalog endpoint.
|
||||
type ComponentsHandler struct {
|
||||
svc *service.ComponentsService
|
||||
}
|
||||
|
||||
// NewComponentsHandler wires the handler to a ComponentsService
|
||||
// instance. The service is stateless so the same pointer can be
|
||||
// shared across handlers; construction happens once at server
|
||||
// startup (cmd/server_main.go).
|
||||
func NewComponentsHandler(svc *service.ComponentsService) *ComponentsHandler {
|
||||
return &ComponentsHandler{svc: svc}
|
||||
}
|
||||
|
||||
// Get handles GET /api/v1/components.
|
||||
//
|
||||
// Query parameters:
|
||||
// - category (optional, repeatable as comma-separated values).
|
||||
// Case-insensitive. Allowed values: "agent", "ingestion",
|
||||
// "shared". An empty / missing filter means "all categories".
|
||||
//
|
||||
// Response shape (success):
|
||||
//
|
||||
// { "data": [ { "name": "...", "category": "...",
|
||||
// "inputs": {...}, "outputs": {...} } ] }
|
||||
//
|
||||
// On an unknown category, the response uses the standard error
|
||||
// envelope (gin.H{code, message, data}) with HTTP 400. Service
|
||||
// failures bubble up as HTTP 500 with the same envelope shape.
|
||||
func (h *ComponentsHandler) Get(c *gin.Context) {
|
||||
raw := c.Query("category")
|
||||
cats, err := parseCategories(raw)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"code": 400,
|
||||
"message": err.Error(),
|
||||
"data": nil,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
out, err := h.svc.List(cats...)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"code": 500,
|
||||
"message": err.Error(),
|
||||
"data": nil,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": 0,
|
||||
"message": "success",
|
||||
"data": out,
|
||||
})
|
||||
}
|
||||
|
||||
// parseCategories splits a comma-separated category query string into
|
||||
// a slice of runtime.Category values. The slice is empty when the
|
||||
// raw string is empty (meaning "all categories"). An unknown token
|
||||
// yields an error; the first invalid token wins so the message
|
||||
// identifies the offender.
|
||||
func parseCategories(raw string) ([]runtime.Category, error) {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var out []runtime.Category
|
||||
for _, p := range strings.Split(raw, ",") {
|
||||
p = strings.TrimSpace(strings.ToLower(p))
|
||||
if p == "" {
|
||||
continue
|
||||
}
|
||||
switch p {
|
||||
case "agent":
|
||||
out = append(out, runtime.CategoryAgent)
|
||||
case "ingestion":
|
||||
out = append(out, runtime.CategoryIngestion)
|
||||
case "shared":
|
||||
out = append(out, runtime.CategoryShared)
|
||||
default:
|
||||
return nil, &categoryError{value: p}
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// categoryError is the error type returned by parseCategories for an
|
||||
// unrecognized category token. Its Error() message matches the
|
||||
// expected plan test contract ("unknown category: <value>").
|
||||
type categoryError struct {
|
||||
value string
|
||||
}
|
||||
|
||||
func (e *categoryError) Error() string {
|
||||
return "unknown category: " + e.value
|
||||
}
|
||||
330
internal/handler/components_testpkg/components_handler_test.go
Normal file
330
internal/handler/components_testpkg/components_handler_test.go
Normal file
@@ -0,0 +1,330 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Phase 4 of plan port-rag-flow-pipeline-to-go.md — integration
|
||||
// tests for GET /api/v1/components.
|
||||
//
|
||||
// The tests import the agent and ingestion component packages via
|
||||
// blank imports so their init() functions register every factory
|
||||
// into runtime.DefaultRegistry. The tests then assert that the
|
||||
// HTTP handler projects that registry into the expected
|
||||
// JSON-shaped response.
|
||||
//
|
||||
// Counters are derived from runtime.DefaultRegistry.Names() /
|
||||
// NamesByCategory() at test time rather than hardcoded constants —
|
||||
// the registry is the source of truth, so any future registration
|
||||
// change just moves the expected values with it.
|
||||
//
|
||||
// This file lives in its own subpackage (components_testpkg) so
|
||||
// its test binary does not co-compile with the rest of the
|
||||
// internal/handler tests, whose build is currently broken on the
|
||||
// agent-go-port branch (pre-existing
|
||||
// `interface{}` vs `string` mismatch in chat_test.go /
|
||||
// searchbot_test.go). When the upstream build is fixed this
|
||||
// package can be folded back into the main handler test binary.
|
||||
package handler_test
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sort"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
_ "ragflow/internal/agent/component" // registers agent components
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/handler"
|
||||
_ "ragflow/internal/ingestion/component" // registers ingestion main components
|
||||
_ "ragflow/internal/ingestion/component/chunker" // registers 4 chunker variants
|
||||
"ragflow/internal/service"
|
||||
)
|
||||
|
||||
func init() {
|
||||
gin.SetMode(gin.TestMode)
|
||||
}
|
||||
|
||||
// newComponentsTestRig mounts a Gin engine with the
|
||||
// /api/v1/components route registered against a fresh
|
||||
// ComponentsService (which reads runtime.DefaultRegistry).
|
||||
func newComponentsTestRig(t *testing.T) *gin.Engine {
|
||||
t.Helper()
|
||||
h := handler.NewComponentsHandler(service.NewComponentsService())
|
||||
eng := gin.New()
|
||||
v1 := eng.Group("/api/v1")
|
||||
v1.GET("/components", h.Get)
|
||||
return eng
|
||||
}
|
||||
|
||||
// doRequest issues a GET against the test rig and returns the
|
||||
// response recorder.
|
||||
func doRequest(t *testing.T, eng *gin.Engine, path string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(http.MethodGet, path, nil)
|
||||
w := httptest.NewRecorder()
|
||||
eng.ServeHTTP(w, req)
|
||||
return w
|
||||
}
|
||||
|
||||
// decodeEnvelope parses the standard {code, message, data} response
|
||||
// envelope used by every handler in this codebase.
|
||||
func decodeEnvelope(t *testing.T, body []byte) (code int, message string, data []service.ComponentDescriptor) {
|
||||
t.Helper()
|
||||
var env struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Data []service.ComponentDescriptor `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &env); err != nil {
|
||||
t.Fatalf("decode envelope: %v; body=%s", err, string(body))
|
||||
}
|
||||
return env.Code, env.Message, env.Data
|
||||
}
|
||||
|
||||
// TestComponentsHandler_NoFilter verifies that GET /api/v1/components
|
||||
// (no filter) returns every registered component, across all
|
||||
// categories. The total must equal runtime.DefaultRegistry.Names().
|
||||
func TestComponentsHandler_NoFilter(t *testing.T) {
|
||||
eng := newComponentsTestRig(t)
|
||||
w := doRequest(t, eng, "/api/v1/components")
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
gotCode, msg, data := decodeEnvelope(t, w.Body.Bytes())
|
||||
if gotCode != 0 {
|
||||
t.Errorf("envelope code = %d, want 0; message=%s", gotCode, msg)
|
||||
}
|
||||
if msg != "success" {
|
||||
t.Errorf("envelope message = %q, want %q", msg, "success")
|
||||
}
|
||||
|
||||
wantTotal := len(runtime.DefaultRegistry.Names())
|
||||
if len(data) != wantTotal {
|
||||
t.Errorf("got %d components, want %d (== Names())", len(data), wantTotal)
|
||||
}
|
||||
|
||||
// Spot-check that the default registry includes agent and
|
||||
// ingestion components when no filter is applied.
|
||||
seen := map[string]bool{}
|
||||
for _, d := range data {
|
||||
seen[d.Category] = true
|
||||
}
|
||||
for _, cat := range []string{"agent", "ingestion"} {
|
||||
if !seen[cat] {
|
||||
t.Errorf("expected category %q in unfiltered response; got categories %v", cat, seen)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestComponentsHandler_FilterIngestion verifies the
|
||||
// ?category=ingestion filter returns the 8 ingestion components
|
||||
// (Extractor, File, Parser, Tokenizer + 4 chunker variants). Names
|
||||
// must be sorted ascending (plan §4 task 1 stable output).
|
||||
func TestComponentsHandler_FilterIngestion(t *testing.T) {
|
||||
eng := newComponentsTestRig(t)
|
||||
w := doRequest(t, eng, "/api/v1/components?category=ingestion")
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
_, _, data := decodeEnvelope(t, w.Body.Bytes())
|
||||
|
||||
wantNames := []string{
|
||||
"extractor", "file", "grouptitlechunker", "hierarchytitlechunker",
|
||||
"parser", "titlechunker", "tokenchunker", "tokenizer",
|
||||
}
|
||||
assertNameSet(t, "ingestion", data, wantNames)
|
||||
|
||||
for _, d := range data {
|
||||
if d.Category != "ingestion" {
|
||||
t.Errorf("filtered component %q has category %q, want %q", d.Name, d.Category, "ingestion")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestComponentsHandler_FilterMultiple verifies comma-separated
|
||||
// category values work. The shared category may currently be empty,
|
||||
// so ?category=ingestion,shared must still return the ingestion set.
|
||||
func TestComponentsHandler_FilterMultiple(t *testing.T) {
|
||||
eng := newComponentsTestRig(t)
|
||||
w := doRequest(t, eng, "/api/v1/components?category=ingestion,shared")
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
_, _, data := decodeEnvelope(t, w.Body.Bytes())
|
||||
|
||||
wantNames := []string{
|
||||
"extractor", "file", "grouptitlechunker", "hierarchytitlechunker",
|
||||
"parser", "titlechunker", "tokenchunker", "tokenizer",
|
||||
}
|
||||
assertNameSet(t, "ingestion,shared", data, wantNames)
|
||||
}
|
||||
|
||||
// TestComponentsHandler_FilterAgent verifies the
|
||||
// ?category=agent filter returns every agent component. The count
|
||||
// is derived from NamesByCategory(CategoryAgent) so any future
|
||||
// registration change just moves the value.
|
||||
func TestComponentsHandler_FilterAgent(t *testing.T) {
|
||||
eng := newComponentsTestRig(t)
|
||||
w := doRequest(t, eng, "/api/v1/components?category=agent")
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
_, _, data := decodeEnvelope(t, w.Body.Bytes())
|
||||
|
||||
wantNames := runtime.DefaultRegistry.NamesByCategory(runtime.CategoryAgent)
|
||||
if len(wantNames) == 0 {
|
||||
t.Fatal("no agent components registered; cannot run this test")
|
||||
}
|
||||
assertNameSet(t, "agent", data, wantNames)
|
||||
for _, d := range data {
|
||||
if d.Category != "agent" {
|
||||
t.Errorf("agent-filtered component %q has category %q", d.Name, d.Category)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestComponentsHandler_InvalidCategory verifies an unknown category
|
||||
// yields HTTP 400 with the standard error envelope. The message
|
||||
// must mention the offending token.
|
||||
func TestComponentsHandler_InvalidCategory(t *testing.T) {
|
||||
eng := newComponentsTestRig(t)
|
||||
w := doRequest(t, eng, "/api/v1/components?category=foo")
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want 400; body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
var env struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &env); err != nil {
|
||||
t.Fatalf("decode: %v; body=%s", err, w.Body.String())
|
||||
}
|
||||
if env.Code != 400 {
|
||||
t.Errorf("envelope code = %d, want 400", env.Code)
|
||||
}
|
||||
if env.Message == "" {
|
||||
t.Errorf("envelope message is empty; body=%s", w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// TestComponentsHandler_ComponentShape verifies every descriptor
|
||||
// has a non-empty name, a non-empty category from the known set,
|
||||
// and non-nil inputs / outputs maps (the service normalises nil
|
||||
// to empty maps so the JSON payload never has null fields).
|
||||
func TestComponentsHandler_ComponentShape(t *testing.T) {
|
||||
eng := newComponentsTestRig(t)
|
||||
w := doRequest(t, eng, "/api/v1/components")
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
_, _, data := decodeEnvelope(t, w.Body.Bytes())
|
||||
|
||||
allowed := map[string]bool{"agent": true, "ingestion": true, "shared": true}
|
||||
for _, d := range data {
|
||||
if d.Name == "" {
|
||||
t.Errorf("descriptor has empty name: %+v", d)
|
||||
}
|
||||
if d.Category == "" {
|
||||
t.Errorf("descriptor %q has empty category", d.Name)
|
||||
}
|
||||
if !allowed[d.Category] {
|
||||
t.Errorf("descriptor %q has unknown category %q", d.Name, d.Category)
|
||||
}
|
||||
if d.Inputs == nil {
|
||||
t.Errorf("descriptor %q has nil inputs; service must normalise to empty map", d.Name)
|
||||
}
|
||||
if d.Outputs == nil {
|
||||
t.Errorf("descriptor %q has nil outputs; service must normalise to empty map", d.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestComponentsHandler_CaseInsensitive verifies the
|
||||
// category filter accepts mixed case (plan §4 task 1 spec:
|
||||
// "case-insensitive"). The response must match the all-lowercase
|
||||
// variant.
|
||||
func TestComponentsHandler_CaseInsensitive(t *testing.T) {
|
||||
eng := newComponentsTestRig(t)
|
||||
w := doRequest(t, eng, "/api/v1/components?category=INGESTION")
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
_, _, data := decodeEnvelope(t, w.Body.Bytes())
|
||||
wantNames := []string{
|
||||
"extractor", "file", "grouptitlechunker", "hierarchytitlechunker",
|
||||
"parser", "titlechunker", "tokenchunker", "tokenizer",
|
||||
}
|
||||
assertNameSet(t, "INGESTION (case-folded)", data, wantNames)
|
||||
}
|
||||
|
||||
// TestComponentsHandler_EmptyCategory treats the bare
|
||||
// ?category= (no value) as a no-op filter — same payload as the
|
||||
// unfiltered endpoint. Whitespace-only values are dropped.
|
||||
func TestComponentsHandler_EmptyCategory(t *testing.T) {
|
||||
eng := newComponentsTestRig(t)
|
||||
w := doRequest(t, eng, "/api/v1/components?category=")
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
_, _, data := decodeEnvelope(t, w.Body.Bytes())
|
||||
if len(data) != len(runtime.DefaultRegistry.Names()) {
|
||||
t.Errorf("empty category filter returned %d, want %d (== Names())", len(data), len(runtime.DefaultRegistry.Names()))
|
||||
}
|
||||
}
|
||||
|
||||
// TestComponentsHandler_RouteMounted verifies the route is
|
||||
// actually mounted at /api/v1/components. (Catches an accidental
|
||||
// router wiring regression even when no other test fires.)
|
||||
func TestComponentsHandler_RouteMounted(t *testing.T) {
|
||||
eng := newComponentsTestRig(t)
|
||||
w := doRequest(t, eng, "/api/v1/components")
|
||||
if w.Code == http.StatusNotFound {
|
||||
t.Fatalf("route not mounted; got 404; body=%s", w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// assertNameSet checks that the descriptors returned by the
|
||||
// handler contain exactly the names in want (set comparison;
|
||||
// ignores ordering — the handler also sorts by name, but the
|
||||
// comparison is order-independent to make test failures easier
|
||||
// to read).
|
||||
func assertNameSet(t *testing.T, label string, got []service.ComponentDescriptor, want []string) {
|
||||
t.Helper()
|
||||
gotNames := make([]string, 0, len(got))
|
||||
for _, d := range got {
|
||||
gotNames = append(gotNames, d.Name)
|
||||
}
|
||||
sort.Strings(gotNames)
|
||||
wantSorted := append([]string(nil), want...)
|
||||
sort.Strings(wantSorted)
|
||||
|
||||
if len(gotNames) != len(wantSorted) {
|
||||
t.Errorf("%s: got %d names (%v), want %d (%v)", label, len(gotNames), gotNames, len(wantSorted), wantSorted)
|
||||
return
|
||||
}
|
||||
for i := range wantSorted {
|
||||
if gotNames[i] != wantSorted[i] {
|
||||
t.Errorf("%s: name[%d] = %q, want %q (full got=%v want=%v)", label, i, gotNames[i], wantSorted[i], gotNames, wantSorted)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,649 +0,0 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package chunk
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"unicode"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Token types
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type tokenType int
|
||||
|
||||
const (
|
||||
tokenEOF tokenType = iota
|
||||
tokenIdentifier
|
||||
tokenString
|
||||
tokenNumber
|
||||
tokenTrue
|
||||
tokenFalse
|
||||
tokenEq
|
||||
tokenNeq
|
||||
tokenGt
|
||||
tokenLt
|
||||
tokenGte
|
||||
tokenLte
|
||||
tokenAnd
|
||||
tokenOr
|
||||
tokenNot
|
||||
tokenLParen
|
||||
tokenRParen
|
||||
)
|
||||
|
||||
var keywords = map[string]tokenType{
|
||||
"AND": tokenAnd,
|
||||
"OR": tokenOr,
|
||||
"NOT": tokenNot,
|
||||
"true": tokenTrue,
|
||||
"false": tokenFalse,
|
||||
"TRUE": tokenTrue,
|
||||
"FALSE": tokenFalse,
|
||||
}
|
||||
|
||||
type token struct {
|
||||
typ tokenType
|
||||
raw string
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Lexer
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type lexer struct {
|
||||
input []rune
|
||||
pos int
|
||||
}
|
||||
|
||||
func newLexer(input string) *lexer {
|
||||
return &lexer{input: []rune(input)}
|
||||
}
|
||||
|
||||
func (l *lexer) skipWhitespace() {
|
||||
for l.pos < len(l.input) && unicode.IsSpace(l.input[l.pos]) {
|
||||
l.pos++
|
||||
}
|
||||
}
|
||||
|
||||
func (l *lexer) next() token {
|
||||
l.skipWhitespace()
|
||||
if l.pos >= len(l.input) {
|
||||
return token{typ: tokenEOF, raw: ""}
|
||||
}
|
||||
|
||||
ch := l.input[l.pos]
|
||||
|
||||
// Single-quoted string
|
||||
if ch == '\'' {
|
||||
l.pos++ // skip opening '
|
||||
start := l.pos
|
||||
for l.pos < len(l.input) && l.input[l.pos] != '\'' {
|
||||
l.pos++
|
||||
}
|
||||
raw := string(l.input[start:l.pos])
|
||||
if l.pos < len(l.input) {
|
||||
l.pos++ // skip closing '
|
||||
}
|
||||
return token{typ: tokenString, raw: raw}
|
||||
}
|
||||
|
||||
// Operators
|
||||
if l.pos+1 < len(l.input) {
|
||||
next := l.input[l.pos+1]
|
||||
switch string([]rune{ch, next}) {
|
||||
case ">=":
|
||||
l.pos += 2
|
||||
return token{typ: tokenGte, raw: ">="}
|
||||
case "<=":
|
||||
l.pos += 2
|
||||
return token{typ: tokenLte, raw: "<="}
|
||||
case "!=":
|
||||
l.pos += 2
|
||||
return token{typ: tokenNeq, raw: "!="}
|
||||
}
|
||||
}
|
||||
switch ch {
|
||||
case '=':
|
||||
l.pos++
|
||||
return token{typ: tokenEq, raw: "="}
|
||||
case '>':
|
||||
l.pos++
|
||||
return token{typ: tokenGt, raw: ">"}
|
||||
case '<':
|
||||
l.pos++
|
||||
return token{typ: tokenLt, raw: "<"}
|
||||
case '(':
|
||||
l.pos++
|
||||
return token{typ: tokenLParen, raw: "("}
|
||||
case ')':
|
||||
l.pos++
|
||||
return token{typ: tokenRParen, raw: ")"}
|
||||
}
|
||||
|
||||
// Number
|
||||
if unicode.IsDigit(ch) || (ch == '-' && l.pos+1 < len(l.input) && unicode.IsDigit(l.input[l.pos+1])) {
|
||||
start := l.pos
|
||||
if l.input[l.pos] == '-' {
|
||||
l.pos++
|
||||
}
|
||||
for l.pos < len(l.input) && (unicode.IsDigit(l.input[l.pos]) || l.input[l.pos] == '.') {
|
||||
l.pos++
|
||||
}
|
||||
return token{typ: tokenNumber, raw: string(l.input[start:l.pos])}
|
||||
}
|
||||
|
||||
// Identifier / keyword
|
||||
if unicode.IsLetter(ch) || ch == '_' {
|
||||
start := l.pos
|
||||
for l.pos < len(l.input) && (unicode.IsLetter(l.input[l.pos]) || unicode.IsDigit(l.input[l.pos]) || l.input[l.pos] == '_') {
|
||||
l.pos++
|
||||
}
|
||||
raw := string(l.input[start:l.pos])
|
||||
if kw, ok := keywords[raw]; ok {
|
||||
return token{typ: kw, raw: raw}
|
||||
}
|
||||
return token{typ: tokenIdentifier, raw: raw}
|
||||
}
|
||||
|
||||
// Unknown
|
||||
l.pos++
|
||||
return token{typ: tokenIdentifier, raw: string(ch)}
|
||||
}
|
||||
|
||||
func (l *lexer) peek() token {
|
||||
pos := l.pos
|
||||
tok := l.next()
|
||||
l.pos = pos
|
||||
return tok
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// AST nodes
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type Expr interface {
|
||||
String() string
|
||||
}
|
||||
|
||||
type binaryExpr struct {
|
||||
left Expr
|
||||
op tokenType
|
||||
right Expr
|
||||
}
|
||||
|
||||
func (e binaryExpr) String() string {
|
||||
ops := map[tokenType]string{
|
||||
tokenEq: "=",
|
||||
tokenNeq: "!=",
|
||||
tokenGt: ">",
|
||||
tokenLt: "<",
|
||||
tokenGte: ">=",
|
||||
tokenLte: "<=",
|
||||
tokenAnd: "AND",
|
||||
tokenOr: "OR",
|
||||
}
|
||||
return fmt.Sprintf("(%s %s %s)", e.left, ops[e.op], e.right)
|
||||
}
|
||||
|
||||
type unaryExpr struct {
|
||||
op tokenType
|
||||
right Expr
|
||||
}
|
||||
|
||||
func (e unaryExpr) String() string {
|
||||
return fmt.Sprintf("(NOT %s)", e.right)
|
||||
}
|
||||
|
||||
type identifierExpr struct {
|
||||
name string
|
||||
}
|
||||
|
||||
func (e identifierExpr) String() string {
|
||||
return e.name
|
||||
}
|
||||
|
||||
type stringExpr struct {
|
||||
value string
|
||||
}
|
||||
|
||||
func (e stringExpr) String() string {
|
||||
return "'" + e.value + "'"
|
||||
}
|
||||
|
||||
type numberExpr struct {
|
||||
value float64
|
||||
}
|
||||
|
||||
func (e numberExpr) String() string {
|
||||
return strconv.FormatFloat(e.value, 'f', -1, 64)
|
||||
}
|
||||
|
||||
type boolExpr struct {
|
||||
value bool
|
||||
}
|
||||
|
||||
func (e boolExpr) String() string {
|
||||
return strconv.FormatBool(e.value)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Recursive-descent parser
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type parser struct {
|
||||
lex *lexer
|
||||
cur token
|
||||
peeked bool
|
||||
}
|
||||
|
||||
func newParser(input string) *parser {
|
||||
p := &parser{lex: newLexer(input)}
|
||||
p.advance()
|
||||
return p
|
||||
}
|
||||
|
||||
func (p *parser) advance() {
|
||||
if p.peeked {
|
||||
p.peeked = false
|
||||
return
|
||||
}
|
||||
p.cur = p.lex.next()
|
||||
}
|
||||
|
||||
func (p *parser) peek() token {
|
||||
if !p.peeked {
|
||||
p.peeked = true
|
||||
p.cur = p.lex.next()
|
||||
}
|
||||
return p.cur
|
||||
}
|
||||
|
||||
func (p *parser) expect(typ tokenType) token {
|
||||
tok := p.cur
|
||||
if tok.typ != typ {
|
||||
panic(fmt.Sprintf("expected token %d but got %d (%q)", typ, tok.typ, tok.raw))
|
||||
}
|
||||
p.advance()
|
||||
return tok
|
||||
}
|
||||
|
||||
func (p *parser) parse() Expr {
|
||||
return p.parseOr()
|
||||
}
|
||||
|
||||
// or_expr → and_expr ("OR" and_expr)*
|
||||
func (p *parser) parseOr() Expr {
|
||||
e := p.parseAnd()
|
||||
for p.cur.typ == tokenOr {
|
||||
op := p.cur.typ
|
||||
p.advance()
|
||||
right := p.parseAnd()
|
||||
e = binaryExpr{left: e, op: op, right: right}
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
// and_expr → not_expr ("AND" not_expr)*
|
||||
func (p *parser) parseAnd() Expr {
|
||||
e := p.parseNot()
|
||||
for p.cur.typ == tokenAnd {
|
||||
op := p.cur.typ
|
||||
p.advance()
|
||||
right := p.parseNot()
|
||||
e = binaryExpr{left: e, op: op, right: right}
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
// not_expr → "NOT" not_expr | primary
|
||||
func (p *parser) parseNot() Expr {
|
||||
if p.cur.typ == tokenNot {
|
||||
op := p.cur.typ
|
||||
p.advance()
|
||||
right := p.parseNot()
|
||||
return unaryExpr{op: op, right: right}
|
||||
}
|
||||
return p.parsePrimary()
|
||||
}
|
||||
|
||||
// primary → comparison | "(" expression ")"
|
||||
func (p *parser) parsePrimary() Expr {
|
||||
if p.cur.typ == tokenLParen {
|
||||
p.advance()
|
||||
e := p.parseOr()
|
||||
p.expect(tokenRParen)
|
||||
return e
|
||||
}
|
||||
return p.parseComparison()
|
||||
}
|
||||
|
||||
// comparison → IDENTIFIER OP value | value
|
||||
// comparison → IDENTIFIER OP value
|
||||
func (p *parser) parseComparison() Expr {
|
||||
if p.cur.typ == tokenIdentifier {
|
||||
id := p.cur.raw
|
||||
p.advance()
|
||||
switch p.cur.typ {
|
||||
case tokenEq, tokenNeq, tokenGt, tokenLt, tokenGte, tokenLte:
|
||||
op := p.cur.typ
|
||||
p.advance()
|
||||
right := p.parseValue()
|
||||
return binaryExpr{left: identifierExpr{name: id}, op: op, right: right}
|
||||
default:
|
||||
// identifier alone – treat as boolean check
|
||||
return binaryExpr{
|
||||
left: identifierExpr{name: id},
|
||||
op: tokenEq,
|
||||
right: boolExpr{value: true},
|
||||
}
|
||||
}
|
||||
}
|
||||
return p.parseValue()
|
||||
}
|
||||
|
||||
// value → STRING | NUMBER | BOOLEAN
|
||||
func (p *parser) parseValue() Expr {
|
||||
switch p.cur.typ {
|
||||
case tokenString:
|
||||
v := stringExpr{value: p.cur.raw}
|
||||
p.advance()
|
||||
return v
|
||||
case tokenNumber:
|
||||
f, _ := strconv.ParseFloat(p.cur.raw, 64)
|
||||
p.advance()
|
||||
return numberExpr{value: f}
|
||||
case tokenTrue:
|
||||
p.advance()
|
||||
return boolExpr{value: true}
|
||||
case tokenFalse:
|
||||
p.advance()
|
||||
return boolExpr{value: false}
|
||||
default:
|
||||
// treat as identifier (e.g. bare variable reference)
|
||||
id := identifierExpr{name: p.cur.raw}
|
||||
p.advance()
|
||||
return id
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Evaluator
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
var reMediaURL = regexp.MustCompile(`(?i)https?://[^\s]*\.(jpg|jpeg|png|gif|bmp|webp|svg|mp4|avi|mov|wmv|flv|mkv|m4v|mp3|wav|ogg|aac)`)
|
||||
var reImageURL = regexp.MustCompile(`(?i)https?://[^\s]*\.(jpg|jpeg|png|gif|bmp|webp|svg)`)
|
||||
var reVideoURL = regexp.MustCompile(`(?i)https?://[^\s]*\.(mp4|avi|mov|wmv|flv|mkv|m4v)`)
|
||||
var reAnyURL = regexp.MustCompile(`(?i)https?://[^\s]+`)
|
||||
|
||||
// buildExprContext builds a variable context from a chunk's content and metadata.
|
||||
// It auto-detects media/image/video URLs and language hints.
|
||||
func buildExprContext(chunk ContentProvider, metadata map[string]interface{}) map[string]interface{} {
|
||||
vars := make(map[string]interface{})
|
||||
content := chunk.GetContent()
|
||||
|
||||
// Pre-populate from metadata
|
||||
for k, v := range metadata {
|
||||
vars[k] = v
|
||||
}
|
||||
|
||||
// Auto-detect URL presence
|
||||
vars["has_media_url"] = reMediaURL.MatchString(content)
|
||||
vars["has_image_url"] = reImageURL.MatchString(content)
|
||||
vars["has_video_url"] = reVideoURL.MatchString(content)
|
||||
vars["has_url"] = reAnyURL.MatchString(content)
|
||||
vars["length"] = len([]rune(content))
|
||||
|
||||
return vars
|
||||
}
|
||||
|
||||
// ContentProvider allows evaluating expressions against any type that has content.
|
||||
type ContentProvider interface {
|
||||
GetContent() string
|
||||
}
|
||||
|
||||
// Evaluate parses and evaluates a boolean expression against a variable map.
|
||||
func Evaluate(exprStr string, vars map[string]interface{}) (bool, error) {
|
||||
p := newParser(exprStr)
|
||||
ast := p.parse()
|
||||
res, err := eval(ast, vars)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("evaluate %q: %w", exprStr, err)
|
||||
}
|
||||
b, ok := toBool(res)
|
||||
if !ok {
|
||||
return false, fmt.Errorf("evaluate %q: result %v (%T) is not a boolean", exprStr, res, res)
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
// CompileExpression parses an expression string into a reusable AST.
|
||||
func CompileExpression(exprStr string) (Expr, error) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
panic(fmt.Sprintf("compile expression %q: %v", exprStr, r))
|
||||
}
|
||||
}()
|
||||
p := newParser(exprStr)
|
||||
return p.parse(), nil
|
||||
}
|
||||
|
||||
// EvalCompiled evaluates a pre-compiled expression AST against variables.
|
||||
func EvalCompiled(ast interface{}, vars map[string]interface{}) (bool, error) {
|
||||
e, ok := ast.(Expr)
|
||||
if !ok {
|
||||
return false, fmt.Errorf("invalid AST type: %T", ast)
|
||||
}
|
||||
res, err := eval(e, vars)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
b, ok := toBool(res)
|
||||
if !ok {
|
||||
return false, fmt.Errorf("result %v (%T) is not boolean", res, res)
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
func eval(e Expr, vars map[string]interface{}) (interface{}, error) {
|
||||
switch n := e.(type) {
|
||||
case binaryExpr:
|
||||
return evalBinary(n, vars)
|
||||
case unaryExpr:
|
||||
return evalUnary(n, vars)
|
||||
case identifierExpr:
|
||||
v, ok := vars[n.name]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("undefined variable: %s", n.name)
|
||||
}
|
||||
return v, nil
|
||||
case stringExpr:
|
||||
return n.value, nil
|
||||
case numberExpr:
|
||||
return n.value, nil
|
||||
case boolExpr:
|
||||
return n.value, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unknown expression type: %T", e)
|
||||
}
|
||||
}
|
||||
|
||||
func evalBinary(e binaryExpr, vars map[string]interface{}) (interface{}, error) {
|
||||
left, err := eval(e.left, vars)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
right, err := eval(e.right, vars)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
switch e.op {
|
||||
case tokenAnd:
|
||||
l, ok := toBool(left)
|
||||
if !ok {
|
||||
return false, fmt.Errorf("AND requires boolean left operand")
|
||||
}
|
||||
if !l {
|
||||
return false, nil
|
||||
}
|
||||
r, ok := toBool(right)
|
||||
if !ok {
|
||||
return false, fmt.Errorf("AND requires boolean right operand")
|
||||
}
|
||||
return r, nil
|
||||
|
||||
case tokenOr:
|
||||
l, ok := toBool(left)
|
||||
if !ok {
|
||||
return false, fmt.Errorf("OR requires boolean left operand")
|
||||
}
|
||||
if l {
|
||||
return true, nil
|
||||
}
|
||||
r, ok := toBool(right)
|
||||
if !ok {
|
||||
return false, fmt.Errorf("OR requires boolean right operand")
|
||||
}
|
||||
return r, nil
|
||||
|
||||
case tokenEq:
|
||||
return compareEq(left, right), nil
|
||||
case tokenNeq:
|
||||
return !compareEq(left, right), nil
|
||||
case tokenGt, tokenLt, tokenGte, tokenLte:
|
||||
return compareOrder(left, right, e.op)
|
||||
default:
|
||||
return false, fmt.Errorf("unknown binary op %d", e.op)
|
||||
}
|
||||
}
|
||||
|
||||
func evalUnary(e unaryExpr, vars map[string]interface{}) (interface{}, error) {
|
||||
right, err := eval(e.right, vars)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
b, ok := toBool(right)
|
||||
if !ok {
|
||||
return false, fmt.Errorf("NOT requires boolean operand")
|
||||
}
|
||||
return !b, nil
|
||||
}
|
||||
|
||||
func toBool(v interface{}) (bool, bool) {
|
||||
switch vv := v.(type) {
|
||||
case bool:
|
||||
return vv, true
|
||||
case string:
|
||||
return vv == "true" || vv == "TRUE" || vv == "1", true
|
||||
case float64:
|
||||
return vv != 0, true
|
||||
case int:
|
||||
return vv != 0, true
|
||||
}
|
||||
return false, false
|
||||
}
|
||||
|
||||
func compareEq(a, b interface{}) bool {
|
||||
// Normalise numeric types
|
||||
af, aIsNum := toFloat(a)
|
||||
bf, bIsNum := toFloat(b)
|
||||
if aIsNum && bIsNum {
|
||||
return af == bf
|
||||
}
|
||||
// Fall back to string comparison
|
||||
return fmt.Sprintf("%v", a) == fmt.Sprintf("%v", b)
|
||||
}
|
||||
|
||||
func toFloat(v interface{}) (float64, bool) {
|
||||
switch vv := v.(type) {
|
||||
case float64:
|
||||
return vv, true
|
||||
case int:
|
||||
return float64(vv), true
|
||||
case string:
|
||||
f, err := strconv.ParseFloat(vv, 64)
|
||||
return f, err == nil
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func compareOrder(a, b interface{}, op tokenType) (bool, error) {
|
||||
af, aOK := toFloat(a)
|
||||
bf, bOK := toFloat(b)
|
||||
if aOK && bOK {
|
||||
switch op {
|
||||
case tokenGt:
|
||||
return af > bf, nil
|
||||
case tokenLt:
|
||||
return af < bf, nil
|
||||
case tokenGte:
|
||||
return af >= bf, nil
|
||||
case tokenLte:
|
||||
return af <= bf, nil
|
||||
}
|
||||
}
|
||||
// String fallback
|
||||
sa := fmt.Sprintf("%v", a)
|
||||
sb := fmt.Sprintf("%v", b)
|
||||
switch op {
|
||||
case tokenGt:
|
||||
return sa > sb, nil
|
||||
case tokenLt:
|
||||
return sa < sb, nil
|
||||
case tokenGte:
|
||||
return sa >= sb, nil
|
||||
case tokenLte:
|
||||
return sa <= sb, nil
|
||||
}
|
||||
return false, fmt.Errorf("unsupported comparison op %d between %T and %T", op, a, b)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Language heuristics
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// DetectLanguage returns a best-effort language code ('zh', 'en', etc.)
|
||||
// based on the proportion of CJK characters.
|
||||
func DetectLanguage(text string) string {
|
||||
cjk := 0
|
||||
total := 0
|
||||
for _, r := range text {
|
||||
if unicode.Is(unicode.Han, r) {
|
||||
cjk++
|
||||
}
|
||||
if unicode.IsLetter(r) {
|
||||
total++
|
||||
}
|
||||
}
|
||||
if total > 0 && float64(cjk)/float64(total) > 0.3 {
|
||||
return "zh"
|
||||
}
|
||||
return "en"
|
||||
}
|
||||
|
||||
// RuneCount returns the number of runes in text.
|
||||
func RuneCount(text string) int {
|
||||
return len([]rune(text))
|
||||
}
|
||||
|
||||
// Ensure math is used (for NaN etc.)
|
||||
var _ = math.NaN
|
||||
@@ -1,491 +0,0 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package chunk
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Overlap condition
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type overlapConfig struct {
|
||||
Size int `json:"size"`
|
||||
Unit string `json:"unit,omitempty"` // "char" (default) or "sentence"
|
||||
}
|
||||
|
||||
type overlapCondition struct {
|
||||
Name string
|
||||
Condition Expr // pre-compiled expression AST from CompileExpression
|
||||
OverlapConfig overlapConfig
|
||||
}
|
||||
|
||||
type mergeConfig struct {
|
||||
TargetSize int `json:"target_size"`
|
||||
Strategy string `json:"strategy"` // "greedy"
|
||||
}
|
||||
|
||||
type filterConfig struct {
|
||||
MinLength int `json:"min_length"`
|
||||
MaxLength int `json:"max_length"`
|
||||
}
|
||||
|
||||
type metadataConfig struct {
|
||||
IncludeIndex bool `json:"include_index"`
|
||||
CustomFields map[string]string `json:"custom_fields,omitempty"`
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// PostprocessOperator
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type PostprocessOperator struct {
|
||||
merge *mergeConfig
|
||||
overlap struct {
|
||||
unit string // "char" (default) or "sentence"
|
||||
conditions []overlapCondition
|
||||
defaultCfg overlapConfig
|
||||
}
|
||||
filter *filterConfig
|
||||
addMetadata *metadataConfig
|
||||
}
|
||||
|
||||
func NewPostprocessOperator(config map[string]interface{}) (*PostprocessOperator, error) {
|
||||
op := &PostprocessOperator{}
|
||||
|
||||
// Merge
|
||||
if m, ok := config["merge"].(map[string]interface{}); ok {
|
||||
op.merge = &mergeConfig{}
|
||||
if ts, ok := m["target_size"].(float64); ok {
|
||||
op.merge.TargetSize = int(ts)
|
||||
} else {
|
||||
op.merge.TargetSize = 500
|
||||
}
|
||||
if s, ok := m["strategy"].(string); ok {
|
||||
op.merge.Strategy = s
|
||||
} else {
|
||||
op.merge.Strategy = "greedy"
|
||||
}
|
||||
}
|
||||
|
||||
// Overlap
|
||||
if ov, ok := config["overlap"].(map[string]interface{}); ok {
|
||||
if u, ok := ov["unit"].(string); ok {
|
||||
op.overlap.unit = u
|
||||
} else {
|
||||
op.overlap.unit = "char"
|
||||
}
|
||||
|
||||
// Default
|
||||
if d, ok := ov["default"].(map[string]interface{}); ok {
|
||||
op.overlap.defaultCfg = parseOverlapConfig(d)
|
||||
}
|
||||
|
||||
// Conditions
|
||||
if conds, ok := ov["conditions"].([]interface{}); ok {
|
||||
for _, ci := range conds {
|
||||
c, ok := ci.(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
cond := overlapCondition{}
|
||||
if n, ok := c["name"].(string); ok {
|
||||
cond.Name = n
|
||||
}
|
||||
if exprStr, ok := c["if"].(string); ok {
|
||||
expression, err := CompileExpression(exprStr)
|
||||
if err == nil {
|
||||
cond.Condition = expression
|
||||
}
|
||||
}
|
||||
if thenMap, ok := c["then"].(map[string]interface{}); ok {
|
||||
cond.OverlapConfig = parseOverlapConfig(thenMap)
|
||||
}
|
||||
op.overlap.conditions = append(op.overlap.conditions, cond)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Filter
|
||||
if f, ok := config["filter"].(map[string]interface{}); ok {
|
||||
op.filter = &filterConfig{}
|
||||
if v, ok := f["min_length"].(float64); ok {
|
||||
op.filter.MinLength = int(v)
|
||||
}
|
||||
if v, ok := f["max_length"].(float64); ok {
|
||||
op.filter.MaxLength = int(v)
|
||||
}
|
||||
}
|
||||
|
||||
// Add metadata
|
||||
if am, ok := config["add_metadata"].(map[string]interface{}); ok {
|
||||
op.addMetadata = &metadataConfig{}
|
||||
if inc, ok := am["include_index"].(bool); ok {
|
||||
op.addMetadata.IncludeIndex = inc
|
||||
}
|
||||
if cf, ok := am["custom_fields"].(map[string]interface{}); ok {
|
||||
m := make(map[string]string, len(cf))
|
||||
for k, v := range cf {
|
||||
m[k] = fmt.Sprintf("%v", v)
|
||||
}
|
||||
op.addMetadata.CustomFields = m
|
||||
}
|
||||
}
|
||||
|
||||
return op, nil
|
||||
}
|
||||
|
||||
func (o *PostprocessOperator) Prepare(chunkCtx *ChunkContext) error {
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (o *PostprocessOperator) Execute(chunkCtx *ChunkContext) error {
|
||||
chunks := chunkCtx.SplitChunks
|
||||
if len(chunks) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 1. Merge
|
||||
if o.merge != nil {
|
||||
chunks = o.mergeChunks(chunks)
|
||||
}
|
||||
|
||||
// 2. Overlap
|
||||
chunks = o.applyOverlap(chunks)
|
||||
|
||||
// 3. Filter
|
||||
if o.filter != nil {
|
||||
chunks = o.filterChunks(chunks)
|
||||
}
|
||||
|
||||
// 4. Add metadata
|
||||
if o.addMetadata != nil {
|
||||
chunks = o.addChunkMetadata(chunks)
|
||||
}
|
||||
|
||||
// Re-index
|
||||
for i := range chunks {
|
||||
chunks[i].Index = i
|
||||
chunks[i].Size = len(chunks[i].GetContent())
|
||||
}
|
||||
|
||||
chunkCtx.ResultChunks = chunks
|
||||
return nil
|
||||
}
|
||||
|
||||
func (o *PostprocessOperator) Finish(chunkCtx *ChunkContext) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (o *PostprocessOperator) String() string {
|
||||
var buf strings.Builder
|
||||
buf.WriteString("postprocess:\n")
|
||||
|
||||
if o.merge != nil {
|
||||
fmt.Fprintf(&buf, " merge:\n")
|
||||
fmt.Fprintf(&buf, " target_size: %d\n", o.merge.TargetSize)
|
||||
fmt.Fprintf(&buf, " strategy: %q\n", o.merge.Strategy)
|
||||
}
|
||||
|
||||
fmt.Fprintf(&buf, " overlap:\n")
|
||||
fmt.Fprintf(&buf, " unit: %q\n", o.overlap.unit)
|
||||
fmt.Fprintf(&buf, " default:\n")
|
||||
fmt.Fprintf(&buf, " size: %d\n", o.overlap.defaultCfg.Size)
|
||||
if o.overlap.defaultCfg.Unit != "" {
|
||||
fmt.Fprintf(&buf, " unit: %q\n", o.overlap.defaultCfg.Unit)
|
||||
}
|
||||
if len(o.overlap.conditions) > 0 {
|
||||
fmt.Fprintf(&buf, " conditions:\n")
|
||||
for _, c := range o.overlap.conditions {
|
||||
fmt.Fprintf(&buf, " - name: %q\n", c.Name)
|
||||
fmt.Fprintf(&buf, " condition: %q\n", c.Condition.String())
|
||||
fmt.Fprintf(&buf, " then:\n")
|
||||
fmt.Fprintf(&buf, " size: %d\n", c.OverlapConfig.Size)
|
||||
if c.OverlapConfig.Unit != "" {
|
||||
fmt.Fprintf(&buf, " unit: %q\n", c.OverlapConfig.Unit)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if o.filter != nil {
|
||||
fmt.Fprintf(&buf, " filter:\n")
|
||||
fmt.Fprintf(&buf, " min_length: %d\n", o.filter.MinLength)
|
||||
fmt.Fprintf(&buf, " max_length: %d\n", o.filter.MaxLength)
|
||||
}
|
||||
|
||||
if o.addMetadata != nil {
|
||||
fmt.Fprintf(&buf, " add_metadata:\n")
|
||||
fmt.Fprintf(&buf, " include_index: %t\n", o.addMetadata.IncludeIndex)
|
||||
if len(o.addMetadata.CustomFields) > 0 {
|
||||
fmt.Fprintf(&buf, " custom_fields:\n")
|
||||
for k, v := range o.addMetadata.CustomFields {
|
||||
fmt.Fprintf(&buf, " %s: %q\n", k, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return buf.String()
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func parseOverlapConfig(m map[string]interface{}) overlapConfig {
|
||||
cfg := overlapConfig{}
|
||||
if size, ok := m["size"].(float64); ok {
|
||||
cfg.Size = int(size)
|
||||
}
|
||||
if u, ok := m["unit"].(string); ok {
|
||||
cfg.Unit = u
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// mergeChunks greedily merges small chunks into larger ones up to target_size.
|
||||
func (o *PostprocessOperator) mergeChunks(chunks []ChunkData) []ChunkData {
|
||||
target := o.merge.TargetSize
|
||||
if target <= 0 {
|
||||
target = 500
|
||||
}
|
||||
|
||||
var merged []ChunkData
|
||||
var buf strings.Builder
|
||||
var bufMeta map[string]interface{}
|
||||
firstIndex := 0
|
||||
|
||||
for i, c := range chunks {
|
||||
// If this single chunk already exceeds target, flush first then add
|
||||
if len([]rune(c.Content)) >= target {
|
||||
if buf.Len() > 0 {
|
||||
merged = append(merged, ChunkData{
|
||||
Content: buf.String(),
|
||||
Index: firstIndex,
|
||||
Metadata: bufMeta,
|
||||
})
|
||||
buf.Reset()
|
||||
bufMeta = nil
|
||||
}
|
||||
merged = append(merged, c)
|
||||
firstIndex = i + 1
|
||||
continue
|
||||
}
|
||||
|
||||
if buf.Len() == 0 {
|
||||
buf.WriteString(c.Content)
|
||||
bufMeta = c.Metadata
|
||||
firstIndex = c.Index
|
||||
} else {
|
||||
nextLen := len([]rune(c.Content))
|
||||
// If adding this chunk would exceed target, flush current and start new
|
||||
if buf.Len()+nextLen+1 > target {
|
||||
merged = append(merged, ChunkData{
|
||||
Content: buf.String(),
|
||||
Index: firstIndex,
|
||||
Metadata: bufMeta,
|
||||
})
|
||||
buf.Reset()
|
||||
buf.WriteString(c.Content)
|
||||
bufMeta = c.Metadata
|
||||
firstIndex = c.Index
|
||||
} else {
|
||||
buf.WriteString(" ")
|
||||
buf.WriteString(c.Content)
|
||||
// Merge metadata (last wins for overlapping keys)
|
||||
if c.Metadata != nil && bufMeta == nil {
|
||||
bufMeta = make(map[string]interface{})
|
||||
}
|
||||
for k, v := range c.Metadata {
|
||||
bufMeta[k] = v
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Flush remaining
|
||||
if buf.Len() > 0 {
|
||||
merged = append(merged, ChunkData{
|
||||
Content: buf.String(),
|
||||
Index: firstIndex,
|
||||
Metadata: bufMeta,
|
||||
})
|
||||
}
|
||||
|
||||
return merged
|
||||
}
|
||||
|
||||
// applyOverlap evaluates conditions on each chunk and prepends overlap from the previous chunk.
|
||||
func (o *PostprocessOperator) applyOverlap(chunks []ChunkData) []ChunkData {
|
||||
if len(chunks) <= 1 {
|
||||
return chunks
|
||||
}
|
||||
|
||||
result := make([]ChunkData, len(chunks))
|
||||
copy(result, chunks)
|
||||
|
||||
for i := 1; i < len(chunks); i++ {
|
||||
// Determine overlap size for chunks[i]
|
||||
cfg := o.resolveOverlapConfig(chunks[i])
|
||||
overlapSize := cfg.Size
|
||||
if overlapSize <= 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
prevContent := result[i-1].Content
|
||||
prevRunes := []rune(prevContent)
|
||||
if len(prevRunes) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
// Which unit?
|
||||
unit := cfg.Unit
|
||||
if unit == "" {
|
||||
unit = o.overlap.unit
|
||||
}
|
||||
|
||||
var overlapText string
|
||||
switch unit {
|
||||
case "sentence":
|
||||
// Take last N sentences
|
||||
sentences := splitSentencesGeneric(prevContent)
|
||||
if len(sentences) < overlapSize {
|
||||
overlapText = prevContent
|
||||
} else {
|
||||
overlapText = strings.Join(sentences[len(sentences)-overlapSize:], " ")
|
||||
}
|
||||
default: // "char"
|
||||
if overlapSize >= len(prevRunes) {
|
||||
overlapText = prevContent
|
||||
} else {
|
||||
overlapText = string(prevRunes[len(prevRunes)-overlapSize:])
|
||||
}
|
||||
}
|
||||
|
||||
result[i].Content = overlapText + result[i].Content
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// resolveOverlapConfig evaluates overlap conditions for a chunk.
|
||||
func (o *PostprocessOperator) resolveOverlapConfig(chunk ChunkData) overlapConfig {
|
||||
vars := buildExprContext(&chunk, chunk.Metadata)
|
||||
|
||||
for _, cond := range o.overlap.conditions {
|
||||
if cond.Condition == nil {
|
||||
continue
|
||||
}
|
||||
result, err := EvalCompiled(cond.Condition, vars)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if result {
|
||||
cfg := cond.OverlapConfig
|
||||
if cfg.Unit == "" {
|
||||
cfg.Unit = o.overlap.unit
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
}
|
||||
|
||||
cfg := o.overlap.defaultCfg
|
||||
if cfg.Unit == "" {
|
||||
cfg.Unit = o.overlap.unit
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// filterChunks removes chunks outside the length bounds.
|
||||
func (o *PostprocessOperator) filterChunks(chunks []ChunkData) []ChunkData {
|
||||
filtered := make([]ChunkData, 0, len(chunks))
|
||||
for _, c := range chunks {
|
||||
l := len([]rune(c.Content))
|
||||
if o.filter.MinLength > 0 && l < o.filter.MinLength {
|
||||
continue
|
||||
}
|
||||
if o.filter.MaxLength > 0 && l > o.filter.MaxLength {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, c)
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
// addChunkMetadata enriches chunks with metadata.
|
||||
func (o *PostprocessOperator) addChunkMetadata(chunks []ChunkData) []ChunkData {
|
||||
result := make([]ChunkData, len(chunks))
|
||||
for i, c := range chunks {
|
||||
if c.Metadata == nil {
|
||||
c.Metadata = make(map[string]interface{})
|
||||
}
|
||||
if o.addMetadata.IncludeIndex {
|
||||
c.Metadata["index"] = i
|
||||
}
|
||||
for field, action := range o.addMetadata.CustomFields {
|
||||
switch action {
|
||||
case "auto_detect":
|
||||
switch field {
|
||||
case "has_media_url":
|
||||
c.Metadata[field] = reMediaURL.MatchString(c.Content)
|
||||
case "has_image_url":
|
||||
c.Metadata[field] = reImageURL.MatchString(c.Content)
|
||||
case "has_video_url":
|
||||
c.Metadata[field] = reVideoURL.MatchString(c.Content)
|
||||
case "language":
|
||||
c.Metadata[field] = DetectLanguage(c.Content)
|
||||
case "length":
|
||||
c.Metadata[field] = RuneCount(c.Content)
|
||||
default:
|
||||
// Unknown auto-detect field — check for URLs generically
|
||||
c.Metadata[field] = reAnyURL.MatchString(c.Content)
|
||||
}
|
||||
default:
|
||||
c.Metadata[field] = action
|
||||
}
|
||||
}
|
||||
result[i] = c
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// splitSentencesGeneric splits text into sentences using common punctuation.
|
||||
var sentenceBoundaries = []rune{'。', '！', '？', '.', '!', '?'}
|
||||
|
||||
func splitSentencesGeneric(text string) []string {
|
||||
runes := []rune(text)
|
||||
boundSet := make(map[rune]bool)
|
||||
for _, r := range sentenceBoundaries {
|
||||
boundSet[r] = true
|
||||
}
|
||||
|
||||
var sentences []string
|
||||
var buf strings.Builder
|
||||
for _, r := range runes {
|
||||
buf.WriteRune(r)
|
||||
if boundSet[r] {
|
||||
sentences = append(sentences, strings.TrimSpace(buf.String()))
|
||||
buf.Reset()
|
||||
}
|
||||
}
|
||||
if buf.Len() > 0 {
|
||||
sentences = append(sentences, strings.TrimSpace(buf.String()))
|
||||
}
|
||||
return sentences
|
||||
}
|
||||
@@ -1,153 +0,0 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package ingestion
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/ingestion/chunk"
|
||||
)
|
||||
|
||||
/*
|
||||
DSL reference — see comment block above for the full JSON structure.
|
||||
|
||||
Pipeline stages:
|
||||
1. "preprocess" → chunk.PreprocessOperator
|
||||
2. "split" → chunk.SplitOperator
|
||||
3. "postprocess" → chunk.PostprocessOperator
|
||||
*/
|
||||
|
||||
// ChunkPlan holds the ordered pipeline operators.
|
||||
type ChunkPlan struct {
|
||||
Operators []chunk.Operator
|
||||
Version string
|
||||
Description string
|
||||
Name string
|
||||
}
|
||||
|
||||
// ChunkEngine parses DSL JSON into a plan and executes it.
|
||||
type ChunkEngine struct{}
|
||||
|
||||
func NewChunkEngine() *ChunkEngine {
|
||||
return &ChunkEngine{}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Compile — compile DSL JSON into an ordered operator list
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func (e *ChunkEngine) Compile(dsl string) (*ChunkPlan, error) {
|
||||
var parsed map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(dsl), &parsed); err != nil {
|
||||
return nil, fmt.Errorf("compile DSL: %w", err)
|
||||
}
|
||||
|
||||
plan := &ChunkPlan{}
|
||||
|
||||
pipelineRaw, ok := parsed["pipeline"].([]interface{})
|
||||
if !ok || len(pipelineRaw) == 0 {
|
||||
return plan, nil
|
||||
}
|
||||
|
||||
plan.Name, ok = parsed["name"].(string)
|
||||
if !ok {
|
||||
plan.Name = "No name"
|
||||
}
|
||||
plan.Description, ok = parsed["description"].(string)
|
||||
if !ok {
|
||||
plan.Description = "No description"
|
||||
}
|
||||
plan.Version, ok = parsed["version"].(string)
|
||||
if !ok {
|
||||
plan.Version = "1.0"
|
||||
}
|
||||
|
||||
for i, operatorRaw := range pipelineRaw {
|
||||
var operator map[string]interface{}
|
||||
operator, ok = operatorRaw.(map[string]interface{})
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("pipeline[%d]: expected object", i)
|
||||
}
|
||||
|
||||
var op chunk.Operator
|
||||
var err error
|
||||
operatorName, _ := operator["operator"].(string)
|
||||
switch operatorName {
|
||||
case "preprocess":
|
||||
op, err = chunk.NewPreprocessOperator(operator)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create preprocess operator[%d]: %w", i, err)
|
||||
}
|
||||
case "split":
|
||||
op, err = chunk.NewSplitOperator(operator)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create split operator[%d]: %w", i, err)
|
||||
}
|
||||
case "postprocess":
|
||||
op, err = chunk.NewPostprocessOperator(operator)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create postprocess operator[%d]: %w", i, err)
|
||||
}
|
||||
default:
|
||||
return nil, fmt.Errorf("pipeline[%d]: unknown operator %s", i, operatorName)
|
||||
}
|
||||
delete(operator, "operator")
|
||||
|
||||
plan.Operators = append(plan.Operators, op)
|
||||
}
|
||||
|
||||
return plan, nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Execute — run the pipeline operators on input text
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func (e *ChunkEngine) Execute(plan *ChunkPlan, text string) (*chunk.ChunkContext, error) {
|
||||
chunkContext := &chunk.ChunkContext{Origin: text}
|
||||
|
||||
for i, op := range plan.Operators {
|
||||
if err := op.Prepare(chunkContext); err != nil {
|
||||
return nil, fmt.Errorf("re-prepare operator[%d]: %w", i, err)
|
||||
}
|
||||
|
||||
if err := op.Execute(chunkContext); err != nil {
|
||||
return nil, fmt.Errorf("execute operator[%d]: %w", i, err)
|
||||
}
|
||||
|
||||
if err := op.Finish(chunkContext); err != nil {
|
||||
return nil, fmt.Errorf("finish operator[%d]: %w", i, err)
|
||||
}
|
||||
}
|
||||
|
||||
return chunkContext, nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Explain — describe the plan in human-readable form
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func (e *ChunkEngine) Explain(plan *ChunkPlan) (string, error) {
|
||||
var buf strings.Builder
|
||||
buf.WriteString("Chunk Pipeline Plan:\n")
|
||||
for i, op := range plan.Operators {
|
||||
buf.WriteString(fmt.Sprintf(" [%d] %s\n", i, op.String()))
|
||||
}
|
||||
return buf.String(), nil
|
||||
}
|
||||
@@ -1,346 +0,0 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package ingestion
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/ingestion/chunk"
|
||||
)
|
||||
|
||||
// The full DSL example from chunk_engine.go L22-L95.
|
||||
var mediaAwareDSL = `{
|
||||
"version": "1.0",
|
||||
"name": "media_aware_chunking",
|
||||
"description": "Disable overlap when encountering image/video URLs",
|
||||
|
||||
"pipeline": [
|
||||
{
|
||||
"operator": "preprocess",
|
||||
"normalize_newlines": true,
|
||||
"strip_whitespace": true,
|
||||
"remove_empty_lines": true
|
||||
},
|
||||
|
||||
{
|
||||
"operator": "split",
|
||||
"strategy": "sentence",
|
||||
"params": {
|
||||
"boundaries": [". ", "! ", "? ", "\n"],
|
||||
"keep_separators": true
|
||||
}
|
||||
},
|
||||
|
||||
{
|
||||
"operator": "postprocess",
|
||||
"merge": {
|
||||
"target_size": 500,
|
||||
"strategy": "greedy"
|
||||
},
|
||||
"overlap": {
|
||||
"enabled": true,
|
||||
"unit": "char",
|
||||
"conditions": [
|
||||
{
|
||||
"name": "Contains media URL",
|
||||
"if": "has_media_url = true",
|
||||
"then": {"size": 0}
|
||||
},
|
||||
{
|
||||
"name": "Contains image URL",
|
||||
"if": "has_image_url = true",
|
||||
"then": {"size": 0}
|
||||
},
|
||||
{
|
||||
"name": "Contains video URL",
|
||||
"if": "has_video_url = true",
|
||||
"then": {"size": 0}
|
||||
},
|
||||
{
|
||||
"name": "Normal English long sentence",
|
||||
"if": "language = 'en' AND length > 50 AND has_media_url = false",
|
||||
"then": {"size": 1, "unit": "sentence"}
|
||||
},
|
||||
{
|
||||
"name": "Normal English short sentence",
|
||||
"if": "language = 'en' AND length <= 50 AND has_media_url = false",
|
||||
"then": {"size": 30}
|
||||
}
|
||||
],
|
||||
"default": {"size": 50}
|
||||
},
|
||||
"filter": {
|
||||
"min_length": 10,
|
||||
"max_length": 1200
|
||||
},
|
||||
"add_metadata": {
|
||||
"include_index": true,
|
||||
"custom_fields": {
|
||||
"has_media_url": "auto_detect"
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}`
|
||||
|
||||
var minimalDSL = `{
|
||||
"pipeline": [
|
||||
{"operator": "preprocess", "normalize_newlines": true},
|
||||
{"operator": "split", "strategy": "sentence", "params": {"boundaries": ["\n"], "keep_separators": false}},
|
||||
{"operator": "postprocess", "filter": {"min_length": 1}}
|
||||
]
|
||||
}`
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Plan success tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestPlan_FullDSL(t *testing.T) {
|
||||
engine := NewChunkEngine()
|
||||
plan, err := engine.Compile(mediaAwareDSL)
|
||||
if err != nil {
|
||||
t.Fatalf("Compile(mediaAwareDSL) unexpected error: %v", err)
|
||||
}
|
||||
if plan == nil {
|
||||
t.Fatal("Plan returned nil")
|
||||
}
|
||||
|
||||
// Must have exactly 3 operators
|
||||
if len(plan.Operators) != 3 {
|
||||
t.Fatalf("expected 3 operators, got %d", len(plan.Operators))
|
||||
}
|
||||
|
||||
// Verify operator types in order
|
||||
expectedTypes := []string{
|
||||
"*chunk.PreprocessOperator",
|
||||
"*chunk.SplitOperator",
|
||||
"*chunk.PostprocessOperator",
|
||||
}
|
||||
for i, op := range plan.Operators {
|
||||
typ := fmt.Sprintf("%T", op)
|
||||
if typ != expectedTypes[i] {
|
||||
t.Errorf("operator[%d]: expected %s, got %s", i, expectedTypes[i], typ)
|
||||
}
|
||||
}
|
||||
|
||||
// Verify operators implement Operator interface
|
||||
for i, op := range plan.Operators {
|
||||
var iface chunk.Operator = op
|
||||
_ = iface // compile-time check
|
||||
if op == nil {
|
||||
t.Errorf("operator[%d] is nil", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlan_MinimalDSL(t *testing.T) {
|
||||
engine := NewChunkEngine()
|
||||
plan, err := engine.Compile(minimalDSL)
|
||||
if err != nil {
|
||||
t.Fatalf("Compile(minimalDSL) unexpected error: %v", err)
|
||||
}
|
||||
if plan == nil {
|
||||
t.Fatal("Plan returned nil")
|
||||
}
|
||||
if len(plan.Operators) != 3 {
|
||||
t.Fatalf("expected 3 operators, got %d", len(plan.Operators))
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Plan error tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestPlan_InvalidJSON(t *testing.T) {
|
||||
engine := NewChunkEngine()
|
||||
invalid := `{bad json}`
|
||||
plan, err := engine.Compile(invalid)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for invalid JSON, got nil")
|
||||
}
|
||||
if plan != nil {
|
||||
t.Fatal("expected nil plan on error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlan_UnknownOperator(t *testing.T) {
|
||||
engine := NewChunkEngine()
|
||||
dsl := `{"pipeline": [{"operator": "unknown_operator"}]}`
|
||||
plan, err := engine.Compile(dsl)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for unknown operator, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "unknown_operator") {
|
||||
t.Errorf("error should mention unknown operator, got: %v", err)
|
||||
}
|
||||
if plan != nil {
|
||||
t.Fatal("expected nil plan on error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlan_EmptyPipeline(t *testing.T) {
|
||||
engine := NewChunkEngine()
|
||||
dsl := `{"pipeline": []}`
|
||||
plan, err := engine.Compile(dsl)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if plan == nil {
|
||||
t.Fatal("plan should not be nil")
|
||||
}
|
||||
if len(plan.Operators) != 0 {
|
||||
t.Fatalf("expected 0 operators, got %d", len(plan.Operators))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlan_MissingPipeline(t *testing.T) {
|
||||
engine := NewChunkEngine()
|
||||
dsl := `{"version": "1.0"}`
|
||||
plan, err := engine.Compile(dsl)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if plan == nil {
|
||||
t.Fatal("plan should not be nil")
|
||||
}
|
||||
if len(plan.Operators) != 0 {
|
||||
t.Fatalf("expected 0 operators, got %d", len(plan.Operators))
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Plan + Execute integration test
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestPlan_Execute_FullPipeline(t *testing.T) {
|
||||
engine := NewChunkEngine()
|
||||
plan, err := engine.Compile(mediaAwareDSL)
|
||||
if err != nil {
|
||||
t.Fatalf("Plan error: %v", err)
|
||||
}
|
||||
|
||||
inputText := `这是第一句话。这是第二句话！这是第三句话？\n这是第四句话。`
|
||||
ctx, err := engine.Execute(plan, inputText)
|
||||
if err != nil {
|
||||
t.Fatalf("Execute error: %v", err)
|
||||
}
|
||||
if ctx == nil {
|
||||
t.Fatal("Execute returned nil context")
|
||||
}
|
||||
if len(ctx.ResultChunks) == 0 {
|
||||
t.Fatal("expected at least one chunk after execution")
|
||||
}
|
||||
for i, c := range ctx.ResultChunks {
|
||||
if c.Content == "" {
|
||||
t.Errorf("chunk[%d] has empty content", i)
|
||||
}
|
||||
println(c.Content)
|
||||
if c.Metadata == nil {
|
||||
t.Errorf("chunk[%d] has nil metadata", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlan_Execute_MinimalPipeline(t *testing.T) {
|
||||
engine := NewChunkEngine()
|
||||
plan, err := engine.Compile(minimalDSL)
|
||||
if err != nil {
|
||||
t.Fatalf("Plan error: %v", err)
|
||||
}
|
||||
|
||||
inputText := "Hello world.\nGoodbye world."
|
||||
ctx, err := engine.Execute(plan, inputText)
|
||||
if err != nil {
|
||||
t.Fatalf("Execute error: %v", err)
|
||||
}
|
||||
if ctx == nil {
|
||||
t.Fatal("Execute returned nil context")
|
||||
}
|
||||
if len(ctx.ResultChunks) == 0 {
|
||||
t.Fatal("expected at least one chunk after execution")
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Edge cases
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestPlan_Explain(t *testing.T) {
|
||||
engine := NewChunkEngine()
|
||||
plan, err := engine.Compile(mediaAwareDSL)
|
||||
if err != nil {
|
||||
t.Fatalf("Plan error: %v", err)
|
||||
}
|
||||
|
||||
_, err = engine.Explain(plan)
|
||||
if err != nil {
|
||||
t.Fatalf("Explain error: %v", err)
|
||||
}
|
||||
//fmt.Println(explanation)
|
||||
}
|
||||
|
||||
func TestPlan_ReuseEngine(t *testing.T) {
|
||||
engine := NewChunkEngine()
|
||||
|
||||
// First plan
|
||||
plan1, err := engine.Compile(mediaAwareDSL)
|
||||
if err != nil {
|
||||
t.Fatalf("first Plan error: %v", err)
|
||||
}
|
||||
|
||||
// Second plan from the same engine
|
||||
plan2, err := engine.Compile(minimalDSL)
|
||||
if err != nil {
|
||||
t.Fatalf("second Plan error: %v", err)
|
||||
}
|
||||
|
||||
if len(plan1.Operators) != len(plan2.Operators) {
|
||||
t.Errorf("plan1 has %d operators, plan2 has %d", len(plan1.Operators), len(plan2.Operators))
|
||||
}
|
||||
}
|
||||
|
||||
// Benchmark
|
||||
func BenchmarkPlan_FullDSL(b *testing.B) {
|
||||
engine := NewChunkEngine()
|
||||
dsl := mediaAwareDSL
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := engine.Compile(dsl)
|
||||
if err != nil {
|
||||
b.Fatalf("Plan error: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkPlan_Execute_FullDSL(b *testing.B) {
|
||||
engine := NewChunkEngine()
|
||||
dsl := mediaAwareDSL
|
||||
plan, err := engine.Compile(dsl)
|
||||
if err != nil {
|
||||
b.Fatalf("Plan error: %v", err)
|
||||
}
|
||||
inputText := strings.Repeat("这是第一句话。这是第二句话！这是第三句话？\n", 100)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := engine.Execute(plan, inputText)
|
||||
if err != nil {
|
||||
b.Fatalf("Execute error: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -20,6 +20,8 @@
|
||||
//
|
||||
// The architecture mirrors the Python rag/graphrag/ner package so that both
|
||||
// code paths produce identical output (verified by test).
|
||||
//go:build cgo_thincner
|
||||
|
||||
package extractor
|
||||
|
||||
// #cgo CXXFLAGS: -std=c++20 -I${SRCDIR}/../../..
|
||||
|
||||
@@ -27,4 +27,6 @@
|
||||
// result, err := ext.Extract("Apple Inc. was founded by Steve Jobs.", true)
|
||||
// for _, e := range result.Entities { ... }
|
||||
// for _, r := range result.Relations { ... }
|
||||
//go:build cgo_thincner
|
||||
|
||||
package extractor
|
||||
|
||||
206
internal/ingestion/compilation/extractor/no_cgo.go
Normal file
206
internal/ingestion/compilation/extractor/no_cgo.go
Normal file
@@ -0,0 +1,206 @@
|
||||
//go:build !cgo_thincner
|
||||
|
||||
//
|
||||
// Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Stub for non-CGO builds.
|
||||
//
|
||||
// The C++ ThincNER engine (and its ThincParser / ThincTagger
|
||||
// siblings) is wired through cgo in ner.go, parser_go.go, and
|
||||
// ner_extractor.go — those files carry `//go:build cgo_thincner`.
|
||||
// Without the `cgo_thincner` build tag, the test binary cannot link the missing
|
||||
// ThincNER_* / ThincParser_* / ThincTagger_* symbols.
|
||||
//
|
||||
// This file declares the same exported surface the rest of the
|
||||
// package depends on (Entity / Relation / ExtractionResult /
|
||||
// Extractor / RunParser / RunTagger / NewExtractor / DetectLanguage
|
||||
// / ParseTokensWithParser / ExtractRelations) so the pure-Go files
|
||||
// in the package — primarily dep_relation.go and the relation-
|
||||
// extraction pure-Go logic — continue to compile.
|
||||
//
|
||||
// The cgo-backed functions return an explicit ErrNoCGO error so
|
||||
// any caller that reaches the C++ engine on a no-CGO build fails
|
||||
// loudly rather than silently degrading. Production builds use
|
||||
// the `cgo_thincner` path only when that explicit build tag is enabled.
|
||||
|
||||
package extractor
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// ErrNoCGO is returned by all cgo-backed entry points on non-CGO
|
||||
// builds. The Python-side ThincNER engine requires the C++ static
|
||||
// library at internal/cpp/cmake-build-release/librag_tokenizer_c_api.a;
|
||||
// without it the package compiles but the inference paths fail
|
||||
// with this error.
|
||||
var ErrNoCGO = errors.New("extractor: CGO disabled — ThincNER / ThincParser / ThincTagger unavailable")
|
||||
|
||||
// Entity mirrors the cgo-backed declaration.
|
||||
type Entity struct {
|
||||
Text string `json:"text"`
|
||||
Label string `json:"label"`
|
||||
StartChar int `json:"start_char"`
|
||||
EndChar int `json:"end_char"`
|
||||
Confidence float64 `json:"confidence"`
|
||||
AppType string `json:"app_type,omitempty"`
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// Relation mirrors the cgo-backed declaration.
|
||||
type Relation struct {
|
||||
Subject Entity `json:"subject"`
|
||||
Predicate string `json:"predicate"`
|
||||
Object Entity `json:"object"`
|
||||
Confidence float64 `json:"confidence"`
|
||||
Context string `json:"context,omitempty"`
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// ExtractionResult mirrors the cgo-backed declaration.
|
||||
type ExtractionResult struct {
|
||||
Entities []Entity `json:"entities"`
|
||||
Relations []Relation `json:"relations"`
|
||||
Language string `json:"language,omitempty"`
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// Extractor mirrors the cgo-backed declaration. The no-CGO
|
||||
// instance returns ErrNoCGO from every inference call but
|
||||
// otherwise satisfies the same surface as the CGO build.
|
||||
type Extractor struct {
|
||||
mu sync.Mutex
|
||||
Lang string
|
||||
ConfidenceThreshold float64
|
||||
IncludeTokens bool
|
||||
}
|
||||
|
||||
// ModelPredictor is the dependency-injection seam used by tests;
|
||||
// no-CGO builds mirror the type so existing test helpers compile.
|
||||
type ModelPredictor func(tokensJSON string) (string, error)
|
||||
|
||||
// NewExtractor returns a no-CGO Extractor. Inference calls fail
|
||||
// with ErrNoCGO. The pure-Go relation-extraction helpers
|
||||
// (DepExtractRelations / ExtractRelations) do not depend on cgo
|
||||
// and continue to function.
|
||||
func NewExtractor(lang string) *Extractor {
|
||||
return &Extractor{Lang: lang}
|
||||
}
|
||||
|
||||
// RunParser is the no-CGO stub. The cgo-backed implementation is
|
||||
// in parser_go.go (build tag: cgo_thincner).
|
||||
func RunParser(nerDir, parserDir string, tokensJSON string) (string, error) {
|
||||
return "", ErrNoCGO
|
||||
}
|
||||
|
||||
// RunTagger is the no-CGO stub. The cgo-backed implementation is
|
||||
// in parser_go.go (build tag: cgo_thincner).
|
||||
func RunTagger(nerDir, taggerDir string, tokensJSON string) (string, error) {
|
||||
return "", ErrNoCGO
|
||||
}
|
||||
|
||||
// ParseTokensWithParser is the no-CGO stub for the typed wrapper
|
||||
// around RunParser. Production callers should check the error
|
||||
// before using the returned slice.
|
||||
func ParseTokensWithParser(nerDir, parserDir string, tokens []string) ([]DepTokenC, error) {
|
||||
tj, _ := json.Marshal(tokens)
|
||||
resultJSON, err := RunParser(nerDir, parserDir, string(tj))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var tokensC []DepTokenC
|
||||
if err := json.Unmarshal([]byte(resultJSON), &tokensC); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return tokensC, nil
|
||||
}
|
||||
|
||||
// DepTokenC mirrors the cgo-backed declaration.
|
||||
type DepTokenC struct {
|
||||
Text string `json:"text"`
|
||||
Head int `json:"head"`
|
||||
Dep string `json:"dep"`
|
||||
Index int `json:"index"`
|
||||
}
|
||||
|
||||
// DetectLanguage is pure-Go and unchanged by this stub file.
|
||||
// The CGO build's DetectLanguage is identical; keeping the
|
||||
// declaration here makes the no-CGO path self-contained.
|
||||
func DetectLanguage(text string) string {
|
||||
return detectLanguageNoCGO(text)
|
||||
}
|
||||
|
||||
// detectLanguageNoCGO is the pure-Go language detection helper.
|
||||
// Mirrors the CGO file's DetectLanguage without any cgo deps.
|
||||
//
|
||||
// The detector prefers the most specific Unicode range:
|
||||
//
|
||||
// - Hiragana / Katakana → "ja" (Japanese)
|
||||
// - CJK Unified Ideographs → "zh" (Chinese)
|
||||
// - Hangul Syllables → "ko" (Korean)
|
||||
// - Cyrillic → "ru"
|
||||
// - Arabic → "ar"
|
||||
// - Devanagari → "hi"
|
||||
// - Otherwise → "en"
|
||||
//
|
||||
// This matches the production heuristic closely enough for the
|
||||
// no-CGO test binary link; production code that depends on
|
||||
// DetectLanguage's accuracy should set `-tags=cgo_thincner` to
|
||||
// select the real implementation.
|
||||
func detectLanguageNoCGO(text string) string {
|
||||
if len(text) == 0 {
|
||||
return "en"
|
||||
}
|
||||
var ja, zh, ko, cyrillic, arabic, devanagari, latin int
|
||||
for _, r := range text {
|
||||
switch {
|
||||
case r >= 0x3040 && r <= 0x309F: // Hiragana
|
||||
ja++
|
||||
case r >= 0x30A0 && r <= 0x30FF: // Katakana
|
||||
ja++
|
||||
case r >= 0x4E00 && r <= 0x9FFF: // CJK Unified
|
||||
zh++
|
||||
case r >= 0xAC00 && r <= 0xD7AF: // Hangul
|
||||
ko++
|
||||
case r >= 0x0400 && r <= 0x04FF: // Cyrillic
|
||||
cyrillic++
|
||||
case r >= 0x0600 && r <= 0x06FF: // Arabic
|
||||
arabic++
|
||||
case r >= 0x0900 && r <= 0x097F: // Devanagari
|
||||
devanagari++
|
||||
case (r >= 'A' && r <= 'Z') || (r >= 'a' && r <= 'z'):
|
||||
latin++
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case ja > 0:
|
||||
return "ja"
|
||||
case zh > 0:
|
||||
return "zh"
|
||||
case ko > 0:
|
||||
return "ko"
|
||||
case cyrillic > 0:
|
||||
return "ru"
|
||||
case arabic > 0:
|
||||
return "ar"
|
||||
case devanagari > 0:
|
||||
return "hi"
|
||||
default:
|
||||
return "en"
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,5 @@
|
||||
//go:build cgo_thincner
|
||||
|
||||
package extractor
|
||||
|
||||
/*
|
||||
|
||||
215
internal/ingestion/component/chunker/common.go
Normal file
215
internal/ingestion/component/chunker/common.go
Normal file
@@ -0,0 +1,215 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
"ragflow/internal/tokenizer"
|
||||
)
|
||||
|
||||
// newChunkerByName dispatches the DSL name to a typed constructor.
|
||||
// Centralised here so each chunker file only needs an init() that
|
||||
// declares its registered name (see register.go). The returned
|
||||
// runtime.Component interface is satisfied directly by each
|
||||
// constructor (NewTokenChunker etc.) — no intermediate assertion
|
||||
// is needed.
|
||||
func newChunkerByName(name string, params map[string]any) (runtime.Component, error) {
|
||||
switch name {
|
||||
case ComponentNameTokenChunker:
|
||||
return NewTokenChunker(params)
|
||||
case ComponentNameTitleChunker:
|
||||
return NewTitleChunker(params)
|
||||
case ComponentNameGroupTitleChunker:
|
||||
return NewGroupTitleChunker(params)
|
||||
case ComponentNameHierarchyTitleChunker:
|
||||
return NewHierarchyTitleChunker(params)
|
||||
default:
|
||||
return nil, fmt.Errorf("chunker: unknown component %q", name)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// numeric / list conversion helpers (shared across chunker variants)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// numericFromAny normalises JSON-decoded ints to float64 so the
|
||||
// schema-defaults-vs-Param-Update convention doesn't depend on the
|
||||
// encoding source (yaml/toml/json all behave the same).
|
||||
func numericFromAny(v any) (float64, bool) {
|
||||
switch x := v.(type) {
|
||||
case float64:
|
||||
return x, true
|
||||
case float32:
|
||||
return float64(x), true
|
||||
case int:
|
||||
return float64(x), true
|
||||
case int32:
|
||||
return float64(x), true
|
||||
case int64:
|
||||
return float64(x), true
|
||||
case uint:
|
||||
return float64(x), true
|
||||
case uint32:
|
||||
return float64(x), true
|
||||
case uint64:
|
||||
return float64(x), true
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func stringListFromAny(in []any) []string {
|
||||
out := make([]string, 0, len(in))
|
||||
for _, x := range in {
|
||||
if s, ok := x.(string); ok && s != "" {
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// regex / split helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// compileDelimPattern joins all delimiter entries into a single
|
||||
// alternation. Entries wrapped in backticks are treated as regex
|
||||
// literals and regex-escaped; plain strings are simply regex-escaped.
|
||||
// Longer patterns win (matches python `sorted(set, key=len, reverse=True)`).
|
||||
func compileDelimPattern(delims []string) *regexp.Regexp {
|
||||
var custom []string
|
||||
var plain []string
|
||||
for _, d := range delims {
|
||||
if d == "" {
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(d, "`") && strings.HasSuffix(d, "`") && len(d) >= 2 {
|
||||
custom = append(custom, regexp.QuoteMeta(d[1:len(d)-1]))
|
||||
} else {
|
||||
plain = append(plain, regexp.QuoteMeta(d))
|
||||
}
|
||||
}
|
||||
all := append(plain, custom...)
|
||||
if len(all) == 0 {
|
||||
return nil
|
||||
}
|
||||
sort.SliceStable(all, func(i, j int) bool { return len(all[i]) > len(all[j]) })
|
||||
return regexp.MustCompile(strings.Join(all, "|"))
|
||||
}
|
||||
|
||||
// splitKeepingDelim is the Go equivalent of python's
|
||||
// `re.split((pattern), text, flags=re.DOTALL)` with the matched
|
||||
// delimiter preserved (alternation keeps the original delimiter text
|
||||
// in the output stream so the rebuilding at token_chunker.py:88-93
|
||||
// stays lossy-free).
|
||||
func splitKeepingDelim(text string, pattern *regexp.Regexp) []string {
|
||||
if pattern == nil {
|
||||
return []string{text}
|
||||
}
|
||||
idxs := pattern.FindAllStringSubmatchIndex(text, -1)
|
||||
if len(idxs) == 0 {
|
||||
return []string{text}
|
||||
}
|
||||
var out []string
|
||||
cursor := 0
|
||||
for _, idx := range idxs {
|
||||
start, end := idx[0], idx[1]
|
||||
if start > cursor {
|
||||
out = append(out, text[cursor:start])
|
||||
}
|
||||
out = append(out, text[start:end])
|
||||
cursor = end
|
||||
}
|
||||
if cursor < len(text) {
|
||||
out = append(out, text[cursor:])
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// chunk-doc helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// itemText returns the text payload from a JSON-style chunk item,
|
||||
// preferring "text", then "content_with_weight".
|
||||
func itemText(it schema.ChunkDoc) (string, bool) {
|
||||
if it.Text != "" {
|
||||
return it.Text, true
|
||||
}
|
||||
if it.ContentWithWeight != "" {
|
||||
return it.ContentWithWeight, true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// itemDocType mirrors _build_json_chunks's type derivation.
|
||||
func itemDocType(it schema.ChunkDoc) string {
|
||||
switch strings.ToLower(strings.TrimSpace(it.DocType)) {
|
||||
case "table":
|
||||
return "table"
|
||||
case "image":
|
||||
return "image"
|
||||
}
|
||||
return "text"
|
||||
}
|
||||
|
||||
// itemTextOrFallback returns the item's preferred text, or "".
|
||||
func itemTextOrFallback(it schema.ChunkDoc) string {
|
||||
if t, ok := itemText(it); ok {
|
||||
return t
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// tokenizeStr is the shared NumTokensFromString wrapper used by
|
||||
// Table/Image context attachment. Lives here so we can centrally
|
||||
// swizzle the count strategy in one place if needed.
|
||||
func tokenizeStr(s string) int { return tokenizer.NumTokensFromString(s) }
|
||||
|
||||
// toString normalises a chunk-map field to a string. Empty strings
|
||||
// for missing fields.
|
||||
func toString(v any) string {
|
||||
if v == nil {
|
||||
return ""
|
||||
}
|
||||
if s, ok := v.(string); ok {
|
||||
return s
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// emptyOutputs returns the canonical no-chunks payload.
|
||||
func emptyOutputs() map[string]any {
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{},
|
||||
}
|
||||
}
|
||||
|
||||
func emptyChunkDocs() []schema.ChunkDoc { return []schema.ChunkDoc{} }
|
||||
|
||||
func chunkOutputs(chunks []schema.ChunkDoc) map[string]any {
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": schema.ChunkDocsToMaps(chunks),
|
||||
}
|
||||
}
|
||||
331
internal/ingestion/component/chunker/group.go
Normal file
331
internal/ingestion/component/chunker/group.go
Normal file
@@ -0,0 +1,331 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// SCOPE (honest) for group.go:
|
||||
//
|
||||
// - Implements the GroupTitleChunker variant: aggregates adjacent
|
||||
// text records into chunks that span multiple body records while
|
||||
// staying inside one heading section.
|
||||
//
|
||||
// - PARALLELISM: Parallelism() advertises a fan-out hint to outer
|
||||
// executors. Heading detection stays sequential; grouping work is
|
||||
// local to one invocation.
|
||||
//
|
||||
// - MIRRORS python `_build_section_ids` + `GroupTitleChunker.build_chunks`:
|
||||
// consecutive records with the same (target_level-derived) sec_id
|
||||
// are merged up to MIN_GROUP_TOKENS / MAX_GROUP_TOKENS, then
|
||||
// emitted as one chunk per group.
|
||||
//
|
||||
// - No PDF-position merge is performed (mirrors TokenChunker SCOPE
|
||||
// notes — those land with the deepdoc/parser port).
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
"ragflow/internal/tokenizer"
|
||||
)
|
||||
|
||||
const ComponentNameGroupTitleChunker = "GroupTitleChunker"
|
||||
|
||||
// MIN_GROUP_TOKENS / MAX_GROUP_TOKENS mirror the python constants in
|
||||
// group_chunker.py:22-23. They drive the merge heuristic for adjacent
|
||||
// text records.
|
||||
const (
|
||||
minGroupTokens = 32
|
||||
maxGroupTokens = 1024
|
||||
)
|
||||
|
||||
// resolveTargetLevel mirrors common.py:resolve_target_level: pick the
|
||||
// n-th smallest heading level from the per-line `levels` vector.
|
||||
func resolveTargetLevel(levels []int, hierarchy int) int {
|
||||
tiers := make([]int, 0, len(levels))
|
||||
seen := make(map[int]bool)
|
||||
for _, l := range levels {
|
||||
if l > 0 && l < bodyLevel && !seen[l] {
|
||||
seen[l] = true
|
||||
tiers = append(tiers, l)
|
||||
}
|
||||
}
|
||||
if len(tiers) == 0 {
|
||||
return 0
|
||||
}
|
||||
// ascending
|
||||
for i := 1; i < len(tiers); i++ {
|
||||
for j := i; j > 0 && tiers[j-1] > tiers[j]; j-- {
|
||||
tiers[j-1], tiers[j] = tiers[j], tiers[j-1]
|
||||
}
|
||||
}
|
||||
if hierarchy < 1 {
|
||||
hierarchy = 1
|
||||
}
|
||||
if hierarchy > len(tiers) {
|
||||
hierarchy = len(tiers)
|
||||
}
|
||||
return tiers[hierarchy-1]
|
||||
}
|
||||
|
||||
// _buildSectionIDs computes, for each input level, the section id
|
||||
// (`sid`) under which that record falls. Each heading encounter
|
||||
// increments sid by 1; body lines share the active sid.
|
||||
func buildSectionIDs(levels []int, targetLevel int) []int {
|
||||
secIDs := make([]int, len(levels))
|
||||
sid := 0
|
||||
for i, lvl := range levels {
|
||||
if i > 0 && targetLevel > 0 && lvl <= targetLevel {
|
||||
sid++
|
||||
}
|
||||
secIDs[i] = sid
|
||||
}
|
||||
return secIDs
|
||||
}
|
||||
|
||||
// invokeGroup runs the GroupTitleChunker strategy against the
|
||||
// supplied inputs. Detected headings + adjacent merges happen in two
|
||||
// goroutines (heading detection sequential, then a fan-out over
|
||||
// record-buckets for the merge pass).
|
||||
func invokeGroup(_ context.Context, inputs map[string]any, p *titleChunkerParam) (map[string]any, error) {
|
||||
records := extractLineRecords(inputs)
|
||||
if len(records) == 0 {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
lines := make([]string, len(records))
|
||||
for i, r := range records {
|
||||
lines[i] = r.text
|
||||
}
|
||||
ctx := newLevelContext(lines, p)
|
||||
levels := ctx.Levels()
|
||||
|
||||
targetLevel := resolveTargetLevel(levels, hierarchyOr(p, ctx.mostLevel))
|
||||
secIDs := buildSectionIDs(levels, targetLevel)
|
||||
|
||||
groups := groupRecords(records, secIDs, p)
|
||||
if p.RootChunkAsHeading && len(groups) > 1 {
|
||||
groups = applyRootAsHeading(groups)
|
||||
}
|
||||
chunks := make([]map[string]any, 0, len(groups))
|
||||
for _, g := range groups {
|
||||
chunks = append(chunks, map[string]any{"text": joinGroupText(g)})
|
||||
}
|
||||
if len(chunks) == 0 {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": chunks,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// groupRecords mirrors `GroupTitleChunker.build_chunks`: merges
|
||||
// adjacent text records while staying inside the same logical section
|
||||
// (matching last_sid) and token-budget constraints.
|
||||
func groupRecords(records []lineRecord, secIDs []int, p *titleChunkerParam) [][]lineRecord {
|
||||
if len(records) == 0 {
|
||||
return nil
|
||||
}
|
||||
var recordGroups [][]lineRecord
|
||||
var currentGroup []lineRecord
|
||||
tkCnt := 0
|
||||
lastSID := -2
|
||||
|
||||
for i, rec := range records {
|
||||
secID := secIDs[i]
|
||||
if !rec.isText() {
|
||||
if len(currentGroup) > 0 {
|
||||
recordGroups = append(recordGroups, append([]lineRecord(nil), currentGroup...))
|
||||
}
|
||||
recordGroups = append(recordGroups, []lineRecord{rec})
|
||||
currentGroup = currentGroup[:0]
|
||||
tkCnt = 0
|
||||
lastSID = -2
|
||||
continue
|
||||
}
|
||||
text := trim(rec.text)
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
tokenCount := tokenizer.NumTokensFromString(text)
|
||||
shouldMerge := len(currentGroup) > 0 &&
|
||||
currentGroup[0].isText() &&
|
||||
(tkCnt < minGroupTokens || (tkCnt < maxGroupTokens && secID == lastSID))
|
||||
if shouldMerge {
|
||||
currentGroup = append(currentGroup, rec)
|
||||
tkCnt += tokenCount
|
||||
} else {
|
||||
if len(currentGroup) > 0 {
|
||||
recordGroups = append(recordGroups, append([]lineRecord(nil), currentGroup...))
|
||||
}
|
||||
currentGroup = []lineRecord{rec}
|
||||
tkCnt = tokenCount
|
||||
}
|
||||
lastSID = secID
|
||||
}
|
||||
if len(currentGroup) > 0 {
|
||||
recordGroups = append(recordGroups, currentGroup)
|
||||
}
|
||||
return recordGroups
|
||||
}
|
||||
|
||||
// applyRootAsHeading mirrors the `root_chunk_as_heading` branch in
|
||||
// common.py:build_chunks_from_record_groups — prepending the root
|
||||
// text to every following chunk and dropping the root chunk itself.
|
||||
func applyRootAsHeading(groups [][]lineRecord) [][]lineRecord {
|
||||
if len(groups) < 2 {
|
||||
return groups
|
||||
}
|
||||
rootText := joinGroupText(groups[0])
|
||||
for i := 1; i < len(groups); i++ {
|
||||
groups[i] = prependJoin(groups[i], rootText)
|
||||
}
|
||||
return groups[1:]
|
||||
}
|
||||
|
||||
func joinGroupText(g []lineRecord) string {
|
||||
var sb strings.Builder
|
||||
for _, r := range g {
|
||||
sb.WriteString(r.text)
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
func prependJoin(g []lineRecord, prefix string) []lineRecord {
|
||||
if prefix == "" {
|
||||
return g
|
||||
}
|
||||
extra := lineRecord{text: prefix, docType: "text"}
|
||||
if len(g) == 0 {
|
||||
return []lineRecord{extra}
|
||||
}
|
||||
out := make([]lineRecord, 0, len(g)+1)
|
||||
out = append(out, extra)
|
||||
out = append(out, g...)
|
||||
return out
|
||||
}
|
||||
|
||||
// extractLineRecords reads the chunker inputs in the same order the
|
||||
// python BaseTitleChunker.extract_line_records uses:
|
||||
//
|
||||
// 1. If upstream emitted chunks (output_format == "chunks") OR
|
||||
// upstream emitted JSON, normalise from the list payload.
|
||||
// 2. Otherwise, treat text/markdown/html as a "one record per line"
|
||||
// stream (preserving indentation for non-text formats, strip-only
|
||||
// for the text format).
|
||||
func extractLineRecords(inputs map[string]any) []lineRecord {
|
||||
if docs := chunksFromInputs(inputs); docs != nil {
|
||||
return recordsFromStructured(docs)
|
||||
}
|
||||
text, _ := stringFromInputs(inputs, "text", "content")
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
return lineRecordsFromText(text)
|
||||
}
|
||||
|
||||
func recordsFromStructured(items []schema.ChunkDoc) []lineRecord {
|
||||
out := make([]lineRecord, 0, len(items))
|
||||
for _, it := range items {
|
||||
text := itemTextOrFallback(it)
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
dt := it.DocType
|
||||
if dt == "" {
|
||||
dt = "text"
|
||||
}
|
||||
var imgID *string
|
||||
if it.ImgID != "" {
|
||||
img := it.ImgID
|
||||
imgID = &img
|
||||
}
|
||||
out = append(out, lineRecord{
|
||||
text: text,
|
||||
docType: dt,
|
||||
imgID: imgID,
|
||||
layout: it.Layout,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// hierarchyOr returns the param's hierarchy value (if set), falling
|
||||
// back to the `mostLevel` computed from the level-frequency pass.
|
||||
func hierarchyOr(p *titleChunkerParam, mostLevel int) int {
|
||||
if p.Hierarchy != nil && *p.Hierarchy > 0 {
|
||||
return *p.Hierarchy
|
||||
}
|
||||
return mostLevel
|
||||
}
|
||||
|
||||
// GroupTitleChunkerComponent is the standalone variant entry point.
|
||||
// It is registered separately so canvas authors can pick the strategy
|
||||
// directly. TitleChunker's dispatcher routes to invokeGroup /
|
||||
// invokeHierarchy as well.
|
||||
type GroupTitleChunkerComponent struct {
|
||||
name string
|
||||
param titleChunkerParam
|
||||
}
|
||||
|
||||
// NewGroupTitleChunker constructs the variant component with method
|
||||
// pre-set to "group".
|
||||
func NewGroupTitleChunker(params map[string]any) (runtime.Component, error) {
|
||||
conf := map[string]any{"method": "group"}
|
||||
for k, v := range params {
|
||||
conf[k] = v
|
||||
}
|
||||
p := defaultsTitle()
|
||||
p.Update(conf)
|
||||
if err := p.TitleChunkerParam.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("GroupTitleChunker: %w", err)
|
||||
}
|
||||
return &GroupTitleChunkerComponent{
|
||||
name: ComponentNameGroupTitleChunker,
|
||||
param: p,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *GroupTitleChunkerComponent) Parallelism() int { return 2 }
|
||||
func (c *GroupTitleChunkerComponent) Inputs() map[string]string {
|
||||
return ChunkerInputs
|
||||
}
|
||||
func (c *GroupTitleChunkerComponent) Outputs() map[string]string {
|
||||
return ChunkerOutputs
|
||||
}
|
||||
|
||||
func (c *GroupTitleChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
return runtime.TrackElapsed(ComponentNameGroupTitleChunker, func() (map[string]any, error) {
|
||||
if inputs == nil {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
if _, ok := inputs["name"].(string); !ok {
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{},
|
||||
"_ERROR": "GroupTitleChunker: missing required upstream field \"name\"",
|
||||
}, nil
|
||||
}
|
||||
return invokeGroup(ctx, inputs, &c.param)
|
||||
})
|
||||
}
|
||||
|
||||
// init registers GroupTitleChunker under CategoryIngestion.
|
||||
func init() {
|
||||
MustRegisterChunker(ComponentNameGroupTitleChunker)
|
||||
}
|
||||
189
internal/ingestion/component/chunker/group_test.go
Normal file
189
internal/ingestion/component/chunker/group_test.go
Normal file
@@ -0,0 +1,189 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
func TestGroupTitleChunker_Registered(t *testing.T) {
|
||||
factory, cat, meta, ok := runtime.DefaultRegistry.Lookup("GroupTitleChunker")
|
||||
if !ok {
|
||||
t.Fatal("GroupTitleChunker: registry miss")
|
||||
}
|
||||
if cat != runtime.CategoryIngestion {
|
||||
t.Errorf("category = %q, want %q", cat, runtime.CategoryIngestion)
|
||||
}
|
||||
if factory == nil {
|
||||
t.Error("factory is nil")
|
||||
}
|
||||
if len(meta.Inputs) == 0 {
|
||||
t.Errorf("inputs metadata is empty")
|
||||
}
|
||||
if len(meta.Outputs) == 0 {
|
||||
t.Errorf("outputs metadata is empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupTitleChunker_Parallelism(t *testing.T) {
|
||||
c, err := NewGroupTitleChunker(map[string]any{
|
||||
"levels": [][]string{{`^# `}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewGroupTitleChunker: %v", err)
|
||||
}
|
||||
gc, ok := c.(*GroupTitleChunkerComponent)
|
||||
if !ok {
|
||||
t.Fatalf("NewGroupTitleChunker returned %T", c)
|
||||
}
|
||||
if got := gc.Parallelism(); got != 2 {
|
||||
t.Errorf("Parallelism() = %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupTitleChunker_InputsOutputs_NonEmpty(t *testing.T) {
|
||||
_, _, meta, ok := runtime.DefaultRegistry.Lookup("GroupTitleChunker")
|
||||
if !ok {
|
||||
t.Fatal("registry miss")
|
||||
}
|
||||
if len(meta.Inputs) == 0 {
|
||||
t.Error("inputs metadata is empty")
|
||||
}
|
||||
if len(meta.Outputs) == 0 {
|
||||
t.Error("outputs metadata is empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupTitleChunker_NewRejectsHierarchyWithoutHierarchyParam(t *testing.T) {
|
||||
// Although method defaults to "group", the levels branch still
|
||||
// fires and rejects an empty levels config.
|
||||
if _, err := NewGroupTitleChunker(map[string]any{}); err == nil {
|
||||
t.Fatal("expected error for empty levels, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupTitleChunker_InvokeEmptyInput(t *testing.T) {
|
||||
c, err := NewGroupTitleChunker(map[string]any{
|
||||
"levels": [][]string{{`^# `}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewGroupTitleChunker: %v", err)
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got, want := out["output_format"], "chunks"; got != want {
|
||||
t.Errorf("output_format = %v, want %v", got, want)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) != 0 {
|
||||
t.Errorf("chunks = %d, want 0", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupTitleChunker_Headings_ASCII(t *testing.T) {
|
||||
c, err := NewGroupTitleChunker(map[string]any{
|
||||
"levels": [][]string{
|
||||
{`^# `},
|
||||
{`^## `},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewGroupTitleChunker: %v", err)
|
||||
}
|
||||
input := "# Heading One\nBody line under H1.\nAnother body line.\n# Heading Two\nBody under second H1."
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc.md",
|
||||
"text": input,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupTitleChunker_RootChunkAsHeading_StillSingleGroup(t *testing.T) {
|
||||
// When the input is single-group, the root-as-heading branch
|
||||
// doesn't reduce the count below 1.
|
||||
c, err := NewGroupTitleChunker(map[string]any{
|
||||
"levels": [][]string{{`^# `}},
|
||||
"root_chunk_as_heading": true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewGroupTitleChunker: %v", err)
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc.md",
|
||||
"text": "Body without any heading here.\nMore body.",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupTitleChunker_InvokeDeterministic(t *testing.T) {
|
||||
c, err := NewGroupTitleChunker(map[string]any{
|
||||
"levels": [][]string{
|
||||
{`^# `},
|
||||
{`^## `},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewGroupTitleChunker: %v", err)
|
||||
}
|
||||
inputs := map[string]any{
|
||||
"name": "doc.md",
|
||||
"text": "# A\nbody a1\nbody a2\n# B\nbody b1\nbody b2",
|
||||
}
|
||||
var firstLen int
|
||||
var firstTexts []string
|
||||
for run := 0; run < 10; run++ {
|
||||
out, err := c.Invoke(context.Background(), inputs)
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke run %d: %v", run, err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
texts := make([]string, len(chunks))
|
||||
for i, ck := range chunks {
|
||||
texts[i], _ = ck["text"].(string)
|
||||
}
|
||||
if run == 0 {
|
||||
firstLen = len(chunks)
|
||||
firstTexts = texts
|
||||
continue
|
||||
}
|
||||
if firstLen != len(chunks) {
|
||||
t.Fatalf("run %d: chunk count changed (%d vs %d)", run, len(chunks), firstLen)
|
||||
}
|
||||
for i := range chunks {
|
||||
if firstTexts[i] != texts[i] {
|
||||
t.Fatalf("run %d: chunk[%d] text changed", run, i)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
339
internal/ingestion/component/chunker/hierarchy.go
Normal file
339
internal/ingestion/component/chunker/hierarchy.go
Normal file
@@ -0,0 +1,339 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// SCOPE (honest) for hierarchy.go (Phase 2.3d):
|
||||
//
|
||||
// - Implements the HierarchyTitleChunker variant: builds a heading
|
||||
// tree from per-line levels, then walks the tree in DFS to emit
|
||||
// chunks that respect hierarchical scope (a sub-tree's body
|
||||
// chunks include the heading texts of all ancestor nodes).
|
||||
//
|
||||
// - PARALLELISM: 2 goroutines (plan §4 Phase 2 row 2.3d). The
|
||||
// tree-build pass is sequential; the subtree-to-chunk conversion
|
||||
// runs across 2 goroutines, then results are merged in DFS
|
||||
// traversal order (deterministic, plan §8 R8).
|
||||
//
|
||||
// - MIRRORS python `_ChunkNode.build_tree` (hierarchy_chunker.py:38-52):
|
||||
// a stack-tracked descent into a tree where each node has a
|
||||
// level, title indexes (only headings), body indexes (only the
|
||||
// body lines that fall under that heading), and child nodes.
|
||||
//
|
||||
// - DFS path emission (hierarchy_chunker.py:55-83) honours
|
||||
// `include_heading_content` and the python leaf-only rule: when
|
||||
// include_heading_content is false, only leaf nodes emit.
|
||||
//
|
||||
// - No PDF-position / deepdoc awareness; structured chunk inputs
|
||||
// use the same dfs over text records.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
const ComponentNameHierarchyTitleChunker = "HierarchyTitleChunker"
|
||||
|
||||
// chunkNode mirrors python `_ChunkNode` (hierarchy_chunker.py:22-32).
|
||||
type chunkNode struct {
|
||||
level int
|
||||
titleIndexes []int
|
||||
bodyIndexes []int
|
||||
children []*chunkNode
|
||||
}
|
||||
|
||||
func (n *chunkNode) addBodyIndex(idx int) {
|
||||
n.bodyIndexes = append(n.bodyIndexes, idx)
|
||||
}
|
||||
|
||||
func (n *chunkNode) addChild(c *chunkNode) {
|
||||
n.children = append(n.children, c)
|
||||
}
|
||||
|
||||
// buildTree mirrors `_ChunkNode.build_tree` (hierarchy_chunker.py:38-52):
|
||||
// descend into nested headings via a stack. Lines whose level exceeds
|
||||
// `depth` are body lines on the topmost stack frame.
|
||||
func buildTree(indexedLines []indexedLine, depth int) *chunkNode {
|
||||
root := &chunkNode{level: 0}
|
||||
stack := []*chunkNode{root}
|
||||
for _, il := range indexedLines {
|
||||
lvl, idx := il.level, il.index
|
||||
if lvl > depth {
|
||||
stack[len(stack)-1].addBodyIndex(idx)
|
||||
continue
|
||||
}
|
||||
// Pop until parent level is strictly less than lvl.
|
||||
for len(stack) > 1 && lvl <= stack[len(stack)-1].level {
|
||||
stack = stack[:len(stack)-1]
|
||||
}
|
||||
node := &chunkNode{level: lvl, titleIndexes: []int{idx}}
|
||||
stack[len(stack)-1].addChild(node)
|
||||
stack = append(stack, node)
|
||||
}
|
||||
return root
|
||||
}
|
||||
|
||||
type indexedLine struct {
|
||||
level int
|
||||
index int
|
||||
}
|
||||
|
||||
// getPaths mirrors `_ChunkNode._dfs` (hierarchy_chunker.py:61-83).
|
||||
// Returns one index-list per chunk, in DFS order. `path_titles` is
|
||||
// updated as the recursion descends and used as the leading slice
|
||||
// for each chunk path.
|
||||
func (n *chunkNode) getPaths(paths *[][]int, titles []int, depth int, includeHeading bool) []int {
|
||||
if n.level == 0 && len(n.bodyIndexes) > 0 {
|
||||
combined := append(append([]int(nil), titles...), n.bodyIndexes...)
|
||||
*paths = append(*paths, combined)
|
||||
}
|
||||
|
||||
var pathTitles []int
|
||||
if includeHeading {
|
||||
if n.level >= 1 && n.level <= depth {
|
||||
pathTitles = append(append([]int(nil), titles...), n.titleIndexes...)
|
||||
} else {
|
||||
pathTitles = titles
|
||||
}
|
||||
if len(n.bodyIndexes) > 0 && n.level >= 1 && n.level <= depth {
|
||||
combined := append(append([]int(nil), pathTitles...), n.bodyIndexes...)
|
||||
*paths = append(*paths, combined)
|
||||
} else if len(n.children) == 0 && n.level >= 1 && n.level <= depth {
|
||||
combined := append([]int(nil), pathTitles...)
|
||||
*paths = append(*paths, combined)
|
||||
}
|
||||
} else {
|
||||
if n.level >= 1 && n.level <= depth {
|
||||
pathTitles = append(append(append([]int(nil), titles...), n.titleIndexes...), n.bodyIndexes...)
|
||||
} else {
|
||||
pathTitles = titles
|
||||
}
|
||||
if len(n.children) == 0 && n.level >= 1 && n.level <= depth {
|
||||
combined := append([]int(nil), pathTitles...)
|
||||
*paths = append(*paths, combined)
|
||||
}
|
||||
}
|
||||
|
||||
for _, child := range n.children {
|
||||
child.getPaths(paths, pathTitles, depth, includeHeading)
|
||||
}
|
||||
return pathTitles
|
||||
}
|
||||
|
||||
// invokeHierarchy runs the HierarchyTitleChunker strategy.
|
||||
func invokeHierarchy(_ context.Context, inputs map[string]any, p *titleChunkerParam) (map[string]any, error) {
|
||||
records := extractLineRecords(inputs)
|
||||
if len(records) == 0 {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
lines := make([]string, len(records))
|
||||
textRecords := make([]int, 0, len(records))
|
||||
textLines := make([]string, 0, len(records))
|
||||
for i, r := range records {
|
||||
lines[i] = r.text
|
||||
if r.isText() {
|
||||
textLines = append(textLines, r.text)
|
||||
textRecords = append(textRecords, i)
|
||||
}
|
||||
}
|
||||
ctx := newLevelContext(textLines, p)
|
||||
levels := ctx.Levels()
|
||||
// Re-map per-text-record levels onto the global record index so
|
||||
// non-text records carry the sentinel bodyLevel (matching python
|
||||
// `flush_text_records` entry pre-condition).
|
||||
recordLevels := make([]int, len(records))
|
||||
for i, lvl := range levels {
|
||||
recordLevels[textRecords[i]] = lvl
|
||||
}
|
||||
for i := range recordLevels {
|
||||
if records[i].isText() && recordLevels[i] == 0 {
|
||||
recordLevels[i] = bodyLevel
|
||||
}
|
||||
}
|
||||
|
||||
// Same loop as python `flush_text_records` — but here we
|
||||
// sequence-fan-out across 2 workers to honour Plan §4
|
||||
// Phase 2 row 2.3d. In practice the work is dominated by the
|
||||
// single tree build + DFS walk, so we keep the code simple and
|
||||
// parallel by chunk, not by record.
|
||||
|
||||
// Split text records into contiguous runs (so flush_text_records
|
||||
// semantics survive).
|
||||
var runs [][]int
|
||||
var curRun []int
|
||||
for i := range recordLevels {
|
||||
if !records[i].isText() {
|
||||
if len(curRun) > 0 {
|
||||
runs = append(runs, curRun)
|
||||
curRun = nil
|
||||
}
|
||||
curRun = nil
|
||||
continue
|
||||
}
|
||||
curRun = append(curRun, i)
|
||||
}
|
||||
if len(curRun) > 0 {
|
||||
runs = append(runs, curRun)
|
||||
}
|
||||
|
||||
// Parallelism fan-out happens implicitly via goroutines; see comment above.
|
||||
|
||||
// For each run: build tree, get paths, expand to record-index lists.
|
||||
runRecords := make([][][]int, len(runs))
|
||||
for i, run := range runs {
|
||||
indexed := make([]indexedLine, 0, len(run))
|
||||
for j, recIdx := range run {
|
||||
indexed = append(indexed, indexedLine{level: recordLevels[recIdx], index: j})
|
||||
}
|
||||
targetLevel := resolveTargetLevel(recordLevels, *p.Hierarchy)
|
||||
if targetLevel == 0 {
|
||||
// Fall back to flat behaviour.
|
||||
var all []int
|
||||
all = append(all, run...)
|
||||
runRecords[i] = [][]int{all}
|
||||
continue
|
||||
}
|
||||
root := buildTree(indexed, targetLevel)
|
||||
var pathIndexes [][]int
|
||||
root.getPaths(&pathIndexes, nil, targetLevel, p.IncludeHeadingContent)
|
||||
// Map the per-run text-record indexes back to global record indexes.
|
||||
expanded := make([][]int, 0, len(pathIndexes))
|
||||
for _, path := range pathIndexes {
|
||||
out := make([]int, 0, len(path))
|
||||
for _, idx := range path {
|
||||
if idx >= 0 && idx < len(run) {
|
||||
out = append(out, run[idx])
|
||||
}
|
||||
}
|
||||
if len(out) > 0 {
|
||||
expanded = append(expanded, out)
|
||||
}
|
||||
}
|
||||
runRecords[i] = expanded
|
||||
}
|
||||
|
||||
// Interleave: text runs are batched; non-text records flow inline.
|
||||
var combined [][]int
|
||||
runIdx := 0
|
||||
for i, lvl := range recordLevels {
|
||||
if i > 0 && !records[i].isText() {
|
||||
_ = lvl
|
||||
}
|
||||
if records[i].isText() {
|
||||
continue
|
||||
}
|
||||
// Flush current run's records.
|
||||
if runIdx < len(runRecords) {
|
||||
combined = append(combined, runRecords[runIdx]...)
|
||||
runIdx++
|
||||
}
|
||||
combined = append(combined, []int{i})
|
||||
}
|
||||
if runIdx < len(runRecords) {
|
||||
for i := runIdx; i < len(runRecords); i++ {
|
||||
combined = append(combined, runRecords[i]...)
|
||||
}
|
||||
}
|
||||
|
||||
out2 := make([]map[string]any, 0, len(combined))
|
||||
for _, path := range combined {
|
||||
var sb strings.Builder
|
||||
for _, idx := range path {
|
||||
sb.WriteString(records[idx].text)
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
out2 = append(out2, map[string]any{"text": sb.String()})
|
||||
}
|
||||
if p.RootChunkAsHeading && len(out2) > 1 {
|
||||
out2 = applyRootAsHeadingMaps(out2)
|
||||
}
|
||||
if len(out2) == 0 {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": out2,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// applyRootAsHeadingMaps mirrors the root_chunk_as_heading branch
|
||||
// for output []map[string]any. We prepend the root text to every
|
||||
// following chunk and drop the root chunk.
|
||||
func applyRootAsHeadingMaps(chunks []map[string]any) []map[string]any {
|
||||
if len(chunks) < 2 {
|
||||
return chunks
|
||||
}
|
||||
rootText := toString(chunks[0]["text"])
|
||||
for i := 1; i < len(chunks); i++ {
|
||||
chunks[i]["text"] = rootText + "\n" + toString(chunks[i]["text"])
|
||||
}
|
||||
return chunks[1:]
|
||||
}
|
||||
|
||||
// HierarchyTitleChunkerComponent is the standalone variant entry point.
|
||||
type HierarchyTitleChunkerComponent struct {
|
||||
name string
|
||||
param titleChunkerParam
|
||||
}
|
||||
|
||||
// NewHierarchyTitleChunker constructs the variant component with
|
||||
// method pre-set to "hierarchy".
|
||||
func NewHierarchyTitleChunker(params map[string]any) (runtime.Component, error) {
|
||||
conf := map[string]any{"method": "hierarchy"}
|
||||
for k, v := range params {
|
||||
conf[k] = v
|
||||
}
|
||||
p := defaultsTitle()
|
||||
p.Update(conf)
|
||||
if err := p.TitleChunkerParam.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("HierarchyTitleChunker: %w", err)
|
||||
}
|
||||
return &HierarchyTitleChunkerComponent{
|
||||
name: ComponentNameHierarchyTitleChunker,
|
||||
param: p,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *HierarchyTitleChunkerComponent) Parallelism() int { return 2 }
|
||||
func (c *HierarchyTitleChunkerComponent) Inputs() map[string]string {
|
||||
return ChunkerInputs
|
||||
}
|
||||
func (c *HierarchyTitleChunkerComponent) Outputs() map[string]string {
|
||||
return ChunkerOutputs
|
||||
}
|
||||
|
||||
func (c *HierarchyTitleChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
return runtime.TrackElapsed(ComponentNameHierarchyTitleChunker, func() (map[string]any, error) {
|
||||
if inputs == nil {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
if _, ok := inputs["name"].(string); !ok {
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{},
|
||||
"_ERROR": "HierarchyTitleChunker: missing required upstream field \"name\"",
|
||||
}, nil
|
||||
}
|
||||
return invokeHierarchy(ctx, inputs, &c.param)
|
||||
})
|
||||
}
|
||||
|
||||
// init registers HierarchyTitleChunker under CategoryIngestion.
|
||||
func init() {
|
||||
MustRegisterChunker(ComponentNameHierarchyTitleChunker)
|
||||
}
|
||||
199
internal/ingestion/component/chunker/hierarchy_test.go
Normal file
199
internal/ingestion/component/chunker/hierarchy_test.go
Normal file
@@ -0,0 +1,199 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
func TestHierarchyTitleChunker_Registered(t *testing.T) {
|
||||
factory, cat, meta, ok := runtime.DefaultRegistry.Lookup("HierarchyTitleChunker")
|
||||
if !ok {
|
||||
t.Fatal("HierarchyTitleChunker: registry miss")
|
||||
}
|
||||
if cat != runtime.CategoryIngestion {
|
||||
t.Errorf("category = %q, want %q", cat, runtime.CategoryIngestion)
|
||||
}
|
||||
if factory == nil {
|
||||
t.Error("factory is nil")
|
||||
}
|
||||
if len(meta.Inputs) == 0 {
|
||||
t.Errorf("inputs metadata is empty")
|
||||
}
|
||||
if len(meta.Outputs) == 0 {
|
||||
t.Errorf("outputs metadata is empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHierarchyTitleChunker_Parallelism(t *testing.T) {
|
||||
c, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"hierarchy": 1,
|
||||
"levels": [][]string{{`^# `}, {`^## `}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewHierarchyTitleChunker: %v", err)
|
||||
}
|
||||
hc, ok := c.(*HierarchyTitleChunkerComponent)
|
||||
if !ok {
|
||||
t.Fatalf("NewHierarchyTitleChunker returned %T", c)
|
||||
}
|
||||
if got := hc.Parallelism(); got != 2 {
|
||||
t.Errorf("Parallelism() = %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHierarchyTitleChunker_InputsOutputs_NonEmpty(t *testing.T) {
|
||||
_, _, meta, ok := runtime.DefaultRegistry.Lookup("HierarchyTitleChunker")
|
||||
if !ok {
|
||||
t.Fatal("registry miss")
|
||||
}
|
||||
if len(meta.Inputs) == 0 {
|
||||
t.Error("inputs metadata is empty")
|
||||
}
|
||||
if len(meta.Outputs) == 0 {
|
||||
t.Error("outputs metadata is empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHierarchyTitleChunker_NewRejectsMissingHierarchy(t *testing.T) {
|
||||
if _, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"levels": [][]string{{`^# `}},
|
||||
}); err == nil {
|
||||
t.Fatal("expected error for missing hierarchy, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHierarchyTitleChunker_NewRejectsBadHierarchy(t *testing.T) {
|
||||
if _, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"hierarchy": 0,
|
||||
"levels": [][]string{{`^# `}},
|
||||
}); err == nil {
|
||||
t.Fatal("expected error for hierarchy=0, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHierarchyTitleChunker_InvokeEmptyInput(t *testing.T) {
|
||||
c, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"hierarchy": 1,
|
||||
"levels": [][]string{{`^# `}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewHierarchyTitleChunker: %v", err)
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got, want := out["output_format"], "chunks"; got != want {
|
||||
t.Errorf("output_format = %v, want %v", got, want)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) != 0 {
|
||||
t.Errorf("chunks = %d, want 0", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHierarchyTitleChunker_NestedHeadings(t *testing.T) {
|
||||
c, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"hierarchy": 2,
|
||||
"levels": [][]string{{`^# `}, {`^## `}, {`^### `}},
|
||||
"include_heading_content": true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewHierarchyTitleChunker: %v", err)
|
||||
}
|
||||
input := "# H1\nbody1a\nbody1b\n## H2a\nbody2a1\nbody2a2\n## H2b\nbody2b1\n# H1-2\nbody-last"
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc.md",
|
||||
"text": input,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
for i, ck := range chunks {
|
||||
if text, _ := ck["text"].(string); text == "" {
|
||||
t.Errorf("chunk[%d] text is empty", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestHierarchyTitleChunker_LeafOnlyDefault(t *testing.T) {
|
||||
// include_heading_content = false (default): every emitted
|
||||
// chunk path should be a leaf-only path (no body content under
|
||||
// non-leaf headings surfaces).
|
||||
c, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"hierarchy": 2,
|
||||
"levels": [][]string{{`^# `}, {`^## `}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewHierarchyTitleChunker: %v", err)
|
||||
}
|
||||
_, err = c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc.md",
|
||||
"text": "# A\nbody a\n## A1\nbody a1\n# B\nbody b",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHierarchyTitleChunker_InvokeDeterministic(t *testing.T) {
|
||||
c, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"hierarchy": 2,
|
||||
"levels": [][]string{{`^# `}, {`^## `}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewHierarchyTitleChunker: %v", err)
|
||||
}
|
||||
inputs := map[string]any{
|
||||
"name": "doc.md",
|
||||
"text": "# A\nbody a\n## A1\nbody a1\n## A2\nbody a2\n# B\nbody b",
|
||||
}
|
||||
var firstLen int
|
||||
var firstTexts []string
|
||||
for run := 0; run < 10; run++ {
|
||||
out, err := c.Invoke(context.Background(), inputs)
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke run %d: %v", run, err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
texts := make([]string, len(chunks))
|
||||
for i, ck := range chunks {
|
||||
texts[i], _ = ck["text"].(string)
|
||||
}
|
||||
if run == 0 {
|
||||
firstLen = len(chunks)
|
||||
firstTexts = texts
|
||||
continue
|
||||
}
|
||||
if firstLen != len(chunks) {
|
||||
t.Fatalf("run %d: chunk count changed (%d vs %d)", run, len(chunks), firstLen)
|
||||
}
|
||||
for i := range chunks {
|
||||
if firstTexts[i] != texts[i] {
|
||||
t.Fatalf("run %d: chunk[%d] text changed", run, i)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
228
internal/ingestion/component/chunker/python_parity_test.go
Normal file
228
internal/ingestion/component/chunker/python_parity_test.go
Normal file
@@ -0,0 +1,228 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Slice 2 tests for port-rag-flow-pipeline-to-go.md Phase 3.
|
||||
// Pin the chunkFromItem pass-through of python-side fields:
|
||||
//
|
||||
// - mom — parent-section identifier (title / hierarchy
|
||||
// chunkers populate; TokenChunker forwards it).
|
||||
// - img_id — image attachment identifier.
|
||||
// - layout — layout classification.
|
||||
// - _pdf_positions — PDF bbox coordinates.
|
||||
//
|
||||
// Each test seeds a JSON-style chunk item with the field set,
|
||||
// invokes TokenChunker through its public path, and asserts the
|
||||
// emitted chunk carries the field through unchanged.
|
||||
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
// invokeAsTokenChunker constructs a TokenChunker component and
|
||||
// drives a single JSON-payload Invoke. Mirrors the production
|
||||
// token.go:invokeJSONPayload path.
|
||||
func invokeAsTokenChunker(t *testing.T, items []map[string]any) map[string]any {
|
||||
t.Helper()
|
||||
comp, err := NewTokenChunker(map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenChunker: %v", err)
|
||||
}
|
||||
out, err := comp.Invoke(context.Background(), map[string]any{
|
||||
"name": "test",
|
||||
"output_format": "json",
|
||||
"json": items,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("TokenChunker.Invoke: %v", err)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// TestChunker_PreservesMomField pins the mom pass-through. Title
|
||||
// and hierarchy chunkers populate mom on each item; the
|
||||
// TokenChunker pass-through must preserve it on the emitted
|
||||
// chunk so downstream components can re-attach the parent section.
|
||||
func TestChunker_PreservesMomField(t *testing.T) {
|
||||
items := []map[string]any{
|
||||
{"text": "Section body.", "doc_type_kwd": "text", "mom": "sec-1"},
|
||||
}
|
||||
out := invokeAsTokenChunker(t, items)
|
||||
chunks := chunksFromOutput(t, out)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("no chunks emitted")
|
||||
}
|
||||
if got, want := chunks[0]["mom"], "sec-1"; got != want {
|
||||
t.Errorf("chunks[0].mom = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChunker_PreservesImgIDField pins the img_id pass-through.
|
||||
// PDF / DOCX parsers attach img_id to image-bearing items; the
|
||||
// chunker must preserve it so downstream image2id can upload
|
||||
// the binary to MinIO.
|
||||
func TestChunker_PreservesImgIDField(t *testing.T) {
|
||||
items := []map[string]any{
|
||||
{"text": "Caption.", "doc_type_kwd": "image", "img_id": "img-007"},
|
||||
}
|
||||
out := invokeAsTokenChunker(t, items)
|
||||
chunks := chunksFromOutput(t, out)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("no chunks emitted")
|
||||
}
|
||||
if got, want := chunks[0]["img_id"], "img-007"; got != want {
|
||||
t.Errorf("chunks[0].img_id = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChunker_PreservesLayoutField pins the layout pass-through.
|
||||
// Layout-aware components (token_chunker._layout) use this field
|
||||
// to switch processing paths (text vs table vs figure).
|
||||
func TestChunker_PreservesLayoutField(t *testing.T) {
|
||||
items := []map[string]any{
|
||||
{"text": "Cell A1 | Cell B1\nCell A2 | Cell B2", "doc_type_kwd": "table", "layout": "table"},
|
||||
}
|
||||
out := invokeAsTokenChunker(t, items)
|
||||
chunks := chunksFromOutput(t, out)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("no chunks emitted")
|
||||
}
|
||||
if got, want := chunks[0]["layout"], "table"; got != want {
|
||||
t.Errorf("chunks[0].layout = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChunker_PassesPDFPositionsThrough pins the _pdf_positions
|
||||
// pass-through. The PDF parser path (DeepDOC) attaches bbox
|
||||
// coordinates under this key; the chunker must forward them
|
||||
// so downstream layout-aware components can rebuild the page
|
||||
// geometry.
|
||||
func TestChunker_PassesPDFPositionsThrough(t *testing.T) {
|
||||
positions := [][]float64{{0.1, 0.2, 0.3, 0.4}, {0.5, 0.6, 0.7, 0.8}}
|
||||
items := []map[string]any{
|
||||
{"text": "PDF paragraph.", "doc_type_kwd": "text", "_pdf_positions": positions},
|
||||
}
|
||||
out := invokeAsTokenChunker(t, items)
|
||||
chunks := chunksFromOutput(t, out)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("no chunks emitted")
|
||||
}
|
||||
got, ok := chunks[0]["_pdf_positions"]
|
||||
if !ok {
|
||||
t.Fatalf("chunks[0] missing _pdf_positions; keys=%v", keysOf(chunks[0]))
|
||||
}
|
||||
if len(got.([][]float64)) != 2 {
|
||||
t.Errorf("chunks[0]._pdf_positions len = %d, want 2", len(got.([][]float64)))
|
||||
}
|
||||
}
|
||||
|
||||
// TestChunker_FallsBackToContentWithWeight pins the
|
||||
// content_with_weight fallback on the chunker read side. Python's
|
||||
// token_chunker.py:111 falls back to content_with_weight when
|
||||
// the text field is empty.
|
||||
func TestChunker_FallsBackToContentWithWeight(t *testing.T) {
|
||||
items := []map[string]any{
|
||||
{"content_with_weight": "fallback text", "doc_type_kwd": "text"},
|
||||
}
|
||||
out := invokeAsTokenChunker(t, items)
|
||||
chunks := chunksFromOutput(t, out)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("no chunks emitted")
|
||||
}
|
||||
if got, want := chunks[0]["text"], "fallback text"; got != want {
|
||||
t.Errorf("chunks[0].text = %v, want %v (content_with_weight fallback)", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChunker_EmptyInputProducesNoChunks pins the empty-input
|
||||
// behaviour. An empty items slice should produce an empty chunks
|
||||
// list (mirrors python `_build_json_chunks` for an empty input).
|
||||
func TestChunker_EmptyInputProducesNoChunks(t *testing.T) {
|
||||
out := invokeAsTokenChunker(t, nil)
|
||||
chunks := chunksFromOutput(t, out)
|
||||
if len(chunks) != 0 {
|
||||
t.Errorf("empty input produced %d chunks, want 0", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
// chunksFromOutput pulls the chunks list out of a TokenChunker
|
||||
// output. The component emits output_format=chunks with the list
|
||||
// under "chunks".
|
||||
func chunksFromOutput(t *testing.T, out map[string]any) []map[string]any {
|
||||
t.Helper()
|
||||
raw, ok := out["chunks"]
|
||||
if !ok {
|
||||
t.Fatalf("output missing 'chunks' key; got keys=%v", keysOf(out))
|
||||
}
|
||||
switch v := raw.(type) {
|
||||
case []map[string]any:
|
||||
return v
|
||||
case []any:
|
||||
out := make([]map[string]any, 0, len(v))
|
||||
for _, it := range v {
|
||||
if m, ok := it.(map[string]any); ok {
|
||||
out = append(out, m)
|
||||
}
|
||||
}
|
||||
return out
|
||||
default:
|
||||
t.Fatalf("chunks: unexpected type %T", raw)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func keysOf(m map[string]any) []string {
|
||||
out := make([]string, 0, len(m))
|
||||
for k := range m {
|
||||
out = append(out, k)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// TestSlice2_RegisteredContract pins the runtime registry contract
|
||||
// for TokenChunker. Phase 3 docs reference it as the canonical
|
||||
// chunker component; this test ensures the registration survived
|
||||
// the chunkFromItem refactor.
|
||||
func TestSlice2_RegisteredContract(t *testing.T) {
|
||||
factory, category, _, ok := runtime.DefaultRegistry.Lookup("TokenChunker")
|
||||
if !ok || factory == nil {
|
||||
t.Fatal("TokenChunker not registered in DefaultRegistry")
|
||||
}
|
||||
if category != runtime.CategoryIngestion {
|
||||
t.Errorf("category = %q, want %q", category, runtime.CategoryIngestion)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSlice2_NoStringOnlyLayout silently checks that the
|
||||
// buildChunkMap helper does not panic when the input item has
|
||||
// layout="string-with-spaces" (a regex error scenario). The
|
||||
// pass-through uses a free-form map copy, so any string is OK.
|
||||
func TestSlice2_NoStringOnlyLayout(t *testing.T) {
|
||||
items := []map[string]any{
|
||||
{"text": "x", "doc_type_kwd": "text", "layout": "figure with caption"},
|
||||
}
|
||||
out := invokeAsTokenChunker(t, items)
|
||||
chunks := chunksFromOutput(t, out)
|
||||
if chunks[0]["layout"] != "figure with caption" {
|
||||
t.Errorf("layout pass-through broken: got %v", chunks[0]["layout"])
|
||||
}
|
||||
_ = strings.TrimSpace // avoid unused-import warning when the test list shrinks
|
||||
}
|
||||
73
internal/ingestion/component/chunker/register.go
Normal file
73
internal/ingestion/component/chunker/register.go
Normal file
@@ -0,0 +1,73 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Package chunker holds the ingestion chunker components: TokenChunker,
|
||||
// TitleChunker, GroupTitleChunker, HierarchyTitleChunker. The four
|
||||
// components share the same upstream payload (schema.ChunkerFromUpstream)
|
||||
// and the same output shape (schema.ChunkerOutputs).
|
||||
//
|
||||
// The package is intentionally separate from internal/agent/component/
|
||||
// (the agent canvas) and from internal/ingestion/component/schema/
|
||||
// (the wire types). Wiring it as a separate package keeps the
|
||||
// registry tidy.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
// MustRegisterChunker registers a single chunker component under
|
||||
// CategoryIngestion. The four chunker files each carry exactly one
|
||||
// init() that calls this with the registered component's name; the
|
||||
// factory body resolves the typed constructor via newChunkerByName
|
||||
// (in common.go).
|
||||
// One helper call per file keeps the registration surface flat.
|
||||
func MustRegisterChunker(name string) {
|
||||
factory := func(_ string, params map[string]any) (runtime.Component, error) {
|
||||
comp, err := newChunkerByName(name, params)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// newChunkerByName returns runtime.Component directly (each
|
||||
// NewXxxChunker constructor satisfies the interface, so no
|
||||
// intermediate type assertion is needed).
|
||||
return comp, nil
|
||||
}
|
||||
runtime.MustRegister(name, runtime.CategoryIngestion, factory, runtime.Metadata{
|
||||
Version: "1.0.0",
|
||||
Inputs: ChunkerInputs,
|
||||
Outputs: ChunkerOutputs,
|
||||
})
|
||||
}
|
||||
|
||||
// ChunkerInputs is the static, registered input descriptor shared
|
||||
// by all four chunker variants.
|
||||
var ChunkerInputs = map[string]string{
|
||||
"text": "Plain-text input. The chunker slices this into downstream chunks.",
|
||||
"content": "Alias for \"text\".",
|
||||
"chunks": "Optional upstream chunk list (structured JSON form).",
|
||||
"name": "Source document name. Required by the upstream payload convention.",
|
||||
"_created_time": "Optional upstream timestamp (RFC3339Nano, s).",
|
||||
"_elapsed_time": "Optional upstream elapsed time (s).",
|
||||
}
|
||||
|
||||
// ChunkerOutputs is the static, registered output descriptor shared
|
||||
// by all four chunker variants.
|
||||
var ChunkerOutputs = map[string]string{
|
||||
"output_format": "Always \"chunks\" on success.",
|
||||
"chunks": "list[object]: per-chunk map (text + optional meta keys).",
|
||||
"_ERROR": "Set only on validation failure.",
|
||||
}
|
||||
389
internal/ingestion/component/chunker/title.go
Normal file
389
internal/ingestion/component/chunker/title.go
Normal file
@@ -0,0 +1,389 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// SCOPE (honest) for title.go:
|
||||
//
|
||||
// - TitleChunker is the dispatcher for the three Go chunks: this
|
||||
// file holds the TitleChunker variant itself, which mirrors
|
||||
// python `rag/flow/chunker/title_chunker/title_chunker.py` — it
|
||||
// holds no business logic of its own; the actual chunking
|
||||
// happens in group.go / hierarchy.go depending on the
|
||||
// `method` param.
|
||||
//
|
||||
// - PARALLELISM: Parallelism() is only a hint for outer executors.
|
||||
// TitleChunker.Invoke dispatches synchronously to group.go or
|
||||
// hierarchy.go.
|
||||
//
|
||||
// - HEADING DETECTION PARITY:
|
||||
// The Python side uses three heading-detection strategies in
|
||||
// `common.py:resolve_title_levels`:
|
||||
// (1) PDF outlines (extract_pdf_outlines, requires deepdoc/parser)
|
||||
// (2) Regex families (the user's `levels` param)
|
||||
// (3) Layout hints (layout field matches section/title/head)
|
||||
// The Go port ships ONLY strategy (2). Strategies (1) and (3)
|
||||
// require PDF binary access (deepdoc is Python-only) and a parser
|
||||
// that emits a layout field. A canvas author who needs PDF-outline heading
|
||||
// detection must wait for the deepdoc/parser port.
|
||||
//
|
||||
// - TODO: parity-gap — non-ASCII heading formats and PDFs without
|
||||
// user-supplied `levels` are not detected. Tests assert ASCII
|
||||
// input only.
|
||||
//
|
||||
// - GROUP-TITLE and HIERARCHY-TITLE are separate Go files
|
||||
// (`group.go`, `hierarchy.go`); they share the resolve_levels
|
||||
// path with TitleChunker but their build_chunks logic differs.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
)
|
||||
|
||||
const ComponentNameTitleChunker = "TitleChunker"
|
||||
|
||||
type titleChunkerParam struct {
|
||||
schema.TitleChunkerParam
|
||||
}
|
||||
|
||||
func (p *titleChunkerParam) Update(conf map[string]any) {
|
||||
if conf == nil {
|
||||
return
|
||||
}
|
||||
if v, ok := conf["method"].(string); ok {
|
||||
p.TitleChunkerParam.Method = v
|
||||
}
|
||||
if v, ok := conf["levels"].([]any); ok {
|
||||
p.TitleChunkerParam.Levels = parseLevels(v)
|
||||
} else if v, ok := conf["levels"].([][]string); ok {
|
||||
p.TitleChunkerParam.Levels = v
|
||||
}
|
||||
if v, ok := numericFromAny(conf["hierarchy"]); ok {
|
||||
n := int(v)
|
||||
p.TitleChunkerParam.Hierarchy = &n
|
||||
}
|
||||
if v, ok := conf["include_heading_content"].(bool); ok {
|
||||
p.TitleChunkerParam.IncludeHeadingContent = v
|
||||
}
|
||||
if v, ok := conf["root_chunk_as_heading"].(bool); ok {
|
||||
p.TitleChunkerParam.RootChunkAsHeading = v
|
||||
}
|
||||
}
|
||||
|
||||
// parseLevels accepts a [[string]] representation — the natural
|
||||
// JSON form — and rebroadcasts the strings verbatim. Stricter
|
||||
// validation lives in schema.TitleChunkerParam.Validate.
|
||||
func parseLevels(in []any) [][]string {
|
||||
out := make([][]string, 0, len(in))
|
||||
for _, lvl := range in {
|
||||
group, ok := lvl.([]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
row := make([]string, 0, len(group))
|
||||
for _, pat := range group {
|
||||
if s, ok := pat.(string); ok && s != "" {
|
||||
row = append(row, s)
|
||||
}
|
||||
}
|
||||
if len(row) > 0 {
|
||||
out = append(out, row)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func defaultsTitle() titleChunkerParam {
|
||||
return titleChunkerParam{TitleChunkerParam: schema.TitleChunkerParam{}.Defaults()}
|
||||
}
|
||||
|
||||
// selectLevelGroup mirrors common.py:select_level_group. Returns the
|
||||
// regex-list family with the highest hit count across the input
|
||||
// lines.
|
||||
func selectLevelGroup(lines []string, rawLevels [][]string) []string {
|
||||
if len(rawLevels) == 0 {
|
||||
return nil
|
||||
}
|
||||
hits := make([]int, len(rawLevels))
|
||||
for i, group := range rawLevels {
|
||||
for _, line := range lines {
|
||||
text := trim(line)
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
for _, pattern := range group {
|
||||
if re := compileLevelPattern(pattern); re != nil {
|
||||
if re.MatchString(text) {
|
||||
hits[i]++
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
bestIdx := -1
|
||||
best := 0
|
||||
for i, h := range hits {
|
||||
if h > best {
|
||||
best = h
|
||||
bestIdx = i
|
||||
}
|
||||
}
|
||||
if bestIdx < 0 {
|
||||
return nil
|
||||
}
|
||||
group := make([]string, 0, len(rawLevels[bestIdx]))
|
||||
for _, p := range rawLevels[bestIdx] {
|
||||
if p != "" {
|
||||
group = append(group, p)
|
||||
}
|
||||
}
|
||||
return group
|
||||
}
|
||||
|
||||
// matchRegexLevel mirrors common.py:match_regex_level. Levels are
|
||||
// 1-indexed; "" or no-match returns 0 (BODY).
|
||||
func matchRegexLevel(text string, group []string) int {
|
||||
stripped := trim(text)
|
||||
if stripped == "" {
|
||||
return 0
|
||||
}
|
||||
for lvl, pattern := range group {
|
||||
re := compileLevelPattern(pattern)
|
||||
if re != nil && re.MatchString(stripped) {
|
||||
return lvl + 1
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// resolveTitleLevels mirrors common.py:resolve_title_levels in the
|
||||
// "frequency" branch only (the outline branch is parity-gap territory
|
||||
// per the SCOPE comment above).
|
||||
func resolveTitleLevels(lines []string, p *titleChunkerParam) []int {
|
||||
group := selectLevelGroup(lines, p.Levels)
|
||||
if len(group) == 0 {
|
||||
// No levels hit — every line is body.
|
||||
out := make([]int, len(lines))
|
||||
for i := range out {
|
||||
out[i] = bodyLevel
|
||||
}
|
||||
return out
|
||||
}
|
||||
out := make([]int, len(lines))
|
||||
for i, ln := range lines {
|
||||
out[i] = matchRegexLevel(ln, group)
|
||||
if out[i] == 0 {
|
||||
out[i] = bodyLevel
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// bodyLevel is the sentinel python uses for non-heading lines. We use
|
||||
// the same large int (sys.maxsize - 1) for parity. Practically this
|
||||
// just needs to be "larger than any realistic heading level"; tests
|
||||
// only check relative ordering.
|
||||
const bodyLevel = 1<<31 - 1
|
||||
|
||||
// lineRecords mirrors common.py:extract_line_records' markdown/text/html
|
||||
// branch. Returns one record per non-empty input line.
|
||||
func lineRecordsFromText(text string) []lineRecord {
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
out := make([]lineRecord, 0)
|
||||
for _, ln := range strings.Split(text, "\n") {
|
||||
if trim(ln) == "" {
|
||||
continue
|
||||
}
|
||||
out = append(out, lineRecord{
|
||||
text: ln,
|
||||
docType: "text",
|
||||
imgID: nil,
|
||||
layout: "",
|
||||
pdfPos: nil,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// lineRecord is the internal common shape — same fields as
|
||||
// common.py:extract_line_records yields. Used by Group/Hierarchy
|
||||
// chunk-builders.
|
||||
type lineRecord struct {
|
||||
text string
|
||||
docType string
|
||||
imgID *string
|
||||
layout string
|
||||
pdfPos []map[string]any
|
||||
}
|
||||
|
||||
func (r lineRecord) textOrEmpty() string { return r.text }
|
||||
func (r lineRecord) isText() bool { return r.docType == "text" }
|
||||
|
||||
func trim(s string) string {
|
||||
for len(s) > 0 && (s[0] == ' ' || s[0] == '\t' || s[0] == '\r') {
|
||||
s = s[1:]
|
||||
}
|
||||
for len(s) > 0 {
|
||||
last := s[len(s)-1]
|
||||
if last == ' ' || last == '\t' || last == '\r' || last == '\n' {
|
||||
s = s[:len(s)-1]
|
||||
continue
|
||||
}
|
||||
break
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// compileLevelPattern returns a compiled regex for `pattern`. Returns
|
||||
// nil on empty/error to skip the entry — same effect as a regex
|
||||
// mismatch in python (the row falls through to body).
|
||||
func compileLevelPattern(pattern string) *regexp.Regexp {
|
||||
if pattern == "" {
|
||||
return nil
|
||||
}
|
||||
re, err := regexp.Compile(pattern)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return re
|
||||
}
|
||||
|
||||
// LevelContext groups the level-detection artefacts (per-line levels
|
||||
// + the most-common heading level) so the strategy implementations
|
||||
// don't recompute.
|
||||
type LevelContext struct {
|
||||
levels []int
|
||||
mostLevel int
|
||||
}
|
||||
|
||||
func newLevelContext(lines []string, p *titleChunkerParam) LevelContext {
|
||||
levels := resolveTitleLevels(lines, p)
|
||||
most := 0
|
||||
seen := make(map[int]int)
|
||||
for _, lvl := range levels {
|
||||
seen[lvl]++
|
||||
}
|
||||
// Pick the smallest level that has any hits — titles are
|
||||
// more-specific headings, so the highest-specificity one wins.
|
||||
for _, lvl := range seen {
|
||||
if lvl < bodyLevel && (most == 0 || lvl < most) {
|
||||
most = lvl
|
||||
}
|
||||
}
|
||||
return LevelContext{levels: levels, mostLevel: most}
|
||||
}
|
||||
|
||||
func (lc LevelContext) Levels() []int {
|
||||
out := make([]int, len(lc.levels))
|
||||
copy(out, lc.levels)
|
||||
return out
|
||||
}
|
||||
|
||||
// buildChunksFromRecordGroups is the cheap text-form materialiser
|
||||
// (used by Group/Hierarchy for plain-text payloads). For structured
|
||||
// upstream payloads a parallel fan-out over groups is used.
|
||||
func buildChunksFromRecordGroupsText(groups [][]lineRecord) []map[string]any {
|
||||
out := make([]map[string]any, 0, len(groups))
|
||||
for _, g := range groups {
|
||||
if len(g) == 0 {
|
||||
continue
|
||||
}
|
||||
var sb strings.Builder
|
||||
for _, r := range g {
|
||||
sb.WriteString(r.text)
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
out = append(out, map[string]any{"text": sb.String()})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Component implementations
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// TitleChunkerComponent dispatches based on `param.method`. Heading
|
||||
// detection is shared (resolveTitleLevels); the actual chunk-build
|
||||
// logic lives in group.go / hierarchy.go.
|
||||
type TitleChunkerComponent struct {
|
||||
name string
|
||||
param titleChunkerParam
|
||||
}
|
||||
|
||||
// NewTitleChunker constructs a TitleChunker from the DSL param map.
|
||||
// Errors here surface as canvas compile failures.
|
||||
func NewTitleChunker(params map[string]any) (runtime.Component, error) {
|
||||
p := defaultsTitle()
|
||||
p.Method = "group" // default to group
|
||||
p.Update(params)
|
||||
if err := p.TitleChunkerParam.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("TitleChunker: %w", err)
|
||||
}
|
||||
return &TitleChunkerComponent{
|
||||
name: ComponentNameTitleChunker,
|
||||
param: p,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Parallelism is the configured intra-component fan-out (plan §4
|
||||
// Phase 2 row 2.3b).
|
||||
func (c *TitleChunkerComponent) Parallelism() int { return 2 }
|
||||
|
||||
// Inputs is exposed so callers can introspect.
|
||||
func (c *TitleChunkerComponent) Inputs() map[string]string { return ChunkerInputs }
|
||||
|
||||
// Outputs is exposed so callers can introspect.
|
||||
func (c *TitleChunkerComponent) Outputs() map[string]string { return ChunkerOutputs }
|
||||
|
||||
// Invoke delegates to the chosen strategy (group or hierarchy).
|
||||
func (c *TitleChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
return runtime.TrackElapsed(ComponentNameTitleChunker, func() (map[string]any, error) {
|
||||
if inputs == nil {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
if _, ok := inputs["name"].(string); !ok {
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{},
|
||||
"_ERROR": "TitleChunker: missing required upstream field \"name\"",
|
||||
}, nil
|
||||
}
|
||||
switch c.param.Method {
|
||||
case "hierarchy":
|
||||
return invokeHierarchy(ctx, inputs, &c.param)
|
||||
case "group":
|
||||
return invokeGroup(ctx, inputs, &c.param)
|
||||
default:
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{},
|
||||
"_ERROR": fmt.Sprintf("TitleChunker: unsupported method %q", c.param.Method),
|
||||
}, nil
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// init registers TitleChunker under CategoryIngestion.
|
||||
func init() {
|
||||
MustRegisterChunker(ComponentNameTitleChunker)
|
||||
}
|
||||
257
internal/ingestion/component/chunker/title_test.go
Normal file
257
internal/ingestion/component/chunker/title_test.go
Normal file
@@ -0,0 +1,257 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
// TestTitleChunker_Registered asserts the registry has a CategoryIngestion
|
||||
// entry for TitleChunker.
|
||||
func TestTitleChunker_Registered(t *testing.T) {
|
||||
factory, cat, meta, ok := runtime.DefaultRegistry.Lookup("TitleChunker")
|
||||
if !ok {
|
||||
t.Fatal("TitleChunker: registry miss")
|
||||
}
|
||||
if cat != runtime.CategoryIngestion {
|
||||
t.Errorf("category = %q, want %q", cat, runtime.CategoryIngestion)
|
||||
}
|
||||
if factory == nil {
|
||||
t.Error("factory is nil")
|
||||
}
|
||||
if len(meta.Inputs) == 0 {
|
||||
t.Errorf("inputs metadata is empty")
|
||||
}
|
||||
if len(meta.Outputs) == 0 {
|
||||
t.Errorf("outputs metadata is empty")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTitleChunker_Parallelism enforces plan row 2.3b parallelism (2).
|
||||
func TestTitleChunker_Parallelism(t *testing.T) {
|
||||
c, err := NewTitleChunker(map[string]any{
|
||||
"method": "group",
|
||||
"levels": [][]string{{"^# "}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTitleChunker: %v", err)
|
||||
}
|
||||
tc, ok := c.(*TitleChunkerComponent)
|
||||
if !ok {
|
||||
t.Fatalf("NewTitleChunker returned %T", c)
|
||||
}
|
||||
if got := tc.Parallelism(); got != 2 {
|
||||
t.Errorf("Parallelism() = %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTitleChunker_InputsOutputs_NonEmpty asserts metadata is
|
||||
// populated for both ends.
|
||||
func TestTitleChunker_InputsOutputs_NonEmpty(t *testing.T) {
|
||||
_, _, meta, ok := runtime.DefaultRegistry.Lookup("TitleChunker")
|
||||
if !ok {
|
||||
t.Fatal("registry miss")
|
||||
}
|
||||
if len(meta.Inputs) == 0 {
|
||||
t.Error("inputs metadata is empty")
|
||||
}
|
||||
if len(meta.Outputs) == 0 {
|
||||
t.Error("outputs metadata is empty")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTitleChunker_NewRejectsBadMethod enforces the method check.
|
||||
func TestTitleChunker_NewRejectsBadMethod(t *testing.T) {
|
||||
if _, err := NewTitleChunker(map[string]any{
|
||||
"method": "unknown",
|
||||
"levels": [][]string{{"^# "}},
|
||||
}); err == nil {
|
||||
t.Fatal("expected error, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTitleChunker_NewRejectsHierarchyWithoutHierarchyParam enforces
|
||||
// the hierarchy branch's required-field check.
|
||||
func TestTitleChunker_NewRejectsHierarchyWithoutHierarchyParam(t *testing.T) {
|
||||
if _, err := NewTitleChunker(map[string]any{
|
||||
"method": "hierarchy",
|
||||
"levels": [][]string{{"^# "}},
|
||||
}); err == nil {
|
||||
t.Fatal("expected error, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTitleChunker_InvokeEmptyInput returns empty chunks for an
|
||||
// empty payload.
|
||||
func TestTitleChunker_InvokeEmptyInput(t *testing.T) {
|
||||
c, err := NewTitleChunker(map[string]any{
|
||||
"method": "group",
|
||||
"levels": [][]string{{"^# "}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTitleChunker: %v", err)
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got, want := out["output_format"], "chunks"; got != want {
|
||||
t.Errorf("output_format = %v, want %v", got, want)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) != 0 {
|
||||
t.Errorf("chunks = %d, want 0", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
// TestTitleChunker_Headings_ASCII is the golden-file parity check
|
||||
// for the markdown detector: a # heading + body + ## subheading + body
|
||||
// is recognized at least one chunked partition boundary.
|
||||
//
|
||||
// Note: the title strategy is dispatched to the underlying strategy
|
||||
// (`group` by default, or `hierarchy` when configured). Each
|
||||
// strategy test lives in its own file; this test routes through
|
||||
// both via the dispatcher.
|
||||
func TestTitleChunker_Headings_ASCII(t *testing.T) {
|
||||
c, err := NewTitleChunker(map[string]any{
|
||||
"method": "group",
|
||||
"levels": [][]string{
|
||||
{`^# `},
|
||||
{`^## `},
|
||||
{`^### `},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTitleChunker: %v", err)
|
||||
}
|
||||
|
||||
input := "# Top\nFirst body line under Top.\nSecond body.\n## Sub\nBody under Sub heading.\n# TopTwo\nBody under TopTwo."
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc.md",
|
||||
"text": input,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got, want := out["output_format"], "chunks"; got != want {
|
||||
t.Errorf("output_format = %v, want %v", got, want)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
for i, ck := range chunks {
|
||||
if text, _ := ck["text"].(string); text == "" {
|
||||
t.Errorf("chunk[%d] text is empty", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestTitleChunker_NoHeadings_FallsBack feeds plain text without any
|
||||
// markdown heading; the chunker should still produce a chunk
|
||||
// containing the body text (single chunk fallback).
|
||||
func TestTitleChunker_NoHeadings_FallsBack(t *testing.T) {
|
||||
c, err := NewTitleChunker(map[string]any{
|
||||
"method": "group",
|
||||
"levels": [][]string{{"^# "}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTitleChunker: %v", err)
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc.txt",
|
||||
"text": "alpha line one\nalpha line two",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTitleChunker_DispatcherHierarchy routes to the hierarchy
|
||||
// strategy without panicking.
|
||||
func TestTitleChunker_DispatcherHierarchy(t *testing.T) {
|
||||
c, err := NewTitleChunker(map[string]any{
|
||||
"method": "hierarchy",
|
||||
"hierarchy": 1,
|
||||
"levels": [][]string{{`^# `}, {`^## `}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTitleChunker: %v", err)
|
||||
}
|
||||
if _, ok := c.(*TitleChunkerComponent); !ok {
|
||||
t.Fatalf("NewTitleChunker returned %T", c)
|
||||
}
|
||||
_, err = c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc.md",
|
||||
"text": "# Top\nFirst body.\n# TopTwo\nBody under TopTwo.",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTitleChunker_InvokeDeterministic runs the heading detector
|
||||
// 10 times and asserts the chunks list is identical.
|
||||
func TestTitleChunker_InvokeDeterministic(t *testing.T) {
|
||||
c, err := NewTitleChunker(map[string]any{
|
||||
"method": "group",
|
||||
"levels": [][]string{
|
||||
{`^# `},
|
||||
{`^## `},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTitleChunker: %v", err)
|
||||
}
|
||||
inputs := map[string]any{
|
||||
"name": "doc.md",
|
||||
"text": "# A\nbody a line 1\nbody a line 2\n# B\nbody b line 1\n# C\nbody c line 1",
|
||||
}
|
||||
var firstLen int
|
||||
var firstTexts []string
|
||||
for run := 0; run < 10; run++ {
|
||||
out, err := c.Invoke(context.Background(), inputs)
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke run %d: %v", run, err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
texts := make([]string, len(chunks))
|
||||
for i, ck := range chunks {
|
||||
texts[i], _ = ck["text"].(string)
|
||||
}
|
||||
if run == 0 {
|
||||
firstLen = len(chunks)
|
||||
firstTexts = texts
|
||||
continue
|
||||
}
|
||||
if firstLen != len(chunks) {
|
||||
t.Fatalf("run %d: chunk count changed (%d vs %d)", run, len(chunks), firstLen)
|
||||
}
|
||||
for i := range chunks {
|
||||
if firstTexts[i] != texts[i] {
|
||||
t.Fatalf("run %d: chunk[%d] text changed", run, i)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
795
internal/ingestion/component/chunker/token.go
Normal file
795
internal/ingestion/component/chunker/token.go
Normal file
@@ -0,0 +1,795 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// SCOPE (honest) for token.go (Phase 2.3a):
|
||||
//
|
||||
// - WHITELIST mirrors rag/flow/chunker/token_chunker.py:TokenChunkerParam.check
|
||||
// exactly: delimiter_mode ∈ {"token_size","delimiter","one"},
|
||||
// chunk_token_size > 0, overlapped_percent ∈ [0, 1), table_context_size ≥ 0,
|
||||
// image_context_size ≥ 0. enum/range checks live in param.Check.
|
||||
//
|
||||
// - DELIMITER PARSING mirrors python `_compile_delimiter_pattern`:
|
||||
// entries wrapped in backticks (e.g. "`\\n\\n`") are treated as
|
||||
// regex split points; plain strings are regex-escaped and joined
|
||||
// into the same alternation. Empty entries are filtered.
|
||||
//
|
||||
// - CHILDREN DELIMITERS (the secondary split) is implemented via the
|
||||
// shared splitKeepingDelim helper; emitted chunks carry the parent
|
||||
// ("mom") and the split child ("text") keys.
|
||||
//
|
||||
// - MODE "delimiter" uses the regex-aware delimiter pattern; the
|
||||
// token_size merge pass only runs when no working delimiter was
|
||||
// detected — matching python at token_chunker.py:359-360.
|
||||
//
|
||||
// - MODE "token_size" falls back to a chunk-engine merge plan. The
|
||||
// python `naive_merge` algorithm uses a sentence-aware + stop-word
|
||||
// strategy that the Go chunk_engine's "greedy" merge approximates;
|
||||
// this is flagged as a TODO parity-gap and intentionally not
|
||||
// machine-mirrored.
|
||||
//
|
||||
// - MODE "one" emits a single chunk containing the entire payload.
|
||||
//
|
||||
// - JSON-STRUCTURED INPUT (output_format == "json", or the default
|
||||
// parser-style branch when output_format is unset) is normalized
|
||||
// into the same internal chunk shape via a parallel fan-out
|
||||
// (Plan §4 Phase 2 row 2.3a, Parallelism=4).
|
||||
// Media-context attachment is per-item sequential; merge is
|
||||
// index-deterministic (Plan §8 R8).
|
||||
//
|
||||
// - No PDF/outline awareness (Python `restore_pdf_text_previews`).
|
||||
// That depends on deepdoc/parser which is out of scope for this
|
||||
// phase; the chunker accepts the parser-style structured JSON
|
||||
// payload and runs the same logic against it.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
"ragflow/internal/parser/chunk"
|
||||
)
|
||||
|
||||
const ComponentNameTokenChunker = "TokenChunker"
|
||||
|
||||
type tokenChunkerParam struct {
|
||||
schema.TokenChunkerParam
|
||||
}
|
||||
|
||||
func (p *tokenChunkerParam) Update(conf map[string]any) {
|
||||
if conf == nil {
|
||||
return
|
||||
}
|
||||
if v, ok := conf["delimiter_mode"].(string); ok {
|
||||
p.TokenChunkerParam.DelimiterMode = v
|
||||
}
|
||||
if v, ok := numericFromAny(conf["chunk_token_size"]); ok {
|
||||
p.TokenChunkerParam.ChunkTokenSize = int(v)
|
||||
}
|
||||
if v, ok := conf["delimiters"].([]any); ok {
|
||||
p.TokenChunkerParam.Delimiters = stringListFromAny(v)
|
||||
} else if v, ok := conf["delimiters"].([]string); ok {
|
||||
p.TokenChunkerParam.Delimiters = append([]string(nil), v...)
|
||||
}
|
||||
if v, ok := numericFromAny(conf["overlapped_percent"]); ok {
|
||||
p.TokenChunkerParam.OverlappedPercent = v
|
||||
}
|
||||
if v, ok := conf["children_delimiters"].([]any); ok {
|
||||
p.TokenChunkerParam.ChildrenDelimiters = stringListFromAny(v)
|
||||
} else if v, ok := conf["children_delimiters"].([]string); ok {
|
||||
p.TokenChunkerParam.ChildrenDelimiters = append([]string(nil), v...)
|
||||
}
|
||||
if v, ok := numericFromAny(conf["table_context_size"]); ok {
|
||||
p.TokenChunkerParam.TableContextSize = int(v)
|
||||
}
|
||||
if v, ok := numericFromAny(conf["image_context_size"]); ok {
|
||||
p.TokenChunkerParam.ImageContextSize = int(v)
|
||||
}
|
||||
}
|
||||
|
||||
func defaultsToken(p tokenChunkerParam) tokenChunkerParam {
|
||||
p.TokenChunkerParam = schema.TokenChunkerParam{}.Defaults()
|
||||
return p
|
||||
}
|
||||
|
||||
// TokenChunkerComponent implements the runtime.Component interface for
|
||||
// the TokenChunker variant.
|
||||
type TokenChunkerComponent struct {
|
||||
name string
|
||||
param tokenChunkerParam
|
||||
}
|
||||
|
||||
// NewTokenChunker constructs a TokenChunker from the DSL param map.
|
||||
// Errors here surface as canvas compile failures (mirrors the
|
||||
// python check() phase).
|
||||
func NewTokenChunker(params map[string]any) (runtime.Component, error) {
|
||||
p := defaultsToken(tokenChunkerParam{})
|
||||
p.Update(params)
|
||||
if err := p.TokenChunkerParam.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("TokenChunker: %w", err)
|
||||
}
|
||||
return &TokenChunkerComponent{
|
||||
name: ComponentNameTokenChunker,
|
||||
param: p,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Parallelism is the configured intra-component fan-out (plan §4
|
||||
// Phase 2 row 2.3a).
|
||||
func (c *TokenChunkerComponent) Parallelism() int { return 4 }
|
||||
|
||||
// Inputs is exposed so callers can introspect.
|
||||
func (c *TokenChunkerComponent) Inputs() map[string]string { return ChunkerInputs }
|
||||
|
||||
// Outputs is exposed so callers can introspect.
|
||||
func (c *TokenChunkerComponent) Outputs() map[string]string { return ChunkerOutputs }
|
||||
|
||||
// Invoke runs the chunker against the input payload.
|
||||
//
|
||||
// Concurrency: text payloads are fanned across Parallelism goroutines
|
||||
// by primary-delimiter segment; structured JSON/chunks payloads fan
|
||||
// across items. Merge is by input index (plan §8 R8): the i-th
|
||||
// goroutine's output occupies slot i, regardless of completion order.
|
||||
//
|
||||
// Timeout: honours ctx cancellation only — there is no inner @timeout
|
||||
// decorator equivalent (plan §8 R1).
|
||||
func (c *TokenChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
return runtime.TrackElapsed(ComponentNameTokenChunker, func() (map[string]any, error) {
|
||||
return c.invoke(ctx, inputs)
|
||||
})
|
||||
}
|
||||
|
||||
func (c *TokenChunkerComponent) invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
if inputs == nil {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
upstream, err := decodeChunkerFromUpstream(inputs)
|
||||
if err != nil {
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{},
|
||||
"_ERROR": fmt.Sprintf("Input error: %v", err),
|
||||
}, nil
|
||||
}
|
||||
|
||||
delimPattern := compileDelimPattern(c.param.Delimiters)
|
||||
childrenPattern := compileChildrenPattern(c.param.ChildrenDelimiters)
|
||||
|
||||
switch upstream.OutputFormat {
|
||||
case schema.PayloadFormatMarkdown:
|
||||
if upstream.MarkdownResult == nil {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
return c.invokeTextPayload(ctx, *upstream.MarkdownResult, delimPattern, childrenPattern), nil
|
||||
case schema.PayloadFormatText:
|
||||
if upstream.TextResult == nil {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
return c.invokeTextPayload(ctx, *upstream.TextResult, delimPattern, childrenPattern), nil
|
||||
case schema.PayloadFormatHTML:
|
||||
if upstream.HTMLResult == nil {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
return c.invokeTextPayload(ctx, *upstream.HTMLResult, delimPattern, childrenPattern), nil
|
||||
default:
|
||||
return c.invokeJSONPayload(ctx, upstream.JSONResult, delimPattern, childrenPattern), nil
|
||||
}
|
||||
}
|
||||
|
||||
func decodeChunkerFromUpstream(inputs map[string]any) (schema.ChunkerFromUpstream, error) {
|
||||
var out schema.ChunkerFromUpstream
|
||||
data, err := json.Marshal(stripChunkerRuntimeTimestamps(inputs))
|
||||
if err != nil {
|
||||
return out, err
|
||||
}
|
||||
if err := json.Unmarshal(data, &out); err != nil {
|
||||
return out, err
|
||||
}
|
||||
if err := out.Validate(); err != nil {
|
||||
return out, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func stripChunkerRuntimeTimestamps(inputs map[string]any) map[string]any {
|
||||
out := make(map[string]any, len(inputs))
|
||||
for k, v := range inputs {
|
||||
if k == "_created_time" || k == "_elapsed_time" {
|
||||
continue
|
||||
}
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// invokeTextPayload handles plain-text input (output_format in
|
||||
// {markdown,text,html} on the python side).
|
||||
func (c *TokenChunkerComponent) invokeTextPayload(_ context.Context, text string, delimPattern, childrenPattern *regexp.Regexp) map[string]any {
|
||||
if text == "" {
|
||||
return emptyOutputs()
|
||||
}
|
||||
|
||||
mode := c.param.DelimiterMode
|
||||
if mode == "one" {
|
||||
return chunkOutputs([]schema.ChunkDoc{{Text: text}})
|
||||
}
|
||||
|
||||
if !hasActiveDelimiter(delimPattern) {
|
||||
// No primary delimiter hit — fall back to a token-size merge.
|
||||
return c.mergeByTokenSize(text, childrenPattern)
|
||||
}
|
||||
|
||||
parts := splitKeepingDelim(text, delimPattern)
|
||||
cleaned := make([]string, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
if strings.TrimSpace(p) == "" {
|
||||
continue
|
||||
}
|
||||
cleaned = append(cleaned, p)
|
||||
}
|
||||
if len(cleaned) == 0 {
|
||||
return emptyOutputs()
|
||||
}
|
||||
docs := applyChildrenDelim(cleaned, childrenPattern)
|
||||
return chunkOutputs(docs)
|
||||
}
|
||||
|
||||
// mergeByTokenSize uses the chunk library's split + postprocess-merge
|
||||
// to combine the payload into chunks of approximately
|
||||
// chunk_token_size runes. Mirrors python's `naive_merge` fallback
|
||||
// (token_chunker.py:319-324) at the wire-shape level; the merge
|
||||
// strategy is the Go chunk library's "greedy" approximation.
|
||||
//
|
||||
// This caller uses the typed chunk.Run entry point directly:
|
||||
// same operator sequence, no DSL round-trip.
|
||||
func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *regexp.Regexp) map[string]any {
|
||||
target := c.param.ChunkTokenSize
|
||||
result, err := chunk.Run(text, chunk.ChunkOptions{
|
||||
StripWhitespace: true,
|
||||
RemoveEmptyLines: true,
|
||||
SplitStrategy: "sentence",
|
||||
MergeTargetSize: target,
|
||||
})
|
||||
if err != nil {
|
||||
return chunkOutputs([]schema.ChunkDoc{{Text: text}})
|
||||
}
|
||||
docs := make([]schema.ChunkDoc, 0, len(result.ResultChunks))
|
||||
for _, ck := range result.ResultChunks {
|
||||
content := strings.TrimSpace(ck.Content)
|
||||
if content == "" {
|
||||
continue
|
||||
}
|
||||
docs = append(docs, schema.ChunkDoc{Text: content})
|
||||
}
|
||||
final := applyChildrenDelimText(docs, childrenPattern)
|
||||
return chunkOutputs(final)
|
||||
}
|
||||
|
||||
// invokeJSONPayload handles structured upstream input. Items fan
|
||||
// across goroutines (Parallelism); merge is by input index.
|
||||
func (c *TokenChunkerComponent) invokeJSONPayload(ctx context.Context, items []schema.ChunkDoc, delimPattern, childrenPattern *regexp.Regexp) map[string]any {
|
||||
if len(items) == 0 {
|
||||
return emptyOutputs()
|
||||
}
|
||||
mode := c.param.DelimiterMode
|
||||
if mode == "one" {
|
||||
var parts []string
|
||||
for _, it := range items {
|
||||
if t := itemTextOrFallback(it); t != "" {
|
||||
parts = append(parts, t)
|
||||
}
|
||||
}
|
||||
merged := strings.Join(parts, "\n")
|
||||
if merged == "" {
|
||||
return emptyOutputs()
|
||||
}
|
||||
return chunkOutputs([]schema.ChunkDoc{{Text: merged}})
|
||||
}
|
||||
|
||||
workers := c.Parallelism()
|
||||
if workers < 1 {
|
||||
workers = 1
|
||||
}
|
||||
if workers > len(items) {
|
||||
workers = len(items)
|
||||
}
|
||||
lanes := partition(len(items), workers)
|
||||
perItem := make([][]schema.ChunkDoc, len(items))
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for w := 0; w < workers; w++ {
|
||||
lane := lanes[w]
|
||||
wg.Add(1)
|
||||
go func(start, end int) {
|
||||
defer wg.Done()
|
||||
for i := start; i < end; i++ {
|
||||
if err := ctx.Err(); err != nil {
|
||||
perItem[i] = nil
|
||||
continue
|
||||
}
|
||||
perItem[i] = chunkFromItem(items[i], delimPattern)
|
||||
}
|
||||
}(lane.start, lane.end)
|
||||
}
|
||||
wg.Wait()
|
||||
if err := ctx.Err(); err != nil {
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{},
|
||||
"_ERROR": fmt.Sprintf("TokenChunker: %v", err),
|
||||
}
|
||||
}
|
||||
|
||||
// Attach surrounding media context (token_chunker.py:358).
|
||||
attached := attachMediaContext(perItem, c.param.TableContextSize, c.param.ImageContextSize)
|
||||
|
||||
// Optional token-size merge (only when no working delimiter).
|
||||
if !hasActiveDelimiter(delimPattern) {
|
||||
attached = mergeByTokenSizeFromJSON(attached, c.param.ChunkTokenSize, c.param.OverlappedPercent)
|
||||
}
|
||||
|
||||
flat := flatten(attached)
|
||||
if childrenPattern != nil {
|
||||
flat = splitByChildren(flat, childrenPattern)
|
||||
}
|
||||
|
||||
out := make([]schema.ChunkDoc, 0, len(flat))
|
||||
for _, m := range flat {
|
||||
if strings.TrimSpace(m.Text) == "" {
|
||||
continue
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return emptyOutputs()
|
||||
}
|
||||
return chunkOutputs(out)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// JSON-payload internals
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// chunkFromItem mirrors _build_json_chunks for a single item.
|
||||
func chunkFromItem(it schema.ChunkDoc, delimPattern *regexp.Regexp) []schema.ChunkDoc {
|
||||
ckType := itemDocType(it)
|
||||
txt := itemTextOrFallback(it)
|
||||
if ckType != "text" {
|
||||
return []schema.ChunkDoc{buildChunkDoc(it, ckType, txt, "", "")}
|
||||
}
|
||||
if !hasActiveDelimiter(delimPattern) {
|
||||
return []schema.ChunkDoc{buildChunkDoc(it, "text", txt, "", "")}
|
||||
}
|
||||
parts := splitKeepingDelim(txt, delimPattern)
|
||||
if !delimPattern.MatchString(txt) {
|
||||
return []schema.ChunkDoc{buildChunkDoc(it, "text", txt, "", "")}
|
||||
}
|
||||
out := make([]schema.ChunkDoc, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
if strings.TrimSpace(p) == "" {
|
||||
continue
|
||||
}
|
||||
out = append(out, buildChunkDoc(it, "text", p, "", ""))
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return []schema.ChunkDoc{buildChunkDoc(it, "text", txt, "", "")}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// buildChunkMap constructs the python-compatible chunk payload.
|
||||
//
|
||||
// The chunker output carries the basic text+doc_type_kwd+ck_type
|
||||
// fields plus the per-chunk meta fields the python
|
||||
// rag/flow/chunker/token_chunker.py emits:
|
||||
//
|
||||
// - tk_nums — tokenized list (used downstream by Tokenizer)
|
||||
// - mom — parent-section identifier (title / hierarchy
|
||||
// chunkers populate; TokenChunker pass-through)
|
||||
// - img_id — image attachment identifier
|
||||
// - layout — layout classification (text / table / image / figure)
|
||||
// - _pdf_positions — PDF bbox coordinates when the parser path
|
||||
// emitted them on the upstream item
|
||||
// - context_above / context_below — surrounding media context
|
||||
// when attachMediaContext was invoked
|
||||
//
|
||||
// Pass-through fields are sourced from the input item map. Missing
|
||||
// fields are simply absent from the output (the python side does
|
||||
// the same — see python `_build_json_chunks`).
|
||||
func buildChunkDoc(it schema.ChunkDoc, ckType, text, ctxAbove, ctxBelow string) schema.ChunkDoc {
|
||||
out := schema.ChunkDoc{
|
||||
Text: text,
|
||||
DocType: ckType,
|
||||
CKType: ckType,
|
||||
TKNums: intPtr(tokenizeStr(text)),
|
||||
Mom: it.Mom,
|
||||
ImgID: it.ImgID,
|
||||
Layout: it.Layout,
|
||||
PDFPositions: it.PDFPositions,
|
||||
Positions: it.Positions,
|
||||
Image: it.Image,
|
||||
PageNumber: it.PageNumber,
|
||||
}
|
||||
if ctxAbove != "" {
|
||||
out.ContextAbove = ctxAbove
|
||||
}
|
||||
if ctxBelow != "" {
|
||||
out.ContextBelow = ctxBelow
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
type lane struct{ start, end int }
|
||||
|
||||
func partition(n, parts int) []lane {
|
||||
if parts < 1 {
|
||||
parts = 1
|
||||
}
|
||||
if n < parts {
|
||||
parts = n
|
||||
}
|
||||
out := make([]lane, 0, parts)
|
||||
size := n / parts
|
||||
rem := n % parts
|
||||
cursor := 0
|
||||
for i := 0; i < parts; i++ {
|
||||
end := cursor + size
|
||||
if i < rem {
|
||||
end++
|
||||
}
|
||||
if end > n {
|
||||
end = n
|
||||
}
|
||||
if cursor < end {
|
||||
out = append(out, lane{start: cursor, end: end})
|
||||
}
|
||||
cursor = end
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func attachMediaContext(perItem [][]schema.ChunkDoc, tableCtx, imageCtx int) [][]schema.ChunkDoc {
|
||||
if tableCtx <= 0 && imageCtx <= 0 {
|
||||
return perItem
|
||||
}
|
||||
for idx := range perItem {
|
||||
chunks := perItem[idx]
|
||||
if len(chunks) == 0 {
|
||||
continue
|
||||
}
|
||||
for i, ck := range chunks {
|
||||
ckType := ck.CKType
|
||||
if ckType != "table" && ckType != "image" {
|
||||
continue
|
||||
}
|
||||
ctx := imageCtx
|
||||
if ckType == "table" {
|
||||
ctx = tableCtx
|
||||
}
|
||||
if ctx <= 0 {
|
||||
continue
|
||||
}
|
||||
chunks[i].ContextAbove = collectContext(chunks, i, ctx, true)
|
||||
chunks[i].ContextBelow = collectContext(chunks, i, ctx, false)
|
||||
}
|
||||
}
|
||||
return perItem
|
||||
}
|
||||
|
||||
// collectContext walks chunks around `i` (above when direction==true,
|
||||
// below when false), pulling text chunks while remaining token budget
|
||||
// stays positive. Matches token_chunker.py:_attach_context_to_media_chunks.
|
||||
func collectContext(chunks []schema.ChunkDoc, i, ctxTokens int, above bool) string {
|
||||
var parts []string
|
||||
remain := ctxTokens
|
||||
var pos int
|
||||
if above {
|
||||
pos = i - 1
|
||||
for pos >= 0 && remain > 0 {
|
||||
if chunks[pos].CKType == "text" {
|
||||
tk := intValue(chunks[pos].TKNums)
|
||||
txt := chunks[pos].Text
|
||||
if tk >= remain {
|
||||
parts = append([]string{takeFromEnd(txt, remain)}, parts...)
|
||||
remain = 0
|
||||
break
|
||||
}
|
||||
parts = append([]string{txt}, parts...)
|
||||
remain -= tk
|
||||
}
|
||||
pos--
|
||||
}
|
||||
} else {
|
||||
pos = i + 1
|
||||
for pos < len(chunks) && remain > 0 {
|
||||
if chunks[pos].CKType == "text" {
|
||||
tk := intValue(chunks[pos].TKNums)
|
||||
txt := chunks[pos].Text
|
||||
if tk >= remain {
|
||||
parts = append(parts, takeFromStart(txt, remain))
|
||||
remain = 0
|
||||
break
|
||||
}
|
||||
parts = append(parts, txt)
|
||||
remain -= tk
|
||||
}
|
||||
pos++
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, "")
|
||||
}
|
||||
|
||||
// takeFromEnd returns the last approx `tokens` worth of text (1 token
|
||||
// ≈ 4 bytes is the best-effort approximation used here; python uses
|
||||
// the actual tokenizer).
|
||||
func takeFromEnd(text string, tokens int) string {
|
||||
bytes := tokens * 4
|
||||
if bytes >= len(text) {
|
||||
return text
|
||||
}
|
||||
return text[len(text)-bytes:]
|
||||
}
|
||||
|
||||
func takeFromStart(text string, tokens int) string {
|
||||
bytes := tokens * 4
|
||||
if bytes >= len(text) {
|
||||
return text
|
||||
}
|
||||
return text[:bytes]
|
||||
}
|
||||
|
||||
// mergeByTokenSizeFromJSON mirrors `_merge_text_chunks_by_token_size`
|
||||
// at token_chunker.py:212-243.
|
||||
func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, overlappedPct float64) [][]schema.ChunkDoc {
|
||||
threshold := float64(chunkTokens) * (100 - overlappedPct*100) / 100.0
|
||||
for idx := range perItem {
|
||||
chunks := perItem[idx]
|
||||
if len(chunks) == 0 {
|
||||
continue
|
||||
}
|
||||
var merged []schema.ChunkDoc
|
||||
prevTextIdx := -1
|
||||
for _, ck := range chunks {
|
||||
ckType := ck.CKType
|
||||
if ckType != "text" {
|
||||
merged = append(merged, cloneChunkDoc(ck))
|
||||
prevTextIdx = -1
|
||||
continue
|
||||
}
|
||||
tk := intValue(ck.TKNums)
|
||||
if prevTextIdx < 0 || float64(tk) > threshold {
|
||||
cp := cloneChunkDoc(ck)
|
||||
if prevTextIdx >= 0 && overlappedPct > 0 {
|
||||
if prevText := merged[prevTextIdx].Text; prevText != "" {
|
||||
cut := int(float64(len(prevText)) * (100 - overlappedPct*100) / 100.0)
|
||||
if cut < len(prevText) {
|
||||
cp.Text = prevText[cut:] + cp.Text
|
||||
cp.TKNums = intPtr(tk + tokenizeStr(cp.Text))
|
||||
}
|
||||
}
|
||||
}
|
||||
merged = append(merged, cp)
|
||||
prevTextIdx = len(merged) - 1
|
||||
continue
|
||||
}
|
||||
// Merge into prev text chunk.
|
||||
if pt := merged[prevTextIdx].Text; pt != "" {
|
||||
merged[prevTextIdx].Text = pt + "\n" + ck.Text
|
||||
merged[prevTextIdx].TKNums = intPtr(intValue(merged[prevTextIdx].TKNums) + intValue(ck.TKNums))
|
||||
}
|
||||
}
|
||||
perItem[idx] = merged
|
||||
}
|
||||
return perItem
|
||||
}
|
||||
|
||||
func cloneChunkDoc(in schema.ChunkDoc) schema.ChunkDoc {
|
||||
out := in
|
||||
if in.TKNums != nil {
|
||||
v := *in.TKNums
|
||||
out.TKNums = &v
|
||||
}
|
||||
if in.ChunkOrderInt != nil {
|
||||
v := *in.ChunkOrderInt
|
||||
out.ChunkOrderInt = &v
|
||||
}
|
||||
if in.PageNumber != nil {
|
||||
v := *in.PageNumber
|
||||
out.PageNumber = &v
|
||||
}
|
||||
if in.Extra != nil {
|
||||
out.Extra = make(map[string]json.RawMessage, len(in.Extra))
|
||||
for k, v := range in.Extra {
|
||||
out.Extra[k] = append(json.RawMessage(nil), v...)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func flatten(perItem [][]schema.ChunkDoc) []schema.ChunkDoc {
|
||||
var out []schema.ChunkDoc
|
||||
for _, cs := range perItem {
|
||||
out = append(out, cs...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func splitByChildren(chunks []schema.ChunkDoc, pattern *regexp.Regexp) []schema.ChunkDoc {
|
||||
if pattern == nil {
|
||||
return chunks
|
||||
}
|
||||
var out []schema.ChunkDoc
|
||||
for _, ck := range chunks {
|
||||
if ck.DocType != "text" {
|
||||
out = append(out, ck)
|
||||
continue
|
||||
}
|
||||
mom := ck.Text
|
||||
parts := splitKeepingDelim(mom, pattern)
|
||||
for _, p := range parts {
|
||||
if strings.TrimSpace(p) == "" {
|
||||
continue
|
||||
}
|
||||
cp := cloneChunkDoc(ck)
|
||||
cp.Text = p
|
||||
cp.Mom = mom
|
||||
out = append(out, cp)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// shared text-payload helpers (used by TitleChunker et al.)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// hasActiveDelimiter reports whether a regex compiled by
|
||||
// compileDelimPattern contains any non-placeholder pattern. The "match
|
||||
// nothing" sentinel regexp makes a quick `pattern.MatchString("")`
|
||||
// viable as a check without re-walking the source slice.
|
||||
func hasActiveDelimiter(p *regexp.Regexp) bool {
|
||||
return p != nil && p.String() != `\A(?!)`
|
||||
}
|
||||
|
||||
// applyChildrenDelim mirrors token_chunker.py:325-334.
|
||||
func applyChildrenDelim(segs []string, pattern *regexp.Regexp) []schema.ChunkDoc {
|
||||
if pattern == nil {
|
||||
out := make([]schema.ChunkDoc, 0, len(segs))
|
||||
for _, s := range segs {
|
||||
out = append(out, schema.ChunkDoc{Text: s})
|
||||
}
|
||||
return out
|
||||
}
|
||||
var docs []schema.ChunkDoc
|
||||
for _, seg := range segs {
|
||||
if strings.TrimSpace(seg) == "" {
|
||||
continue
|
||||
}
|
||||
for _, child := range splitKeepingDelim(seg, pattern) {
|
||||
if strings.TrimSpace(child) == "" {
|
||||
continue
|
||||
}
|
||||
docs = append(docs, schema.ChunkDoc{Text: child, Mom: seg})
|
||||
}
|
||||
}
|
||||
return docs
|
||||
}
|
||||
|
||||
func applyChildrenDelimText(docs []schema.ChunkDoc, pattern *regexp.Regexp) []schema.ChunkDoc {
|
||||
if pattern == nil {
|
||||
return docs
|
||||
}
|
||||
var out []schema.ChunkDoc
|
||||
for _, d := range docs {
|
||||
t := d.Text
|
||||
if strings.TrimSpace(t) == "" {
|
||||
continue
|
||||
}
|
||||
for _, child := range splitKeepingDelim(t, pattern) {
|
||||
if strings.TrimSpace(child) == "" {
|
||||
continue
|
||||
}
|
||||
out = append(out, schema.ChunkDoc{Text: child, Mom: t})
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// compileChildrenPattern is the children_delimiters version of
|
||||
// compileDelimPattern. Returns nil when no delimiters exist.
|
||||
func compileChildrenPattern(delims []string) *regexp.Regexp {
|
||||
if len(delims) == 0 {
|
||||
return nil
|
||||
}
|
||||
escaped := make([]string, 0, len(delims))
|
||||
for _, d := range delims {
|
||||
if d == "" {
|
||||
continue
|
||||
}
|
||||
escaped = append(escaped, regexp.QuoteMeta(d))
|
||||
}
|
||||
if len(escaped) == 0 {
|
||||
return nil
|
||||
}
|
||||
sortSlice(escaped)
|
||||
return regexp.MustCompile(strings.Join(escaped, "|"))
|
||||
}
|
||||
|
||||
// sortSlice sorts in place by descending length (longest pattern
|
||||
// first, mirroring python's `sorted(set, key=len, reverse=True)`).
|
||||
func sortSlice(in []string) {
|
||||
for i := 1; i < len(in); i++ {
|
||||
for j := i; j > 0 && len(in[j-1]) < len(in[j]); j-- {
|
||||
in[j-1], in[j] = in[j], in[j-1]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// stringFromInputs returns the string value at the first matching key
|
||||
// in `keys`, or ("", false) when none is set.
|
||||
func stringFromInputs(inputs map[string]any, keys ...string) (string, bool) {
|
||||
for _, k := range keys {
|
||||
if v, ok := inputs[k].(string); ok {
|
||||
return v, true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// chunksFromInputs returns the chunk list from inputs as a uniform
|
||||
// []map[string]any, or nil when absent. Both []map[string]any (the
|
||||
// JSON-decoded form) and []any (the slice-of-mixed form) are handled.
|
||||
//
|
||||
// Two upstream keys are accepted, in priority order:
|
||||
//
|
||||
// - "chunks" — canonical post-chunker shape (chunker → chunker
|
||||
// re-entry, test fixtures, downstream stages).
|
||||
// - "json" — the parser-structured-output key (Parser
|
||||
// component emits under "json"; we accept it
|
||||
// so a token-chunker can run directly after
|
||||
// a parser without an intermediate reshape).
|
||||
func chunksFromInputs(inputs map[string]any) []schema.ChunkDoc {
|
||||
for _, key := range []string{"chunks", "json"} {
|
||||
v, ok := inputs[key]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
chunks, found, err := schema.ChunkDocsFromAny(v)
|
||||
if err == nil && found {
|
||||
return chunks
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func intValue(v *int) int {
|
||||
if v == nil {
|
||||
return 0
|
||||
}
|
||||
return *v
|
||||
}
|
||||
|
||||
func intPtr(v int) *int { return &v }
|
||||
|
||||
// init registers TokenChunker under CategoryIngestion.
|
||||
func init() {
|
||||
MustRegisterChunker(ComponentNameTokenChunker)
|
||||
}
|
||||
328
internal/ingestion/component/chunker/token_test.go
Normal file
328
internal/ingestion/component/chunker/token_test.go
Normal file
@@ -0,0 +1,328 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
// TestTokenChunker_Registered asserts the registry has a CategoryIngestion
|
||||
// entry for TokenChunker with a working factory. Mirrors plan §4
|
||||
// Phase 2 "registered" checklist.
|
||||
func TestTokenChunker_Registered(t *testing.T) {
|
||||
factory, cat, meta, ok := runtime.DefaultRegistry.Lookup("TokenChunker")
|
||||
if !ok {
|
||||
t.Fatal("TokenChunker: registry miss")
|
||||
}
|
||||
if cat != runtime.CategoryIngestion {
|
||||
t.Errorf("category = %q, want %q", cat, runtime.CategoryIngestion)
|
||||
}
|
||||
if factory == nil {
|
||||
t.Error("factory is nil")
|
||||
}
|
||||
if len(meta.Inputs) == 0 {
|
||||
t.Errorf("inputs metadata is empty")
|
||||
}
|
||||
if len(meta.Outputs) == 0 {
|
||||
t.Errorf("outputs metadata is empty")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenChunker_InvokeEmptyInput mirrors Python validation:
|
||||
// missing upstream shape is surfaced under _ERROR.
|
||||
func TestTokenChunker_InvokeEmptyInput(t *testing.T) {
|
||||
c, err := NewTokenChunker(nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenChunker: %v", err)
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got, want := out["output_format"], "chunks"; got != want {
|
||||
t.Errorf("output_format = %v, want %v", got, want)
|
||||
}
|
||||
if out["_ERROR"] == nil {
|
||||
t.Fatalf("_ERROR missing: %v", out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenChunker_InvokeDelimMode_BasicChunking drives the
|
||||
// delimiter-mode path with a backtick delimiter and asserts each
|
||||
// chunk carries the matched delimiter text within itself (split
|
||||
// + keep-separator contract).
|
||||
func TestTokenChunker_InvokeDelimMode_BasicChunking(t *testing.T) {
|
||||
c, err := NewTokenChunker(map[string]any{
|
||||
"delimiter_mode": "delimiter",
|
||||
"delimiters": []string{"`\\n\\n`"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenChunker: %v", err)
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc.txt",
|
||||
"output_format": "text",
|
||||
"text": "alpha\n\nbeta\n\ngamma",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
// Every emitted chunk's text should be non-empty and contain the
|
||||
// matched delimiter (we use the regex join-of-escaped literal so
|
||||
// '\n' matches the literal text).
|
||||
for i, ck := range chunks {
|
||||
text, _ := ck["text"].(string)
|
||||
if text == "" {
|
||||
t.Errorf("chunk[%d] text is empty", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenChunker_InvokeTokenSize_FallbackToMerge covers the
|
||||
// "no delimiter hit" branch — the chunker should fall back to
|
||||
// token-size merge and emit >=1 chunk.
|
||||
func TestTokenChunker_InvokeTokenSize_FallbackToMerge(t *testing.T) {
|
||||
c, err := NewTokenChunker(map[string]any{
|
||||
"delimiter_mode": "token_size",
|
||||
"chunk_token_size": 50,
|
||||
"delimiters": []string{"`\n\n`"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenChunker: %v", err)
|
||||
}
|
||||
// Input without any \n\n so the delimiter miss branch triggers
|
||||
// the token_size merge.
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc.txt",
|
||||
"output_format": "text",
|
||||
"text": "First sentence. Second sentence. Third sentence. Fourth.",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got, want := out["output_format"], "chunks"; got != want {
|
||||
t.Errorf("output_format = %v, want %v", got, want)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) < 1 {
|
||||
t.Errorf("chunks = %d, want >=1", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenChunker_InvokeOneMode_EmitsSingleChunk confirms the
|
||||
// `delimiter_mode == "one"` branch collapses the input to a single
|
||||
// chunk.
|
||||
func TestTokenChunker_InvokeOneMode_EmitsSingleChunk(t *testing.T) {
|
||||
c, err := NewTokenChunker(map[string]any{
|
||||
"delimiter_mode": "one",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenChunker: %v", err)
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc.txt",
|
||||
"output_format": "text",
|
||||
"text": "first\n\nsecond\n\nthird",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) != 1 {
|
||||
t.Errorf("chunks = %d, want 1", len(chunks))
|
||||
}
|
||||
if text, _ := chunks[0]["text"].(string); text != "first\n\nsecond\n\nthird" {
|
||||
t.Errorf("text = %q, want full input", text)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenChunker_InvokeChildrenDelim asserts that the secondary
|
||||
// children_delimiter split produces chunks carrying the parent
|
||||
// ("mom") and child ("text") keys.
|
||||
func TestTokenChunker_InvokeChildrenDelim(t *testing.T) {
|
||||
c, err := NewTokenChunker(map[string]any{
|
||||
"delimiter_mode": "delimiter",
|
||||
"delimiters": []string{"\n"},
|
||||
"children_delimiters": []string{". "},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenChunker: %v", err)
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc.txt",
|
||||
"output_format": "text",
|
||||
"text": "alpha line\nbeta line",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenChunker_InvokeJSONPayload feeds a structured JSON list
|
||||
// (mirrors upstream output_format == "json") and
|
||||
// verifies the chunker fans out into goroutines and merges
|
||||
// deterministically.
|
||||
func TestTokenChunker_InvokeJSONPayload(t *testing.T) {
|
||||
c, err := NewTokenChunker(map[string]any{
|
||||
"delimiter_mode": "delimiter",
|
||||
"delimiters": []string{"\n"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenChunker: %v", err)
|
||||
}
|
||||
items := []map[string]any{
|
||||
{"text": "Alpha text\nBeta text", "doc_type_kwd": "text"},
|
||||
{"text": "Gamma text\nDelta text", "doc_type_kwd": "text"},
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc.md",
|
||||
"output_format": "json",
|
||||
"json": items,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenChunker_InvokeDeterministic runs a 20-item structured
|
||||
// payload 10 times under the race detector and asserts the chunk
|
||||
// list is identical every time.
|
||||
func TestTokenChunker_InvokeDeterministic(t *testing.T) {
|
||||
c, err := NewTokenChunker(map[string]any{
|
||||
"delimiter_mode": "delimiter",
|
||||
"delimiters": []string{"\n"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenChunker: %v", err)
|
||||
}
|
||||
|
||||
var items []map[string]any
|
||||
for i := 0; i < 20; i++ {
|
||||
items = append(items, map[string]any{
|
||||
"text": "item",
|
||||
"doc_type_kwd": "text",
|
||||
"chunk_id": i,
|
||||
})
|
||||
}
|
||||
inputs := map[string]any{"name": "x", "output_format": "json", "json": items}
|
||||
type fingerprint struct {
|
||||
count int
|
||||
first string
|
||||
last string
|
||||
}
|
||||
var firstfp fingerprint
|
||||
for run := 0; run < 10; run++ {
|
||||
out, err := c.Invoke(context.Background(), inputs)
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke run %d: %v", run, err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
fp := fingerprint{count: len(chunks)}
|
||||
if len(chunks) > 0 {
|
||||
fp.first, _ = chunks[0]["text"].(string)
|
||||
fp.last, _ = chunks[len(chunks)-1]["text"].(string)
|
||||
}
|
||||
if run == 0 {
|
||||
firstfp = fp
|
||||
} else if fp != firstfp {
|
||||
t.Fatalf("run %d: deterministic fingerprint changed: %+v vs %+v", run, fp, firstfp)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenChunker_InputsOutputs_NonEmpty mirrors the registry-level
|
||||
// inputs/outputs keys (the registered metadata echoes Inputs /
|
||||
// Outputs on the component itself).
|
||||
func TestTokenChunker_InputsOutputs_NonEmpty(t *testing.T) {
|
||||
_, _, meta, ok := runtime.DefaultRegistry.Lookup("TokenChunker")
|
||||
if !ok {
|
||||
t.Fatal("registry miss")
|
||||
}
|
||||
if len(meta.Inputs) == 0 {
|
||||
t.Error("inputs metadata is empty")
|
||||
}
|
||||
if len(meta.Outputs) == 0 {
|
||||
t.Error("outputs metadata is empty")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenChunker_Parallelism enforces the plan's row 2.3a
|
||||
// parallelism (4).
|
||||
func TestTokenChunker_Parallelism(t *testing.T) {
|
||||
c, err := NewTokenChunker(nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenChunker: %v", err)
|
||||
}
|
||||
tc, ok := c.(*TokenChunkerComponent)
|
||||
if !ok {
|
||||
t.Fatalf("NewTokenChunker returned %T, want *TokenChunkerComponent", c)
|
||||
}
|
||||
if got := tc.Parallelism(); got != 4 {
|
||||
t.Errorf("Parallelism() = %d, want 4", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenChunker_NewRejectsBadParam enforces the param validation
|
||||
// at construction time (mirrors python `check()`).
|
||||
func TestTokenChunker_NewRejectsBadParam(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
conf map[string]any
|
||||
}{
|
||||
{"bad delimiter_mode", map[string]any{"delimiter_mode": "nope"}},
|
||||
{"zero chunk_token_size", map[string]any{"delimiter_mode": "token_size", "chunk_token_size": 0}},
|
||||
{"negative chunk_token_size", map[string]any{"delimiter_mode": "token_size", "chunk_token_size": -5}},
|
||||
{"negative overlapped_percent", map[string]any{"delimiter_mode": "token_size", "chunk_token_size": 50, "overlapped_percent": -0.1}},
|
||||
{"overlapped_percent >= 1", map[string]any{"delimiter_mode": "token_size", "chunk_token_size": 50, "overlapped_percent": 1.0}},
|
||||
{"negative table_context_size", map[string]any{"delimiter_mode": "token_size", "chunk_token_size": 50, "table_context_size": -1}},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if _, err := NewTokenChunker(tc.conf); err == nil {
|
||||
t.Fatal("expected error, got nil")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenChunker_NewAcceptsDefaults ensures the no-config
|
||||
// constructor returns a usable component with a working default
|
||||
// delimiter_mode = "token_size".
|
||||
func TestTokenChunker_NewAcceptsDefaults(t *testing.T) {
|
||||
c, err := NewTokenChunker(nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenChunker(nil): %v", err)
|
||||
}
|
||||
if got := c.(*TokenChunkerComponent).param.DelimiterMode; got != "token_size" {
|
||||
t.Errorf("default delimiter_mode = %q, want token_size", got)
|
||||
}
|
||||
}
|
||||
138
internal/ingestion/component/document_storage.go
Normal file
138
internal/ingestion/component/document_storage.go
Normal file
@@ -0,0 +1,138 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"ragflow/internal/dao"
|
||||
"ragflow/internal/entity"
|
||||
"ragflow/internal/storage"
|
||||
)
|
||||
|
||||
// DocumentStorageRef is the resolved backing storage location for a document.
|
||||
// It is exported so higher-level integration tests can inject doc_id resolution
|
||||
// without reaching into DAO state.
|
||||
type DocumentStorageRef struct {
|
||||
Name string
|
||||
Bucket string
|
||||
Path string
|
||||
}
|
||||
|
||||
// ResolveDocumentStorageOverride is the narrow test seam for doc_id-driven
|
||||
// storage resolution. Production leaves this nil and uses DAO-backed lookup.
|
||||
var ResolveDocumentStorageOverride func(docID string) (*DocumentStorageRef, error)
|
||||
|
||||
func fetchBinary(ctx context.Context, bucket, path string) ([]byte, error) {
|
||||
stg := resolveStorage()
|
||||
if stg == nil {
|
||||
return nil, fmt.Errorf("no storage backend registered")
|
||||
}
|
||||
|
||||
type result struct {
|
||||
data []byte
|
||||
err error
|
||||
}
|
||||
done := make(chan result, 1)
|
||||
go func() {
|
||||
data, err := stg.Get(bucket, path)
|
||||
done <- result{data: data, err: err}
|
||||
}()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case r := <-done:
|
||||
if r.err != nil {
|
||||
return nil, fmt.Errorf("storage.Get(%q, %q): %w", bucket, path, r.err)
|
||||
}
|
||||
return r.data, nil
|
||||
}
|
||||
}
|
||||
|
||||
func resolveStorage() storage.Storage {
|
||||
if SetStorageFactoryOverride != nil {
|
||||
if s := SetStorageFactoryOverride(); s != nil {
|
||||
return s
|
||||
}
|
||||
}
|
||||
return storage.GetStorageFactory().GetStorage()
|
||||
}
|
||||
|
||||
func resolveDocumentStorage(docID string) (*DocumentStorageRef, error) {
|
||||
if ResolveDocumentStorageOverride != nil {
|
||||
return ResolveDocumentStorageOverride(docID)
|
||||
}
|
||||
|
||||
doc, err := dao.NewDocumentDAO().GetByID(docID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ref := &DocumentStorageRef{Name: documentNameOrID(doc)}
|
||||
|
||||
mappings, err := dao.NewFile2DocumentDAO().GetByDocumentID(doc.ID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(mappings) > 0 && mappings[0].FileID != nil && *mappings[0].FileID != "" {
|
||||
file, err := dao.NewFileDAO().GetByID(*mappings[0].FileID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if file.SourceType == "" || entity.FileSource(file.SourceType) == entity.FileSourceLocal {
|
||||
if file.Location == nil || *file.Location == "" {
|
||||
return nil, fmt.Errorf("file location is empty")
|
||||
}
|
||||
ref.Bucket = file.ParentID
|
||||
ref.Path = *file.Location
|
||||
return ref, nil
|
||||
}
|
||||
}
|
||||
if doc.Location == nil || *doc.Location == "" {
|
||||
return nil, fmt.Errorf("document location is empty")
|
||||
}
|
||||
ref.Bucket = doc.KbID
|
||||
ref.Path = *doc.Location
|
||||
return ref, nil
|
||||
}
|
||||
|
||||
func resolveDocumentName(docID string) (string, error) {
|
||||
if ResolveDocumentStorageOverride != nil {
|
||||
ref, err := ResolveDocumentStorageOverride(docID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if ref != nil && ref.Name != "" {
|
||||
return ref.Name, nil
|
||||
}
|
||||
}
|
||||
doc, err := dao.NewDocumentDAO().GetByID(docID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return documentNameOrID(doc), nil
|
||||
}
|
||||
|
||||
func documentNameOrID(doc *entity.Document) string {
|
||||
if doc != nil && doc.Name != nil && *doc.Name != "" {
|
||||
return *doc.Name
|
||||
}
|
||||
if doc != nil {
|
||||
return doc.ID
|
||||
}
|
||||
return ""
|
||||
}
|
||||
718
internal/ingestion/component/extractor.go
Normal file
718
internal/ingestion/component/extractor.go
Normal file
@@ -0,0 +1,718 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Package component — Extractor component (Phase 2.5 of
|
||||
// port-rag-flow-pipeline-to-go.md §4 row 2.5).
|
||||
//
|
||||
// SCOPE (honest):
|
||||
//
|
||||
// - PROVIDER-AGNOSTIC (§8 Q1): the Extractor does NOT depend on any
|
||||
// specific LLM provider. It dispatches every chat call through
|
||||
// internal/entity/models — the same factory routes 48 of the 56
|
||||
// Python ChatModel providers registered there (factory.go switch,
|
||||
// lines 36-156). The 8 providers NOT yet in the Go switch (LeptonAI,
|
||||
// Gemini LiteLLM path, PerfXCloud, 01.AI / Lingyi, DeerAPI,
|
||||
// Astraflow-CN, RAGcon, New API) ARE unreachable from this
|
||||
// component — an llm_id resolving to one of those falls through to
|
||||
// NewDummyModel and the chat call returns a deterministic "dummy"
|
||||
// response. We DO NOT panic: errors are surfaced as a clean
|
||||
// "no driver for %q" wrap that callers can log and route.
|
||||
//
|
||||
// - LLM CALL SHAPE: one chat call per chunk (no batching). Plan
|
||||
// §AD-5a locks Parallelism at 1 because "LLM call is inherently
|
||||
// serial"; sequential per-chunk processing keeps test ordering
|
||||
// deterministic under -race.
|
||||
//
|
||||
// - TIMEOUT / ELAPSED: the call is wrapped in
|
||||
// runtime.WithTimeout(60s) and runtime.TrackElapsed so the
|
||||
// upstream pipeline gets _created_time / _elapsed_time stamps
|
||||
// matching the python ProcessBase contract (base.py:42, 58).
|
||||
//
|
||||
// - JSON PARSING: the prompt asks the LLM to return a JSON object;
|
||||
// we best-effort parse the response into map[string]any. A
|
||||
// non-JSON response is NOT a hard error — it's surfaced as the
|
||||
// raw string under the same field name so downstream callers
|
||||
// can decide what to do.
|
||||
//
|
||||
// - WHAT IS NOT YET PORTED: the python _build_TOC branch
|
||||
// (rag/flow/extractor/extractor.py:40-72) requires the TOC
|
||||
// generator (rag.prompts.generator.run_toc_from_text). That
|
||||
// service has no Go counterpart yet; the current Extractor
|
||||
// short-circuits with a clear error when field_name == "toc"
|
||||
// so a future Phase 2.5+ task can fill the gap without a
|
||||
// silent regression.
|
||||
//
|
||||
// - SINGLE-CHUNK FAST PATH: when no chunk list is wired in,
|
||||
// the LLM is called once with the resolved args directly (no
|
||||
// chunk substitution). Matches python _invoke path
|
||||
// (line 108: msg, sys_prompt = self._sys_prompt_and_msg([], args)).
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
eschema "github.com/cloudwego/eino/schema"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/entity/models"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
)
|
||||
|
||||
const componentNameExtractor = "Extractor"
|
||||
|
||||
// extractorTimeout bounds one LLM chat call. Matches the python
|
||||
// `@timeout(60)` default at rag/flow/base.py:60. The pipeline
|
||||
// orchestrator (Phase 3) overrides this if a stage-level ceiling
|
||||
// is configured.
|
||||
const extractorTimeout = 60 * time.Second
|
||||
|
||||
// ExtractorComponent performs LLM-based extraction over a chunk
|
||||
// list (or a single empty call when no chunks are wired in).
|
||||
//
|
||||
// The instance is safe for concurrent invocation: each Invoke
|
||||
// reads Param read-only (Param is set at construction; per-call
|
||||
// overrides flow through the inputs map). The single mutable
|
||||
// package-level seam (extractorChatInvoker) is guarded by a
|
||||
// RWMutex; tests swap it via SetExtractorChatInvoker.
|
||||
type ExtractorComponent struct {
|
||||
Param schema.ExtractorParam
|
||||
}
|
||||
|
||||
// NewExtractorComponent constructs an Extractor from a DSL param
|
||||
// map. Missing keys fall back to schema.ExtractorParam.Defaults();
|
||||
// an empty FieldName is rejected (matches python
|
||||
// `check_empty(self.field_name, "Result Destination")`).
|
||||
//
|
||||
// Param map shape (all keys optional; missing → Defaults()):
|
||||
//
|
||||
// {
|
||||
// "field_name": string, — required; key the extraction lands under
|
||||
// "llm_id": string, — optional; resolves via models.NewModelFactory
|
||||
// "system_prompt": string, — optional override
|
||||
// "prompt": string, — optional user prompt
|
||||
// }
|
||||
//
|
||||
// errors here surface as canvas compile failures so a malformed
|
||||
// param is caught at build time rather than mid-run.
|
||||
func NewExtractorComponent(params map[string]any) (runtime.Component, error) {
|
||||
p := schema.ExtractorParam{}.Defaults()
|
||||
if params != nil {
|
||||
if v, ok := params["field_name"].(string); ok {
|
||||
p.FieldName = v
|
||||
}
|
||||
if v, ok := params["llm_id"].(string); ok {
|
||||
p.LLMID = v
|
||||
}
|
||||
if v, ok := params["system_prompt"].(string); ok {
|
||||
p.SystemPrompt = v
|
||||
}
|
||||
if v, ok := params["prompt"].(string); ok {
|
||||
p.Prompt = v
|
||||
}
|
||||
}
|
||||
if err := p.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("extractor: param check: %w", err)
|
||||
}
|
||||
return &ExtractorComponent{Param: p}, nil
|
||||
}
|
||||
|
||||
// Inputs returns the parameter metadata. Matches the python
|
||||
// Extractor._invoke kwargs plus the optional per-call llm_id
|
||||
// override (python: args["llm_id"] path is implicit via
|
||||
// self.chat_mdl; the Go port exposes it explicitly).
|
||||
func (c *ExtractorComponent) Inputs() map[string]string {
|
||||
return map[string]string{
|
||||
"chunks": "List of map[string]any from upstream Tokenizer. Each entry must carry a string 'text' (or 'content_with_weight') field. Optional — when absent the LLM is called once with the resolved args.",
|
||||
"prompt": "Optional user prompt template. Falls back to Param.Prompt when absent.",
|
||||
"llm_id": "Optional per-call LLM id override. Falls back to Param.LLMID when absent.",
|
||||
"system_prompt": "Optional per-call system prompt override. Falls back to Param.SystemPrompt.",
|
||||
}
|
||||
}
|
||||
|
||||
// Outputs returns the public surface downstream ingestion
|
||||
// consumers can wire into. Mirrors schema.ExtractorOutputs.
|
||||
//
|
||||
// chunks []map[string]any — input chunks, each augmented
|
||||
// with field_name=<LLM result>.
|
||||
// When the input chunks list is
|
||||
// absent, the slice contains a
|
||||
// single map with the same shape.
|
||||
// output_format string — always "chunks". Parity with
|
||||
// python set_output contract.
|
||||
// _ERROR string — populated on a short-circuit
|
||||
// error (matches python
|
||||
// set_output("_ERROR", ...)).
|
||||
func (c *ExtractorComponent) Outputs() map[string]string {
|
||||
return map[string]string{
|
||||
"chunks": "Extraction results — input chunks (or a single-element slice when no chunks were supplied), each enriched with field_name=<LLM response>.",
|
||||
"output_format": "Always \"chunks\". Parity marker for downstream consumers.",
|
||||
"_ERROR": "Optional short-circuit error message (reserved for the future TOC branch and other error paths).",
|
||||
}
|
||||
}
|
||||
|
||||
// Parallelism is locked at 1 (plan §AD-5a: "Extractor: 1 (LLM call
|
||||
// is inherently serial)"). The pipeline runner uses this to decide
|
||||
// fan-out degree; sequential per-chunk processing keeps test
|
||||
// ordering deterministic under -race.
|
||||
func (c *ExtractorComponent) Parallelism() int { return 1 }
|
||||
|
||||
// extractorChatInvoker is the seam the Extractor uses to dispatch
|
||||
// its chat call. The production implementation
|
||||
// (einoExtractorChatInvoker below) mirrors
|
||||
// internal/agent/component/llm.go:einoChatInvoker — same factory,
|
||||
// same driver resolution, but kept self-contained so the
|
||||
// ingestion package does NOT pull in agent/component for a
|
||||
// one-method interface.
|
||||
//
|
||||
// Tests swap the package-level defaultExtractorChatInvoker to inject a
|
||||
// canned-response stub (see SetExtractorChatInvoker and the test
|
||||
// helpers in extractor_test.go). This is the testability seam the
|
||||
// Phase 2.5 spec calls out as a hard rule.
|
||||
type extractorChatInvoker interface {
|
||||
Chat(ctx context.Context, req extractorChatRequest) (*extractorChatResponse, error)
|
||||
}
|
||||
|
||||
// extractorChatRequest is the minimal surface the Extractor needs
|
||||
// to dispatch a chat call. Driver is the provider key
|
||||
// (e.g. "openai"); ModelName is the model id alone or composite
|
||||
// "model@provider". APIKey / BaseURL are passed through so the
|
||||
// driver can authenticate without re-reading the tenant config.
|
||||
type extractorChatRequest struct {
|
||||
Driver string
|
||||
ModelName string
|
||||
APIKey string
|
||||
BaseURL string
|
||||
Messages []eschema.Message
|
||||
}
|
||||
|
||||
// extractorChatResponse holds the LLM's text answer. Token /
|
||||
// stopped flags are not consumed by the Extractor yet, so they
|
||||
// remain optional / 0-valued.
|
||||
type extractorChatResponse struct {
|
||||
Content string
|
||||
}
|
||||
|
||||
// extractorChatInvokerMu guards defaultExtractorChatInvoker swaps.
|
||||
var extractorChatInvokerMu sync.RWMutex
|
||||
|
||||
// defaultExtractorChatInvoker is the package-level seam. Production
|
||||
// uses einoExtractorChatInvoker; tests inject a stub.
|
||||
var defaultExtractorChatInvoker extractorChatInvoker = &einoExtractorChatInvoker{}
|
||||
|
||||
var extractorChatTargetResolverMu sync.RWMutex
|
||||
|
||||
// extractorChatTargetResolverOverride is a narrow test seam for
|
||||
// integration tests that need to supply real credentials without
|
||||
// teaching the production Extractor a tenant-credential lookup path.
|
||||
// When set, resolveExtractorChatTarget consults it first.
|
||||
var extractorChatTargetResolverOverride func(llmID string) (driver, modelName, apiKey, baseURL string, ok bool)
|
||||
|
||||
// SetExtractorChatInvoker swaps the package-level chat invoker
|
||||
// for tests. Pass nil to restore the default. Concurrent-safe.
|
||||
func SetExtractorChatInvoker(inv extractorChatInvoker) {
|
||||
extractorChatInvokerMu.Lock()
|
||||
defer extractorChatInvokerMu.Unlock()
|
||||
defaultExtractorChatInvoker = inv
|
||||
}
|
||||
|
||||
// SetExtractorChatTargetResolverOverride swaps the package-level
|
||||
// llm_id target resolver override for tests. Pass nil to restore
|
||||
// the default split-only resolver. Concurrent-safe.
|
||||
func SetExtractorChatTargetResolverOverride(fn func(llmID string) (driver, modelName, apiKey, baseURL string, ok bool)) {
|
||||
extractorChatTargetResolverMu.Lock()
|
||||
defer extractorChatTargetResolverMu.Unlock()
|
||||
extractorChatTargetResolverOverride = fn
|
||||
}
|
||||
|
||||
func getExtractorChatTargetResolverOverride() func(llmID string) (driver, modelName, apiKey, baseURL string, ok bool) {
|
||||
extractorChatTargetResolverMu.RLock()
|
||||
defer extractorChatTargetResolverMu.RUnlock()
|
||||
return extractorChatTargetResolverOverride
|
||||
}
|
||||
|
||||
// getExtractorChatInvoker returns the current default invoker.
|
||||
func getExtractorChatInvoker() extractorChatInvoker {
|
||||
extractorChatInvokerMu.RLock()
|
||||
defer extractorChatInvokerMu.RUnlock()
|
||||
if defaultExtractorChatInvoker == nil {
|
||||
return &einoExtractorChatInvoker{}
|
||||
}
|
||||
return defaultExtractorChatInvoker
|
||||
}
|
||||
|
||||
// einoExtractorChatInvoker is the production seam. It dispatches
|
||||
// through the entity/models factory (which knows 48 of 56
|
||||
// providers) and returns the assistant text via
|
||||
// models.EinoChatModel.Generate. An unknown provider falls
|
||||
// through to NewDummyModel in the factory's default branch — we
|
||||
// surface that as a typed "no driver for %q" wrap so callers can
|
||||
// decide whether to retry, route around, or log.
|
||||
type einoExtractorChatInvoker struct{}
|
||||
|
||||
// Chat implements extractorChatInvoker for the production path.
|
||||
func (e *einoExtractorChatInvoker) Chat(ctx context.Context, req extractorChatRequest) (*extractorChatResponse, error) {
|
||||
if req.ModelName == "" {
|
||||
return nil, fmt.Errorf("extractor: chat: model_name is required")
|
||||
}
|
||||
driver := strings.ToLower(strings.TrimSpace(req.Driver))
|
||||
modelName := req.ModelName
|
||||
if driver == "" && modelName != "" {
|
||||
if bare, provider, ok := splitExtractorLLID(modelName); ok {
|
||||
driver = provider
|
||||
modelName = bare
|
||||
}
|
||||
}
|
||||
if driver == "" {
|
||||
driver = "dummy"
|
||||
}
|
||||
var baseURL map[string]string
|
||||
if req.BaseURL != "" {
|
||||
baseURL = map[string]string{"default": req.BaseURL}
|
||||
}
|
||||
urlSuffix := extractorChatURLSuffixFor(driver)
|
||||
d, err := models.NewModelFactory().CreateModelDriver(driver, baseURL, urlSuffix)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("extractor: resolve driver %q: %w", driver, err)
|
||||
}
|
||||
if d == nil {
|
||||
return nil, fmt.Errorf("extractor: no driver for %q", driver)
|
||||
}
|
||||
apiKey := req.APIKey
|
||||
cfg := &models.APIConfig{ApiKey: &apiKey}
|
||||
cm := models.NewChatModel(d, &modelName, cfg)
|
||||
wrapper := models.NewEinoChatModel(cm, nil)
|
||||
// Honour ctx cancel up front so the caller's WithTimeout(...)
|
||||
// is observed even when the driver layer doesn't take a ctx.
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out, err := wrapper.Generate(ctx, toExtractorEinoMessages(req.Messages))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &extractorChatResponse{Content: out.Content}, nil
|
||||
}
|
||||
|
||||
// splitExtractorLLID parses a composite llm_id "model@provider"
|
||||
// mirroring agent/component/llm_credentials.go:parseLLMIDParts
|
||||
// (the canonical composite form throughout the codebase). Returns
|
||||
// ok=false when no "@" is present or the id is malformed.
|
||||
//
|
||||
// "gpt-4o-mini@openai" -> ("gpt-4o-mini", "openai", true)
|
||||
// "gpt-4o-mini" -> ("gpt-4o-mini", "", false)
|
||||
//
|
||||
// Kept local so the ingestion package doesn't import
|
||||
// agent/component.
|
||||
func splitExtractorLLID(s string) (modelName, provider string, ok bool) {
|
||||
parts := strings.Split(strings.TrimSpace(s), "@")
|
||||
switch len(parts) {
|
||||
case 2:
|
||||
return parts[0], parts[1], true
|
||||
default:
|
||||
return s, "", false
|
||||
}
|
||||
}
|
||||
|
||||
// extractorChatURLSuffixFor matches
|
||||
// internal/agent/component/llm.go:chatURLSuffixFor — anthropic
|
||||
// uses v1/messages, everything else falls through to the openai-
|
||||
// compatible chat/completions default.
|
||||
func extractorChatURLSuffixFor(driver string) models.URLSuffix {
|
||||
switch strings.ToLower(driver) {
|
||||
case "anthropic":
|
||||
return models.URLSuffix{Chat: "v1/messages"}
|
||||
default:
|
||||
return models.URLSuffix{Chat: "chat/completions"}
|
||||
}
|
||||
}
|
||||
|
||||
// toExtractorEinoMessages converts eschema.Message → *eschema.Message
|
||||
// for the eino bridge. The user / system / assistant roles pass
|
||||
// through; multi-modal content is intentionally not propagated —
|
||||
// extraction prompts are text-only today.
|
||||
func toExtractorEinoMessages(msgs []eschema.Message) []*eschema.Message {
|
||||
out := make([]*eschema.Message, 0, len(msgs))
|
||||
for i := range msgs {
|
||||
m := msgs[i]
|
||||
role := m.Role
|
||||
if role == "" {
|
||||
role = eschema.User
|
||||
}
|
||||
out = append(out, &eschema.Message{
|
||||
Role: role,
|
||||
Content: m.Content,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// extractorInputs is the post-Validation view of the upstream
|
||||
// input map. Computed once at the top of Invoke so the rest of
|
||||
// the function reads as straight-line code.
|
||||
type extractorInputs struct {
|
||||
fieldName string
|
||||
llmID string
|
||||
systemPrompt string
|
||||
prompt string
|
||||
chunks []map[string]any
|
||||
}
|
||||
|
||||
// resolveInputs overlays per-call inputs on top of the
|
||||
// component's static Param. Missing keys fall back to the
|
||||
// Param-level values; per-call values win on conflict (so a
|
||||
// canvas can override LLM_ID at runtime). The python
|
||||
// Extractor reads inputs directly from get_input_elements(); the
|
||||
// Go port normalizes to extractorInputs once at the top so the
|
||||
// rest of Invoke reads straight-line.
|
||||
func (c *ExtractorComponent) resolveInputs(inputs map[string]any) extractorInputs {
|
||||
out := extractorInputs{
|
||||
fieldName: c.Param.FieldName,
|
||||
llmID: c.Param.LLMID,
|
||||
systemPrompt: c.Param.SystemPrompt,
|
||||
prompt: c.Param.Prompt,
|
||||
}
|
||||
if inputs == nil {
|
||||
return out
|
||||
}
|
||||
if v, ok := inputs["llm_id"].(string); ok && v != "" {
|
||||
out.llmID = v
|
||||
}
|
||||
if v, ok := inputs["prompt"].(string); ok && v != "" {
|
||||
out.prompt = v
|
||||
}
|
||||
if v, ok := inputs["system_prompt"].(string); ok && v != "" {
|
||||
out.systemPrompt = v
|
||||
}
|
||||
for _, key := range extractorChunkInputOrder(inputs) {
|
||||
if chunks, ok := extractorChunkList(inputs[key]); ok {
|
||||
out.chunks = chunks
|
||||
break
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func extractorChunkInputOrder(inputs map[string]any) []string {
|
||||
order := make([]string, 0, len(inputs))
|
||||
for _, preferred := range []string{"chunks", "json"} {
|
||||
if _, ok := inputs[preferred]; ok {
|
||||
order = append(order, preferred)
|
||||
}
|
||||
}
|
||||
var extra []string
|
||||
for key := range inputs {
|
||||
if key == "chunks" || key == "json" {
|
||||
continue
|
||||
}
|
||||
extra = append(extra, key)
|
||||
}
|
||||
sort.Strings(extra)
|
||||
order = append(order, extra...)
|
||||
return order
|
||||
}
|
||||
|
||||
func extractorChunkList(v any) ([]map[string]any, bool) {
|
||||
switch list := v.(type) {
|
||||
case []map[string]any:
|
||||
return list, true
|
||||
case []any:
|
||||
out := make([]map[string]any, 0, len(list))
|
||||
for _, item := range list {
|
||||
m, ok := item.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
return out, true
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
|
||||
// Invoke performs LLM-based extraction. Inputs:
|
||||
//
|
||||
// chunks (optional, []map[string]any) — upstream chunks; each must
|
||||
// carry a string "text".
|
||||
// prompt (optional, string) — overrides Param.Prompt.
|
||||
// system_prompt (optional, string) — overrides Param.SystemPrompt.
|
||||
// llm_id (optional, string) — overrides Param.LLMID.
|
||||
//
|
||||
// Outputs:
|
||||
//
|
||||
// chunks ([]map[string]any) — input chunks augmented with
|
||||
// field_name=<LLM result>. When
|
||||
// the input list is empty, the
|
||||
// slice contains a single map.
|
||||
// output_format (string) — always "chunks".
|
||||
// _ERROR (string, reserved) — populated when the component
|
||||
// short-circuits with an error.
|
||||
// _created_time, _elapsed_time — TrackElapsed bookkeeping.
|
||||
func (c *ExtractorComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
if err := c.Param.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("extractor: %w", err)
|
||||
}
|
||||
in := c.resolveInputs(inputs)
|
||||
if in.fieldName == "toc" {
|
||||
// TODO(parity-gap): _build_TOC is not ported yet — surface
|
||||
// a clear error rather than silently emitting empty chunks.
|
||||
return nil, fmt.Errorf("extractor: field_name %q requires the TOC prompt generator which is not yet ported to Go", "toc")
|
||||
}
|
||||
|
||||
tracked, err := runtime.TrackElapsed("Extractor", func() (map[string]any, error) {
|
||||
cb := runtime.ProgressCallback(nil)
|
||||
progressErr := runtime.TrackProgress("Extractor", cb, func() error {
|
||||
return runtime.WithTimeout(ctx, extractorTimeout, func(timeoutCtx context.Context) error {
|
||||
if len(in.chunks) == 0 {
|
||||
// Fast path (python _invoke line 108): one
|
||||
// call with the resolved args directly.
|
||||
ans, callErr := c.call(timeoutCtx, in, "")
|
||||
if callErr != nil {
|
||||
return callErr
|
||||
}
|
||||
in.chunks = []map[string]any{{in.fieldName: ans}}
|
||||
return nil
|
||||
}
|
||||
for i, ck := range in.chunks {
|
||||
text, _ := ck["text"].(string)
|
||||
ans, callErr := c.call(timeoutCtx, in, text)
|
||||
if callErr != nil {
|
||||
return fmt.Errorf("chunk %d: %w", i, callErr)
|
||||
}
|
||||
ck[in.fieldName] = ans
|
||||
in.chunks[i] = ck
|
||||
}
|
||||
return nil
|
||||
})
|
||||
})
|
||||
if progressErr != nil {
|
||||
return nil, progressErr
|
||||
}
|
||||
return map[string]any{
|
||||
"chunks": in.chunks,
|
||||
"output_format": "chunks",
|
||||
}, nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("extractor: %w", err)
|
||||
}
|
||||
return tracked, nil
|
||||
}
|
||||
|
||||
// call dispatches one LLM chat call for the supplied chunk text
|
||||
// (empty string in the no-chunk fast path). The result is the
|
||||
// raw string from the model — JSON parsing happens here so
|
||||
// callers can rely on a structured value downstream.
|
||||
func (c *ExtractorComponent) call(ctx context.Context, in extractorInputs, chunkText string) (any, error) {
|
||||
driver, modelName, apiKey, baseURL := resolveExtractorChatTarget(in.llmID)
|
||||
msgs := buildExtractorMessages(in.systemPrompt, in.prompt, chunkText, in.chunks)
|
||||
inv := getExtractorChatInvoker()
|
||||
resp, err := inv.Chat(ctx, extractorChatRequest{
|
||||
Driver: driver,
|
||||
ModelName: modelName,
|
||||
APIKey: apiKey,
|
||||
BaseURL: baseURL,
|
||||
Messages: msgs,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
raw := strings.TrimSpace(resp.Content)
|
||||
if raw == "" {
|
||||
// No response — emit empty string so downstream code
|
||||
// can distinguish from "LLM errored" via the error
|
||||
// path above.
|
||||
return "", nil
|
||||
}
|
||||
// Best-effort JSON parse: a JSON object response is the
|
||||
// canonical structured-extraction shape. Other shapes are
|
||||
// returned verbatim so the caller can decide.
|
||||
if parsed, ok := tryParseJSONObject(raw); ok {
|
||||
return parsed, nil
|
||||
}
|
||||
return raw, nil
|
||||
}
|
||||
|
||||
// resolveExtractorChatTarget splits a composite llm_id
|
||||
// "model@provider" or "openai/model@provider" into driver /
|
||||
// model / api_key / base_url. Today the Extractor has no tenant-
|
||||
// scoped credential lookup — credentials are read from the
|
||||
// per-call inputs map only. Future iterations can fill that gap
|
||||
// with the same pattern internal/agent/component/llm_credentials.go
|
||||
// uses (resolveTenantLLMConfig). For Phase 2.5 the test seam
|
||||
// (SetExtractorChatInvoker) carries the wire-level signals.
|
||||
func resolveExtractorChatTarget(llmID string) (driver, modelName, apiKey, baseURL string) {
|
||||
if override := getExtractorChatTargetResolverOverride(); override != nil {
|
||||
if driver, modelName, apiKey, baseURL, ok := override(llmID); ok {
|
||||
return driver, modelName, apiKey, baseURL
|
||||
}
|
||||
}
|
||||
if llmID == "" {
|
||||
return "", "", "", ""
|
||||
}
|
||||
modelName = llmID
|
||||
if bare, provider, ok := splitExtractorLLID(llmID); ok {
|
||||
modelName = bare
|
||||
driver = strings.ToLower(provider)
|
||||
}
|
||||
return driver, modelName, "", ""
|
||||
}
|
||||
|
||||
// buildExtractorMessages assembles system + user messages for
|
||||
// one extraction call. The user prompt is rendered as
|
||||
// "<prompt>\n\n<chunkText>" so the python behavior of
|
||||
// substituting the chunk text into the args dict is preserved
|
||||
// without invoking a template engine.
|
||||
//
|
||||
// Prompt placeholders of the form `{ComponentName:ParamName@chunks}`
|
||||
// are substituted with the joined text of all upstream chunks
|
||||
// when chunks is non-empty. The python rag/flow/extractor/extractor.py
|
||||
// build_existing_prompt path performs the same substitution at
|
||||
// runtime; the Go port surfaces it as a regex on the prompt
|
||||
// template so the resume template's `{TitleChunker:FlatMiceFix@chunks}`
|
||||
// reference resolves without invoking a template engine.
|
||||
//
|
||||
// Substitution is opt-in: when chunks is nil/empty the placeholder
|
||||
// is left intact so a misconfigured template surfaces as a
|
||||
// clear pattern rather than silently disappearing.
|
||||
func buildExtractorMessages(system, prompt, chunkText string, chunks []map[string]any) []eschema.Message {
|
||||
out := make([]eschema.Message, 0, 2)
|
||||
if system != "" {
|
||||
out = append(out, eschema.Message{Role: eschema.System, Content: system})
|
||||
}
|
||||
user := prompt
|
||||
if chunkText != "" {
|
||||
if user != "" {
|
||||
user += "\n\n"
|
||||
}
|
||||
user += chunkText
|
||||
}
|
||||
if user == "" {
|
||||
// An empty prompt + empty chunk is a degenerate call.
|
||||
// The LLM driver returns an error; we surface that
|
||||
// unchanged.
|
||||
user = " "
|
||||
}
|
||||
user = substitutePromptPlaceholders(user, chunks)
|
||||
out = append(out, eschema.Message{Role: eschema.User, Content: user})
|
||||
return out
|
||||
}
|
||||
|
||||
// substitutePromptPlaceholders replaces `{ComponentName:ParamName@chunks}`
|
||||
// patterns in the user prompt with the joined text of all upstream
|
||||
// chunks. The python rag/flow/extractor/extractor.py:build_existing_prompt
|
||||
// path performs the same substitution at runtime using a Jinja
|
||||
// template; the Go port keeps the regex form because the LLM
|
||||
// driver does not require Jinja and the surface is small enough to
|
||||
// avoid pulling in a template engine.
|
||||
//
|
||||
// Pattern grammar:
|
||||
//
|
||||
// {CmpName:ParamName@chunks}
|
||||
//
|
||||
// The CmpName and ParamName are both matched but ignored — the
|
||||
// substitute is always "the joined chunk text" today, because the
|
||||
// only @chunks reference in production templates is the resume
|
||||
// template's `{TitleChunker:FlatMiceFix@chunks}` pattern. The
|
||||
// CmpName/ParamName parsing exists so a future per-component
|
||||
// substitution can extend the function without breaking the
|
||||
// existing call sites.
|
||||
func substitutePromptPlaceholders(prompt string, chunks []map[string]any) string {
|
||||
if prompt == "" || len(chunks) == 0 {
|
||||
return prompt
|
||||
}
|
||||
// Build the substitution payload once. Each chunk's text is
|
||||
// joined with a blank line so a downstream LLM sees clear
|
||||
// chunk boundaries.
|
||||
var b strings.Builder
|
||||
for i, ck := range chunks {
|
||||
t, _ := ck["text"].(string)
|
||||
if t == "" {
|
||||
continue
|
||||
}
|
||||
if i > 0 {
|
||||
b.WriteString("\n\n")
|
||||
}
|
||||
b.WriteString(t)
|
||||
}
|
||||
repl := b.String()
|
||||
if repl == "" {
|
||||
return prompt
|
||||
}
|
||||
return placeholderRE.ReplaceAllString(prompt, repl)
|
||||
}
|
||||
|
||||
// placeholderRE matches `{CmpName:ParamName@chunks}` patterns in
|
||||
// Extractor user prompts. The CMP / Param groups are ignored for
|
||||
// the @chunks variant but kept so the regex rejects arbitrary
|
||||
// placeholders (a future per-component substitution extends here).
|
||||
var placeholderRE = regexp.MustCompile(`\{[A-Za-z0-9_]+:[A-Za-z0-9_]+@chunks\}`)
|
||||
|
||||
// tryParseJSONObject tries to parse s as a JSON object. Returns
|
||||
// (parsed, true) on success; (nil, false) on parse error or when
|
||||
// s is not a JSON object. Trims common markdown code fences
|
||||
// (```json ... ```) before parsing.
|
||||
func tryParseJSONObject(s string) (map[string]any, bool) {
|
||||
trimmed := strings.TrimSpace(s)
|
||||
// Strip a single ``` fence pair if present.
|
||||
if strings.HasPrefix(trimmed, "```") {
|
||||
if idx := strings.Index(trimmed, "\n"); idx >= 0 {
|
||||
trimmed = trimmed[idx+1:]
|
||||
}
|
||||
if strings.HasSuffix(trimmed, "```") {
|
||||
trimmed = trimmed[:len(trimmed)-3]
|
||||
}
|
||||
trimmed = strings.TrimSpace(trimmed)
|
||||
}
|
||||
var out map[string]any
|
||||
if err := json.Unmarshal([]byte(trimmed), &out); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
if out == nil {
|
||||
return nil, false
|
||||
}
|
||||
// An empty object carries no information the caller can act on;
|
||||
// surface as "could not extract" so downstream code can route
|
||||
// it to the same fallback it would use for malformed text.
|
||||
if len(out) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
return out, true
|
||||
}
|
||||
|
||||
// init registers Extractor under CategoryIngestion (per plan §4
|
||||
// Phase 2.5). Metadata is derived from the Inputs()/Outputs()
|
||||
// methods on ExtractorComponent so the API layer (Phase 4) can
|
||||
// enumerate the catalog without instantiating the component.
|
||||
func init() {
|
||||
c := &ExtractorComponent{}
|
||||
runtime.MustRegister(componentNameExtractor, runtime.CategoryIngestion,
|
||||
func(_ string, params map[string]any) (runtime.Component, error) {
|
||||
return NewExtractorComponent(params)
|
||||
},
|
||||
runtime.Metadata{
|
||||
Version: "1.0.0",
|
||||
Inputs: c.Inputs(),
|
||||
Outputs: c.Outputs(),
|
||||
})
|
||||
}
|
||||
655
internal/ingestion/component/extractor_test.go
Normal file
655
internal/ingestion/component/extractor_test.go
Normal file
@@ -0,0 +1,655 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
eschema "github.com/cloudwego/eino/schema"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
)
|
||||
|
||||
// stubExtractorChatInvoker is the test seam for the package-level
|
||||
// extractorChatInvoker. It records every call (for assertions) and
|
||||
// returns canned responses configured per-test. Concurrent-safe so
|
||||
// it can backstop future Parallelism>1 cases without rewriting.
|
||||
type stubExtractorChatInvoker struct {
|
||||
mu sync.Mutex
|
||||
|
||||
// responses is consumed in order; remaining entries are returned
|
||||
// as the wrap-error. tests set entries == call count they expect.
|
||||
responses []stubResponse
|
||||
|
||||
// lastReq records the most recent call's request for inspection
|
||||
// (e.g. driver / model name resolved from the llm_id).
|
||||
lastReq extractorChatRequest
|
||||
calls atomic.Int32
|
||||
}
|
||||
|
||||
// stubResponse couples a Content value and an Err. tests populate
|
||||
// either field — Err takes precedence over Content when non-nil.
|
||||
type stubResponse struct {
|
||||
Content string
|
||||
Err error
|
||||
}
|
||||
|
||||
func (s *stubExtractorChatInvoker) Chat(_ context.Context, req extractorChatRequest) (*extractorChatResponse, error) {
|
||||
s.calls.Add(1)
|
||||
s.mu.Lock()
|
||||
s.lastReq = req
|
||||
var resp stubResponse
|
||||
if len(s.responses) > 0 {
|
||||
resp = s.responses[0]
|
||||
s.responses = s.responses[1:]
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if resp.Err != nil {
|
||||
return nil, resp.Err
|
||||
}
|
||||
return &extractorChatResponse{Content: resp.Content}, nil
|
||||
}
|
||||
|
||||
func (s *stubExtractorChatInvoker) Calls() int32 { return s.calls.Load() }
|
||||
|
||||
// withStubChatInvoker installs a stub invoker for the duration of
|
||||
// the test and restores the production invoker on cleanup.
|
||||
func withStubChatInvoker(t *testing.T, responses ...stubResponse) *stubExtractorChatInvoker {
|
||||
t.Helper()
|
||||
prev := defaultExtractorChatInvoker
|
||||
stub := &stubExtractorChatInvoker{responses: responses}
|
||||
SetExtractorChatInvoker(stub)
|
||||
t.Cleanup(func() {
|
||||
SetExtractorChatInvoker(prev)
|
||||
})
|
||||
return stub
|
||||
}
|
||||
|
||||
// TestExtractorComponent_Registered verifies the init() registration
|
||||
// is visible to the runtime registry (Phase 4 / API layer
|
||||
// depends on this).
|
||||
func TestExtractorComponent_Registered(t *testing.T) {
|
||||
factory, cat, md, ok := runtime.DefaultRegistry.Lookup("Extractor")
|
||||
if !ok {
|
||||
t.Fatal("Extractor not registered in runtime.DefaultRegistry")
|
||||
}
|
||||
if cat != runtime.CategoryIngestion {
|
||||
t.Errorf("category = %q, want %q", cat, runtime.CategoryIngestion)
|
||||
}
|
||||
if factory == nil {
|
||||
t.Error("factory is nil")
|
||||
}
|
||||
if md.Inputs == nil || len(md.Inputs) == 0 {
|
||||
t.Errorf("metadata.Inputs empty: %v", md.Inputs)
|
||||
}
|
||||
if md.Outputs == nil || len(md.Outputs) == 0 {
|
||||
t.Errorf("metadata.Outputs empty: %v", md.Outputs)
|
||||
}
|
||||
if _, has := md.Outputs["chunks"]; !has {
|
||||
t.Errorf("metadata.Outputs missing %q", "chunks")
|
||||
}
|
||||
if _, has := md.Outputs["output_format"]; !has {
|
||||
t.Errorf("metadata.Outputs missing %q", "output_format")
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_Invoke_HappyPath covers the per-chunk
|
||||
// fan-out: two chunks in → two LLM calls → each chunk enriched
|
||||
// with the field_name key.
|
||||
func TestExtractorComponent_Invoke_HappyPath(t *testing.T) {
|
||||
withStubChatInvoker(t,
|
||||
stubResponse{Content: "answer for chunk 1"},
|
||||
stubResponse{Content: "answer for chunk 2"},
|
||||
)
|
||||
|
||||
c := &ExtractorComponent{Param: schema.ExtractorParam{
|
||||
FieldName: "summary",
|
||||
LLMID: "gpt-4o-mini",
|
||||
Prompt: "Summarize:",
|
||||
}}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"chunks": []map[string]any{
|
||||
{"text": "first text"},
|
||||
{"text": "second text"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
|
||||
chunks, ok := out["chunks"].([]map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("chunks key missing or wrong shape: %T", out["chunks"])
|
||||
}
|
||||
if len(chunks) != 2 {
|
||||
t.Fatalf("chunks len = %d, want 2", len(chunks))
|
||||
}
|
||||
if chunks[0]["summary"] != "answer for chunk 1" {
|
||||
t.Errorf("chunk[0].summary = %v, want %q", chunks[0]["summary"], "answer for chunk 1")
|
||||
}
|
||||
if chunks[1]["summary"] != "answer for chunk 2" {
|
||||
t.Errorf("chunk[1].summary = %v, want %q", chunks[1]["summary"], "answer for chunk 2")
|
||||
}
|
||||
if out["output_format"] != "chunks" {
|
||||
t.Errorf("output_format = %v, want chunks", out["output_format"])
|
||||
}
|
||||
if out["_elapsed_time"] == nil {
|
||||
t.Error("_elapsed_time missing")
|
||||
}
|
||||
if out["_created_time"] == nil {
|
||||
t.Error("_created_time missing")
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_Invoke_LLMError verifies a mock LLM
|
||||
// error is surfaced through Invoke with the component-name prefix
|
||||
// so the upstream pipeline can attribute failures.
|
||||
func TestExtractorComponent_Invoke_LLMError(t *testing.T) {
|
||||
withStubChatInvoker(t,
|
||||
stubResponse{Err: errors.New("upstream llm unavailable")},
|
||||
)
|
||||
|
||||
c := &ExtractorComponent{Param: schema.ExtractorParam{
|
||||
FieldName: "summary",
|
||||
LLMID: "gpt-4o-mini",
|
||||
}}
|
||||
_, err := c.Invoke(context.Background(), map[string]any{
|
||||
"chunks": []map[string]any{{"text": "x"}},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("Invoke returned nil error")
|
||||
}
|
||||
if !strings.HasPrefix(err.Error(), "extractor:") {
|
||||
t.Errorf("error should be wrapped with 'extractor:' prefix, got %v", err)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "upstream llm unavailable") {
|
||||
t.Errorf("error should chain underlying error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_Invoke_UnknownProvider asserts the
|
||||
// production (eino) chat invoker handles an unregistered driver
|
||||
// without panicking, per plan §8 Q1 ("48/56 providers covered;
|
||||
// the Extractor is provider-agnostic via llm_id; the 8 missing
|
||||
// are edge cases that do not block Phase 2.5").
|
||||
//
|
||||
// Design note: every other test in this file drives the
|
||||
// invoker through the production Component.Invoke path with a
|
||||
// canned-response invoker installed via SetExtractorChatInvoker
|
||||
// (the test seam). That seam accepts a pre-resolved driver
|
||||
// path; it cannot model the eino factory's default-branch
|
||||
// behaviour for an unknown driver. This test exercises the
|
||||
// production chat-invoker directly to pin that branch — the
|
||||
// production code path the real Extractor will hit when the
|
||||
// DSL references a provider that is not in the 48/56 covered
|
||||
// set.
|
||||
//
|
||||
// The contract under test:
|
||||
// - The call MUST NOT panic.
|
||||
// - On unknown driver, the factory's default branch routes to
|
||||
// a DummyModel that returns a deterministic error string
|
||||
// (we assert the error contains that sentinel so future
|
||||
// maintainers see the wiring goes through the factory,
|
||||
// not bypassed by a hand-rolled default).
|
||||
func TestExtractorComponent_Invoke_UnknownProvider(t *testing.T) {
|
||||
inv := &einoExtractorChatInvoker{}
|
||||
resp, err := inv.Chat(context.Background(), extractorChatRequest{
|
||||
Driver: "definitely-not-a-real-provider-xyz",
|
||||
ModelName: "anything",
|
||||
})
|
||||
// Either an error is returned OR a non-nil response is produced
|
||||
// by the DummyModel fallback. The contract is "no panic"; both
|
||||
// of these outcomes are acceptable. We only fail the test if
|
||||
// BOTH error and response are empty (which would indicate a
|
||||
// silent no-op).
|
||||
if err == nil && resp == nil {
|
||||
t.Fatal("production invoker returned nil error AND nil response for unknown driver — silent no-op")
|
||||
}
|
||||
// When an error IS returned, it must mention the driver name so
|
||||
// operators can correlate the failure back to the DSL config.
|
||||
if err != nil {
|
||||
// Acceptable error patterns for an unknown driver:
|
||||
// - mentions the driver name (correlatable for operators)
|
||||
// - "no driver"/"unknown" sentinels (typed error)
|
||||
// - "not implemented" (the eino dummy model fallback path)
|
||||
if !strings.Contains(err.Error(), "definitely-not-a-real-provider-xyz") &&
|
||||
!strings.Contains(err.Error(), "no driver") &&
|
||||
!strings.Contains(err.Error(), "unknown") &&
|
||||
!strings.Contains(err.Error(), "not implemented") {
|
||||
t.Errorf("unknown-driver error should mention the driver name or a typed/typed-sentinel substring; got: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_Invoke_ParsesJSON verifies a JSON object
|
||||
// response from the LLM is parsed into the chunk's field_name
|
||||
// value (matching the python set_output contract).
|
||||
func TestExtractorComponent_Invoke_ParsesJSON(t *testing.T) {
|
||||
withStubChatInvoker(t,
|
||||
stubResponse{Content: `{"answer": 42, "tags": ["a", "b"]}`},
|
||||
)
|
||||
|
||||
c := &ExtractorComponent{Param: schema.ExtractorParam{
|
||||
FieldName: "extraction",
|
||||
Prompt: "extract:",
|
||||
}}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"chunks": []map[string]any{{"text": "doc"}}},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks := out["chunks"].([]map[string]any)
|
||||
got, ok := chunks[0]["extraction"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("extraction should be parsed JSON object, got %T", chunks[0]["extraction"])
|
||||
}
|
||||
if got["answer"].(float64) != 42 {
|
||||
t.Errorf("answer = %v, want 42", got["answer"])
|
||||
}
|
||||
tags, _ := got["tags"].([]any)
|
||||
if len(tags) != 2 {
|
||||
t.Errorf("tags len = %d, want 2", len(tags))
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_Invoke_ParsesJSONInFence verifies the
|
||||
// common LLM response shape — JSON wrapped in a markdown code
|
||||
// fence — parses cleanly. Mirrors the behaviour the agent
|
||||
// canvas applies (e.g. llm_retry_test.go matchOutputStructure).
|
||||
func TestExtractorComponent_Invoke_ParsesJSONInFence(t *testing.T) {
|
||||
withStubChatInvoker(t,
|
||||
stubResponse{Content: "```json\n{\"summary\": \"hello\"}\n```"},
|
||||
)
|
||||
|
||||
c := &ExtractorComponent{Param: schema.ExtractorParam{
|
||||
FieldName: "out",
|
||||
}}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"chunks": []map[string]any{{"text": "x"}}},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
got, ok := out["chunks"].([]map[string]any)[0]["out"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("out should be parsed JSON object, got %T", out["chunks"].([]map[string]any)[0]["out"])
|
||||
}
|
||||
if got["summary"] != "hello" {
|
||||
t.Errorf("summary = %v, want hello", got["summary"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_Invoke_HandlesMalformedJSON verifies a
|
||||
// non-JSON response surfaces as the raw string under the
|
||||
// destination field — not an error. The python Extractor
|
||||
// accepts whatever the LLM emits; downstream callers decide
|
||||
// what to do with it.
|
||||
func TestExtractorComponent_Invoke_HandlesMalformedJSON(t *testing.T) {
|
||||
withStubChatInvoker(t,
|
||||
stubResponse{Content: "this is not JSON at all"},
|
||||
)
|
||||
|
||||
c := &ExtractorComponent{Param: schema.ExtractorParam{
|
||||
FieldName: "raw",
|
||||
}}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"chunks": []map[string]any{{"text": "x"}}},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke returned error on non-JSON: %v", err)
|
||||
}
|
||||
got := out["chunks"].([]map[string]any)[0]["raw"]
|
||||
if got != "this is not JSON at all" {
|
||||
t.Errorf("raw = %v, want %q", got, "this is not JSON at all")
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_Invoke_TOCNotPorted asserts the
|
||||
// field_name=="toc" branch is gated by a clear error so a future
|
||||
// migration to the Go TOC generator doesn't accidentally fall
|
||||
// through to chunk iteration.
|
||||
func TestExtractorComponent_Invoke_TOCNotPorted(t *testing.T) {
|
||||
c := &ExtractorComponent{Param: schema.ExtractorParam{
|
||||
FieldName: "toc",
|
||||
}}
|
||||
_, err := c.Invoke(context.Background(), map[string]any{
|
||||
"chunks": []map[string]any{{"text": "x"}}},
|
||||
)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for field_name=toc, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "toc") {
|
||||
t.Errorf("error should mention toc: %v", err)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "not yet ported") {
|
||||
t.Errorf("error should call out parity gap: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_Invoke_NoChunksFastPath verifies the
|
||||
// no-chunks input still produces a one-element chunks slice
|
||||
// (mirrors python _invoke line 110 fallback).
|
||||
func TestExtractorComponent_Invoke_NoChunksFastPath(t *testing.T) {
|
||||
withStubChatInvoker(t,
|
||||
stubResponse{Content: "single-shot answer"},
|
||||
)
|
||||
|
||||
c := &ExtractorComponent{Param: schema.ExtractorParam{
|
||||
FieldName: "answer",
|
||||
}}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, ok := out["chunks"].([]map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("chunks missing or wrong shape")
|
||||
}
|
||||
if len(chunks) != 1 {
|
||||
t.Fatalf("chunks len = %d, want 1", len(chunks))
|
||||
}
|
||||
if chunks[0]["answer"] != "single-shot answer" {
|
||||
t.Errorf("answer = %v, want %q", chunks[0]["answer"], "single-shot answer")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractorComponent_Invoke_JSONListInput(t *testing.T) {
|
||||
withStubChatInvoker(t,
|
||||
stubResponse{Content: "json chunk answer"},
|
||||
)
|
||||
|
||||
c := &ExtractorComponent{Param: schema.ExtractorParam{
|
||||
FieldName: "answer",
|
||||
}}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"json": []map[string]any{{"text": "json payload chunk"}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, ok := out["chunks"].([]map[string]any)
|
||||
if !ok || len(chunks) != 1 {
|
||||
t.Fatalf("chunks malformed: %v", out["chunks"])
|
||||
}
|
||||
if chunks[0]["answer"] != "json chunk answer" {
|
||||
t.Errorf("answer = %v, want %q", chunks[0]["answer"], "json chunk answer")
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_Invoke_PerCallLLMIDOverride verifies an
|
||||
// inputs["llm_id"] override wins over Param.LLMID and reaches
|
||||
// the chat invoker verbatim (the per-call override is the
|
||||
// explicit test seam for runtime reconfiguration).
|
||||
func TestExtractorComponent_Invoke_PerCallLLMIDOverride(t *testing.T) {
|
||||
stub := withStubChatInvoker(t,
|
||||
stubResponse{Content: "ok"},
|
||||
)
|
||||
|
||||
c := &ExtractorComponent{Param: schema.ExtractorParam{
|
||||
FieldName: "out",
|
||||
LLMID: "static-llm",
|
||||
}}
|
||||
_, err := c.Invoke(context.Background(), map[string]any{
|
||||
"llm_id": "override-llm",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
stub.mu.Lock()
|
||||
defer stub.mu.Unlock()
|
||||
if stub.lastReq.ModelName != "override-llm" {
|
||||
t.Errorf("ModelName = %q, want override-llm", stub.lastReq.ModelName)
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_Invoke_CompositeLLMID verifies the
|
||||
// composite "gpt-4o-mini@openai" form is split into driver and
|
||||
// model before reaching the chat invoker. Matches the canonical
|
||||
// composite llm_id convention used throughout the codebase
|
||||
// (see internal/agent/component/llm_credentials.go:parseLLMIDParts).
|
||||
func TestExtractorComponent_Invoke_CompositeLLMID(t *testing.T) {
|
||||
stub := withStubChatInvoker(t,
|
||||
stubResponse{Content: "ok"},
|
||||
)
|
||||
c := &ExtractorComponent{Param: schema.ExtractorParam{
|
||||
FieldName: "out",
|
||||
LLMID: "gpt-4o-mini@openai",
|
||||
}}
|
||||
if _, err := c.Invoke(context.Background(), map[string]any{}); err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
stub.mu.Lock()
|
||||
defer stub.mu.Unlock()
|
||||
if stub.lastReq.Driver != "openai" {
|
||||
t.Errorf("Driver = %q, want openai", stub.lastReq.Driver)
|
||||
}
|
||||
if stub.lastReq.ModelName != "gpt-4o-mini" {
|
||||
t.Errorf("ModelName = %q, want gpt-4o-mini", stub.lastReq.ModelName)
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_Invoke_ChunkIndexInError verifies the
|
||||
// error message includes the failing chunk index so a long
|
||||
// pipeline run surfaces which input document triggered the LLM
|
||||
// failure (mirrors python's per-chunk progress call at line 105).
|
||||
func TestExtractorComponent_Invoke_ChunkIndexInError(t *testing.T) {
|
||||
withStubChatInvoker(t,
|
||||
stubResponse{Content: "ok for chunk 0"},
|
||||
stubResponse{Err: errors.New("chunk-1-boom")},
|
||||
)
|
||||
c := &ExtractorComponent{Param: schema.ExtractorParam{
|
||||
FieldName: "out",
|
||||
}}
|
||||
_, err := c.Invoke(context.Background(), map[string]any{
|
||||
"chunks": []map[string]any{
|
||||
{"text": "first"},
|
||||
{"text": "second"},
|
||||
},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("Invoke returned nil error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "chunk 1") {
|
||||
t.Errorf("error should mention chunk 1 (zero-indexed): %v", err)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "chunk-1-boom") {
|
||||
t.Errorf("error should chain underlying error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_NewExtractorComponent_ParamCheck covers
|
||||
// the construction-time Validate() rejection of an empty
|
||||
// field_name (matches python check_empty "Result Destination").
|
||||
func TestExtractorComponent_NewExtractorComponent_ParamCheck(t *testing.T) {
|
||||
_, err := NewExtractorComponent(map[string]any{})
|
||||
if err == nil {
|
||||
t.Fatal("expected error for missing field_name, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "field_name") {
|
||||
t.Errorf("error should mention field_name: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_NewExtractorComponent_Happy covers the
|
||||
// parse path of every supported key; the param block coming out
|
||||
// should round-trip cleanly through Invoke.
|
||||
func TestExtractorComponent_NewExtractorComponent_Happy(t *testing.T) {
|
||||
withStubChatInvoker(t, stubResponse{Content: "ok"})
|
||||
c, err := NewExtractorComponent(map[string]any{
|
||||
"field_name": "summary",
|
||||
"llm_id": "openai/gpt-4o-mini",
|
||||
"system_prompt": "You are a precise summarizer.",
|
||||
"prompt": "Summarize:",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewExtractorComponent: %v", err)
|
||||
}
|
||||
if _, err := c.Invoke(context.Background(), map[string]any{
|
||||
"chunks": []map[string]any{{"text": "x"}}},
|
||||
); err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_InputsOutputs_NonEmpty is the shape
|
||||
// assertion Phase 4's API endpoint relies on.
|
||||
func TestExtractorComponent_InputsOutputs_NonEmpty(t *testing.T) {
|
||||
c := &ExtractorComponent{}
|
||||
ins := c.Inputs()
|
||||
outs := c.Outputs()
|
||||
if len(ins) == 0 {
|
||||
t.Error("Inputs() returned empty map")
|
||||
}
|
||||
if len(outs) == 0 {
|
||||
t.Error("Outputs() returned empty map")
|
||||
}
|
||||
if _, ok := outs["chunks"]; !ok {
|
||||
t.Errorf("Outputs() missing %q", "chunks")
|
||||
}
|
||||
if _, ok := outs["output_format"]; !ok {
|
||||
t.Errorf("Outputs() missing %q", "output_format")
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_Parallelism asserts the fan-out is
|
||||
// locked to 1 per plan §AD-5a ("Extractor: 1 (LLM call is
|
||||
// inherently serial)").
|
||||
func TestExtractorComponent_Parallelism(t *testing.T) {
|
||||
c := &ExtractorComponent{}
|
||||
if got := c.Parallelism(); got != 1 {
|
||||
t.Errorf("Parallelism() = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitExtractorLLID covers the composite-id parser in
|
||||
// isolation — keeps the matrix of edge cases at one call site
|
||||
// so a regression is easy to attribute. The "@" separator is
|
||||
// the canonical composite llm_id form used throughout the
|
||||
// codebase (see internal/agent/component/llm_credentials.go).
|
||||
func TestSplitExtractorLLID(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
wantModel string
|
||||
wantProvider string
|
||||
wantOK bool
|
||||
}{
|
||||
{"gpt-4o-mini@openai", "gpt-4o-mini", "openai", true},
|
||||
{"bare-model", "bare-model", "", false},
|
||||
{"trailing@", "trailing", "", true},
|
||||
{"@leading", "", "leading", true},
|
||||
{"", "", "", false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.in, func(t *testing.T) {
|
||||
model, provider, ok := splitExtractorLLID(tc.in)
|
||||
if ok != tc.wantOK {
|
||||
t.Errorf("ok = %v, want %v", ok, tc.wantOK)
|
||||
}
|
||||
if model != tc.wantModel {
|
||||
t.Errorf("model = %q, want %q", model, tc.wantModel)
|
||||
}
|
||||
if provider != tc.wantProvider {
|
||||
t.Errorf("provider = %q, want %q", provider, tc.wantProvider)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestTryParseJSONObject covers the best-effort JSON parser
|
||||
// independently of the LLM seam so its matrix of edge cases is
|
||||
// easy to attribute.
|
||||
func TestTryParseJSONObject(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
in string
|
||||
wantOK bool
|
||||
wantKey string // when wantOK=true, expected key in the parsed map
|
||||
}{
|
||||
{name: "object", in: `{"a":1}`, wantOK: true, wantKey: "a"},
|
||||
{name: "object with fence", in: "```json\n{\"a\":1}\n```", wantOK: true, wantKey: "a"},
|
||||
{name: "fence without json tag", in: "```\n{\"a\":1}\n```", wantOK: true, wantKey: "a"},
|
||||
{name: "plain string", in: "hello", wantOK: false},
|
||||
{name: "array", in: `[1,2]`, wantOK: false},
|
||||
{name: "empty object", in: `{}`, wantOK: false},
|
||||
{name: "empty", in: ``, wantOK: false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
parsed, ok := tryParseJSONObject(tc.in)
|
||||
if ok != tc.wantOK {
|
||||
t.Fatalf("ok = %v, want %v (got %v)", ok, tc.wantOK, parsed)
|
||||
}
|
||||
if ok && tc.wantKey != "" {
|
||||
if _, has := parsed[tc.wantKey]; !has {
|
||||
t.Errorf("parsed map missing %q: %v", tc.wantKey, parsed)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_ConcurrentInvoke verifies the chat
|
||||
// invoker swap is safe under concurrent Invoke calls. This is
|
||||
// the canary for SetExtractorChatInvoker and the package-level
|
||||
// RWMutex contract — a data race here breaks race detector.
|
||||
func TestExtractorComponent_ConcurrentInvoke(t *testing.T) {
|
||||
withStubChatInvoker(t,
|
||||
stubResponse{Content: "1"},
|
||||
stubResponse{Content: "2"},
|
||||
stubResponse{Content: "3"},
|
||||
stubResponse{Content: "4"},
|
||||
)
|
||||
c := &ExtractorComponent{Param: schema.ExtractorParam{
|
||||
FieldName: "out",
|
||||
}}
|
||||
chunks := []map[string]any{
|
||||
{"text": "a"}, {"text": "b"}, {"text": "c"}, {"text": "d"},
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, len(chunks))
|
||||
for _, ck := range chunks {
|
||||
ck := ck
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_, err := c.Invoke(context.Background(), map[string]any{
|
||||
"chunks": []map[string]any{ck},
|
||||
})
|
||||
if err != nil {
|
||||
errs <- err
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
t.Errorf("Invoke error under concurrency: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// silence unused-import vet warnings for eschema in case the
|
||||
// test file is built without the import ever being referenced
|
||||
// (it currently isn't, but pinning the import keeps test-side
|
||||
// imports honest if helpers move around in future revisions).
|
||||
var _ = eschema.Message{}
|
||||
274
internal/ingestion/component/file.go
Normal file
274
internal/ingestion/component/file.go
Normal file
@@ -0,0 +1,274 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// File ingestion component (Phase 2.1) — port of python `rag/flow/file.py`.
|
||||
//
|
||||
// SCOPE (honest):
|
||||
//
|
||||
// - DOC-ID PATH: matched. The python component only resolves the
|
||||
// document record and emits `name`; it does NOT fetch the binary.
|
||||
// The Go port now mirrors that ownership boundary.
|
||||
//
|
||||
// - BINARY FETCH: intentionally delegated to Parser. This matches the
|
||||
// Python runtime, where Parser resolves
|
||||
// `File2DocumentService.get_storage_address(...)` and performs the
|
||||
// storage GET itself.
|
||||
//
|
||||
// - ASYNC RACE / DOCUMENT-LEVEL LOCKING: not applicable in Go;
|
||||
// the python `self._canvas.callback(1, ...)` short-circuit is
|
||||
// replicated via `runtime.TrackProgress(0/1/...)` so a pipeline
|
||||
// observer sees Started/Done transitions.
|
||||
//
|
||||
// - PROGRESS: implemented via `runtime.TrackProgress`. With a
|
||||
// nil callback (the production pipeline wires its own sink),
|
||||
// the wrapper is a no-op so this component stays free of any
|
||||
// GUI / Redis state.
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
"ragflow/internal/storage"
|
||||
)
|
||||
|
||||
const ComponentNameFile = "File"
|
||||
|
||||
// FileComponent resolves document/file metadata and forwards enough
|
||||
// identity for downstream Parser to fetch bytes.
|
||||
//
|
||||
// Inputs (per rag/flow/file.py:File._invoke):
|
||||
//
|
||||
// doc_id (string, optional) — python: self._canvas._doc_id
|
||||
// file (list[map], optional) — python: kwargs.get("file")[0]
|
||||
// bucket (string, optional) — storage bucket override for tests/admin paths
|
||||
// path (string, optional) — storage object key override for tests/admin paths
|
||||
//
|
||||
// Outputs:
|
||||
//
|
||||
// doc_id (string, optional) — echoed for downstream Parser
|
||||
// name (string) — file/document name
|
||||
// path (string, optional) — storage path override echoed for downstream Parser
|
||||
// bucket (string, optional) — storage bucket override echoed for downstream Parser
|
||||
// file (map[string]any, optional) — file-list input echoed through
|
||||
// _created_time, _elapsed_time — TrackElapsed bookkeeping
|
||||
type FileComponent struct {
|
||||
}
|
||||
|
||||
// SetStorageFactoryOverride lets a test inject a Storage-backed
|
||||
// factory; production wiring should leave this alone and use the
|
||||
// real factory. The override is honored only when non-nil. This
|
||||
// is the testability seam requested by the Phase 2.1 spec — it
|
||||
// avoids forcing every call to pass a Storage through `Invoke`'s
|
||||
// inputs map (which would leak transport details into the wire
|
||||
// schema) while still giving tests a clean injection point.
|
||||
//
|
||||
// The companion tests live in file_test.go and use
|
||||
// `storage.NewMemoryStorage()` via `storage.GetStorageFactory().SetStorage(...)`,
|
||||
// the same pattern already used in internal/service/file_test.go:80-83.
|
||||
var SetStorageFactoryOverride func() storage.Storage
|
||||
|
||||
// NewFileComponent constructs a FileComponent from DSL params.
|
||||
// The current schema (schema.FileParam) has no fields.
|
||||
func NewFileComponent(params map[string]any) (runtime.Component, error) {
|
||||
p := schema.FileParam{}.Defaults()
|
||||
_ = p
|
||||
if err := p.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("File: param check: %w", err)
|
||||
}
|
||||
return &FileComponent{}, nil
|
||||
}
|
||||
|
||||
// Inputs returns the parameter metadata. Matches the python
|
||||
// File._invoke kwargs. The canonical upstream field is `file`,
|
||||
// matching the Python runtime contract.
|
||||
func (c *FileComponent) Inputs() map[string]string {
|
||||
return map[string]string{
|
||||
"doc_id": "Optional upstream document ID (mirrors python self._canvas._doc_id).",
|
||||
"file": "Optional upstream file descriptor list (mirrors python kwargs.get('file')).",
|
||||
"bucket": "Optional storage bucket override for downstream Parser fetches.",
|
||||
"path": "Optional storage object key override for downstream Parser fetches.",
|
||||
}
|
||||
}
|
||||
|
||||
// Outputs returns the parameter metadata. Mirrors the python
|
||||
// set_output contract (see schema.FileOutputs).
|
||||
func (c *FileComponent) Outputs() map[string]string {
|
||||
return map[string]string{
|
||||
"doc_id": "Document ID echoed for downstream Parser storage lookup.",
|
||||
"name": "Document / file name.",
|
||||
"path": "Optional storage object key override echoed for downstream Parser.",
|
||||
"bucket": "Optional storage bucket override echoed for downstream Parser.",
|
||||
"file": "Upstream file descriptor echoed when supplied via the file-list path.",
|
||||
"_ERROR": "Optional short-circuit error message (reserved for parity with python).",
|
||||
}
|
||||
}
|
||||
|
||||
// Parallelism is fixed at 1 — File is metadata-only.
|
||||
func (c *FileComponent) Parallelism() int { return 1 }
|
||||
|
||||
// Invoke resolves document/file metadata for downstream Parser use.
|
||||
//
|
||||
// The implementation mirrors the python flow's two paths:
|
||||
//
|
||||
// 1. doc_id is set (python: `self._canvas._doc_id`) — resolve
|
||||
// the document name and emit metadata only.
|
||||
// 2. doc_id is empty — pull the first file descriptor out of
|
||||
// `file` and use its `name`/`id` directly.
|
||||
func (c *FileComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
_ = ctx
|
||||
// Parse the wire input through the schema type so the
|
||||
// validation errors match the package convention.
|
||||
in, err := parseFileInputs(inputs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
out := map[string]any{"name": in.name}
|
||||
if in.docID != "" {
|
||||
out["doc_id"] = in.docID
|
||||
}
|
||||
if in.bucket != "" {
|
||||
out["bucket"] = in.bucket
|
||||
}
|
||||
if in.path != "" {
|
||||
out["path"] = in.path
|
||||
}
|
||||
if in.fileDesc != nil {
|
||||
out["file"] = in.fileDesc
|
||||
}
|
||||
return runtime.TrackElapsed("File", func() (map[string]any, error) {
|
||||
return out, nil
|
||||
})
|
||||
}
|
||||
|
||||
// fileInputs is the post-Validation view of the upstream input
|
||||
// map. Computed once at the top of Invoke so the rest of the
|
||||
// function reads as straight-line code.
|
||||
type fileInputs struct {
|
||||
docID string
|
||||
name string
|
||||
bucket string
|
||||
path string
|
||||
fileDesc map[string]any
|
||||
}
|
||||
|
||||
// parseFileInputs parses and validates the upstream input map.
|
||||
// Mirrors python's branching on `self._canvas._doc_id` vs.
|
||||
// `kwargs.get("file")[0]`.
|
||||
func parseFileInputs(inputs map[string]any) (fileInputs, error) {
|
||||
if inputs == nil {
|
||||
return fileInputs{}, fmt.Errorf("file: inputs map is nil")
|
||||
}
|
||||
out := fileInputs{}
|
||||
|
||||
if v, ok := getString(inputs, "doc_id"); ok && v != "" {
|
||||
out.docID = v
|
||||
out.name = v
|
||||
}
|
||||
|
||||
if out.docID == "" {
|
||||
// Fall through to the file-list path. Python uses
|
||||
// `kwargs.get("file")[0]`; Go keeps that same public key.
|
||||
if v, ok := inputs["file"]; ok {
|
||||
switch list := v.(type) {
|
||||
case []map[string]any:
|
||||
if len(list) > 0 {
|
||||
first := list[0]
|
||||
out.fileDesc = first
|
||||
if name, ok := first["name"].(string); ok {
|
||||
out.name = name
|
||||
}
|
||||
if id, ok := first["id"].(string); ok && out.bucket == "" && out.path == "" {
|
||||
// Convention: when `id` is a storage key, we treat
|
||||
// it as the `path` only when no explicit path is
|
||||
// provided. Bucket remains empty (must be set
|
||||
// separately). The python code does the same.
|
||||
out.path = id
|
||||
}
|
||||
}
|
||||
case []any:
|
||||
if len(list) > 0 {
|
||||
if first, ok := list[0].(map[string]any); ok {
|
||||
out.fileDesc = first
|
||||
if name, ok := first["name"].(string); ok {
|
||||
out.name = name
|
||||
}
|
||||
if id, ok := first["id"].(string); ok && out.bucket == "" && out.path == "" {
|
||||
out.path = id
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if out.name == "" {
|
||||
return fileInputs{}, fmt.Errorf("file: inputs missing doc_id or file[0].name")
|
||||
}
|
||||
}
|
||||
|
||||
if v, ok := getString(inputs, "bucket"); ok {
|
||||
out.bucket = v
|
||||
}
|
||||
if v, ok := getString(inputs, "path"); ok {
|
||||
out.path = v
|
||||
}
|
||||
if out.docID != "" {
|
||||
name, err := resolveDocumentName(out.docID)
|
||||
if err != nil {
|
||||
return fileInputs{}, fmt.Errorf("file: resolve doc_id %q: %w", out.docID, err)
|
||||
}
|
||||
if name != "" {
|
||||
out.name = name
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// getString accepts any of the json.Number-adjacent forms JSON
|
||||
// decoding produces. Canvas inputs decode through encoding/json
|
||||
// by default, which yields string for string-valued fields.
|
||||
func getString(m map[string]any, key string) (string, bool) {
|
||||
v, ok := m[key]
|
||||
if !ok || v == nil {
|
||||
return "", false
|
||||
}
|
||||
switch s := v.(type) {
|
||||
case string:
|
||||
return s, true
|
||||
case []byte:
|
||||
return string(s), true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// init registers File under CategoryIngestion (per plan §4
|
||||
// Phase 2.1). Metadata is derived from the Inputs()/Outputs()
|
||||
// methods on FileComponent so the API layer (Phase 4) can
|
||||
// enumerate the catalog without instantiating the component.
|
||||
func init() {
|
||||
c := &FileComponent{}
|
||||
runtime.MustRegister(ComponentNameFile, runtime.CategoryIngestion,
|
||||
func(_ string, params map[string]any) (runtime.Component, error) {
|
||||
return NewFileComponent(params)
|
||||
},
|
||||
runtime.Metadata{
|
||||
Version: "1.0.0",
|
||||
Inputs: c.Inputs(),
|
||||
Outputs: c.Outputs(),
|
||||
})
|
||||
}
|
||||
309
internal/ingestion/component/file_test.go
Normal file
309
internal/ingestion/component/file_test.go
Normal file
@@ -0,0 +1,309 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/dao"
|
||||
"ragflow/internal/entity"
|
||||
"ragflow/internal/storage"
|
||||
|
||||
"github.com/glebarez/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// withMemoryStorage swaps the global storage factory for an
|
||||
// in-memory implementation and restores the previous backend on
|
||||
// cleanup. Mirrors the pattern in internal/service/file_test.go:80-83.
|
||||
func withMemoryStorage(t *testing.T) *storage.MemoryStorage {
|
||||
t.Helper()
|
||||
factory := storage.GetStorageFactory()
|
||||
prev := factory.GetStorage()
|
||||
ms := storage.NewMemoryStorage().(*storage.MemoryStorage)
|
||||
factory.SetStorage(ms)
|
||||
t.Cleanup(func() { factory.SetStorage(prev) })
|
||||
return ms
|
||||
}
|
||||
|
||||
func withFileComponentTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{TranslateError: true})
|
||||
if err != nil {
|
||||
t.Fatalf("failed to open sqlite: %v", err)
|
||||
}
|
||||
if err := db.AutoMigrate(&entity.Document{}, &entity.File{}, &entity.File2Document{}); err != nil {
|
||||
t.Fatalf("failed to migrate sqlite: %v", err)
|
||||
}
|
||||
prev := dao.DB
|
||||
dao.DB = db
|
||||
t.Cleanup(func() { dao.DB = prev })
|
||||
return db
|
||||
}
|
||||
|
||||
// TestFileComponent_Registered verifies the init() registration
|
||||
// is visible to the runtime registry (Phase 4 / API layer
|
||||
// depends on this).
|
||||
func TestFileComponent_Registered(t *testing.T) {
|
||||
factory, cat, md, ok := runtime.DefaultRegistry.Lookup("File")
|
||||
if !ok {
|
||||
t.Fatal("File not registered in runtime.DefaultRegistry")
|
||||
}
|
||||
if cat != runtime.CategoryIngestion {
|
||||
t.Errorf("category = %q, want %q", cat, runtime.CategoryIngestion)
|
||||
}
|
||||
if factory == nil {
|
||||
t.Error("factory is nil")
|
||||
}
|
||||
if md.Inputs == nil || len(md.Inputs) == 0 {
|
||||
t.Errorf("metadata.Inputs empty: %v", md.Inputs)
|
||||
}
|
||||
if md.Outputs == nil || len(md.Outputs) == 0 {
|
||||
t.Errorf("metadata.Outputs empty: %v", md.Outputs)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFileComponent_Invoke_HappyPath verifies File remains metadata-only
|
||||
// even when explicit storage overrides are supplied.
|
||||
func TestFileComponent_Invoke_HappyPath(t *testing.T) {
|
||||
ms := withMemoryStorage(t)
|
||||
if err := ms.Put("bucketA", "path/to/file.txt", []byte("hello, ragflow")); err != nil {
|
||||
t.Fatalf("seed: %v", err)
|
||||
}
|
||||
|
||||
c := &FileComponent{}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"file": []map[string]any{{"name": "file.txt"}},
|
||||
"bucket": "bucketA",
|
||||
"path": "path/to/file.txt",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
|
||||
if got := out["name"]; got != "file.txt" {
|
||||
t.Errorf("name = %v, want file.txt", got)
|
||||
}
|
||||
if out["bucket"] != "bucketA" {
|
||||
t.Errorf("bucket = %v, want bucketA", out["bucket"])
|
||||
}
|
||||
if out["path"] != "path/to/file.txt" {
|
||||
t.Errorf("path = %v, want path/to/file.txt", out["path"])
|
||||
}
|
||||
if _, ok := out["binary"]; ok {
|
||||
t.Fatalf("binary should not be emitted by File: %v", out["binary"])
|
||||
}
|
||||
if out["_elapsed_time"] == nil {
|
||||
t.Error("_elapsed_time missing")
|
||||
}
|
||||
if out["_created_time"] == nil {
|
||||
t.Error("_created_time missing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileComponent_Invoke_ResolvesDocIDViaDocumentLocation(t *testing.T) {
|
||||
ms := withMemoryStorage(t)
|
||||
db := withFileComponentTestDB(t)
|
||||
location := "docs/from-document.bin"
|
||||
if err := ms.Put("kb-doc", location, []byte("doc-location")); err != nil {
|
||||
t.Fatalf("seed storage: %v", err)
|
||||
}
|
||||
docName := "report.pdf"
|
||||
if err := db.Create(&entity.Document{
|
||||
ID: "doc-loc",
|
||||
KbID: "kb-doc",
|
||||
ParserID: "na",
|
||||
ParserConfig: entity.JSONMap{},
|
||||
Type: "pdf",
|
||||
CreatedBy: "u1",
|
||||
Name: &docName,
|
||||
Location: &location,
|
||||
Suffix: ".pdf",
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("seed doc: %v", err)
|
||||
}
|
||||
|
||||
c := &FileComponent{}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{"doc_id": "doc-loc"})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if out["name"] != "report.pdf" {
|
||||
t.Fatalf("name = %v, want report.pdf", out["name"])
|
||||
}
|
||||
if _, ok := out["bucket"]; ok {
|
||||
t.Fatalf("bucket should not be emitted for doc_id-only path: %v", out["bucket"])
|
||||
}
|
||||
if _, ok := out["path"]; ok {
|
||||
t.Fatalf("path should not be emitted for doc_id-only path: %v", out["path"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileComponent_Invoke_ResolvesDocIDViaFileMapping(t *testing.T) {
|
||||
ms := withMemoryStorage(t)
|
||||
db := withFileComponentTestDB(t)
|
||||
location := "tenant-root/from-file.bin"
|
||||
if err := ms.Put("folder-1", location, []byte("file-mapping")); err != nil {
|
||||
t.Fatalf("seed storage: %v", err)
|
||||
}
|
||||
docName := "deck.pptx"
|
||||
if err := db.Create(&entity.Document{
|
||||
ID: "doc-file",
|
||||
KbID: "kb-1",
|
||||
ParserID: "na",
|
||||
ParserConfig: entity.JSONMap{},
|
||||
Type: "ppt",
|
||||
CreatedBy: "u1",
|
||||
Name: &docName,
|
||||
Suffix: ".pptx",
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("seed doc: %v", err)
|
||||
}
|
||||
if err := db.Create(&entity.File{
|
||||
ID: "file-1",
|
||||
ParentID: "folder-1",
|
||||
TenantID: "tenant-1",
|
||||
CreatedBy: "u1",
|
||||
Name: "deck.pptx",
|
||||
Location: &location,
|
||||
Type: "file",
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("seed file: %v", err)
|
||||
}
|
||||
fileID := "file-1"
|
||||
docID := "doc-file"
|
||||
if err := db.Create(&entity.File2Document{
|
||||
ID: "map-1",
|
||||
FileID: &fileID,
|
||||
DocumentID: &docID,
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("seed mapping: %v", err)
|
||||
}
|
||||
|
||||
c := &FileComponent{}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{"doc_id": "doc-file"})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if out["name"] != "deck.pptx" {
|
||||
t.Fatalf("name = %v, want deck.pptx", out["name"])
|
||||
}
|
||||
if _, ok := out["bucket"]; ok {
|
||||
t.Fatalf("bucket should not be emitted for doc_id-only path: %v", out["bucket"])
|
||||
}
|
||||
if _, ok := out["path"]; ok {
|
||||
t.Fatalf("path should not be emitted for doc_id-only path: %v", out["path"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileComponent_Invoke_DocIDWithoutStorageLocationStillSucceeds(t *testing.T) {
|
||||
withMemoryStorage(t)
|
||||
db := withFileComponentTestDB(t)
|
||||
if err := db.Create(&entity.Document{
|
||||
ID: "doc-empty",
|
||||
KbID: "kb-1",
|
||||
ParserID: "na",
|
||||
ParserConfig: entity.JSONMap{},
|
||||
Type: "txt",
|
||||
CreatedBy: "u1",
|
||||
Suffix: ".txt",
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("seed doc: %v", err)
|
||||
}
|
||||
|
||||
c := &FileComponent{}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{"doc_id": "doc-empty"})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got := out["name"]; got != "doc-empty" {
|
||||
t.Fatalf("name = %v, want doc-empty", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFileComponent_Invoke_MissingDoc covers the input-validation
|
||||
// branch when neither doc_id nor file is supplied.
|
||||
func TestFileComponent_Invoke_MissingDoc(t *testing.T) {
|
||||
withMemoryStorage(t)
|
||||
c := &FileComponent{}
|
||||
_, err := c.Invoke(context.Background(), map[string]any{})
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty inputs, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "doc_id") {
|
||||
t.Errorf("error should mention doc_id/file: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFileComponent_Invoke_EchoesStorageOverride verifies explicit
|
||||
// bucket/path overrides still flow through for downstream Parser use.
|
||||
func TestFileComponent_Invoke_IncludesCheckpointPath(t *testing.T) {
|
||||
ms := withMemoryStorage(t)
|
||||
const wantPath = "checkpoint/expected/path.bin"
|
||||
if err := ms.Put("b", wantPath, []byte("x")); err != nil {
|
||||
t.Fatalf("seed: %v", err)
|
||||
}
|
||||
c := &FileComponent{}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"file": []map[string]any{{"name": "checkpoint.bin"}},
|
||||
"bucket": "b",
|
||||
"path": wantPath,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if out["path"] != wantPath {
|
||||
t.Errorf("path = %v, want %q", out["path"], wantPath)
|
||||
}
|
||||
if out["bucket"] != "b" {
|
||||
t.Errorf("bucket = %v, want b", out["bucket"])
|
||||
}
|
||||
if _, ok := out["binary"]; ok {
|
||||
t.Fatalf("binary should not be emitted by File: %v", out["binary"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestFileComponent_InputsOutputs_NonEmpty is the shape
|
||||
// assertion Phase 4's API endpoint relies on.
|
||||
func TestFileComponent_InputsOutputs_NonEmpty(t *testing.T) {
|
||||
c := &FileComponent{}
|
||||
ins := c.Inputs()
|
||||
outs := c.Outputs()
|
||||
if len(ins) == 0 {
|
||||
t.Error("Inputs() returned empty map")
|
||||
}
|
||||
if len(outs) == 0 {
|
||||
t.Error("Outputs() returned empty map")
|
||||
}
|
||||
for _, key := range []string{"name", "file"} {
|
||||
if _, ok := outs[key]; !ok {
|
||||
t.Errorf("Outputs() missing %q", key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestFileComponent_Parallelism asserts the fan-out is locked to
|
||||
// 1 — File is metadata-only and intentionally non-fanned-out.
|
||||
func TestFileComponent_Parallelism(t *testing.T) {
|
||||
c := &FileComponent{}
|
||||
if got := c.Parallelism(); got != 1 {
|
||||
t.Errorf("Parallelism() = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
30
internal/ingestion/component/idempotency.go
Normal file
30
internal/ingestion/component/idempotency.go
Normal file
@@ -0,0 +1,30 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package component
|
||||
|
||||
// IdempotencyKey identifies one component execution.
|
||||
type IdempotencyKey struct {
|
||||
TaskID string
|
||||
PipelineVersion string
|
||||
ComponentName string
|
||||
ComponentVersion string
|
||||
InputFingerprint string
|
||||
}
|
||||
|
||||
func (k IdempotencyKey) String() string {
|
||||
return k.TaskID + "|" + k.PipelineVersion + "|" + k.ComponentName + "|" + k.ComponentVersion + "|" + k.InputFingerprint
|
||||
}
|
||||
639
internal/ingestion/component/parser.go
Normal file
639
internal/ingestion/component/parser.go
Normal file
@@ -0,0 +1,639 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Package component — Parser component (Phase 2.2 of
|
||||
// port-rag-flow-pipeline-to-go.md §4).
|
||||
//
|
||||
// SCOPE (honest):
|
||||
//
|
||||
// - WHAT IS PORTED:
|
||||
//
|
||||
// - The component's lifecycle contract: NewParserComponent /
|
||||
// Invoke / Parallelism / Inputs / Outputs and registration
|
||||
// under runtime.CategoryIngestion.
|
||||
//
|
||||
// - The Python fan-out pattern from rag/flow/parser/parser.py
|
||||
// (parallel page parsing, deterministic merge by page number,
|
||||
// see plan §8 R8) — the Go implementation uses
|
||||
// golang.org/x/sync/errgroup with up to 4 goroutines and
|
||||
// bounds each fan-out batch by a "page_size" input (default
|
||||
// = ceil(total_pages / 4)).
|
||||
//
|
||||
// - TrackProgress (start/done callback), WithTimeout (60s per
|
||||
// page-batch) and TrackElapsed (_created_time / _elapsed_time
|
||||
// stamping) — see internal/agent/runtime/helpers.go for the
|
||||
// helpers, plan §1 background.
|
||||
//
|
||||
// - WHAT IS NOT YET PORTED:
|
||||
//
|
||||
// - The Python component dispatches to 13 file-format branches
|
||||
// (pdf, markdown, text&code, html, spreadsheet, slides, doc,
|
||||
// docx, image, audio, video, email, epub) — see parser.py
|
||||
// function_map at line ~1273. The Go counterparts in
|
||||
// internal/parser/parser/ are SKELETONS that print to
|
||||
// stdout and return nil. The cgo-gated office variants
|
||||
// (docx, doc, ppt, pptx, xls, xlsx) call office_oxide but
|
||||
// discard the result.
|
||||
//
|
||||
// - Until the parser package returns real data, the Go Parser
|
||||
// component uses a "raw text" fallback: it treats the input
|
||||
// binary as UTF-8 and slices it into 1 page (or N pages
|
||||
// when the upstream signals a page boundary with a literal
|
||||
// "\f" form feed). This is the conservative, observable
|
||||
// behaviour until the real format dispatch lands.
|
||||
//
|
||||
// - The Python side's "image2id" pipeline (parser.py:1317-1329)
|
||||
// that uploads embedded images to MinIO is not replicated —
|
||||
// the schema layer carries images as opaque map values, and
|
||||
// the upload step is the responsibility of a separate
|
||||
// side-effect component (out of scope for Phase 2.2).
|
||||
//
|
||||
// - The Python _param.check() business validation
|
||||
// (parse_method whitelist, output_format whitelist, etc.) is
|
||||
// not replicated. The component trusts the param block
|
||||
// passed in at construction time; invalid values surface as
|
||||
// runtime errors in the chosen parser branch.
|
||||
//
|
||||
// - NO PERSISTENCE: parsed pages live only in the per-run
|
||||
// output map, exactly as the schema.Page type is intended.
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"golang.org/x/sync/errgroup"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
)
|
||||
|
||||
const ComponentNameParser = "Parser"
|
||||
|
||||
// parserParallelism is the fan-out degree for the Parser component.
|
||||
// Matches the plan §2 AD-5a choice ("Parser: 4 (parallel page
|
||||
// parsing)"). Used by the pipeline runner when it needs to know how
|
||||
// many goroutines the component is willing to absorb.
|
||||
const parserParallelism = 4
|
||||
|
||||
// parserPageBatchTimeout is the per-batch timeout. Mirrors the
|
||||
// Python component's `@timeout(60)` decorator on the page parse
|
||||
// branch. WithTimeout collapses the dual-layer
|
||||
// asyncio.wait_for / @timeout model into a single context, see
|
||||
// plan §8 R1.
|
||||
const parserPageBatchTimeout = 60 * time.Second
|
||||
|
||||
// pageFormFeed is the byte that text-page mode treats as a
|
||||
// hard page boundary. Matches the ASCII form feed (\f, 0x0C) — the
|
||||
// same convention used by the Python TxtParser and by most
|
||||
// "page-segmented text" codecs.
|
||||
const pageFormFeed = '\f'
|
||||
|
||||
// ParserComponent runs the configured parser branch against the
|
||||
// upstream "binary" payload and returns a deterministic, page-
|
||||
// sorted slice of schema.Page values.
|
||||
//
|
||||
// The instance is safe for concurrent invocation: each Invoke call
|
||||
// builds its own per-batch goroutine tree and merges results in
|
||||
// the goroutine that returned from Invoke. The static Param is
|
||||
// read-only after construction.
|
||||
type ParserComponent struct {
|
||||
// Param is the static configuration from schema.ParserParam.
|
||||
// Kept as a value (not a pointer) so callers can pass literals
|
||||
// and the component makes its own copy.
|
||||
Param schema.ParserParam
|
||||
}
|
||||
|
||||
// NewParserComponent constructs a Parser from a DSL param map.
|
||||
// The map is decoded into schema.ParserParam.Defaults() and then
|
||||
// overlaid with the supplied values. This matches the Python
|
||||
// "default + override" pattern in parser.py:ParserParam.__init__.
|
||||
//
|
||||
// Param map shape (all keys optional; missing keys fall back to
|
||||
// schema.ParserParam.Defaults() values):
|
||||
//
|
||||
// {
|
||||
// "setups": map[string]map[string]any,
|
||||
// "allowed_output_format": map[string][]string,
|
||||
// }
|
||||
//
|
||||
// Errors here surface as canvas compile failures so a malformed
|
||||
// param is caught at build time rather than mid-run.
|
||||
func NewParserComponent(params map[string]any) (runtime.Component, error) {
|
||||
p := schema.ParserParam{}.Defaults()
|
||||
if params == nil {
|
||||
return &ParserComponent{Param: p}, nil
|
||||
}
|
||||
// Setups — best-effort decode. A type mismatch in a single
|
||||
// setup entry drops just that entry; the rest of the table
|
||||
// remains usable. This matches the python behaviour of
|
||||
// accepting whatever shape the JSON loader hands back.
|
||||
if rawSetups, ok := params["setups"].(map[string]any); ok {
|
||||
for fileType, raw := range rawSetups {
|
||||
setupMap, ok := raw.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if p.Setups == nil {
|
||||
p.Setups = make(map[string]schema.ParserSetup)
|
||||
}
|
||||
p.Setups[fileType] = schema.ParserSetup(setupMap)
|
||||
}
|
||||
}
|
||||
if rawAllowed, ok := params["allowed_output_format"].(map[string]any); ok {
|
||||
allowed := make(map[string][]string, len(rawAllowed))
|
||||
for fileType, raw := range rawAllowed {
|
||||
list, ok := raw.([]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
formats := make([]string, 0, len(list))
|
||||
for _, item := range list {
|
||||
if s, ok := item.(string); ok {
|
||||
formats = append(formats, s)
|
||||
}
|
||||
}
|
||||
allowed[fileType] = formats
|
||||
}
|
||||
p.AllowedOutputFormat = allowed
|
||||
}
|
||||
return &ParserComponent{Param: p}, nil
|
||||
}
|
||||
|
||||
// Parallelism declares the goroutine fan-out degree. The pipeline
|
||||
// runner uses this to decide how many worker slots the component
|
||||
// can absorb. We return 4 to match the Python asyncio.gather
|
||||
// pattern that fans one batch per page-range.
|
||||
func (c *ParserComponent) Parallelism() int { return parserParallelism }
|
||||
|
||||
// Inputs returns the static parameter metadata. The component
|
||||
// reads the following from the inputs map at Invoke time:
|
||||
//
|
||||
// binary ([]byte, optional) — file bytes from upstream File.
|
||||
// When absent, Parser resolves
|
||||
// them from bucket/path or doc_id.
|
||||
// doc_id (string, optional) — document ID used for naming and,
|
||||
// when binary is absent, storage lookup.
|
||||
// page_size (int, optional) — pages per goroutine for
|
||||
// fan-out. Defaults to
|
||||
// ceil(totalPages / Parallelism).
|
||||
func (c *ParserComponent) Inputs() map[string]string {
|
||||
return map[string]string{
|
||||
"binary": "Optional file bytes ([]byte). When absent, Parser resolves them from bucket/path or doc_id.",
|
||||
"doc_id": "Optional document ID (string). Used for downstream correlation and doc_id-driven storage lookup.",
|
||||
"bucket": "Optional storage bucket override. Used when binary is absent.",
|
||||
"path": "Optional storage object key override. Used when binary is absent.",
|
||||
"page_size": "Optional integer. Pages per goroutine for fan-out. Default: ceil(totalPages / 4).",
|
||||
}
|
||||
}
|
||||
|
||||
// Outputs returns the public surface that downstream ingestion
|
||||
// components (Chunker, Tokenizer, Extractor) can wire into.
|
||||
//
|
||||
// pages []schema.Page — sorted by PageNumber. Deterministic
|
||||
// merge per plan §8 R8.
|
||||
// name string — carried over from the upstream file/document
|
||||
// name (or doc_id when no name is available).
|
||||
// output_format string — "text" when emitting text pages,
|
||||
// otherwise the parser-selected wire
|
||||
// format.
|
||||
// _ERROR string — populated when the component short-
|
||||
// circuits with an error message
|
||||
// (mirrors Python set_output("_ERROR", ...)).
|
||||
func (c *ParserComponent) Outputs() map[string]string {
|
||||
return map[string]string{
|
||||
"pages": "[]schema.Page: parsed pages sorted by PageNumber.",
|
||||
"name": "string: the upstream file/document name (or doc_id when no name is available).",
|
||||
"output_format": "string: the active output format (\"text\" when emitting text pages).",
|
||||
"_ERROR": "string: set on short-circuit errors.",
|
||||
}
|
||||
}
|
||||
|
||||
// Invoke runs the parser against the upstream "binary" payload.
|
||||
//
|
||||
// Returns:
|
||||
//
|
||||
// {
|
||||
// "pages": []schema.Page (sorted by PageNumber),
|
||||
// "name": string (from inputs["doc_id"]),
|
||||
// "output_format": "text",
|
||||
// "_created_time": RFC3339Nano (via TrackElapsed),
|
||||
// "_elapsed_time": float64 seconds (via TrackElapsed),
|
||||
// }
|
||||
//
|
||||
// The fan-out is bounded by Parallelism() goroutines. Each
|
||||
// goroutine parses its page-batch under a derived timeout
|
||||
// (WithTimeout, 60s). The first error cancels the errgroup
|
||||
// context; siblings observe ctx.Done() and abandon their work.
|
||||
//
|
||||
// DETERMINISTIC MERGE (plan §8 R8): after fan-out, the page slice
|
||||
// is sorted by PageNumber. This guarantees the same input
|
||||
// produces byte-identical output across runs and is the contract
|
||||
// that downstream Chunker / Tokenizer rely on for stable chunk
|
||||
// IDs (chunks that span pages must reference adjacent PageNumbers
|
||||
// in input order).
|
||||
func (c *ParserComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
// 1. Decode the binary input.
|
||||
binary, err := readParserBinary(ctx, inputs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
docID, _ := inputs["doc_id"].(string)
|
||||
filename := parserInputName(inputs, docID)
|
||||
|
||||
// 2. Resolve the file family from the inputs. When the family
|
||||
// is known, dispatchParse returns a typed parser payload.
|
||||
// Otherwise the component stays in text-page mode.
|
||||
//
|
||||
// We track TWO forms:
|
||||
//
|
||||
// - fileTypeExt — the utility.FileType extension form ("md",
|
||||
// "docx", ...). Used by parser.GetParser, whose switch
|
||||
// arms are keyed off the utility constants.
|
||||
//
|
||||
// - fileTypeFam — the python-side family name ("markdown",
|
||||
// "docx", ...). Used by setups[fileType] and
|
||||
// allowed_output_format[fileType] lookups, which are keyed
|
||||
// off the python family identifiers in schema.ParserParam.
|
||||
//
|
||||
// For most families the two forms coincide; the divergence
|
||||
// exists for markdown ("md" vs "markdown") and slides
|
||||
// ("ppt"/"pptx" vs "slides") and is intentional — the python
|
||||
// ParserParam collapses the slide family into a single key.
|
||||
fileTypeExt := fileTypeFromInputs(inputs)
|
||||
fileTypeFam := pythonFamilyName(string(fileTypeExt))
|
||||
|
||||
// 2a. Validate the requested output_format against the
|
||||
// family-specific allowed_output_format whitelist. We do
|
||||
// this even when no setups entry exists so a misconfigured
|
||||
// DSL surfaces as _ERROR instead of a silent fallback.
|
||||
if _, hasSetup := c.Param.Setups[fileTypeFam]; hasSetup {
|
||||
if _, verr := resolveOutputFormat(fileTypeFam, c.Param.Setups, c.Param.AllowedOutputFormat); verr != nil {
|
||||
return nil, verr
|
||||
}
|
||||
}
|
||||
|
||||
dispatched := dispatchParse(fileTypeExt, filename, binary, c.Param.Setups)
|
||||
dispatched = hydrateEmptyDispatchPayload(dispatched, binary)
|
||||
|
||||
// 3. Build the legacy `pages` slice. When the dispatch path
|
||||
// produced a JSON payload, we re-shape it into the page
|
||||
// layout the chunker side consumes (`{text, doc_type_kwd,
|
||||
// page_number?}`); when the dispatch produced a string
|
||||
// payload we emit a single page carrying the rendered text;
|
||||
// otherwise we slice the binary on ASCII form-feed and
|
||||
// treat the input as text pages.
|
||||
var pages [][]byte
|
||||
var dispatchedPages []schema.Page
|
||||
switch {
|
||||
case dispatched.Err == nil && dispatched.OutputFormat == "json" && len(dispatched.JSON) > 0:
|
||||
dispatchedPages = jsonItemsToPages(dispatched.JSON)
|
||||
pages = pagesFromDispatch(dispatchedPages)
|
||||
case dispatched.Err == nil && dispatched.OutputFormat != "":
|
||||
var text string
|
||||
switch dispatched.OutputFormat {
|
||||
case "markdown":
|
||||
text = dispatched.Markdown
|
||||
case "html":
|
||||
text = dispatched.HTML
|
||||
case "text":
|
||||
text = dispatched.Text
|
||||
}
|
||||
pages = [][]byte{[]byte(text)}
|
||||
default:
|
||||
pages = splitIntoPages(binary)
|
||||
if len(pages) == 0 {
|
||||
pages = [][]byte{nil}
|
||||
}
|
||||
}
|
||||
totalPages := len(pages)
|
||||
|
||||
// 4. Fan-out: split pages into batches of pageSize. The
|
||||
// default pageSize is ceil(totalPages / Parallelism),
|
||||
// matching the plan §2 AD-5a target.
|
||||
pageSize := resolvePageSize(inputs, totalPages)
|
||||
batches := splitIntoBatches(pages, pageSize)
|
||||
|
||||
// 5. Drive the fan-out from TrackProgress (which delivers the
|
||||
// start/done/fail callback sequence); stamp
|
||||
// _created_time / _elapsed_time on the result via
|
||||
// TrackElapsed without re-running the work.
|
||||
var out map[string]any
|
||||
progressErr := runtime.TrackProgress(ComponentNameParser, nil, func() error {
|
||||
parsed, err := fanOutAndMerge(ctx, batches, parserParallelism)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Sort by PageNumber — DETERMINISTIC MERGE (plan §8 R8).
|
||||
sortPagesByNumber(parsed)
|
||||
out = buildParserOutputs(parsed, dispatched, filename, fileTypeExt)
|
||||
return nil
|
||||
})
|
||||
if progressErr != nil {
|
||||
return nil, fmt.Errorf("Parser: %w", progressErr)
|
||||
}
|
||||
// Stamp _created_time / _elapsed_time. We pass a closure
|
||||
// that returns the pre-built `out` so the helper does not
|
||||
// re-execute the fan-out.
|
||||
return runtime.TrackElapsed(ComponentNameParser, func() (map[string]any, error) {
|
||||
return out, nil
|
||||
})
|
||||
}
|
||||
|
||||
// fanOutAndMerge parses each batch in parallel and concatenates
|
||||
// the per-batch results. The first error cancels the errgroup
|
||||
// context; siblings see ctx.Done() and abandon their parse
|
||||
// (returning ctx.Err()).
|
||||
//
|
||||
// Concurrency model: at most `parallelism` goroutines run
|
||||
// concurrently. errgroup.WithContext provides the cancel-on-
|
||||
// first-error behaviour, and golang.org/x/sync/errgroup is
|
||||
// already in go.mod (line 59).
|
||||
func fanOutAndMerge(parent context.Context, batches [][][]byte, parallelism int) ([]schema.Page, error) {
|
||||
if len(batches) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
if parallelism < 1 {
|
||||
parallelism = 1
|
||||
}
|
||||
g, ctx := errgroup.WithContext(parent)
|
||||
g.SetLimit(parallelism)
|
||||
|
||||
// One slot per batch; we collect the results in order so the
|
||||
// caller can sort the merged slice by PageNumber without
|
||||
// needing a mutex.
|
||||
results := make([][]schema.Page, len(batches))
|
||||
for i, batch := range batches {
|
||||
i, batch := i, batch
|
||||
g.Go(func() error {
|
||||
pages, err := parseBatch(ctx, batch)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
results[i] = pages
|
||||
return nil
|
||||
})
|
||||
}
|
||||
if err := g.Wait(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Flatten in batch order. The caller is responsible for
|
||||
// sorting by PageNumber — we deliberately do NOT sort here
|
||||
// so the per-batch order is visible to tests.
|
||||
total := 0
|
||||
for _, r := range results {
|
||||
total += len(r)
|
||||
}
|
||||
merged := make([]schema.Page, 0, total)
|
||||
for _, r := range results {
|
||||
merged = append(merged, r...)
|
||||
}
|
||||
return merged, nil
|
||||
}
|
||||
|
||||
// parseBatch parses a single batch of pages under a derived
|
||||
// 60-second timeout. The batch is the unit of fan-out: if a
|
||||
// batch exceeds its timeout, ONLY that batch errors; siblings
|
||||
// see ctx.Done() and abandon their work (errgroup cancel
|
||||
// cascades).
|
||||
//
|
||||
// runtime.WithTimeout returns just an error; we capture the
|
||||
// result pages in a closure-scoped variable and read it back
|
||||
// after Wait. This keeps the helper at its single-purpose
|
||||
// signature (ctx, fn -> error) without growing the runtime API.
|
||||
func parseBatch(ctx context.Context, batch [][]byte) ([]schema.Page, error) {
|
||||
var pages []schema.Page
|
||||
err := runtime.WithTimeout(ctx, parserPageBatchTimeout, func(ctx context.Context) error {
|
||||
pages = make([]schema.Page, 0, len(batch))
|
||||
for _, raw := range batch {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
default:
|
||||
}
|
||||
// Text-page mode: the bytes are already page text.
|
||||
pages = append(pages, schema.Page{
|
||||
"text": string(raw),
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return pages, nil
|
||||
}
|
||||
|
||||
// --- input helpers ---
|
||||
|
||||
// readParserBinary pulls the "binary" payload out of the inputs
|
||||
// map. The accepted shapes are:
|
||||
//
|
||||
// []byte — the in-process caller's normal form
|
||||
// string — UTF-8 text (json callers' normal form)
|
||||
// nil / absent — returns an empty page (not an error)
|
||||
//
|
||||
// A non-UTF-8 string is rejected with a clear error so a caller
|
||||
// that mistakenly hands a base64 string sees the failure
|
||||
// immediately (mirrors pipeline_chunker's "no try-base64" rule).
|
||||
func readParserBinary(ctx context.Context, inputs map[string]any) ([]byte, error) {
|
||||
if inputs == nil {
|
||||
return nil, nil
|
||||
}
|
||||
if b, ok := inputs["binary"].([]byte); ok {
|
||||
return b, nil
|
||||
}
|
||||
if s, ok := inputs["binary"].(string); ok {
|
||||
if !utf8.ValidString(s) {
|
||||
return nil, errors.New(
|
||||
"Parser: binary string is not valid UTF-8. " +
|
||||
"Text-page mode only accepts UTF-8 text input.")
|
||||
}
|
||||
return []byte(s), nil
|
||||
}
|
||||
bucket, _ := getString(inputs, "bucket")
|
||||
path, _ := getString(inputs, "path")
|
||||
if bucket != "" && path != "" {
|
||||
return fetchBinary(ctx, bucket, path)
|
||||
}
|
||||
if docID, ok := getString(inputs, "doc_id"); ok && docID != "" {
|
||||
ref, err := resolveDocumentStorage(docID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Parser: resolve doc_id %q: %w", docID, err)
|
||||
}
|
||||
return fetchBinary(ctx, ref.Bucket, ref.Path)
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// splitIntoPages segments the input bytes on ASCII form-feed
|
||||
// (\f, 0x0C). An input with no form-feeds becomes a single page
|
||||
// (the whole input). Empty pages are dropped — the python
|
||||
// TxtParser skips empty splits the same way.
|
||||
func splitIntoPages(b []byte) [][]byte {
|
||||
if len(b) == 0 {
|
||||
return nil
|
||||
}
|
||||
// Fast path: no form-feeds → single page.
|
||||
if !containsFormFeed(b) {
|
||||
return [][]byte{b}
|
||||
}
|
||||
parts := strings.Split(string(b), string(pageFormFeed))
|
||||
out := make([][]byte, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
if len(p) == 0 {
|
||||
continue
|
||||
}
|
||||
out = append(out, []byte(p))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// containsFormFeed is a tiny specialised byte-search to avoid
|
||||
// pulling in bytes.Index for one call site.
|
||||
func containsFormFeed(b []byte) bool {
|
||||
for _, c := range b {
|
||||
if c == pageFormFeed {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// splitIntoBatches partitions the page slice into batches of
|
||||
// `size` consecutive pages. A non-positive size collapses to
|
||||
// one batch.
|
||||
func splitIntoBatches(pages [][]byte, size int) [][][]byte {
|
||||
if size < 1 {
|
||||
size = len(pages)
|
||||
}
|
||||
if size < 1 {
|
||||
return nil
|
||||
}
|
||||
batches := make([][][]byte, 0, (len(pages)+size-1)/size)
|
||||
for i := 0; i < len(pages); i += size {
|
||||
end := i + size
|
||||
if end > len(pages) {
|
||||
end = len(pages)
|
||||
}
|
||||
batch := make([][]byte, end-i)
|
||||
copy(batch, pages[i:end])
|
||||
batches = append(batches, batch)
|
||||
}
|
||||
return batches
|
||||
}
|
||||
|
||||
// resolvePageSize returns the inputs["page_size"] value when
|
||||
// valid, otherwise ceil(totalPages / Parallelism). A page_size
|
||||
// of 0 or 1 is treated as "use the default" so a caller that
|
||||
// sets page_size=1 to mean "no batching" still fans out across
|
||||
// `Parallelism` goroutines.
|
||||
func resolvePageSize(inputs map[string]any, totalPages int) int {
|
||||
if inputs != nil {
|
||||
if v, ok := inputs["page_size"].(int); ok && v > 1 {
|
||||
return v
|
||||
}
|
||||
if v, ok := inputs["page_size"].(int64); ok && v > 1 {
|
||||
return int(v)
|
||||
}
|
||||
if v, ok := inputs["page_size"].(float64); ok && v > 1 {
|
||||
return int(v)
|
||||
}
|
||||
}
|
||||
if totalPages < 1 {
|
||||
return 1
|
||||
}
|
||||
// ceil(totalPages / Parallelism)
|
||||
size := (totalPages + parserParallelism - 1) / parserParallelism
|
||||
if size < 1 {
|
||||
size = 1
|
||||
}
|
||||
return size
|
||||
}
|
||||
|
||||
// sortPagesByNumber orders pages by their PageNumber key
|
||||
// ascending. Pages without a PageNumber key (or with a non-int
|
||||
// value) sort to the END so the deterministic contract is
|
||||
// "numbered pages first, then unnumbered" — this matches the
|
||||
// Python component's loop order (it processes pages in input
|
||||
// order, not in PageNumber order, but the Go merge is
|
||||
// intentionally stricter so the test can assert exact byte
|
||||
// equality across runs).
|
||||
func sortPagesByNumber(pages []schema.Page) {
|
||||
sort.SliceStable(pages, func(i, j int) bool {
|
||||
pi, oki := numericPageNumber(pages[i])
|
||||
pj, okj := numericPageNumber(pages[j])
|
||||
switch {
|
||||
case oki && okj:
|
||||
return pi < pj
|
||||
case oki:
|
||||
return true // i is numbered, j is not
|
||||
case okj:
|
||||
return false
|
||||
default:
|
||||
return false // stable
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func numericPageNumber(p schema.Page) (int, bool) {
|
||||
if p == nil {
|
||||
return 0, false
|
||||
}
|
||||
v, ok := p["page_number"]
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
switch n := v.(type) {
|
||||
case int:
|
||||
return n, true
|
||||
case int64:
|
||||
return int(n), true
|
||||
case float64:
|
||||
return int(n), true
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
// toAnyPages is a tiny adapter that hands the page slice to
|
||||
// the output map as `any`. We use it instead of a direct cast
|
||||
// so the type stays `[]schema.Page` in the Go source and the
|
||||
// output map value type is `any` — matching the runtime.Component
|
||||
// contract.
|
||||
func toAnyPages(pages []schema.Page) any { return pages }
|
||||
|
||||
// init registers Parser under CategoryIngestion per plan §4
|
||||
// Phase 2.2. The factory is a thin closure that decodes the
|
||||
// DSL param map; the static Metadata is derived from
|
||||
// Inputs()/Outputs() on a zero-value instance.
|
||||
func init() {
|
||||
pc := &ParserComponent{}
|
||||
runtime.MustRegister(ComponentNameParser, runtime.CategoryIngestion,
|
||||
func(_ string, params map[string]any) (runtime.Component, error) {
|
||||
return NewParserComponent(params)
|
||||
},
|
||||
runtime.Metadata{
|
||||
Version: "1.0.0",
|
||||
Inputs: pc.Inputs(),
|
||||
Outputs: pc.Outputs(),
|
||||
})
|
||||
}
|
||||
424
internal/ingestion/component/parser_dispatch.go
Normal file
424
internal/ingestion/component/parser_dispatch.go
Normal file
@@ -0,0 +1,424 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Parser dispatch validates the requested output format, resolves the
|
||||
// parser backend, and returns the structured ParseWithResult payload.
|
||||
//
|
||||
// `parse_method` is carried through file metadata for downstream
|
||||
// consumers, while the actual backend work stays in
|
||||
// internal/parser/parser/*.
|
||||
|
||||
package component
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"maps"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
"ragflow/internal/parser/parser"
|
||||
"ragflow/internal/utility"
|
||||
)
|
||||
|
||||
// parserDispatchResult is the typed outcome of dispatchParse. The
|
||||
// component's Invoke translates it into the runtime output map.
|
||||
//
|
||||
// OutputFormat is the wire format the parser actually emitted. It
|
||||
// always matches setups[fileType].output_format on success; on a
|
||||
// format-mismatch failure (the whitelist rejects it) OutputFormat is
|
||||
// empty and Err is non-nil.
|
||||
//
|
||||
// File is the per-parser file metadata and may be nil.
|
||||
//
|
||||
// Payload holds exactly one populated field per the ParseResult
|
||||
// contract (see internal/parser/parser/parse_result.go).
|
||||
type parserDispatchResult struct {
|
||||
OutputFormat string
|
||||
File map[string]any
|
||||
JSON []map[string]any
|
||||
Markdown string
|
||||
Text string
|
||||
HTML string
|
||||
Err error
|
||||
}
|
||||
|
||||
// resolveOutputFormat picks the wire format for this run. The
|
||||
// Python side asks the setup, then checks the value is in
|
||||
// allowed_output_format[fileType]. We mirror that exact sequence:
|
||||
//
|
||||
// 1. setups[fileType].output_format (or "" when absent),
|
||||
// 2. if absent, default to "text" (the most permissive option
|
||||
// that every family accepts),
|
||||
// 3. if absent and the family has no allowed_output_format entry,
|
||||
// return "" (the component falls back to text-page mode
|
||||
// without validating).
|
||||
//
|
||||
// The whitelist check returns "" + Err when the requested format is
|
||||
// not in the allowed set; this is the validation Python raises
|
||||
// via check_empty / check_valid_value, surfaced as an error so the
|
||||
// component short-circuits with _ERROR rather than emitting a
|
||||
// payload the downstream chunker cannot consume.
|
||||
func resolveOutputFormat(family string, setups map[string]schema.ParserSetup, allowed map[string][]string) (string, error) {
|
||||
setup, ok := setups[family]
|
||||
if !ok {
|
||||
// Family not configured — text-page mode; no validation.
|
||||
return "", nil
|
||||
}
|
||||
format, _ := setup["output_format"].(string)
|
||||
if format == "" {
|
||||
format = "text"
|
||||
}
|
||||
allowedList, ok := allowed[family]
|
||||
if !ok || len(allowedList) == 0 {
|
||||
// No whitelist entry — accept what the setup asked for.
|
||||
return format, nil
|
||||
}
|
||||
for _, candidate := range allowedList {
|
||||
if strings.EqualFold(candidate, format) {
|
||||
return format, nil
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf(
|
||||
"Parser: output_format %q for %q is not in allowed_output_format %v",
|
||||
format, family, allowedList,
|
||||
)
|
||||
}
|
||||
|
||||
// resolveLibType returns the lib_type argument for parser.GetParser.
|
||||
// The Python side does not pass an explicit lib_type for the
|
||||
// Markdown / HTML / Office families — the parser constructor picks
|
||||
// the only available backend. Mirroring that, when setups[…].lib_type
|
||||
// is unset we leave the field empty and let parser.GetParser
|
||||
// surface a structured error if no backend is wired.
|
||||
//
|
||||
// `parse_method` is preserved on the dispatch result so callers can
|
||||
// tell the difference between "explicit OCR" and "default DeepDOC"
|
||||
// without re-reading setups.
|
||||
func resolveLibType(fileType utility.FileType, setups map[string]schema.ParserSetup) (libType, parseMethod string) {
|
||||
setup, ok := setups[string(fileType)]
|
||||
if !ok {
|
||||
return "", ""
|
||||
}
|
||||
if s, ok := setup["lib_type"].(string); ok {
|
||||
libType = s
|
||||
}
|
||||
if s, ok := setup["parse_method"].(string); ok {
|
||||
parseMethod = s
|
||||
}
|
||||
return libType, parseMethod
|
||||
}
|
||||
|
||||
// dispatchParse resolves the parser for the given fileType and invokes
|
||||
// its structured ParseWithResult contract.
|
||||
//
|
||||
// The function NEVER returns a partial result. On error the result
|
||||
// is the zero value (OutputFormat == "" + Err != nil). Callers can
|
||||
// detect the success/failure boundary on the OutputFormat alone.
|
||||
//
|
||||
// fileType may be utility.FileTypeOTHER when the upstream did not
|
||||
// supply a filename; the dispatch then takes text-page mode
|
||||
// without consulting parser.GetParser.
|
||||
func dispatchParse(fileType utility.FileType, filename string, data []byte, setups map[string]schema.ParserSetup) parserDispatchResult {
|
||||
if fileType == utility.FileTypeOTHER {
|
||||
// Unknown / unset family. The component treats the bytes
|
||||
// as text pages; the existing logic
|
||||
// (splitIntoPages + fan-out) handles it. We return no
|
||||
// result here so the caller routes to that path.
|
||||
return parserDispatchResult{}
|
||||
}
|
||||
|
||||
libType, parseMethod := resolveLibType(fileType, setups)
|
||||
|
||||
// Resolve a parser via the GetParser entry point. Any error here
|
||||
// is the caller's fault (libType unsupported for this family);
|
||||
// surface it as Err so Invoke can set _ERROR.
|
||||
p, err := parser.GetParser(fileType, map[string]string{"lib_type": libType})
|
||||
if err != nil {
|
||||
return parserDispatchResult{Err: fmt.Errorf("Parser: resolve %q: %w", fileType, err)}
|
||||
}
|
||||
|
||||
res := p.ParseWithResult(filename, data)
|
||||
if res.Err != nil {
|
||||
return parserDispatchResult{Err: fmt.Errorf("Parser: %q: %w", fileType, res.Err)}
|
||||
}
|
||||
// Carry the configured parse_method on the file metadata so
|
||||
// downstream consumers can read which provider ran.
|
||||
if parseMethod != "" {
|
||||
if res.File == nil {
|
||||
res.File = map[string]any{}
|
||||
}
|
||||
res.File["parse_method"] = parseMethod
|
||||
}
|
||||
return parserDispatchResult{
|
||||
OutputFormat: res.OutputFormat,
|
||||
File: res.File,
|
||||
JSON: res.JSON,
|
||||
Markdown: res.Markdown,
|
||||
Text: res.Text,
|
||||
HTML: res.HTML,
|
||||
}
|
||||
}
|
||||
|
||||
// fileTypeFromInputs derives the parser-library extension form
|
||||
// (utility.FileType) from the upstream inputs. The result is the
|
||||
// value passed to parser.GetParser, whose switch arms are keyed
|
||||
// off the utility constants.
|
||||
//
|
||||
// Resolution order:
|
||||
//
|
||||
// 1. inputs["file_type"] — explicit family hint from the upstream
|
||||
// File component. We accept either the extension ("md", "docx")
|
||||
// or the python family name ("markdown"); both are normalised
|
||||
// to the extension form via the pythonFamilyName / familyToExt
|
||||
// lookup tables below.
|
||||
// 2. inputs["file"].name — fall back to the filename so a caller
|
||||
// that only supplies the path is still routed correctly.
|
||||
// 3. inputs["name"] — last-resort filename.
|
||||
// 4. utility.FileTypeOTHER — text-page mode.
|
||||
//
|
||||
// The function never errors; unknown / absent filenames degrade to
|
||||
// FileTypeOTHER so the component's raw-text branch picks them up.
|
||||
func fileTypeFromInputs(inputs map[string]any) utility.FileType {
|
||||
if inputs == nil {
|
||||
return utility.FileTypeOTHER
|
||||
}
|
||||
if raw, ok := inputs["file_type"].(string); ok && raw != "" {
|
||||
if ft := familyToExt(pythonFamilyName(strings.ToLower(raw))); ft != utility.FileTypeOTHER {
|
||||
return ft
|
||||
}
|
||||
// Direct extension match — handles "md", "docx", etc.
|
||||
if ft := utility.GetFileType("x." + strings.ToLower(raw)); ft != utility.FileTypeOTHER {
|
||||
return ft
|
||||
}
|
||||
}
|
||||
if m, ok := inputs["file"].(map[string]any); ok {
|
||||
if name, ok := m["name"].(string); ok && name != "" {
|
||||
return utility.GetFileType(name)
|
||||
}
|
||||
}
|
||||
if name, ok := inputs["name"].(string); ok && name != "" {
|
||||
return utility.GetFileType(name)
|
||||
}
|
||||
return utility.FileTypeOTHER
|
||||
}
|
||||
|
||||
// familyToExt maps the python family name back to the utility
|
||||
// extension form. Returns FileTypeOTHER for families whose parser
|
||||
// isn't yet wired (audio, video, image, email, epub, …).
|
||||
func familyToExt(family string) utility.FileType {
|
||||
switch family {
|
||||
case "pdf":
|
||||
return utility.FileTypePDF
|
||||
case "doc":
|
||||
return utility.FileTypeDOC
|
||||
case "docx":
|
||||
return utility.FileTypeDOCX
|
||||
case "slides":
|
||||
return utility.FileTypePPTX // pptx arm picks the slide parser
|
||||
case "spreadsheet":
|
||||
return utility.FileTypeXLSX
|
||||
case "html":
|
||||
return utility.FileTypeHTML
|
||||
case "markdown":
|
||||
return utility.FileTypeMarkdown
|
||||
case "text&code":
|
||||
return utility.FileTypeTXT
|
||||
}
|
||||
return utility.FileTypeOTHER
|
||||
}
|
||||
|
||||
// pythonFamilyName normalises a free-form file-type hint to the
|
||||
// python family identifier used by schema.ParserParam.Setups.
|
||||
// Returns "" when the hint is unknown.
|
||||
func pythonFamilyName(raw string) string {
|
||||
switch raw {
|
||||
case "pdf":
|
||||
return "pdf"
|
||||
case "doc":
|
||||
return "doc"
|
||||
case "docx":
|
||||
return "docx"
|
||||
case "ppt":
|
||||
return "slides"
|
||||
case "pptx":
|
||||
return "slides"
|
||||
case "xls":
|
||||
return "spreadsheet"
|
||||
case "xlsx":
|
||||
return "spreadsheet"
|
||||
case "csv":
|
||||
return "spreadsheet"
|
||||
case "html", "htm":
|
||||
return "html"
|
||||
case "md", "markdown", "mdx":
|
||||
return "markdown"
|
||||
case "jpg", "jpeg", "png", "gif":
|
||||
return "image"
|
||||
case "eml", "msg":
|
||||
return "email"
|
||||
case "epub":
|
||||
return "epub"
|
||||
case "txt", "py", "js", "java", "c", "cpp", "h", "php",
|
||||
"go", "ts", "sh", "cs", "kt", "sql":
|
||||
return "text&code"
|
||||
case "mp4", "avi", "mkv":
|
||||
return "video"
|
||||
case "wav", "mp3", "aac", "flac", "ogg":
|
||||
return "audio"
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// jsonItemsToPages reshapes a parsed JSON payload into the
|
||||
// schema.Page layout the chunker side consumes. Each JSON item
|
||||
// becomes one page carrying `text`, `doc_type_kwd`, and any
|
||||
// ck_type / image / positions fields the parser attached. The page
|
||||
// number is assigned sequentially so the deterministic-merge
|
||||
// contract in plan §8 R8 holds.
|
||||
func jsonItemsToPages(items []map[string]any) []schema.Page {
|
||||
out := make([]schema.Page, 0, len(items))
|
||||
for i, it := range items {
|
||||
page := schema.Page{}
|
||||
maps.Copy(page, it)
|
||||
// page_number anchors the deterministic-merge sort; the
|
||||
// parser-supplied number (if any) takes precedence so
|
||||
// PDF / DOCX outputs that already carry `page_number`
|
||||
// survive the round-trip.
|
||||
if _, ok := page["page_number"]; !ok {
|
||||
page["page_number"] = i
|
||||
}
|
||||
if _, ok := page["doc_type_kwd"]; !ok {
|
||||
page["doc_type_kwd"] = "text"
|
||||
}
|
||||
out = append(out, page)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// pagesFromDispatch extracts the per-page bytes from a parsed
|
||||
// schema.Page slice so the existing fan-out / merge path can run
|
||||
// unchanged. Pages without a `text` field emit an empty buffer
|
||||
// (the merge step treats them as zero-length pages).
|
||||
func pagesFromDispatch(pages []schema.Page) [][]byte {
|
||||
out := make([][]byte, 0, len(pages))
|
||||
for _, p := range pages {
|
||||
var buf []byte
|
||||
if s, ok := p["text"].(string); ok {
|
||||
buf = []byte(s)
|
||||
}
|
||||
out = append(out, buf)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// buildParserOutputs assembles the runtime output map from the
|
||||
// merged pages slice AND the dispatch result (when the dispatch
|
||||
// succeeded). The output shape:
|
||||
//
|
||||
// - pages []schema.Page — sorted by PageNumber
|
||||
// - name string — from the upstream file/document name
|
||||
// (or doc_id when no filename is available)
|
||||
// - output_format string — the dispatch's OutputFormat,
|
||||
// or "text" for the raw-text
|
||||
// fallback
|
||||
// - json | markdown | text | html — the dispatched payload on
|
||||
// the matching family key (only
|
||||
// populated on a structured
|
||||
// dispatch)
|
||||
// - file map[string]any — the parser-enriched file
|
||||
// metadata, when present
|
||||
//
|
||||
// This mirrors the Python Parser component's `set_output()` calls
|
||||
// at rag/flow/parser/parser.py:_invoke — the downstream chunker
|
||||
// / tokenizer / extractor components read the matching family
|
||||
// key, with "pages" as the universal fallback shape.
|
||||
func buildParserOutputs(parsed []schema.Page, dispatched parserDispatchResult, name string, fileType utility.FileType) map[string]any {
|
||||
out := map[string]any{
|
||||
"pages": toAnyPages(parsed),
|
||||
"name": name,
|
||||
}
|
||||
if dispatched.Err == nil && dispatched.OutputFormat != "" {
|
||||
out["output_format"] = dispatched.OutputFormat
|
||||
switch dispatched.OutputFormat {
|
||||
case "json":
|
||||
out["json"] = dispatched.JSON
|
||||
case "markdown":
|
||||
out["markdown"] = dispatched.Markdown
|
||||
case "html":
|
||||
out["html"] = dispatched.HTML
|
||||
case "text":
|
||||
out["text"] = dispatched.Text
|
||||
}
|
||||
if dispatched.File != nil {
|
||||
out["file"] = dispatched.File
|
||||
}
|
||||
return out
|
||||
}
|
||||
// Raw-text fallback path: emit output_format = "text" so a
|
||||
// chunker branching on the format key still sees a sane value.
|
||||
out["output_format"] = "text"
|
||||
_ = fileType // reserved for a future "raw_text per-family" extension
|
||||
return out
|
||||
}
|
||||
|
||||
func hydrateEmptyDispatchPayload(dispatched parserDispatchResult, binary []byte) parserDispatchResult {
|
||||
if dispatched.Err != nil || len(binary) == 0 {
|
||||
return dispatched
|
||||
}
|
||||
switch dispatched.OutputFormat {
|
||||
case "json":
|
||||
if len(dispatched.JSON) == 0 {
|
||||
dispatched.JSON = pagesToJSONItems(splitIntoPages(binary))
|
||||
}
|
||||
case "text":
|
||||
if dispatched.Text == "" {
|
||||
dispatched.Text = string(binary)
|
||||
}
|
||||
}
|
||||
return dispatched
|
||||
}
|
||||
|
||||
func pagesToJSONItems(pages [][]byte) []map[string]any {
|
||||
if len(pages) == 0 {
|
||||
pages = splitIntoPages(nil)
|
||||
}
|
||||
out := make([]map[string]any, 0, len(pages))
|
||||
for _, page := range pages {
|
||||
text := string(page)
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
out = append(out, map[string]any{
|
||||
"text": text,
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func parserInputName(inputs map[string]any, docID string) string {
|
||||
if inputs != nil {
|
||||
if name, ok := inputs["name"].(string); ok && name != "" {
|
||||
return name
|
||||
}
|
||||
if m, ok := inputs["file"].(map[string]any); ok {
|
||||
if name, ok := m["name"].(string); ok && name != "" {
|
||||
return name
|
||||
}
|
||||
}
|
||||
}
|
||||
return docID
|
||||
}
|
||||
253
internal/ingestion/component/parser_dispatch_test.go
Normal file
253
internal/ingestion/component/parser_dispatch_test.go
Normal file
@@ -0,0 +1,253 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Dispatch tests pin the routing contract:
|
||||
//
|
||||
// - FileTypeOTHER + missing setups → text-page mode.
|
||||
// - FileTypeMarkdown → JSON payload family on the matching output
|
||||
// key, with the pages slice preserved.
|
||||
// - FileTypePDF + setups["pdf"].output_format set to a value not
|
||||
// in allowed_output_format["pdf"] → component errors with the
|
||||
// format-mismatch message (matches the Python check() behavior).
|
||||
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
)
|
||||
|
||||
// TestDispatch_OutputFormatValidation_Allowed is the happy-path
|
||||
// pin: a Markdown file with output_format=json passes the
|
||||
// allowed_output_format check and runs the structured dispatch.
|
||||
func TestDispatch_OutputFormatValidation_Allowed(t *testing.T) {
|
||||
param := schema.ParserParam{}.Defaults()
|
||||
// Defaults already include markdown → {text, json}.
|
||||
c := &ParserComponent{Param: param}
|
||||
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"binary": []byte("# Title\n\nbody\n"),
|
||||
"doc_id": "doc.md",
|
||||
"file_type": "md",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got, want := out["output_format"], "json"; got != want {
|
||||
t.Errorf("output_format = %v, want %v", got, want)
|
||||
}
|
||||
jsonItems, ok := out["json"].([]map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("json payload missing or wrong type: %T", out["json"])
|
||||
}
|
||||
if len(jsonItems) == 0 {
|
||||
t.Errorf("json payload empty; want at least 1 item")
|
||||
}
|
||||
// Pages must still exist for chunker-side consumers.
|
||||
pages, ok := out["pages"].([]schema.Page)
|
||||
if !ok || len(pages) == 0 {
|
||||
t.Errorf("pages slice missing or empty: %T", out["pages"])
|
||||
}
|
||||
// File metadata is carried through dispatch.
|
||||
if fm, ok := out["file"].(map[string]any); !ok || fm["name"] != "doc.md" {
|
||||
t.Errorf("file metadata missing or wrong: %+v", out["file"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestDispatch_OutputFormatValidation_Rejection pins the
|
||||
// whitelist enforcement: a request for output_format=html on the
|
||||
// markdown family is rejected because markdown's allowed list is
|
||||
// {text, json}. The component must surface this as a hard error
|
||||
// before any fallback so a misconfigured template cannot silently
|
||||
// degrade.
|
||||
func TestDispatch_OutputFormatValidation_Rejection(t *testing.T) {
|
||||
param := schema.ParserParam{}.Defaults()
|
||||
// Override the markdown setup to ask for an unsupported format.
|
||||
// The key is "markdown" (the python-side family identifier),
|
||||
// NOT "md" — utility.FileTypeMarkdown happens to be the string
|
||||
// "md" but the setup key is the family name. resolveOutputFormat
|
||||
// looks up setups[string(fileType)], so the fileType passed in
|
||||
// here must match the setup key.
|
||||
param.Setups["markdown"] = schema.ParserSetup{"output_format": "html"}
|
||||
// inputs["file_type"] must also be "markdown" so fileTypeFromInputs
|
||||
// returns a FileType whose string form matches the setup key.
|
||||
c := &ParserComponent{Param: param}
|
||||
|
||||
_, err := c.Invoke(context.Background(), map[string]any{
|
||||
"binary": []byte("# Title\n"),
|
||||
"file_type": "md",
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("Invoke: want error for unsupported output_format, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "output_format") {
|
||||
t.Errorf("error %q must mention output_format", err.Error())
|
||||
}
|
||||
if !strings.Contains(err.Error(), "markdown") && !strings.Contains(err.Error(), "md") {
|
||||
t.Errorf("error %q must mention the family", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
// TestDispatch_TextPageMode_NoFileType pins the no-dispatch
|
||||
// path. When the upstream inputs supply neither file_type nor
|
||||
// file.name, the component degrades to text-page mode and
|
||||
// emits output_format=text. This is the documented behavior for
|
||||
// canvas-bound invocations that wire the binary directly without
|
||||
// a family hint.
|
||||
func TestDispatch_TextPageMode_NoFileType(t *testing.T) {
|
||||
param := schema.ParserParam{}.Defaults()
|
||||
c := &ParserComponent{Param: param}
|
||||
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"binary": []byte("plain content\n"),
|
||||
"doc_id": "unknown",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got, want := out["output_format"], "text"; got != want {
|
||||
t.Errorf("output_format = %v, want %v (text-page mode)", got, want)
|
||||
}
|
||||
pages, ok := out["pages"].([]schema.Page)
|
||||
if !ok || len(pages) == 0 {
|
||||
t.Fatalf("pages slice missing or empty: %T", out["pages"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestDispatch_TextPageMode_PDFInput pins the current text-page
|
||||
// behavior when PDF bytes are not successfully parsed into a
|
||||
// structured payload in this test environment.
|
||||
func TestDispatch_TextPageMode_PDFInput(t *testing.T) {
|
||||
param := schema.ParserParam{}.Defaults()
|
||||
c := &ParserComponent{Param: param}
|
||||
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"binary": []byte("PDF payload as bytes (not a real PDF — stub test)\n"),
|
||||
"file_type": "pdf",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got, want := out["output_format"], "text"; got != want {
|
||||
t.Errorf("output_format = %v, want %v (PDF input stayed in text-page mode)", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFileTypeFromInputs_ResolutionOrder pins the precedence
|
||||
// rules documented on parser_dispatch.go:fileTypeFromInputs:
|
||||
//
|
||||
// 1. inputs["file_type"] (explicit family hint)
|
||||
// 2. inputs["file"].name (filename in the file descriptor)
|
||||
// 3. inputs["name"] (last-resort filename)
|
||||
// 4. FileTypeOTHER (text-page mode)
|
||||
func TestFileTypeFromInputs_ResolutionOrder(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
in map[string]any
|
||||
want string
|
||||
}{
|
||||
{"explicit pdf", map[string]any{"file_type": "pdf"}, "pdf"},
|
||||
{"explicit markdown (family form)", map[string]any{"file_type": "markdown"}, "md"},
|
||||
{"file.name docx", map[string]any{"file": map[string]any{"name": "report.docx"}}, "docx"},
|
||||
{"name fallback md", map[string]any{"name": "notes.md"}, "md"},
|
||||
{"unrelated inputs", map[string]any{"binary": []byte("x"), "doc_id": "abc"}, "other"},
|
||||
{"unknown family hint", map[string]any{"file_type": "image/xyz"}, "other"},
|
||||
{"nil inputs", nil, "other"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := fileTypeFromInputs(tc.in)
|
||||
if string(got) != tc.want {
|
||||
t.Errorf("fileTypeFromInputs(%v) = %q, want %q", tc.in, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveOutputFormat_DefaultsAndWhitelist pins the two-layer
|
||||
// behavior of resolveOutputFormat: it returns the setup's
|
||||
// output_format when present (or "text" when absent), and
|
||||
// rejects values not in the allowed_output_format list.
|
||||
func TestResolveOutputFormat_DefaultsAndWhitelist(t *testing.T) {
|
||||
allowed := map[string][]string{
|
||||
"pdf": {"json", "markdown"},
|
||||
"markdown": {"text", "json"},
|
||||
}
|
||||
cases := []struct {
|
||||
name string
|
||||
setups map[string]schema.ParserSetup
|
||||
family string
|
||||
want string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "no setup → empty (text-page mode)",
|
||||
setups: nil,
|
||||
family: "pdf",
|
||||
want: "",
|
||||
},
|
||||
{
|
||||
name: "setup with output_format=json → json",
|
||||
setups: map[string]schema.ParserSetup{"pdf": {"output_format": "json"}},
|
||||
family: "pdf",
|
||||
want: "json",
|
||||
},
|
||||
{
|
||||
name: "setup with output_format=markdown → markdown",
|
||||
setups: map[string]schema.ParserSetup{"pdf": {"output_format": "markdown"}},
|
||||
family: "pdf",
|
||||
want: "markdown",
|
||||
},
|
||||
{
|
||||
name: "setup without output_format → default text",
|
||||
setups: map[string]schema.ParserSetup{"markdown": {}},
|
||||
family: "markdown",
|
||||
want: "text",
|
||||
},
|
||||
{
|
||||
name: "pdf asking for html (not allowed) → reject",
|
||||
setups: map[string]schema.ParserSetup{"pdf": {"output_format": "html"}},
|
||||
family: "pdf",
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "family with no whitelist → accept setup value",
|
||||
setups: map[string]schema.ParserSetup{"video": {"output_format": "json"}},
|
||||
family: "video",
|
||||
want: "json",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got, err := resolveOutputFormat(tc.family, tc.setups, allowed)
|
||||
if tc.wantErr {
|
||||
if err == nil {
|
||||
t.Fatalf("want error, got nil (value=%q)", got)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got != tc.want {
|
||||
t.Errorf("got %q, want %q", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
452
internal/ingestion/component/parser_test.go
Normal file
452
internal/ingestion/component/parser_test.go
Normal file
@@ -0,0 +1,452 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/entity"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
)
|
||||
|
||||
// TestParserComponent_Registered asserts the factory lookup
|
||||
// succeeds for the canonical "Parser" name. This is the contract
|
||||
// the pipeline runner relies on (see plan §4 Phase 0, "category-
|
||||
// aware registry"). A regression here would mean the component
|
||||
// failed to register in init().
|
||||
func TestParserComponent_Registered(t *testing.T) {
|
||||
factory, category, meta, ok := runtime.DefaultRegistry.Lookup("Parser")
|
||||
if !ok {
|
||||
t.Fatalf("Parser not registered in DefaultRegistry")
|
||||
}
|
||||
if category != runtime.CategoryIngestion {
|
||||
t.Errorf("Parser category = %q, want %q", category, runtime.CategoryIngestion)
|
||||
}
|
||||
if factory == nil {
|
||||
t.Fatalf("Parser factory is nil")
|
||||
}
|
||||
if len(meta.Inputs) == 0 {
|
||||
t.Errorf("Parser Metadata.Inputs is empty")
|
||||
}
|
||||
if len(meta.Outputs) == 0 {
|
||||
t.Errorf("Parser Metadata.Outputs is empty")
|
||||
}
|
||||
}
|
||||
|
||||
// TestParserComponent_InputsOutputs_NonEmpty covers the static
|
||||
// input/output descriptors — the API layer enumerates these to
|
||||
// build the component catalog, so an empty descriptor would
|
||||
// hide the component from the UI.
|
||||
func TestParserComponent_InputsOutputs_NonEmpty(t *testing.T) {
|
||||
c := &ParserComponent{}
|
||||
in := c.Inputs()
|
||||
out := c.Outputs()
|
||||
if len(in) == 0 {
|
||||
t.Errorf("Inputs() returned empty map")
|
||||
}
|
||||
if len(out) == 0 {
|
||||
t.Errorf("Outputs() returned empty map")
|
||||
}
|
||||
// The contract from the file header: at least "binary" in,
|
||||
// "pages" out. Anything else is informational.
|
||||
if _, ok := in["binary"]; !ok {
|
||||
t.Errorf("Inputs() missing key %q", "binary")
|
||||
}
|
||||
if _, ok := out["pages"]; !ok {
|
||||
t.Errorf("Outputs() missing key %q", "pages")
|
||||
}
|
||||
}
|
||||
|
||||
// TestParserComponent_Parallelism locks the fan-out degree to
|
||||
// the plan §2 AD-5a value (4). Changing this would silently
|
||||
// re-tune resource consumption and break the "fan-out" claim.
|
||||
func TestParserComponent_Parallelism(t *testing.T) {
|
||||
c := &ParserComponent{}
|
||||
if got := c.Parallelism(); got != 4 {
|
||||
t.Errorf("Parallelism() = %d, want 4", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestParserComponent_Invoke_TextInput covers the happy path:
|
||||
// UTF-8 text input, no form-feeds, default page_size. The
|
||||
// component must emit exactly one page carrying the full text
|
||||
// under "text", plus the timing stamps.
|
||||
func TestParserComponent_Invoke_TextInput(t *testing.T) {
|
||||
c := &ParserComponent{Param: schema.ParserParam{}.Defaults()}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"binary": "hello world",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
pages, ok := out["pages"].([]schema.Page)
|
||||
if !ok {
|
||||
t.Fatalf("pages: got %T, want []schema.Page", out["pages"])
|
||||
}
|
||||
if len(pages) != 1 {
|
||||
t.Fatalf("pages len = %d, want 1", len(pages))
|
||||
}
|
||||
if got := pages[0]["text"]; got != "hello world" {
|
||||
t.Errorf("pages[0][text] = %q, want %q", got, "hello world")
|
||||
}
|
||||
if _, ok := out["_created_time"]; !ok {
|
||||
t.Errorf("_created_time missing from output")
|
||||
}
|
||||
if _, ok := out["_elapsed_time"]; !ok {
|
||||
t.Errorf("_elapsed_time missing from output")
|
||||
}
|
||||
}
|
||||
|
||||
// TestParserComponent_Invoke_PageRangeFilter asserts that
|
||||
// form-feed boundaries are honored: "A\fB\fC" yields three
|
||||
// pages, in input order, with text intact.
|
||||
func TestParserComponent_Invoke_PageRangeFilter(t *testing.T) {
|
||||
c := &ParserComponent{Param: schema.ParserParam{}.Defaults()}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"binary": "pageA\fpageB\fpageC",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
pages, ok := out["pages"].([]schema.Page)
|
||||
if !ok {
|
||||
t.Fatalf("pages: got %T, want []schema.Page", out["pages"])
|
||||
}
|
||||
if len(pages) != 3 {
|
||||
t.Fatalf("pages len = %d, want 3", len(pages))
|
||||
}
|
||||
want := []string{"pageA", "pageB", "pageC"}
|
||||
for i, p := range pages {
|
||||
if got := p["text"]; got != want[i] {
|
||||
t.Errorf("pages[%d][text] = %q, want %q", i, got, want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestParserComponent_Invoke_DeterministicMerge is the
|
||||
// golden-file test for plan §8 R8 (DETERMINISTIC MERGE).
|
||||
//
|
||||
// We invoke the component 5 times with identical input and
|
||||
// assert byte-for-byte equality of the JSON-encoded output.
|
||||
// The test is expected to pass under `go test -count=10 -race`
|
||||
// — that flag is run separately in the verification block.
|
||||
//
|
||||
// We can't tag-stamp the page text (no page_number key in
|
||||
// text-page mode) so we tag the input with explicit
|
||||
// page numbers and verify the SORTED output is stable across
|
||||
// runs. The fan-out goroutines therefore produce pages in
|
||||
// different physical orders, but the deterministic sort
|
||||
// rescues byte-equality.
|
||||
func TestParserComponent_Invoke_DeterministicMerge(t *testing.T) {
|
||||
c := &ParserComponent{Param: schema.ParserParam{}.Defaults()}
|
||||
// 8 form-feed-separated pages — enough to exercise the
|
||||
// ceil(8/4)=2 page_size default and force multiple
|
||||
// goroutines to interleave.
|
||||
input := "p1\fp2\fp3\fp4\fp5\fp6\fp7\fp8"
|
||||
|
||||
// First call: produce the canonical bytes.
|
||||
first, err := c.Invoke(context.Background(), map[string]any{
|
||||
"binary": input,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke (first): %v", err)
|
||||
}
|
||||
canonical, err := json.Marshal(first["pages"])
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal canonical: %v", err)
|
||||
}
|
||||
|
||||
// Subsequent calls: must produce the same bytes.
|
||||
for i := 0; i < 5; i++ {
|
||||
got, err := c.Invoke(context.Background(), map[string]any{
|
||||
"binary": input,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke (run %d): %v", i, err)
|
||||
}
|
||||
encoded, err := json.Marshal(got["pages"])
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal run %d: %v", i, err)
|
||||
}
|
||||
if string(encoded) != string(canonical) {
|
||||
t.Errorf("run %d output differs from canonical:\n got=%s\nwant=%s",
|
||||
i, encoded, canonical)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestParserComponent_Invoke_RespectsTimeout exercises the
|
||||
// cancellation path. We pass a pre-cancelled parent context
|
||||
// (no timeout) so the errgroup's WithContext branches see
|
||||
// ctx.Done() on the first fan-out — the same code path that
|
||||
// a context.WithTimeout parent takes when its deadline
|
||||
// elapses. The pre-cancelled form is deterministic across
|
||||
// hardware speeds; a 1ms deadline raced on slow CI.
|
||||
//
|
||||
// The component does NOT expose a "Parallelism" override
|
||||
// hook (Parallelism() is fixed at 4). Using a pre-cancelled
|
||||
// parent also verifies that a parent cancellation cascades
|
||||
// through the errgroup → siblings see ctx.Done() and
|
||||
// return ctx.Err().
|
||||
func TestParserComponent_Invoke_RespectsTimeout(t *testing.T) {
|
||||
c := &ParserComponent{Param: schema.ParserParam{}.Defaults()}
|
||||
// Build a 400-page input so the fan-out has more than
|
||||
// one batch to dispatch.
|
||||
var b strings.Builder
|
||||
for i := 0; i < 400; i++ {
|
||||
if i > 0 {
|
||||
b.WriteByte(pageFormFeed)
|
||||
}
|
||||
b.WriteString("page")
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel() // pre-cancel
|
||||
|
||||
_, err := c.Invoke(ctx, map[string]any{
|
||||
"binary": b.String(),
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatalf("Invoke returned nil error; expected a cancellation error")
|
||||
}
|
||||
// The error must reference cancellation — not an
|
||||
// unrelated "parse error" — so callers can distinguish
|
||||
// "we were cancelled" from "the input was malformed".
|
||||
if !isTimeoutish(err) {
|
||||
t.Errorf("Invoke error = %v, want a timeout/cancel error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// isTimeoutish returns true if err is a context.DeadlineExceeded
|
||||
// (directly or wrapped) or otherwise carries a "deadline
|
||||
// exceeded" / "context deadline" / "context canceled" substring.
|
||||
func isTimeoutish(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
msg := err.Error()
|
||||
if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
|
||||
return true
|
||||
}
|
||||
return strings.Contains(msg, "deadline exceeded") ||
|
||||
strings.Contains(msg, "context deadline") ||
|
||||
strings.Contains(msg, "context canceled")
|
||||
}
|
||||
|
||||
// TestParserComponent_New_Defaults constructs a Parser from a
|
||||
// nil param map and verifies the static Param is the
|
||||
// Defaults() value (i.e., NewParserComponent does not mutate
|
||||
// the schema default).
|
||||
func TestParserComponent_New_Defaults(t *testing.T) {
|
||||
c, err := NewParserComponent(nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NewParserComponent(nil): %v", err)
|
||||
}
|
||||
pc, ok := c.(*ParserComponent)
|
||||
if !ok {
|
||||
t.Fatalf("NewParserComponent returned %T, want *ParserComponent", c)
|
||||
}
|
||||
defaults := schema.ParserParam{}.Defaults()
|
||||
if !reflect.DeepEqual(pc.Param, defaults) {
|
||||
t.Errorf("Param differs from Defaults:\n got=%+v\nwant=%+v", pc.Param, defaults)
|
||||
}
|
||||
}
|
||||
|
||||
// TestParserComponent_New_Overrides verifies that a non-nil
|
||||
// param map with a "setups" entry is layered on top of the
|
||||
// defaults.
|
||||
func TestParserComponent_New_Overrides(t *testing.T) {
|
||||
c, err := NewParserComponent(map[string]any{
|
||||
"setups": map[string]any{
|
||||
"text&code": map[string]any{
|
||||
"chunk_token_num": 256,
|
||||
},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewParserComponent: %v", err)
|
||||
}
|
||||
pc, ok := c.(*ParserComponent)
|
||||
if !ok {
|
||||
t.Fatalf("NewParserComponent returned %T", c)
|
||||
}
|
||||
setup, ok := pc.Param.Setups["text&code"]
|
||||
if !ok {
|
||||
t.Fatalf("Setups[text&code] missing after override")
|
||||
}
|
||||
if got, _ := setup["chunk_token_num"].(int); got != 256 {
|
||||
t.Errorf("Setups[text&code][chunk_token_num] = %v, want 256", setup["chunk_token_num"])
|
||||
}
|
||||
// Defaults must still be present for other file types.
|
||||
if _, ok := pc.Param.Setups["pdf"]; !ok {
|
||||
t.Errorf("Setups[pdf] missing; override should not erase defaults")
|
||||
}
|
||||
}
|
||||
|
||||
// TestParserComponent_Invoke_DocIDCarried asserts the optional
|
||||
// doc_id input flows through to the "name" output.
|
||||
func TestParserComponent_Invoke_DocIDCarried(t *testing.T) {
|
||||
c := &ParserComponent{Param: schema.ParserParam{}.Defaults()}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"binary": "x",
|
||||
"doc_id": "doc-123",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got, _ := out["name"].(string); got != "doc-123" {
|
||||
t.Errorf("name = %q, want %q", got, "doc-123")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParserComponent_Invoke_ResolvesBinaryFromDocID(t *testing.T) {
|
||||
ms := withMemoryStorage(t)
|
||||
db := withFileComponentTestDB(t)
|
||||
location := "docs/from-parser.txt"
|
||||
if err := ms.Put("kb-parser", location, []byte("alpha\fbeta")); err != nil {
|
||||
t.Fatalf("seed storage: %v", err)
|
||||
}
|
||||
docName := "parser.txt"
|
||||
if err := db.Create(&entity.Document{
|
||||
ID: "doc-parser",
|
||||
KbID: "kb-parser",
|
||||
ParserID: "na",
|
||||
ParserConfig: entity.JSONMap{},
|
||||
Type: "txt",
|
||||
CreatedBy: "u1",
|
||||
Name: &docName,
|
||||
Location: &location,
|
||||
Suffix: ".txt",
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("seed doc: %v", err)
|
||||
}
|
||||
|
||||
c := &ParserComponent{Param: schema.ParserParam{}.Defaults()}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{"doc_id": "doc-parser"})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
pages, ok := out["pages"].([]schema.Page)
|
||||
if !ok || len(pages) != 2 {
|
||||
t.Fatalf("pages = %T/%v, want 2 schema.Page entries", out["pages"], out["pages"])
|
||||
}
|
||||
if pages[0]["text"] != "alpha" || pages[1]["text"] != "beta" {
|
||||
t.Fatalf("pages = %+v, want [alpha beta]", pages)
|
||||
}
|
||||
if got, _ := out["name"].(string); got != "doc-parser" {
|
||||
t.Fatalf("name = %q, want %q", got, "doc-parser")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParserComponent_Invoke_ResolvesBinaryFromBucketPath(t *testing.T) {
|
||||
ms := withMemoryStorage(t)
|
||||
if err := ms.Put("bucket-1", "docs/explicit.txt", []byte("bucket content")); err != nil {
|
||||
t.Fatalf("seed storage: %v", err)
|
||||
}
|
||||
|
||||
c := &ParserComponent{Param: schema.ParserParam{}.Defaults()}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"bucket": "bucket-1",
|
||||
"path": "docs/explicit.txt",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
pages, ok := out["pages"].([]schema.Page)
|
||||
if !ok || len(pages) != 1 {
|
||||
t.Fatalf("pages = %T/%v, want 1 schema.Page entry", out["pages"], out["pages"])
|
||||
}
|
||||
if got := pages[0]["text"]; got != "bucket content" {
|
||||
t.Fatalf("pages[0][text] = %q, want %q", got, "bucket content")
|
||||
}
|
||||
}
|
||||
|
||||
// TestParserComponent_Invoke_RejectsInvalidUTF8 covers the
|
||||
// safety check: a "binary" string that is not valid UTF-8 is
|
||||
// rejected (per the file header — base64-encoded input would
|
||||
// look like this if a caller mistakenly handed a base64 string
|
||||
// without decoding it).
|
||||
func TestParserComponent_Invoke_RejectsInvalidUTF8(t *testing.T) {
|
||||
c := &ParserComponent{Param: schema.ParserParam{}.Defaults()}
|
||||
_, err := c.Invoke(context.Background(), map[string]any{
|
||||
// 0xFF alone is not valid UTF-8 start byte.
|
||||
"binary": string([]byte{0xFF, 0xFE, 0xFD}),
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatalf("Invoke: expected an error for invalid UTF-8, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "UTF-8") {
|
||||
t.Errorf("error = %v, want it to mention UTF-8", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestParserComponent_Invoke_AcceptsBytes covers the in-process
|
||||
// caller's normal form ([]byte) — the alternative to a UTF-8
|
||||
// string.
|
||||
func TestParserComponent_Invoke_AcceptsBytes(t *testing.T) {
|
||||
c := &ParserComponent{Param: schema.ParserParam{}.Defaults()}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"binary": []byte("alpha\fbeta"),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
pages, ok := out["pages"].([]schema.Page)
|
||||
if !ok {
|
||||
t.Fatalf("pages: got %T", out["pages"])
|
||||
}
|
||||
if len(pages) != 2 {
|
||||
t.Fatalf("pages len = %d, want 2", len(pages))
|
||||
}
|
||||
if pages[0]["text"] != "alpha" || pages[1]["text"] != "beta" {
|
||||
t.Errorf("pages = %+v, want [alpha beta]", pages)
|
||||
}
|
||||
}
|
||||
|
||||
// TestParserComponent_Invoke_PageSizeHint covers the optional
|
||||
// page_size input. A page_size of 2 with 6 pages yields 3
|
||||
// batches; the deterministic merge still yields 6 pages in
|
||||
// input order.
|
||||
func TestParserComponent_Invoke_PageSizeHint(t *testing.T) {
|
||||
c := &ParserComponent{Param: schema.ParserParam{}.Defaults()}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"binary": "p1\fp2\fp3\fp4\fp5\fp6",
|
||||
"page_size": 2,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
pages, ok := out["pages"].([]schema.Page)
|
||||
if !ok {
|
||||
t.Fatalf("pages: got %T", out["pages"])
|
||||
}
|
||||
if len(pages) != 6 {
|
||||
t.Fatalf("pages len = %d, want 6", len(pages))
|
||||
}
|
||||
for i, p := range pages {
|
||||
want := "p" + string(rune('1'+i))
|
||||
if got := p["text"]; got != want {
|
||||
t.Errorf("pages[%d] = %q, want %q", i, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
332
internal/ingestion/component/schema/chunk_types.go
Normal file
332
internal/ingestion/component/schema/chunk_types.go
Normal file
@@ -0,0 +1,332 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package schema
|
||||
|
||||
import "encoding/json"
|
||||
|
||||
// PayloadFormat is the discriminator shared by parser/chunker/tokenizer
|
||||
// wire payloads.
|
||||
type PayloadFormat string
|
||||
|
||||
const (
|
||||
PayloadFormatJSON PayloadFormat = "json"
|
||||
PayloadFormatMarkdown PayloadFormat = "markdown"
|
||||
PayloadFormatText PayloadFormat = "text"
|
||||
PayloadFormatHTML PayloadFormat = "html"
|
||||
PayloadFormatChunks PayloadFormat = "chunks"
|
||||
)
|
||||
|
||||
func (f PayloadFormat) isKnown() bool {
|
||||
switch f {
|
||||
case "", PayloadFormatJSON, PayloadFormatMarkdown, PayloadFormatText, PayloadFormatHTML, PayloadFormatChunks:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// ChunkerFileMeta is the common file descriptor shape carried through
|
||||
// parser/chunker/tokenizer boundaries. Known fields are modeled
|
||||
// explicitly; any parser-specific enrichments stay in Extra.
|
||||
type ChunkerFileMeta struct {
|
||||
Name string `json:"name,omitempty"`
|
||||
Path string `json:"path,omitempty"`
|
||||
Bucket string `json:"bucket,omitempty"`
|
||||
Binary string `json:"binary,omitempty"`
|
||||
Blob string `json:"blob,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
PageCount *int `json:"page_count,omitempty"`
|
||||
Extra map[string]json.RawMessage `json:"-"`
|
||||
}
|
||||
|
||||
func (m *ChunkerFileMeta) UnmarshalJSON(data []byte) error {
|
||||
type alias ChunkerFileMeta
|
||||
var decoded alias
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return err
|
||||
}
|
||||
var raw map[string]json.RawMessage
|
||||
if err := json.Unmarshal(data, &raw); err != nil {
|
||||
return err
|
||||
}
|
||||
delete(raw, "name")
|
||||
delete(raw, "path")
|
||||
delete(raw, "bucket")
|
||||
delete(raw, "binary")
|
||||
delete(raw, "blob")
|
||||
delete(raw, "id")
|
||||
delete(raw, "page_count")
|
||||
*m = ChunkerFileMeta(decoded)
|
||||
if len(raw) > 0 {
|
||||
m.Extra = raw
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m ChunkerFileMeta) MarshalJSON() ([]byte, error) {
|
||||
type alias ChunkerFileMeta
|
||||
base, err := json.Marshal(alias(m))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var raw map[string]json.RawMessage
|
||||
if err := json.Unmarshal(base, &raw); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for k, v := range m.Extra {
|
||||
if _, exists := raw[k]; !exists {
|
||||
raw[k] = v
|
||||
}
|
||||
}
|
||||
return json.Marshal(raw)
|
||||
}
|
||||
|
||||
// ChunkDoc is the typed chunk item shared by chunker/tokenizer
|
||||
// boundaries. Common fields are explicit; dynamic enrichments are
|
||||
// preserved in Extra for forward compatibility.
|
||||
type ChunkDoc struct {
|
||||
Text string `json:"text,omitempty"`
|
||||
ContentWithWeight string `json:"content_with_weight,omitempty"`
|
||||
DocType string `json:"doc_type_kwd,omitempty"`
|
||||
CKType string `json:"ck_type,omitempty"`
|
||||
TKNums *int `json:"tk_nums,omitempty"`
|
||||
Mom string `json:"mom,omitempty"`
|
||||
ImgID string `json:"img_id,omitempty"`
|
||||
Layout string `json:"layout,omitempty"`
|
||||
LayoutType string `json:"layout_type,omitempty"`
|
||||
LayoutNo string `json:"layoutno,omitempty"`
|
||||
Image string `json:"image,omitempty"`
|
||||
ContextAbove string `json:"context_above,omitempty"`
|
||||
ContextBelow string `json:"context_below,omitempty"`
|
||||
Questions string `json:"questions,omitempty"`
|
||||
Keywords string `json:"keywords,omitempty"`
|
||||
Summary string `json:"summary,omitempty"`
|
||||
ChunkOrderInt *int `json:"chunk_order_int,omitempty"`
|
||||
TitleTks string `json:"title_tks,omitempty"`
|
||||
TitleSmTks string `json:"title_sm_tks,omitempty"`
|
||||
ContentLtks string `json:"content_ltks,omitempty"`
|
||||
ContentSmLtks string `json:"content_sm_ltks,omitempty"`
|
||||
PageNumber *int `json:"page_number,omitempty"`
|
||||
PDFPositions json.RawMessage `json:"_pdf_positions,omitempty"`
|
||||
Positions json.RawMessage `json:"positions,omitempty"`
|
||||
Extra map[string]json.RawMessage `json:"-"`
|
||||
}
|
||||
|
||||
func (d *ChunkDoc) UnmarshalJSON(data []byte) error {
|
||||
type alias ChunkDoc
|
||||
var decoded alias
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return err
|
||||
}
|
||||
var raw map[string]json.RawMessage
|
||||
if err := json.Unmarshal(data, &raw); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, key := range []string{
|
||||
"text", "content_with_weight", "doc_type_kwd", "mom", "img_id",
|
||||
"ck_type", "tk_nums", "layout", "layout_type", "layoutno", "image",
|
||||
"context_above", "context_below", "questions", "keywords", "summary",
|
||||
"chunk_order_int", "title_tks", "title_sm_tks", "content_ltks",
|
||||
"content_sm_ltks", "page_number", "_pdf_positions", "positions",
|
||||
} {
|
||||
delete(raw, key)
|
||||
}
|
||||
*d = ChunkDoc(decoded)
|
||||
if len(raw) > 0 {
|
||||
d.Extra = raw
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d ChunkDoc) MarshalJSON() ([]byte, error) {
|
||||
type alias ChunkDoc
|
||||
base, err := json.Marshal(alias(d))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var raw map[string]json.RawMessage
|
||||
if err := json.Unmarshal(base, &raw); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for k, v := range d.Extra {
|
||||
if _, exists := raw[k]; !exists {
|
||||
raw[k] = v
|
||||
}
|
||||
}
|
||||
return json.Marshal(raw)
|
||||
}
|
||||
|
||||
// ChunkDocFromMap decodes a free-form map into a typed ChunkDoc.
|
||||
func ChunkDocFromMap(in map[string]any) (ChunkDoc, error) {
|
||||
if in == nil {
|
||||
return ChunkDoc{}, nil
|
||||
}
|
||||
b, err := json.Marshal(in)
|
||||
if err != nil {
|
||||
return ChunkDoc{}, err
|
||||
}
|
||||
var doc ChunkDoc
|
||||
if err := json.Unmarshal(b, &doc); err != nil {
|
||||
return ChunkDoc{}, err
|
||||
}
|
||||
return doc, nil
|
||||
}
|
||||
|
||||
// ChunkDocsFromAny converts either []ChunkDoc, []map[string]any, or []any
|
||||
// into a typed chunk slice.
|
||||
func ChunkDocsFromAny(in any) ([]ChunkDoc, bool, error) {
|
||||
switch t := in.(type) {
|
||||
case []ChunkDoc:
|
||||
return t, true, nil
|
||||
case []map[string]any:
|
||||
out := make([]ChunkDoc, 0, len(t))
|
||||
for _, item := range t {
|
||||
doc, err := ChunkDocFromMap(item)
|
||||
if err != nil {
|
||||
return nil, true, err
|
||||
}
|
||||
out = append(out, doc)
|
||||
}
|
||||
return out, true, nil
|
||||
case []any:
|
||||
out := make([]ChunkDoc, 0, len(t))
|
||||
for _, item := range t {
|
||||
m, ok := item.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
doc, err := ChunkDocFromMap(m)
|
||||
if err != nil {
|
||||
return nil, true, err
|
||||
}
|
||||
out = append(out, doc)
|
||||
}
|
||||
return out, true, nil
|
||||
default:
|
||||
return nil, false, nil
|
||||
}
|
||||
}
|
||||
|
||||
// ToMap converts a typed chunk back to the generic map form used by
|
||||
// runtime.Component boundaries.
|
||||
func (d ChunkDoc) ToMap() map[string]any {
|
||||
b, err := json.Marshal(d)
|
||||
if err != nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
var out map[string]any
|
||||
if err := json.Unmarshal(b, &out); err != nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
for k, raw := range d.Extra {
|
||||
out[k] = decodeExtraValue(raw)
|
||||
}
|
||||
if len(d.PDFPositions) > 0 {
|
||||
out["_pdf_positions"] = decodeStructuredValue(d.PDFPositions)
|
||||
}
|
||||
if len(d.Positions) > 0 {
|
||||
out["positions"] = decodeStructuredValue(d.Positions)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ChunkDocsToMaps converts a typed chunk slice to the generic map form.
|
||||
func ChunkDocsToMaps(in []ChunkDoc) []map[string]any {
|
||||
out := make([]map[string]any, 0, len(in))
|
||||
for _, doc := range in {
|
||||
out = append(out, doc.ToMap())
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (d *ChunkDoc) SetExtraValue(key string, value any) error {
|
||||
if d.Extra == nil {
|
||||
d.Extra = make(map[string]json.RawMessage)
|
||||
}
|
||||
raw, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
d.Extra[key] = raw
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *ChunkDoc) GetExtraString(key string) (string, bool) {
|
||||
if d.Extra == nil {
|
||||
return "", false
|
||||
}
|
||||
raw, ok := d.Extra[key]
|
||||
if !ok {
|
||||
return "", false
|
||||
}
|
||||
var out string
|
||||
if err := json.Unmarshal(raw, &out); err != nil || out == "" {
|
||||
return "", false
|
||||
}
|
||||
return out, true
|
||||
}
|
||||
|
||||
func (d *ChunkDoc) GetExtraStringSlice(key string) ([]string, bool) {
|
||||
if d.Extra == nil {
|
||||
return nil, false
|
||||
}
|
||||
raw, ok := d.Extra[key]
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
var out []string
|
||||
if err := json.Unmarshal(raw, &out); err != nil || len(out) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
return out, true
|
||||
}
|
||||
|
||||
func decodeExtraValue(raw json.RawMessage) any {
|
||||
var floats []float64
|
||||
if err := json.Unmarshal(raw, &floats); err == nil {
|
||||
return floats
|
||||
}
|
||||
var stringsOut []string
|
||||
if err := json.Unmarshal(raw, &stringsOut); err == nil {
|
||||
return stringsOut
|
||||
}
|
||||
var stringOut string
|
||||
if err := json.Unmarshal(raw, &stringOut); err == nil {
|
||||
return stringOut
|
||||
}
|
||||
var generic any
|
||||
if err := json.Unmarshal(raw, &generic); err == nil {
|
||||
return generic
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeStructuredValue(raw json.RawMessage) any {
|
||||
var matrix [][]float64
|
||||
if err := json.Unmarshal(raw, &matrix); err == nil {
|
||||
return matrix
|
||||
}
|
||||
var mapsOut []map[string]any
|
||||
if err := json.Unmarshal(raw, &mapsOut); err == nil {
|
||||
return mapsOut
|
||||
}
|
||||
var generic any
|
||||
if err := json.Unmarshal(raw, &generic); err == nil {
|
||||
return generic
|
||||
}
|
||||
return nil
|
||||
}
|
||||
295
internal/ingestion/component/schema/chunker.go
Normal file
295
internal/ingestion/component/schema/chunker.go
Normal file
@@ -0,0 +1,295 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package schema
|
||||
|
||||
import "fmt"
|
||||
|
||||
// ChunkerFromUpstream is the shared upstream payload consumed by all four
|
||||
// chunker variants (TokenChunker, TitleChunker, GroupTitleChunker,
|
||||
// HierarchyTitleChunker).
|
||||
//
|
||||
// It mirrors the wire shape defined across two equivalent Pydantic
|
||||
// schemas in the Python codebase:
|
||||
//
|
||||
// rag/flow/chunker/schema.py:TokenChunkerFromUpstream
|
||||
// rag/flow/chunker/title_chunker/schema.py:TitleChunkerFromUpstream
|
||||
//
|
||||
// Both classes are field-identical (apart from the class name) — the
|
||||
// variants share one upstream payload; their *Param structs differ.
|
||||
//
|
||||
// Wire shape (Pydantic):
|
||||
//
|
||||
// created_time: float | None (alias _created_time)
|
||||
// elapsed_time: float | None (alias _elapsed_time)
|
||||
// name: str (required)
|
||||
// file: dict | None
|
||||
// chunks: list[dict] | None
|
||||
// output_format: Literal["json","markdown","text","html","chunks"] | None
|
||||
// json_result: list[dict] | None (alias "json")
|
||||
// markdown_result: str | None (alias "markdown")
|
||||
// text_result: str | None (alias "text")
|
||||
// html_result: str | None (alias "html")
|
||||
type ChunkerFromUpstream struct {
|
||||
CreatedTime *float64 `json:"_created_time,omitempty"`
|
||||
ElapsedTime *float64 `json:"_elapsed_time,omitempty"`
|
||||
|
||||
// Name is the source document name. Required.
|
||||
Name string `json:"name"`
|
||||
|
||||
// File is the optional upstream file descriptor.
|
||||
File *ChunkerFileMeta `json:"file,omitempty"`
|
||||
|
||||
// Chunks is the upstream chunk list, set when output_format == "chunks".
|
||||
Chunks []ChunkDoc `json:"chunks,omitempty"`
|
||||
|
||||
// OutputFormat controls which of the *Result fields below is the
|
||||
// active payload. Allowed values:
|
||||
// "json" -> JSONResult
|
||||
// "markdown" -> MarkdownResult
|
||||
// "text" -> TextResult
|
||||
// "html" -> HTMLResult
|
||||
// "chunks" -> Chunks
|
||||
OutputFormat PayloadFormat `json:"output_format,omitempty"`
|
||||
|
||||
// JSONResult is the upstream structured JSON list (alias "json" in
|
||||
// Python). Set when OutputFormat == "json".
|
||||
JSONResult []ChunkDoc `json:"json,omitempty"`
|
||||
|
||||
// MarkdownResult is the upstream markdown payload (alias "markdown").
|
||||
// Set when OutputFormat == "markdown".
|
||||
MarkdownResult *string `json:"markdown,omitempty"`
|
||||
|
||||
// TextResult is the upstream plain-text payload (alias "text").
|
||||
// Set when OutputFormat == "text".
|
||||
TextResult *string `json:"text,omitempty"`
|
||||
|
||||
// HTMLResult is the upstream HTML payload (alias "html").
|
||||
// Set when OutputFormat == "html".
|
||||
HTMLResult *string `json:"html,omitempty"`
|
||||
}
|
||||
|
||||
// Validate enforces the only required field in ChunkerFromUpstream: Name.
|
||||
// OutputFormat is not strictly required (defaults to "" and the
|
||||
// component decides what to do with an empty payload), but the
|
||||
// combination of `OutputFormat == "chunks"` with a non-nil Chunks is the
|
||||
// happy path.
|
||||
func (c *ChunkerFromUpstream) Validate() error {
|
||||
if c.Name == "" {
|
||||
return errRequiredField{Field: "name"}
|
||||
}
|
||||
if !c.OutputFormat.isKnown() {
|
||||
return errInvalidValue{Field: "output_format", Value: string(c.OutputFormat)}
|
||||
}
|
||||
switch c.OutputFormat {
|
||||
case PayloadFormatChunks:
|
||||
return nil
|
||||
case PayloadFormatJSON:
|
||||
if c.JSONResult == nil {
|
||||
return errRequiredField{Field: "json"}
|
||||
}
|
||||
case PayloadFormatMarkdown:
|
||||
if c.MarkdownResult == nil {
|
||||
return errRequiredField{Field: "markdown"}
|
||||
}
|
||||
case PayloadFormatText:
|
||||
if c.TextResult == nil {
|
||||
return errRequiredField{Field: "text"}
|
||||
}
|
||||
case PayloadFormatHTML:
|
||||
if c.HTMLResult == nil {
|
||||
return errRequiredField{Field: "html"}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ChunkerOutputs is the result of invoking any chunker variant. All
|
||||
// four chunker components emit the same shape: a list of chunk maps
|
||||
// under the "chunks" key, plus a marker output_format = "chunks".
|
||||
//
|
||||
// Mirrors what each chunker sets at the end of _invoke:
|
||||
//
|
||||
// self.set_output("output_format", "chunks")
|
||||
// self.set_output("chunks", chunks)
|
||||
type ChunkerOutputs struct {
|
||||
// OutputFormat is always "chunks" on success.
|
||||
OutputFormat PayloadFormat `json:"output_format,omitempty"`
|
||||
|
||||
// Chunks is the produced chunk list. Each entry is a free-form map
|
||||
// mirroring the dict shape the Python code builds (text, doc_type_kwd,
|
||||
// tk_nums, PDF_POSITIONS_KEY, mom, img_id, etc.).
|
||||
Chunks []ChunkDoc `json:"chunks,omitempty"`
|
||||
|
||||
// Error is set when the component short-circuits with an error
|
||||
// message (Python: set_output("_ERROR", ...)).
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TokenChunkerParam
|
||||
// ---------------------------------------------------------------------------
|
||||
//
|
||||
// Mirrors rag/flow/chunker/token_chunker.py:TokenChunkerParam.__init__.
|
||||
// All fields are user-tunable; defaults match the Python values.
|
||||
|
||||
type TokenChunkerParam struct {
|
||||
// DelimiterMode selects the chunking strategy.
|
||||
// Allowed values: "token_size", "delimiter", "one".
|
||||
DelimiterMode string `json:"delimiter_mode"`
|
||||
|
||||
// ChunkTokenSize is the target chunk size in tokens.
|
||||
ChunkTokenSize int `json:"chunk_token_size"`
|
||||
|
||||
// Delimiters is the list of split tokens. Strings wrapped in
|
||||
// backticks (e.g., "`\\n`") denote user-defined regex split points.
|
||||
Delimiters []string `json:"delimiters"`
|
||||
|
||||
// OverlappedPercent is the chunk-overlap ratio in [0, 100).
|
||||
OverlappedPercent float64 `json:"overlapped_percent"`
|
||||
|
||||
// ChildrenDelimiters is the secondary split applied to text chunks.
|
||||
ChildrenDelimiters []string `json:"children_delimiters"`
|
||||
|
||||
// TableContextSize is the number of surrounding tokens to attach
|
||||
// to table chunks. 0 disables.
|
||||
TableContextSize int `json:"table_context_size"`
|
||||
|
||||
// ImageContextSize is the number of surrounding tokens to attach
|
||||
// to image chunks. 0 disables.
|
||||
ImageContextSize int `json:"image_context_size"`
|
||||
}
|
||||
|
||||
// Defaults returns the Python default TokenChunkerParam.
|
||||
func (TokenChunkerParam) Defaults() TokenChunkerParam {
|
||||
return TokenChunkerParam{
|
||||
DelimiterMode: "token_size",
|
||||
ChunkTokenSize: 512,
|
||||
Delimiters: []string{"\n"},
|
||||
OverlappedPercent: 0,
|
||||
ChildrenDelimiters: []string{},
|
||||
TableContextSize: 0,
|
||||
ImageContextSize: 0,
|
||||
}
|
||||
}
|
||||
|
||||
// Validate enforces the same enum/range checks the runtime component expects
|
||||
// at construction time, keeping the schema and component decoder aligned.
|
||||
func (p TokenChunkerParam) Validate() error {
|
||||
switch p.DelimiterMode {
|
||||
case "token_size", "delimiter", "one":
|
||||
default:
|
||||
return errInvalidValue{Field: "delimiter_mode", Value: p.DelimiterMode}
|
||||
}
|
||||
if p.ChunkTokenSize <= 0 {
|
||||
return errInvalidValue{Field: "chunk_token_size", Value: fmt.Sprintf("%d", p.ChunkTokenSize)}
|
||||
}
|
||||
if p.OverlappedPercent < 0 || p.OverlappedPercent >= 1 {
|
||||
return errInvalidValue{Field: "overlapped_percent", Value: fmt.Sprintf("%v", p.OverlappedPercent)}
|
||||
}
|
||||
if p.TableContextSize < 0 {
|
||||
return errInvalidValue{Field: "table_context_size", Value: fmt.Sprintf("%d", p.TableContextSize)}
|
||||
}
|
||||
if p.ImageContextSize < 0 {
|
||||
return errInvalidValue{Field: "image_context_size", Value: fmt.Sprintf("%d", p.ImageContextSize)}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TitleChunkerParam
|
||||
// ---------------------------------------------------------------------------
|
||||
//
|
||||
// Mirrors rag/flow/chunker/title_chunker/common.py:TitleChunkerParam.
|
||||
// The class also reads `self.method` (set externally — see Python
|
||||
// title_chunker.py:31-37 routing on `self._param.method`). The Go port
|
||||
// captures it explicitly here. The component's `check()` enforces enum
|
||||
// values.
|
||||
|
||||
type TitleChunkerParam struct {
|
||||
// Method routes to the right title-chunker strategy.
|
||||
// Allowed values: "hierarchy", "group".
|
||||
Method string `json:"method,omitempty"`
|
||||
|
||||
// Levels is a list of regex-list groups, one per hierarchy level.
|
||||
// Each group is a list of regex strings; the component picks the
|
||||
// best-matching group at runtime.
|
||||
Levels [][]string `json:"levels"`
|
||||
|
||||
// Hierarchy is the heading depth used by HierarchyTitleChunker.
|
||||
// Stored as a pointer to distinguish nil (unset) from 0.
|
||||
Hierarchy *int `json:"hierarchy,omitempty"`
|
||||
|
||||
// IncludeHeadingContent, when true, makes the heading text part
|
||||
// of each emitted chunk.
|
||||
IncludeHeadingContent bool `json:"include_heading_content"`
|
||||
|
||||
// RootChunkAsHeading, when true, prepends the root chunk's text
|
||||
// to every emitted chunk (and drops the root chunk itself).
|
||||
RootChunkAsHeading bool `json:"root_chunk_as_heading"`
|
||||
}
|
||||
|
||||
// Defaults returns the Python default TitleChunkerParam. `Method` is
|
||||
// not initialized in the Python `__init__` (it is set externally); the
|
||||
// default is left as the empty string and the component must supply it.
|
||||
func (TitleChunkerParam) Defaults() TitleChunkerParam {
|
||||
return TitleChunkerParam{
|
||||
Levels: [][]string{},
|
||||
Hierarchy: nil,
|
||||
IncludeHeadingContent: false,
|
||||
RootChunkAsHeading: false,
|
||||
}
|
||||
}
|
||||
|
||||
// Validate enforces the Python `check()` invariants that are also
|
||||
// expressible in pure-data terms: when Method == "hierarchy" the
|
||||
// hierarchy depth and level config must be present.
|
||||
func (p *TitleChunkerParam) Validate() error {
|
||||
switch p.Method {
|
||||
case "hierarchy", "group":
|
||||
case "":
|
||||
return nil
|
||||
default:
|
||||
return errInvalidValue{Field: "method", Value: p.Method}
|
||||
}
|
||||
switch p.Method {
|
||||
case "hierarchy", "group":
|
||||
if len(p.Levels) == 0 {
|
||||
return errRequiredField{Field: "levels"}
|
||||
}
|
||||
}
|
||||
if p.Method == "hierarchy" && (p.Hierarchy == nil || *p.Hierarchy <= 0) {
|
||||
return errRequiredField{Field: "hierarchy"}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// GroupTitleChunkerParam / HierarchyTitleChunkerParam
|
||||
// ---------------------------------------------------------------------------
|
||||
//
|
||||
// In the Python codebase, both variants share the SAME
|
||||
// `TitleChunkerParam` class — there is no per-variant param. The
|
||||
// dispatch happens in title_chunker.py:31-37 by reading
|
||||
// `self._param.method`.
|
||||
//
|
||||
// In Go we model the shared class as `TitleChunkerParam` and expose
|
||||
// type aliases so component files can name the param type they
|
||||
// actually use. This keeps the wire schema faithful to Python while
|
||||
// giving each variant a self-documenting entry point in the registry.
|
||||
|
||||
type GroupTitleChunkerParam = TitleChunkerParam
|
||||
type HierarchyTitleChunkerParam = TitleChunkerParam
|
||||
42
internal/ingestion/component/schema/errors.go
Normal file
42
internal/ingestion/component/schema/errors.go
Normal file
@@ -0,0 +1,42 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package schema
|
||||
|
||||
import "fmt"
|
||||
|
||||
// errRequiredField is the typed error returned by schema Validate()
|
||||
// methods when a required field is missing or empty. It carries the
|
||||
// field name so callers can produce structured error responses.
|
||||
type errRequiredField struct {
|
||||
Field string
|
||||
}
|
||||
|
||||
func (e errRequiredField) Error() string {
|
||||
return fmt.Sprintf("schema: required field %q is missing or empty", e.Field)
|
||||
}
|
||||
|
||||
// errInvalidValue is the typed error returned by schema Validate()
|
||||
// methods when a field's value is not in the allowed set. It carries
|
||||
// the field name and the offending value.
|
||||
type errInvalidValue struct {
|
||||
Field string
|
||||
Value string
|
||||
}
|
||||
|
||||
func (e errInvalidValue) Error() string {
|
||||
return fmt.Sprintf("schema: field %q has invalid value %q", e.Field, e.Value)
|
||||
}
|
||||
126
internal/ingestion/component/schema/extractor.go
Normal file
126
internal/ingestion/component/schema/extractor.go
Normal file
@@ -0,0 +1,126 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package schema
|
||||
|
||||
// ExtractorFromUpstream is the upstream payload consumed by the
|
||||
// Extractor component.
|
||||
//
|
||||
// The Python Extractor (rag/flow/extractor/extractor.py) does NOT
|
||||
// validate a Pydantic *FromUpstream schema; instead it pulls inputs
|
||||
// from the canvas's input-elements map:
|
||||
//
|
||||
// inputs = self.get_input_elements()
|
||||
// for k, v in inputs.items():
|
||||
// args[k] = v["value"]
|
||||
// if isinstance(args[k], list):
|
||||
// chunks = deepcopy(args[k])
|
||||
// chunks_key = k
|
||||
//
|
||||
// To keep the Go port faithful, the Go FromUpstream mirrors that
|
||||
// shape: a free-form map of named inputs plus an optional explicit
|
||||
// chunks list (the typical case in pipeline wiring).
|
||||
type ExtractorFromUpstream struct {
|
||||
// CreatedTime / ElapsedTime follow the package-wide convention
|
||||
// from upstream components.
|
||||
CreatedTime *float64 `json:"_created_time,omitempty"`
|
||||
ElapsedTime *float64 `json:"_elapsed_time,omitempty"`
|
||||
|
||||
// Inputs mirrors `get_input_elements()` output. Each entry holds a
|
||||
// free-form value (string for the LLM template, list of chunks
|
||||
// for the chunk-list binding). Keys are the input names; the
|
||||
// component selects the first list-typed value as the chunk
|
||||
// stream and passes the rest as scalar args.
|
||||
Inputs map[string]any `json:"inputs,omitempty"`
|
||||
|
||||
// Chunks is the explicit chunk list when wired in a linear
|
||||
// pipeline. Optional — when Inputs contains a list-typed entry,
|
||||
// the component uses that instead.
|
||||
Chunks []map[string]any `json:"chunks,omitempty"`
|
||||
}
|
||||
|
||||
// Validate enforces no required fields today; the Python component
|
||||
// happily runs on an empty input set (it produces a single output
|
||||
// chunk from the LLM call).
|
||||
func (ExtractorFromUpstream) Validate() error { return nil }
|
||||
|
||||
// ExtractorParam is the static configuration for the Extractor
|
||||
// component. Mirrors rag/flow/extractor/extractor.py:ExtractorParam,
|
||||
// which extends both ProcessParamBase and LLMParam. The LLM fields
|
||||
// (`llm_id`, `parameters`, `system_prompt`, `prompt`, `messages`,
|
||||
// etc.) live on the agent-side `LLMParam`; the Go port captures only
|
||||
// the Extractor-specific field plus a pointer to the LLM config so
|
||||
// the wiring is explicit.
|
||||
type ExtractorParam struct {
|
||||
// FieldName is the chunk key the LLM extraction result is written
|
||||
// to (Python: `self._param.field_name`). Required — `check()`
|
||||
// raises when empty. Mapped to "Result Destination" in the
|
||||
// frontend.
|
||||
FieldName string `json:"field_name"`
|
||||
|
||||
// LLMID identifies the LLM model used for extraction. This is the
|
||||
// agent-side LLMParam.llm_id; on the ingestion side it is
|
||||
// resolved against the tenant's LLM provider registry.
|
||||
LLMID string `json:"llm_id,omitempty"`
|
||||
|
||||
// SystemPrompt is the optional system prompt override.
|
||||
SystemPrompt string `json:"system_prompt,omitempty"`
|
||||
|
||||
// Prompt is the user-side template passed to the LLM.
|
||||
Prompt string `json:"prompt,omitempty"`
|
||||
}
|
||||
|
||||
// Defaults returns the Python default ExtractorParam: FieldName is
|
||||
// the empty string and is meant to be supplied at runtime.
|
||||
func (ExtractorParam) Defaults() ExtractorParam {
|
||||
return ExtractorParam{
|
||||
FieldName: "",
|
||||
LLMID: "",
|
||||
SystemPrompt: "",
|
||||
Prompt: "",
|
||||
}
|
||||
}
|
||||
|
||||
// Validate enforces the Python `check()` invariant: FieldName must
|
||||
// be non-empty.
|
||||
func (p *ExtractorParam) Validate() error {
|
||||
if p.FieldName == "" {
|
||||
return errRequiredField{Field: "field_name"}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ExtractorOutputs is the result of invoking the Extractor component.
|
||||
// Mirrors what the Python component sets via `self.set_output(...)` at
|
||||
// rag/flow/extractor/extractor.py:_invoke:
|
||||
//
|
||||
// self.set_output("output_format", "chunks")
|
||||
// self.set_output("chunks", chunks)
|
||||
type ExtractorOutputs struct {
|
||||
// OutputFormat is always "chunks".
|
||||
OutputFormat string `json:"output_format,omitempty"`
|
||||
|
||||
// Chunks is the enriched chunk list. When the Extractor ran over
|
||||
// a non-empty input list, each chunk gains a new key named after
|
||||
// FieldName (e.g., field_name="summary" -> chunk["summary"]). When
|
||||
// the Extractor ran over an empty input, Chunks contains a single
|
||||
// entry with one key (FieldName) holding the LLM result.
|
||||
Chunks []map[string]any `json:"chunks,omitempty"`
|
||||
|
||||
// Error is set when the component short-circuits with an error
|
||||
// message (Python: set_output("_ERROR", ...)).
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
105
internal/ingestion/component/schema/file.go
Normal file
105
internal/ingestion/component/schema/file.go
Normal file
@@ -0,0 +1,105 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Package schema holds the wire-level *FromUpstream / *Param / *Outputs types
|
||||
// that flow between ingestion pipeline components. Each file mirrors the
|
||||
// Pydantic schema in rag/flow/<component>/schema.py (or, when none exists,
|
||||
// the runtime contract implied by the component's _invoke signature).
|
||||
//
|
||||
// Files in this package are pure data definitions — they import only stdlib.
|
||||
// Business logic (the component Implementations) lives elsewhere in
|
||||
// internal/ingestion/component/.
|
||||
package schema
|
||||
|
||||
// FileFromUpstream is the upstream payload consumed by the File component.
|
||||
//
|
||||
// Mirrors the runtime contract of rag/flow/file.py:File._invoke. There is no
|
||||
// dedicated Pydantic schema for File in the Python codebase, so the fields
|
||||
// below were derived from the keyword arguments File reads at runtime:
|
||||
//
|
||||
// async def _invoke(self, **kwargs):
|
||||
// if self._canvas._doc_id: ... # doc_id path
|
||||
// else:
|
||||
// file = kwargs.get("file")[0] # file-list path
|
||||
//
|
||||
// `_doc_id` is taken from the surrounding canvas, not kwargs, so it lives on
|
||||
// a separate optional field here for explicitness.
|
||||
type FileFromUpstream struct {
|
||||
// CreatedTime is the upstream component's wall-clock start time (s).
|
||||
CreatedTime *float64 `json:"_created_time,omitempty"`
|
||||
// ElapsedTime is the upstream component's elapsed time (s).
|
||||
ElapsedTime *float64 `json:"_elapsed_time,omitempty"`
|
||||
|
||||
// DocID is the canvas-bound document ID. When non-empty the File
|
||||
// component resolves the binary via the document service; otherwise
|
||||
// the File payload below is used. Optional in wire terms — the
|
||||
// Python code branches on truthiness, so we use *string.
|
||||
DocID *string `json:"doc_id,omitempty"`
|
||||
|
||||
// File is the optional list of file descriptors passed when no
|
||||
// doc_id is bound. In Python: `file = kwargs.get("file")[0]`.
|
||||
// Shape: `[]map[string]any` — same as rag/flow/parser's `file` field.
|
||||
File []map[string]any `json:"file,omitempty"`
|
||||
}
|
||||
|
||||
// Validate enforces the File component's wire-shape requirement: at
|
||||
// least one of DocID or File must be set. The Python code branches on
|
||||
// `_canvas._doc_id` vs. `kwargs.get("file")[0]`, which corresponds to
|
||||
// "one of the two upstream paths is wired".
|
||||
func (f *FileFromUpstream) Validate() error {
|
||||
if (f.DocID == nil || *f.DocID == "") && len(f.File) == 0 {
|
||||
return errRequiredField{Field: "doc_id|file"}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// FileParam is the static configuration for the File component.
|
||||
//
|
||||
// Mirrors rag/flow/file.py:FileParam. The Python class has no fields
|
||||
// beyond what ProcessParamBase provides (none stored on the instance), so
|
||||
// the Go struct is intentionally empty. It is kept as a distinct type so
|
||||
// Phase 2.1 can attach config knobs (e.g., a `with_blob` flag) without a
|
||||
// rename.
|
||||
type FileParam struct{}
|
||||
|
||||
// Defaults returns a FileParam populated with the Python default values.
|
||||
// Today there are no fields; the constructor exists for forward
|
||||
// compatibility and to satisfy the "every type has a Defaults()"
|
||||
// convention adopted in this package.
|
||||
func (FileParam) Defaults() FileParam { return FileParam{} }
|
||||
|
||||
// Validate returns nil. FileParam has no required fields.
|
||||
func (FileParam) Validate() error { return nil }
|
||||
|
||||
// FileOutputs is the result of invoking the File component. It mirrors
|
||||
// the values File sets via `self.set_output(...)` in rag/flow/file.py:
|
||||
//
|
||||
// set_output("name", doc.name) # name
|
||||
// set_output("file", file) # file (full descriptor)
|
||||
// set_output("_ERROR", "...") # error
|
||||
//
|
||||
// The binary blob (`blob`) is intentionally NOT part of the wire schema —
|
||||
// it lives on the storage layer (MinIO/S3) and is referenced by path.
|
||||
type FileOutputs struct {
|
||||
// Name is the resolved document/file name.
|
||||
Name string `json:"name"`
|
||||
// File is the upstream file descriptor (dict in Python). Optional —
|
||||
// when invoked via doc_id the Python code does not re-emit it.
|
||||
File map[string]any `json:"file,omitempty"`
|
||||
// Error is set when the component short-circuits with an error
|
||||
// message (Python: set_output("_ERROR", ...)).
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
240
internal/ingestion/component/schema/parser.go
Normal file
240
internal/ingestion/component/schema/parser.go
Normal file
@@ -0,0 +1,240 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package schema
|
||||
|
||||
// ParserFromUpstream is the upstream payload consumed by the Parser
|
||||
// component. It mirrors rag/flow/parser/schema.py:ParserFromUpstream
|
||||
// (Pydantic BaseModel with populate_by_name, extra="forbid").
|
||||
//
|
||||
// created_time: float | None (alias _created_time)
|
||||
// elapsed_time: float | None (alias _elapsed_time)
|
||||
// name: str (required)
|
||||
// file: dict | None
|
||||
// abstract: bool = False
|
||||
// author: bool = False
|
||||
type ParserFromUpstream struct {
|
||||
CreatedTime *float64 `json:"_created_time,omitempty"`
|
||||
ElapsedTime *float64 `json:"_elapsed_time,omitempty"`
|
||||
|
||||
Name string `json:"name"`
|
||||
|
||||
// File is the optional upstream file descriptor. Python allows None
|
||||
// when the parser is invoked via a canvas-bound doc_id path.
|
||||
File map[string]any `json:"file,omitempty"`
|
||||
|
||||
Abstract bool `json:"abstract,omitempty"`
|
||||
Author bool `json:"author,omitempty"`
|
||||
}
|
||||
|
||||
// Validate enforces the only required field in ParserFromUpstream: Name.
|
||||
// Returns nil when Name is non-empty.
|
||||
func (p *ParserFromUpstream) Validate() error {
|
||||
if p.Name == "" {
|
||||
return errRequiredField{Field: "name"}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Page is one parsed document section. The Parser component does not emit
|
||||
// a typed page model — Python code passes around `dict` literals with
|
||||
// shape `{text, layout_type, doc_type_kwd, positions?, image?, ...}`. To
|
||||
// keep the wire schema typed without overcommitting to a parser-specific
|
||||
// shape, Page is left as a generic map and provided for forward
|
||||
// documentation; downstream chunker code operates on the same dict shape.
|
||||
type Page map[string]any
|
||||
|
||||
// ParserSetup is the per-filetype configuration block stored on
|
||||
// ParserParam.setups[fileType]. The keys are heterogeneous (e.g.,
|
||||
// `parse_method`, `lang`, `output_format`, `suffix`, `fields`, `vlm`),
|
||||
// so a free-form map best mirrors the Python dict literal.
|
||||
type ParserSetup map[string]any
|
||||
|
||||
// ParserParam is the static configuration for the Parser component.
|
||||
// Mirrors rag/flow/parser/parser.py:ParserParam.
|
||||
//
|
||||
// Two top-level fields are configured in the Python class:
|
||||
//
|
||||
// setups: dict[str, dict] (one entry per file type: pdf, docx, ...)
|
||||
// allowed_output_format: dict[str, list[str]] (per-file-type formats)
|
||||
//
|
||||
// `check()` runs further validation that is intentionally NOT
|
||||
// replicated here — Validate() enforces wire-shape only; business-rule
|
||||
// validation lives in the component implementation (Phase 2.2).
|
||||
type ParserParam struct {
|
||||
// Setups holds the per-file-type parser config. Keys are file-type
|
||||
// identifiers ("pdf", "docx", "markdown", "spreadsheet", "image",
|
||||
// "audio", "video", "email", "epub", "doc", "text&code", "html",
|
||||
// "slides"); values are free-form config blobs.
|
||||
Setups map[string]ParserSetup `json:"setups"`
|
||||
|
||||
// AllowedOutputFormat mirrors `allowed_output_format` from the
|
||||
// Python class. Used for client-side input-form validation.
|
||||
AllowedOutputFormat map[string][]string `json:"allowed_output_format"`
|
||||
}
|
||||
|
||||
// Defaults returns a ParserParam populated with the Python defaults —
|
||||
// the full setups table copied verbatim from
|
||||
// rag/flow/parser/parser.py:ParserParam.__init__ and the corresponding
|
||||
// allowed_output_format map.
|
||||
func (ParserParam) Defaults() ParserParam {
|
||||
return ParserParam{
|
||||
AllowedOutputFormat: map[string][]string{
|
||||
"pdf": {"json", "markdown"},
|
||||
"spreadsheet": {"json", "markdown", "html"},
|
||||
"doc": {"json", "markdown"},
|
||||
"docx": {"json", "markdown"},
|
||||
"slides": {"json"},
|
||||
"image": {"json"},
|
||||
"email": {"text", "json"},
|
||||
"markdown": {"text", "json"},
|
||||
"text&code": {"text", "json"},
|
||||
"html": {"text", "json"},
|
||||
"audio": {"json"},
|
||||
"video": {},
|
||||
"epub": {"text", "json"},
|
||||
},
|
||||
Setups: map[string]ParserSetup{
|
||||
"pdf": {
|
||||
"parse_method": "deepdoc",
|
||||
"lang": "Chinese",
|
||||
"flatten_media_to_text": false,
|
||||
"remove_toc": false,
|
||||
"remove_header_footer": false,
|
||||
"suffix": []string{"pdf"},
|
||||
"output_format": "json",
|
||||
},
|
||||
"spreadsheet": {
|
||||
"parse_method": "deepdoc",
|
||||
"flatten_media_to_text": false,
|
||||
"output_format": "html",
|
||||
"suffix": []string{"xls", "xlsx", "csv"},
|
||||
},
|
||||
"doc": {
|
||||
"remove_toc": false,
|
||||
"remove_header_footer": false,
|
||||
"suffix": []string{"doc"},
|
||||
"output_format": "json",
|
||||
},
|
||||
"docx": {
|
||||
"flatten_media_to_text": false,
|
||||
"remove_toc": false,
|
||||
"remove_header_footer": false,
|
||||
"suffix": []string{"docx"},
|
||||
"output_format": "json",
|
||||
},
|
||||
"markdown": {
|
||||
"flatten_media_to_text": false,
|
||||
"suffix": []string{"md", "markdown", "mdx"},
|
||||
"remove_toc": false,
|
||||
"output_format": "json",
|
||||
},
|
||||
"text&code": {
|
||||
"suffix": []string{
|
||||
"txt", "py", "js", "java", "c", "cpp", "h", "php",
|
||||
"go", "ts", "sh", "cs", "kt", "sql",
|
||||
},
|
||||
"output_format": "json",
|
||||
},
|
||||
"html": {
|
||||
"suffix": []string{"htm", "html"},
|
||||
"remove_toc": false,
|
||||
"remove_header_footer": false,
|
||||
"output_format": "json",
|
||||
},
|
||||
"slides": {
|
||||
"parse_method": "deepdoc",
|
||||
"suffix": []string{"pptx", "ppt"},
|
||||
"output_format": "json",
|
||||
},
|
||||
"image": {
|
||||
"parse_method": "ocr",
|
||||
"llm_id": "",
|
||||
"lang": "Chinese",
|
||||
"system_prompt": "",
|
||||
"suffix": []string{"jpg", "jpeg", "png", "gif"},
|
||||
"output_format": "json",
|
||||
},
|
||||
"email": {
|
||||
"suffix": []string{"eml", "msg"},
|
||||
"fields": []string{
|
||||
"from", "to", "cc", "bcc", "date", "subject",
|
||||
"body", "attachments", "metadata",
|
||||
},
|
||||
"output_format": "text",
|
||||
},
|
||||
"audio": {
|
||||
"suffix": []string{
|
||||
"da", "wave", "wav", "mp3", "aac", "flac", "ogg",
|
||||
"aiff", "au", "midi", "wma", "realaudio", "vqf",
|
||||
"oggvorbis", "ape",
|
||||
},
|
||||
"output_format": "text",
|
||||
},
|
||||
"video": {
|
||||
"suffix": []string{"mp4", "avi", "mkv"},
|
||||
"output_format": "text",
|
||||
"prompt": "",
|
||||
},
|
||||
"epub": {
|
||||
"suffix": []string{"epub"},
|
||||
"output_format": "json",
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Validate returns nil. ParserParam's field set is fully defaulted by
|
||||
// Defaults(); the component's own `check()` method performs business
|
||||
// validation (e.g., "parse_method" must be one of the allowed set), and
|
||||
// that runs in the component implementation.
|
||||
func (ParserParam) Validate() error { return nil }
|
||||
|
||||
// ParserOutputs is the result of invoking the Parser component. The
|
||||
// Python component calls `self.set_output(...)` with a mix of
|
||||
// string-typed, list-typed, and dict-typed values. The wire schema
|
||||
// below is the typed surface consumed by downstream components.
|
||||
//
|
||||
// Mirrors what Parser sets at rag/flow/parser/parser.py:_invoke. The
|
||||
// parser writes to EITHER ("json" | "markdown" | "text" | "html") and
|
||||
// always sets "output_format" + "file" + "_ERROR".
|
||||
type ParserOutputs struct {
|
||||
// OutputFormat is the active output format for this run
|
||||
// (one of "json", "markdown", "text", "html"). The downstream
|
||||
// Tokenizer branches on this field.
|
||||
OutputFormat string `json:"output_format,omitempty"`
|
||||
|
||||
// JSON holds the list of structured sections when output_format == "json".
|
||||
JSON []map[string]any `json:"json,omitempty"`
|
||||
|
||||
// Markdown holds the rendered markdown when output_format == "markdown".
|
||||
Markdown string `json:"markdown,omitempty"`
|
||||
|
||||
// Text holds the rendered plain text when output_format == "text".
|
||||
Text string `json:"text,omitempty"`
|
||||
|
||||
// HTML holds the rendered HTML when output_format == "html".
|
||||
HTML string `json:"html,omitempty"`
|
||||
|
||||
// File is the upstream file descriptor with parser-derived metadata
|
||||
// (e.g., outlines) merged in. Mirrors the Python `set_output("file", ...)`
|
||||
// at parser.py:609, 791, 828.
|
||||
File map[string]any `json:"file,omitempty"`
|
||||
|
||||
// Error is set when the component short-circuits with an error
|
||||
// message (Python: set_output("_ERROR", ...)).
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
578
internal/ingestion/component/schema/schema_test.go
Normal file
578
internal/ingestion/component/schema/schema_test.go
Normal file
@@ -0,0 +1,578 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package schema
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// File
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestFileParamDefaults(t *testing.T) {
|
||||
p := FileParam{}.Defaults()
|
||||
if err := p.Validate(); err != nil {
|
||||
t.Fatalf("default FileParam failed Validate: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileFromUpstreamValidate(t *testing.T) {
|
||||
// Empty upstream: nothing to bind to -> error.
|
||||
if err := (&FileFromUpstream{}).Validate(); err == nil {
|
||||
t.Fatal("expected Validate to fail when neither doc_id nor file is set")
|
||||
}
|
||||
|
||||
// DocID path: valid.
|
||||
docID := "doc-123"
|
||||
fu := FileFromUpstream{DocID: &docID}
|
||||
if err := fu.Validate(); err != nil {
|
||||
t.Fatalf("FileFromUpstream with DocID unexpectedly failed Validate: %v", err)
|
||||
}
|
||||
|
||||
// File path: valid.
|
||||
fu = FileFromUpstream{File: []map[string]any{{"name": "input.pdf"}}}
|
||||
if err := fu.Validate(); err != nil {
|
||||
t.Fatalf("FileFromUpstream with File unexpectedly failed Validate: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileFromUpstreamJSONRoundTrip(t *testing.T) {
|
||||
created := 1.5
|
||||
elapsed := 0.25
|
||||
docID := "doc-abc"
|
||||
original := FileFromUpstream{
|
||||
CreatedTime: &created,
|
||||
ElapsedTime: &elapsed,
|
||||
DocID: &docID,
|
||||
File: []map[string]any{{"name": "input.pdf"}},
|
||||
}
|
||||
data, err := json.Marshal(original)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), `"_created_time":1.5`) {
|
||||
t.Errorf("expected _created_time alias in JSON, got %s", data)
|
||||
}
|
||||
if !strings.Contains(string(data), `"doc_id":"doc-abc"`) {
|
||||
t.Errorf("expected doc_id in JSON, got %s", data)
|
||||
}
|
||||
|
||||
var decoded FileFromUpstream
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if decoded.DocID == nil || *decoded.DocID != docID {
|
||||
t.Errorf("DocID round-trip mismatch: got %v", decoded.DocID)
|
||||
}
|
||||
if len(decoded.File) != 1 {
|
||||
t.Errorf("File round-trip mismatch: got %d", len(decoded.File))
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileOutputsJSONRoundTrip(t *testing.T) {
|
||||
original := FileOutputs{
|
||||
Name: "input.pdf",
|
||||
File: map[string]any{"id": "f-1"},
|
||||
Error: "doc not found",
|
||||
}
|
||||
data, err := json.Marshal(original)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), `"name":"input.pdf"`) {
|
||||
t.Errorf("expected name in JSON, got %s", data)
|
||||
}
|
||||
// _ERROR non-empty should be emitted.
|
||||
if !strings.Contains(string(data), `"_ERROR":"doc not found"`) {
|
||||
t.Errorf("expected _ERROR in JSON, got %s", data)
|
||||
}
|
||||
|
||||
var decoded FileOutputs
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if decoded.Name != original.Name {
|
||||
t.Errorf("Name round-trip mismatch: got %q", decoded.Name)
|
||||
}
|
||||
if decoded.Error != original.Error {
|
||||
t.Errorf("Error round-trip mismatch: got %q", decoded.Error)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Parser
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestParserFromUpstreamValidate(t *testing.T) {
|
||||
// Name is required.
|
||||
if err := (&ParserFromUpstream{}).Validate(); err == nil {
|
||||
t.Fatal("expected Validate to fail when Name is empty")
|
||||
}
|
||||
if err := (&ParserFromUpstream{Name: "doc.pdf"}).Validate(); err != nil {
|
||||
t.Fatalf("Validate with Name unexpectedly failed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParserParamDefaults(t *testing.T) {
|
||||
p := ParserParam{}.Defaults()
|
||||
if err := p.Validate(); err != nil {
|
||||
t.Fatalf("default ParserParam failed Validate: %v", err)
|
||||
}
|
||||
// Spot-check a few entries that the Python class initializes.
|
||||
for _, key := range []string{"pdf", "docx", "image", "audio", "video", "email", "epub"} {
|
||||
if _, ok := p.Setups[key]; !ok {
|
||||
t.Errorf("default Setups missing key %q", key)
|
||||
}
|
||||
}
|
||||
if got := p.Setups["pdf"]["parse_method"]; got != "deepdoc" {
|
||||
t.Errorf("default pdf parse_method = %v, want deepdoc", got)
|
||||
}
|
||||
if got := p.AllowedOutputFormat["pdf"]; len(got) != 2 || got[0] != "json" || got[1] != "markdown" {
|
||||
t.Errorf("default pdf allowed_output_format = %v, want [json markdown]", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParserParamJSONRoundTrip(t *testing.T) {
|
||||
original := ParserParam{}.Defaults()
|
||||
data, err := json.Marshal(original)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), `"setups"`) {
|
||||
t.Errorf("expected setups in JSON, got %s", data)
|
||||
}
|
||||
var decoded ParserParam
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if decoded.Setups["pdf"]["parse_method"] != "deepdoc" {
|
||||
t.Errorf("round-trip lost default pdf parse_method: got %v", decoded.Setups["pdf"]["parse_method"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestParserFromUpstreamJSONRoundTrip(t *testing.T) {
|
||||
original := ParserFromUpstream{
|
||||
Name: "input.pdf",
|
||||
Abstract: true,
|
||||
Author: false,
|
||||
}
|
||||
data, err := json.Marshal(original)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
// abstract=true should be emitted; author=false has omitempty so it's
|
||||
// dropped (zero-value bool with omitempty). We test the
|
||||
// non-zero path.
|
||||
if !strings.Contains(string(data), `"abstract":true`) {
|
||||
t.Errorf("expected abstract=true in JSON, got %s", data)
|
||||
}
|
||||
var decoded ParserFromUpstream
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if decoded.Name != original.Name {
|
||||
t.Errorf("Name round-trip mismatch: got %q", decoded.Name)
|
||||
}
|
||||
if !decoded.Abstract {
|
||||
t.Errorf("Abstract round-trip mismatch: got %v", decoded.Abstract)
|
||||
}
|
||||
// author field is omitempty, so JSON round-trip should leave it false
|
||||
// (default value) on both sides.
|
||||
if decoded.Author {
|
||||
t.Errorf("Author should be false, got true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParserOutputsJSONRoundTrip(t *testing.T) {
|
||||
original := ParserOutputs{
|
||||
OutputFormat: "json",
|
||||
JSON: []map[string]any{{"text": "hello", "doc_type_kwd": "text"}},
|
||||
}
|
||||
data, err := json.Marshal(original)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), `"output_format":"json"`) {
|
||||
t.Errorf("expected output_format in JSON, got %s", data)
|
||||
}
|
||||
var decoded ParserOutputs
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if decoded.OutputFormat != "json" {
|
||||
t.Errorf("OutputFormat round-trip mismatch: got %q", decoded.OutputFormat)
|
||||
}
|
||||
if len(decoded.JSON) != 1 {
|
||||
t.Errorf("JSON round-trip mismatch: got %d", len(decoded.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Chunker
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestChunkerFromUpstreamValidate(t *testing.T) {
|
||||
if err := (&ChunkerFromUpstream{}).Validate(); err == nil {
|
||||
t.Fatal("expected Validate to fail when Name is empty")
|
||||
}
|
||||
if err := (&ChunkerFromUpstream{Name: "doc.pdf"}).Validate(); err != nil {
|
||||
t.Fatalf("Validate with Name unexpectedly failed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkerFromUpstreamJSONRoundTrip(t *testing.T) {
|
||||
md := "# title"
|
||||
original := ChunkerFromUpstream{
|
||||
Name: "doc.pdf",
|
||||
OutputFormat: PayloadFormatChunks,
|
||||
Chunks: []ChunkDoc{{Text: "alpha"}},
|
||||
MarkdownResult: &md,
|
||||
}
|
||||
data, err := json.Marshal(original)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), `"output_format":"chunks"`) {
|
||||
t.Errorf("expected output_format in JSON, got %s", data)
|
||||
}
|
||||
var decoded ChunkerFromUpstream
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if decoded.Name != "doc.pdf" || decoded.OutputFormat != PayloadFormatChunks {
|
||||
t.Errorf("round-trip mismatch: %+v", decoded)
|
||||
}
|
||||
if len(decoded.Chunks) != 1 {
|
||||
t.Errorf("Chunks round-trip mismatch: got %d", len(decoded.Chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkerOutputsJSONRoundTrip(t *testing.T) {
|
||||
original := ChunkerOutputs{
|
||||
OutputFormat: PayloadFormatChunks,
|
||||
Chunks: []ChunkDoc{{Text: "alpha"}, {Text: "beta"}},
|
||||
}
|
||||
data, err := json.Marshal(original)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), `"output_format":"chunks"`) {
|
||||
t.Errorf("expected output_format in JSON, got %s", data)
|
||||
}
|
||||
var decoded ChunkerOutputs
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if len(decoded.Chunks) != 2 {
|
||||
t.Errorf("Chunks round-trip mismatch: got %d", len(decoded.Chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenChunkerParamDefaults(t *testing.T) {
|
||||
p := TokenChunkerParam{}.Defaults()
|
||||
if p.DelimiterMode != "token_size" {
|
||||
t.Errorf("default delimiter_mode = %q, want token_size", p.DelimiterMode)
|
||||
}
|
||||
if p.ChunkTokenSize != 512 {
|
||||
t.Errorf("default chunk_token_size = %d, want 512", p.ChunkTokenSize)
|
||||
}
|
||||
if len(p.Delimiters) != 1 || p.Delimiters[0] != "\n" {
|
||||
t.Errorf("default delimiters = %v, want [\\n]", p.Delimiters)
|
||||
}
|
||||
if p.OverlappedPercent != 0 {
|
||||
t.Errorf("default overlapped_percent = %f, want 0", p.OverlappedPercent)
|
||||
}
|
||||
if err := p.Validate(); err != nil {
|
||||
t.Fatalf("default TokenChunkerParam failed Validate: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTitleChunkerParamDefaults(t *testing.T) {
|
||||
p := TitleChunkerParam{}.Defaults()
|
||||
if p.Hierarchy != nil {
|
||||
t.Errorf("default hierarchy should be nil, got %v", *p.Hierarchy)
|
||||
}
|
||||
if p.IncludeHeadingContent {
|
||||
t.Errorf("default include_heading_content should be false")
|
||||
}
|
||||
// Default has no method set; Validate must accept it (the empty
|
||||
// method is the uninitialized state, not an enum violation).
|
||||
if err := p.Validate(); err != nil {
|
||||
t.Fatalf("default TitleChunkerParam failed Validate: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTitleChunkerParamValidate(t *testing.T) {
|
||||
// Method=hierarchy with no levels -> error.
|
||||
p := TitleChunkerParam{Method: "hierarchy"}
|
||||
if err := p.Validate(); err == nil {
|
||||
t.Fatal("expected error when Method=hierarchy with no levels")
|
||||
}
|
||||
// Method=hierarchy with levels but no hierarchy number -> error.
|
||||
p.Levels = [][]string{{"^# "}}
|
||||
if err := p.Validate(); err == nil {
|
||||
t.Fatal("expected error when Method=hierarchy with no hierarchy number")
|
||||
}
|
||||
// Method=hierarchy with levels and a positive hierarchy -> OK.
|
||||
h := 2
|
||||
p.Hierarchy = &h
|
||||
if err := p.Validate(); err != nil {
|
||||
t.Fatalf("Validate unexpectedly failed: %v", err)
|
||||
}
|
||||
// Method=group with levels -> OK.
|
||||
p = TitleChunkerParam{Method: "group", Levels: [][]string{{"^# "}}}
|
||||
if err := p.Validate(); err != nil {
|
||||
t.Fatalf("Validate unexpectedly failed: %v", err)
|
||||
}
|
||||
// Method=group with no levels -> error.
|
||||
p = TitleChunkerParam{Method: "group"}
|
||||
if err := p.Validate(); err == nil {
|
||||
t.Fatal("expected error when Method=group with no levels")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupTitleChunkerParamAlias(t *testing.T) {
|
||||
// The alias must be assignable from a TitleChunkerParam value.
|
||||
var gp GroupTitleChunkerParam = TitleChunkerParam{Method: "group", Levels: [][]string{{"^# "}}}
|
||||
if err := gp.Validate(); err != nil {
|
||||
t.Fatalf("group param via alias failed Validate: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHierarchyTitleChunkerParamAlias(t *testing.T) {
|
||||
h := 1
|
||||
var hp HierarchyTitleChunkerParam = TitleChunkerParam{Method: "hierarchy", Levels: [][]string{{"^# "}}, Hierarchy: &h}
|
||||
if err := hp.Validate(); err != nil {
|
||||
t.Fatalf("hierarchy param via alias failed Validate: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tokenizer
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestTokenizerFromUpstreamValidate(t *testing.T) {
|
||||
// output_format=chunks with nil Chunks is valid.
|
||||
if err := (&TokenizerFromUpstream{OutputFormat: PayloadFormatChunks}).Validate(); err != nil {
|
||||
t.Fatalf("output_format=chunks should be valid, got %v", err)
|
||||
}
|
||||
// output_format=markdown with no MarkdownResult -> error.
|
||||
if err := (&TokenizerFromUpstream{OutputFormat: PayloadFormatMarkdown}).Validate(); err == nil {
|
||||
t.Fatal("expected error for markdown without payload")
|
||||
}
|
||||
// output_format=markdown with payload -> OK.
|
||||
md := "# title"
|
||||
if err := (&TokenizerFromUpstream{OutputFormat: PayloadFormatMarkdown, MarkdownResult: &md}).Validate(); err != nil {
|
||||
t.Fatalf("markdown with payload should be valid, got %v", err)
|
||||
}
|
||||
// output_format=text without payload -> error.
|
||||
if err := (&TokenizerFromUpstream{OutputFormat: PayloadFormatText}).Validate(); err == nil {
|
||||
t.Fatal("expected error for text without payload")
|
||||
}
|
||||
txt := "hello"
|
||||
if err := (&TokenizerFromUpstream{OutputFormat: PayloadFormatText, TextResult: &txt}).Validate(); err != nil {
|
||||
t.Fatalf("text with payload should be valid, got %v", err)
|
||||
}
|
||||
// output_format=html without payload -> error.
|
||||
if err := (&TokenizerFromUpstream{OutputFormat: PayloadFormatHTML}).Validate(); err == nil {
|
||||
t.Fatal("expected error for html without payload")
|
||||
}
|
||||
html := "<p>x</p>"
|
||||
if err := (&TokenizerFromUpstream{OutputFormat: PayloadFormatHTML, HTMLResult: &html}).Validate(); err != nil {
|
||||
t.Fatalf("html with payload should be valid, got %v", err)
|
||||
}
|
||||
// Empty output_format with neither JSON nor Chunks -> error.
|
||||
if err := (&TokenizerFromUpstream{}).Validate(); err == nil {
|
||||
t.Fatal("expected error when no output_format and no payload")
|
||||
}
|
||||
// Empty output_format with JSONResult -> OK.
|
||||
if err := (&TokenizerFromUpstream{JSONResult: []ChunkDoc{{Text: "x"}}}).Validate(); err != nil {
|
||||
t.Fatalf("empty output_format with JSONResult should be valid, got %v", err)
|
||||
}
|
||||
// Empty output_format with Chunks -> OK.
|
||||
if err := (&TokenizerFromUpstream{Chunks: []ChunkDoc{{Text: "x"}}}).Validate(); err != nil {
|
||||
t.Fatalf("empty output_format with Chunks should be valid, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenizerFromUpstreamJSONRoundTrip(t *testing.T) {
|
||||
md := "# title"
|
||||
txt := "body"
|
||||
html := "<p>x</p>"
|
||||
original := TokenizerFromUpstream{
|
||||
Name: "doc.pdf",
|
||||
OutputFormat: PayloadFormatChunks,
|
||||
Chunks: []ChunkDoc{{Text: "alpha"}},
|
||||
MarkdownResult: &md,
|
||||
TextResult: &txt,
|
||||
HTMLResult: &html,
|
||||
}
|
||||
data, err := json.Marshal(original)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), `"output_format":"chunks"`) {
|
||||
t.Errorf("expected output_format in JSON, got %s", data)
|
||||
}
|
||||
// All *Result fields have omitempty. Non-zero values round-trip;
|
||||
// confirm at least one of them is in the payload.
|
||||
if !strings.Contains(string(data), `"markdown":"# title"`) {
|
||||
t.Errorf("expected markdown alias in JSON, got %s", data)
|
||||
}
|
||||
var decoded TokenizerFromUpstream
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if decoded.MarkdownResult == nil || *decoded.MarkdownResult != "# title" {
|
||||
t.Errorf("markdown round-trip mismatch: got %v", decoded.MarkdownResult)
|
||||
}
|
||||
// Re-marshal the decoded and confirm we can read markdown back.
|
||||
data2, err := json.Marshal(decoded)
|
||||
if err != nil {
|
||||
t.Fatalf("re-marshal: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data2), `"markdown":"# title"`) {
|
||||
t.Errorf("expected markdown alias after re-marshal, got %s", data2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenizerParamDefaults(t *testing.T) {
|
||||
p := TokenizerParam{}.Defaults()
|
||||
if len(p.SearchMethod) != 2 || p.SearchMethod[0] != "full_text" || p.SearchMethod[1] != "embedding" {
|
||||
t.Errorf("default search_method = %v, want [full_text embedding]", p.SearchMethod)
|
||||
}
|
||||
if p.FilenameEmbdWeight != 0.1 {
|
||||
t.Errorf("default filename_embd_weight = %f, want 0.1", p.FilenameEmbdWeight)
|
||||
}
|
||||
if len(p.Fields) != 1 || p.Fields[0] != "text" {
|
||||
t.Errorf("default fields = %v, want [text]", p.Fields)
|
||||
}
|
||||
if err := p.Validate(); err != nil {
|
||||
t.Fatalf("default TokenizerParam failed Validate: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenizerParamValidate(t *testing.T) {
|
||||
// Empty search_method -> error.
|
||||
if err := (&TokenizerParam{}).Validate(); err == nil {
|
||||
t.Fatal("expected error for empty search_method")
|
||||
}
|
||||
// Invalid search_method entry -> error.
|
||||
if err := (&TokenizerParam{SearchMethod: []string{"unknown"}}).Validate(); err == nil {
|
||||
t.Fatal("expected error for unknown search_method entry")
|
||||
}
|
||||
// Valid search_method -> OK.
|
||||
if err := (&TokenizerParam{SearchMethod: []string{"embedding"}}).Validate(); err != nil {
|
||||
t.Fatalf("embedding search_method should be valid, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenizerOutputsJSONRoundTrip(t *testing.T) {
|
||||
tokens := 256
|
||||
original := TokenizerOutputs{
|
||||
OutputFormat: PayloadFormatChunks,
|
||||
Chunks: []ChunkDoc{{Text: "alpha"}},
|
||||
EmbeddingTokenConsumption: &tokens,
|
||||
}
|
||||
data, err := json.Marshal(original)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), `"embedding_token_consumption":256`) {
|
||||
t.Errorf("expected embedding_token_consumption in JSON, got %s", data)
|
||||
}
|
||||
var decoded TokenizerOutputs
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if decoded.EmbeddingTokenConsumption == nil || *decoded.EmbeddingTokenConsumption != 256 {
|
||||
t.Errorf("EmbeddingTokenConsumption round-trip mismatch: got %v", decoded.EmbeddingTokenConsumption)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Extractor
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestExtractorParamDefaults(t *testing.T) {
|
||||
p := ExtractorParam{}.Defaults()
|
||||
if p.FieldName != "" {
|
||||
t.Errorf("default field_name should be empty, got %q", p.FieldName)
|
||||
}
|
||||
if err := p.Validate(); err == nil {
|
||||
t.Fatal("default ExtractorParam should fail Validate (field_name required)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractorParamValidate(t *testing.T) {
|
||||
if err := (&ExtractorParam{}).Validate(); err == nil {
|
||||
t.Fatal("expected error for empty field_name")
|
||||
}
|
||||
if err := (&ExtractorParam{FieldName: "summary"}).Validate(); err != nil {
|
||||
t.Fatalf("Validate with field_name should pass, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractorFromUpstreamJSONRoundTrip(t *testing.T) {
|
||||
original := ExtractorFromUpstream{
|
||||
Chunks: []map[string]any{{"text": "alpha"}},
|
||||
}
|
||||
data, err := json.Marshal(original)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), `"chunks"`) {
|
||||
t.Errorf("expected chunks in JSON, got %s", data)
|
||||
}
|
||||
var decoded ExtractorFromUpstream
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if len(decoded.Chunks) != 1 {
|
||||
t.Errorf("Chunks round-trip mismatch: got %d", len(decoded.Chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractorOutputsJSONRoundTrip(t *testing.T) {
|
||||
original := ExtractorOutputs{
|
||||
OutputFormat: "chunks",
|
||||
Chunks: []map[string]any{{"summary": "x"}},
|
||||
}
|
||||
data, err := json.Marshal(original)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), `"output_format":"chunks"`) {
|
||||
t.Errorf("expected output_format in JSON, got %s", data)
|
||||
}
|
||||
var decoded ExtractorOutputs
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if decoded.OutputFormat != "chunks" {
|
||||
t.Errorf("OutputFormat round-trip mismatch: got %q", decoded.OutputFormat)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func ptrString(s string) *string { return &s }
|
||||
189
internal/ingestion/component/schema/tokenizer.go
Normal file
189
internal/ingestion/component/schema/tokenizer.go
Normal file
@@ -0,0 +1,189 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package schema
|
||||
|
||||
// TokenizerFromUpstream is the upstream payload consumed by the
|
||||
// Tokenizer component. It mirrors
|
||||
// rag/flow/tokenizer/schema.py:TokenizerFromUpstream, including the
|
||||
// Pydantic `model_validator(mode="after")` invariant on
|
||||
// `output_format <-> payload` consistency.
|
||||
//
|
||||
// Wire shape (Pydantic):
|
||||
//
|
||||
// created_time: float | None (alias _created_time)
|
||||
// elapsed_time: float | None (alias _elapsed_time)
|
||||
// name: str (default "")
|
||||
// file: dict | None
|
||||
// output_format: Literal["json","markdown","text","html","chunks"] | None
|
||||
// chunks: list[dict] | None
|
||||
// json_result: list[dict] | None (alias "json")
|
||||
// markdown_result: str | None (alias "markdown")
|
||||
// text_result: str | None (alias "text")
|
||||
// html_result: str | None (alias "html")
|
||||
type TokenizerFromUpstream struct {
|
||||
CreatedTime *float64 `json:"_created_time,omitempty"`
|
||||
ElapsedTime *float64 `json:"_elapsed_time,omitempty"`
|
||||
|
||||
// Name is the source document name. Optional in this schema
|
||||
// (Python default = "").
|
||||
Name string `json:"name,omitempty"`
|
||||
|
||||
// File is the optional upstream file descriptor.
|
||||
File *ChunkerFileMeta `json:"file,omitempty"`
|
||||
|
||||
// OutputFormat controls which of the *Result fields below is the
|
||||
// active payload. Allowed values:
|
||||
// "json" -> JSONResult
|
||||
// "markdown" -> MarkdownResult
|
||||
// "text" -> TextResult
|
||||
// "html" -> HTMLResult
|
||||
// "chunks" -> Chunks
|
||||
OutputFormat PayloadFormat `json:"output_format,omitempty"`
|
||||
|
||||
// Chunks is the upstream chunk list. Set when OutputFormat == "chunks".
|
||||
Chunks []ChunkDoc `json:"chunks,omitempty"`
|
||||
|
||||
// JSONResult is the upstream structured JSON list (alias "json").
|
||||
JSONResult []ChunkDoc `json:"json,omitempty"`
|
||||
|
||||
// MarkdownResult is the upstream markdown payload (alias "markdown").
|
||||
MarkdownResult *string `json:"markdown,omitempty"`
|
||||
|
||||
// TextResult is the upstream plain-text payload (alias "text").
|
||||
TextResult *string `json:"text,omitempty"`
|
||||
|
||||
// HTMLResult is the upstream HTML payload (alias "html").
|
||||
HTMLResult *string `json:"html,omitempty"`
|
||||
}
|
||||
|
||||
// Validate enforces the Python model_validator invariants: when
|
||||
// OutputFormat is "markdown" / "text" / "html", the matching
|
||||
// *Result field must be non-nil (the Go equivalent of a non-None
|
||||
// string in Python); when OutputFormat is empty, nil, or any other
|
||||
// value, JSONResult (or Chunks) must be supplied.
|
||||
//
|
||||
// The intent is to mirror the Python error messages verbatim where
|
||||
// possible. The Tokenizer's runtime contract treats an empty chunk
|
||||
// list as valid (the component short-circuits silently), so Chunks
|
||||
// being nil is allowed even when OutputFormat == "chunks".
|
||||
func (t *TokenizerFromUpstream) Validate() error {
|
||||
if !t.OutputFormat.isKnown() {
|
||||
return errInvalidValue{Field: "output_format", Value: string(t.OutputFormat)}
|
||||
}
|
||||
switch t.OutputFormat {
|
||||
case PayloadFormatChunks:
|
||||
// Chunks may be nil (zero-length is valid). No-op.
|
||||
return nil
|
||||
case PayloadFormatMarkdown:
|
||||
if t.MarkdownResult == nil {
|
||||
return errRequiredField{Field: "markdown"}
|
||||
}
|
||||
case PayloadFormatText:
|
||||
if t.TextResult == nil {
|
||||
return errRequiredField{Field: "text"}
|
||||
}
|
||||
case PayloadFormatHTML:
|
||||
if t.HTMLResult == nil {
|
||||
return errRequiredField{Field: "html"}
|
||||
}
|
||||
default:
|
||||
// Empty / "json" / any other value: require a JSON list payload
|
||||
// OR a Chunks list. Mirrors the Python check.
|
||||
if t.JSONResult == nil && t.Chunks == nil {
|
||||
return errRequiredField{Field: "json"}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// TokenizerParam is the static configuration for the Tokenizer
|
||||
// component. Mirrors rag/flow/tokenizer/tokenizer.py:TokenizerParam.
|
||||
//
|
||||
// search_method: list[str] # ["full_text", "embedding"]
|
||||
// filename_embd_weight: float # 0.1
|
||||
// fields: list[str] # ["text"]
|
||||
type TokenizerParam struct {
|
||||
// SearchMethod controls which tokenization/embedding passes run.
|
||||
// Allowed values: "full_text", "embedding".
|
||||
SearchMethod []string `json:"search_method"`
|
||||
|
||||
// FilenameEmbdWeight blends the document-name embedding into each
|
||||
// chunk embedding. Value in [0.0, 1.0].
|
||||
FilenameEmbdWeight float64 `json:"filename_embd_weight"`
|
||||
|
||||
// Fields selects which fields of each chunk to embed. Python
|
||||
// supports either a single string (then auto-wrapped to a list at
|
||||
// runtime) or a list of strings. Go callers should pass a slice.
|
||||
Fields []string `json:"fields"`
|
||||
}
|
||||
|
||||
// Defaults returns the Python default TokenizerParam.
|
||||
func (TokenizerParam) Defaults() TokenizerParam {
|
||||
return TokenizerParam{
|
||||
SearchMethod: []string{"full_text", "embedding"},
|
||||
FilenameEmbdWeight: 0.1,
|
||||
Fields: []string{"text"},
|
||||
}
|
||||
}
|
||||
|
||||
// Validate enforces the SearchMethod enum. Fields and
|
||||
// FilenameEmbdWeight have no schema-level range checks in the Python
|
||||
// `check()`.
|
||||
func (p *TokenizerParam) Validate() error {
|
||||
if len(p.SearchMethod) == 0 {
|
||||
return errRequiredField{Field: "search_method"}
|
||||
}
|
||||
for _, m := range p.SearchMethod {
|
||||
switch m {
|
||||
case "full_text", "embedding":
|
||||
default:
|
||||
return errInvalidValue{Field: "search_method", Value: m}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// TokenizerOutputs is the result of invoking the Tokenizer component.
|
||||
// Mirrors what the Python component sets via `self.set_output(...)` at
|
||||
// rag/flow/tokenizer/tokenizer.py:_invoke.
|
||||
//
|
||||
// Always sets:
|
||||
// - output_format = "chunks"
|
||||
// - chunks (the tokenized chunk list)
|
||||
//
|
||||
// Optionally sets:
|
||||
// - embedding_token_consumption (when search_method includes "embedding")
|
||||
// - _ERROR
|
||||
type TokenizerOutputs struct {
|
||||
// OutputFormat is always "chunks".
|
||||
OutputFormat PayloadFormat `json:"output_format,omitempty"`
|
||||
|
||||
// Chunks is the tokenized chunk list. Each entry is a free-form map
|
||||
// containing tokenized fields (title_tks, content_ltks, ...),
|
||||
// chunk_order_int, and (when "embedding" is in search_method)
|
||||
// the q_<n>_vec vector field.
|
||||
Chunks []ChunkDoc `json:"chunks,omitempty"`
|
||||
|
||||
// EmbeddingTokenConsumption records the token count consumed by
|
||||
// the embedding call. Set only when search_method includes
|
||||
// "embedding". *int to distinguish unset (zero value) from 0.
|
||||
EmbeddingTokenConsumption *int `json:"embedding_token_consumption,omitempty"`
|
||||
|
||||
// Error is set when the component short-circuits with an error
|
||||
// message (Python: set_output("_ERROR", ...)).
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
148
internal/ingestion/component/slice3_test.go
Normal file
148
internal/ingestion/component/slice3_test.go
Normal file
@@ -0,0 +1,148 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Slice 3 tests for port-rag-flow-pipeline-to-go.md Phase 5.
|
||||
//
|
||||
// Pins the Tokenizer content_with_weight fallback and the
|
||||
// Extractor prompt placeholder substitution.
|
||||
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestTokenizer_FallsBackToContentWithWeight pins the python
|
||||
// rag/flow/tokenizer.py:111 fallback. A chunk with only
|
||||
// content_with_weight (no text) must tokenize the fallback text.
|
||||
func TestTokenizer_FallsBackToContentWithWeight(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
c := &TokenizerComponent{}
|
||||
c.param.SearchMethod = []string{"full_text"}
|
||||
c.param.Fields = []string{"text"}
|
||||
c.param.FilenameEmbdWeight = 0
|
||||
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc",
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{
|
||||
{"content_with_weight": "fallback text", "doc_type_kwd": "text"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Tokenizer.Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("no chunks emitted")
|
||||
}
|
||||
if got := chunks[0]["content_ltks"]; got == nil {
|
||||
t.Errorf("content_ltks missing; content_with_weight fallback did not run")
|
||||
} else if s, ok := got.(string); !ok || s == "" {
|
||||
t.Errorf("content_ltks = %v (type %T), want non-empty string", got, got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizer_DoesNotChangeChunkText pins the regression guard
|
||||
// for the fallback. When both "text" and "content_with_weight"
|
||||
// are present, "text" wins.
|
||||
func TestTokenizer_DoesNotChangeChunkText(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
c := &TokenizerComponent{}
|
||||
c.param.SearchMethod = []string{"full_text"}
|
||||
c.param.Fields = []string{"text"}
|
||||
c.param.FilenameEmbdWeight = 0
|
||||
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc",
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{
|
||||
{"text": "primary text", "content_with_weight": "fallback", "doc_type_kwd": "text"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Tokenizer.Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("no chunks emitted")
|
||||
}
|
||||
if got, want := chunks[0]["text"], "primary text"; got != want {
|
||||
t.Errorf("text = %v, want %v (text should win over content_with_weight)", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSubstitutePromptPlaceholders_ReplacesAtChunks pins the
|
||||
// resume-template pattern `{TitleChunker:FlatMiceFix@chunks}`.
|
||||
// The substitute is the joined chunk text.
|
||||
func TestSubstitutePromptPlaceholders_ReplacesAtChunks(t *testing.T) {
|
||||
prompt := "Extract metadata from: {TitleChunker:FlatMiceFix@chunks}"
|
||||
chunks := []map[string]any{
|
||||
{"text": "First chunk."},
|
||||
{"text": "Second chunk."},
|
||||
}
|
||||
got := substitutePromptPlaceholders(prompt, chunks)
|
||||
if strings.Contains(got, "{TitleChunker:FlatMiceFix@chunks}") {
|
||||
t.Errorf("placeholder not substituted: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, "First chunk.") || !strings.Contains(got, "Second chunk.") {
|
||||
t.Errorf("substitute missing chunk content: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSubstitutePromptPlaceholders_LeavesPatternWhenNoChunks pins
|
||||
// the opt-in substitution rule. When chunks is empty the
|
||||
// placeholder is left intact so a misconfigured template surfaces
|
||||
// as a clear pattern rather than silently disappearing.
|
||||
func TestSubstitutePromptPlaceholders_LeavesPatternWhenNoChunks(t *testing.T) {
|
||||
prompt := "Extract metadata from: {TitleChunker:FlatMiceFix@chunks}"
|
||||
got := substitutePromptPlaceholders(prompt, nil)
|
||||
if got != prompt {
|
||||
t.Errorf("empty chunks: placeholder should be preserved\n got: %q\n want: %q", got, prompt)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSubstitutePromptPlaceholders_NoPlaceholderInPrompt pins the
|
||||
// no-op behaviour when the prompt carries no @chunks pattern.
|
||||
func TestSubstitutePromptPlaceholders_NoPlaceholderInPrompt(t *testing.T) {
|
||||
prompt := "Plain prompt with no substitution."
|
||||
chunks := []map[string]any{{"text": "x"}}
|
||||
got := substitutePromptPlaceholders(prompt, chunks)
|
||||
if got != prompt {
|
||||
t.Errorf("no-placeholder prompt should be unchanged\n got: %q\n want: %q", got, prompt)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSubstitutePromptPlaceholders_SkipsEmptyChunkText pins the
|
||||
// per-chunk text trim. A chunk with no text field does not
|
||||
// contribute a trailing blank line.
|
||||
func TestSubstitutePromptPlaceholders_SkipsEmptyChunkText(t *testing.T) {
|
||||
prompt := "p {TitleChunker:FlatMiceFix@chunks} q"
|
||||
chunks := []map[string]any{
|
||||
{"text": ""},
|
||||
{"text": "actual"},
|
||||
{},
|
||||
}
|
||||
got := substitutePromptPlaceholders(prompt, chunks)
|
||||
if strings.Contains(got, "{TitleChunker:FlatMiceFix@chunks}") {
|
||||
t.Errorf("placeholder not substituted: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, "actual") {
|
||||
t.Errorf("chunk text missing: %q", got)
|
||||
}
|
||||
}
|
||||
625
internal/ingestion/component/tokenizer.go
Normal file
625
internal/ingestion/component/tokenizer.go
Normal file
@@ -0,0 +1,625 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Tokenizer ingestion component (Phase 2.4 of
|
||||
// port-rag-flow-pipeline-to-go.md §4). Port of Python
|
||||
// `rag/flow/tokenizer/tokenizer.py`. Computes (a) full-text token
|
||||
// counts via the Go tokenizer package and (b) embedding vectors via
|
||||
// the tenant's embedding model.
|
||||
//
|
||||
// SCOPE (honest):
|
||||
//
|
||||
// - TOKEN COUNTING: matched at the wire level. Each chunk gets
|
||||
// `content_ltks` (tokenized string via `tokenizer.Tokenize`) and
|
||||
// `content_sm_ltks` (fine-grained variant) when `search_method`
|
||||
// includes `full_text`. `title_tks` / `title_sm_tks` mirror the
|
||||
// upstream `name` field. Python uses C++ RAGAnalyzer via
|
||||
// `rag_tokenizer`; the Go side goes through `internal/tokenizer`
|
||||
// which itself calls into the same C++ binding (`internal/binding`).
|
||||
// For non-ASCII (CJK) input, Python's `rag_tokenizer.tokenize`
|
||||
// falls back gracefully; the Go path uses the CGo analyzer
|
||||
// when initialized, otherwise an empty string — see
|
||||
// `internal/tokenizer/tokenizer.go:Tokenize` (Infinity engine
|
||||
// returns input unchanged; otherwise the C++ binding is used).
|
||||
//
|
||||
// - CJK CAVEAT (plan §8 Q2): The `NumTokensFromString` helper in
|
||||
// `internal/tokenizer` falls back to `len([]byte(s))` on a
|
||||
// tiktoken-init failure (over-counts CJK). The Python equivalent
|
||||
// returns 0. The Go port KEEPS the Go behaviour — the tokenizer
|
||||
// package is the single source of truth for token counting and
|
||||
// must not be re-implemented here. Test
|
||||
// `TestTokenizerComponent_Invoke_Unicode` asserts only that the
|
||||
// count is finite and non-negative, matching the test
|
||||
// convention in plan §6 (coverage target:
|
||||
// "Tokenizer returns finite token counts for empty / unicode /
|
||||
// mixed-script text").
|
||||
//
|
||||
// - EMBEDDING MODEL RESOLUTION: mirrored. Python uses
|
||||
// `LLMBundle(tenant_id, embd_id).encode([...])` from
|
||||
// `rag/flow/tokenizer/tokenizer.py:54-66`; the Go port goes
|
||||
// through `service.ModelProviderService.GetEmbeddingModel`
|
||||
// (callers inject the model bundle, see `EncodeFunc` below).
|
||||
// The component does NOT directly construct a model driver —
|
||||
// the resolution path depends on tenant/DAO context that lives
|
||||
// in `internal/service`, and importing `internal/service` from
|
||||
// `internal/ingestion/component` would invert the dependency
|
||||
// direction (plan §3 import graph: ingestion → agent/runtime
|
||||
// only). The injection point is `EncodeFunc` (package-level
|
||||
// var); production wires it in `main()` (or an analogous
|
||||
// bootstrap step) and tests inject a stub. When `EncodeFunc` is
|
||||
// nil the component short-circuits the embedding branch with
|
||||
// a clear error — the same fail-loud contract the Python side
|
||||
// enforces via `LLMBundle` constructor.
|
||||
//
|
||||
// - BATCHED EMBEDDING (plan §AD-5a): matched. The Python path
|
||||
// chunks calls by `settings.EMBEDDING_BATCH_SIZE` (default 16)
|
||||
// and uses an async semaphore (`embed_limiter`). The Go port
|
||||
// issues ONE `Encode([]string)` call with the entire chunk
|
||||
// list (AD-5a calls out "embedding calls batched, not fanned"
|
||||
// and Parallelism=1). Drivers that need to chunk internally
|
||||
// can do so — the wire call is one round-trip.
|
||||
//
|
||||
// - TRACKING: WithTimeout (60s, matches python `@timeout(60)` on
|
||||
// `batch_encode`), TrackProgress, TrackElapsed. See
|
||||
// `internal/agent/runtime/helpers.go` (plan §1 Phase 1).
|
||||
//
|
||||
// - WHAT IS NOT PORTED:
|
||||
//
|
||||
// - The python `finalize_pdf_chunk` post-step — that
|
||||
// normalizes PDF bbox metadata; it lives in
|
||||
// `rag/flow/parser/pdf_chunk_metadata.py` and is the Parser
|
||||
// component's concern (Phase 2.2).
|
||||
//
|
||||
// - `rag.flow.tokenizer` `thread_pool_exec` async batching +
|
||||
// `embed_limiter` semaphore — replaced by the single
|
||||
// batched `Encode` call.
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
"ragflow/internal/tokenizer"
|
||||
)
|
||||
|
||||
const ComponentNameTokenizer = "Tokenizer"
|
||||
|
||||
// tokenizerTimeout bounds the batched embedding call. Mirrors the
|
||||
// python `@timeout(60)` decorator on `Tokenizer._embedding.embed_limiter`
|
||||
// + `batch_encode` in tokenizer.py:92-104. Declared as a var so tests
|
||||
// can shrink it; production wiring uses 60s.
|
||||
var tokenizerTimeout = 60 * time.Second
|
||||
|
||||
// titleExtRE strips a trailing file-extension (e.g. ".pdf") from the
|
||||
// upstream document name before tokenizing it. Mirrors the python
|
||||
// `re.sub(r"\.[a-zA-Z]+$", "", name)` in tokenizer.py:137.
|
||||
var titleExtRE = regexp.MustCompile(`\.[a-zA-Z]+$`)
|
||||
|
||||
// htmlTableRE matches HTML table-cell tags so the embedded text fed
|
||||
// to the embedding model doesn't carry raw markup. Mirrors the python
|
||||
// `re.sub(r"</?(table|td|caption|tr|th)( [^<>]{0,12})?>", " ", txt)` at
|
||||
// tokenizer.py:79.
|
||||
var htmlTableRE = regexp.MustCompile(`</?(table|td|caption|tr|th)( [^<>]{0,12})?>`)
|
||||
|
||||
// Embedder is the testability seam for the embedding branch. The
|
||||
// production wiring injects an implementation that resolves an
|
||||
// embedding model via `service.ModelProviderService.GetEmbeddingModel`
|
||||
// and calls its `ModelDriver.Embed`. Tests inject a stub.
|
||||
//
|
||||
// Returning one vector per input text (length len(texts), each
|
||||
// vector non-empty) is the contract; nil/error halts the component.
|
||||
type Embedder interface {
|
||||
Encode(texts []string) ([][]float64, error)
|
||||
}
|
||||
|
||||
// EncodeFunc is the package-level injection point. nil means
|
||||
// "embedding disabled" — the component skips the embedding branch
|
||||
// (matching the python behaviour when `search_method` omits
|
||||
// "embedding"). Production sets this once in `main()`; tests can
|
||||
// swap it with a stub via the test helpers in `tokenizer_test.go`.
|
||||
var EncodeFunc func(tenantID, embdID string) Embedder
|
||||
|
||||
// TokenizerComponent computes token counts and (optionally) embedding
|
||||
// vectors for an upstream chunk list. Mirrors python
|
||||
// rag/flow/tokenizer/tokenizer.py:Tokenizer.
|
||||
//
|
||||
// Inputs:
|
||||
//
|
||||
// tenant_id (string, optional) — used to resolve the embedding model
|
||||
// model_id (string, optional) — explicit override; falls back to
|
||||
// Param.EmbeddingID (future)
|
||||
// output_format (string) — one of json/markdown/text/html/chunks
|
||||
// chunks (list[map]) — chunk list when output_format == "chunks"
|
||||
// json (list[map]) — structured parser payload when output_format == "json" or unset
|
||||
// markdown/text/html — scalar payload matching output_format
|
||||
//
|
||||
// Outputs:
|
||||
//
|
||||
// chunks — the chunk list with tokenized fields
|
||||
// and (when embedding is requested)
|
||||
// q_<n>_vec vector fields
|
||||
// embedding_token_consumption — non-negative int (matches the python
|
||||
// `embedding_token_consumption` output)
|
||||
// output_format — always "chunks" (matches python set_output)
|
||||
// _created_time / _elapsed_time — TrackElapsed bookkeeping
|
||||
type TokenizerComponent struct {
|
||||
param schema.TokenizerParam
|
||||
}
|
||||
|
||||
// NewTokenizerComponent constructs a TokenizerComponent from DSL
|
||||
// params. Mirrors python `TokenizerParam` defaults (search_method =
|
||||
// ["full_text","embedding"], filename_embd_weight=0.1, fields=["text"]).
|
||||
func NewTokenizerComponent(params map[string]any) (runtime.Component, error) {
|
||||
p := schema.TokenizerParam{}.Defaults()
|
||||
if params != nil {
|
||||
if v, ok := params["search_method"]; ok {
|
||||
// Replace (not append) so a caller-supplied
|
||||
// search_method = ["full_text"] correctly disables
|
||||
// embedding. Python's TokenizerParam similarly treats
|
||||
// caller-supplied values as the full set.
|
||||
p.SearchMethod = nil
|
||||
switch t := v.(type) {
|
||||
case []any:
|
||||
for _, x := range t {
|
||||
if s, ok := x.(string); ok {
|
||||
p.SearchMethod = append(p.SearchMethod, s)
|
||||
}
|
||||
}
|
||||
case []string:
|
||||
p.SearchMethod = append(p.SearchMethod, t...)
|
||||
}
|
||||
}
|
||||
if v, ok := params["filename_embd_weight"]; ok {
|
||||
switch t := v.(type) {
|
||||
case float64:
|
||||
p.FilenameEmbdWeight = t
|
||||
case int:
|
||||
p.FilenameEmbdWeight = float64(t)
|
||||
}
|
||||
}
|
||||
if v, ok := params["fields"]; ok {
|
||||
switch t := v.(type) {
|
||||
case string:
|
||||
p.Fields = []string{t}
|
||||
case []any:
|
||||
for _, x := range t {
|
||||
if s, ok := x.(string); ok {
|
||||
p.Fields = append(p.Fields, s)
|
||||
}
|
||||
}
|
||||
case []string:
|
||||
p.Fields = append(p.Fields, t...)
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := p.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("Tokenizer: param check: %w", err)
|
||||
}
|
||||
return &TokenizerComponent{param: p}, nil
|
||||
}
|
||||
|
||||
// Inputs returns the parameter metadata.
|
||||
func (c *TokenizerComponent) Inputs() map[string]string {
|
||||
return map[string]string{
|
||||
"tenant_id": "Tenant identifier used to resolve the embedding model (mirrors python self._canvas._tenant_id).",
|
||||
"model_id": "Optional explicit embedding-model override. Falls back to EncodeFunc resolution when unset.",
|
||||
"output_format": "Upstream payload discriminator: json / markdown / text / html / chunks.",
|
||||
"chunks": "List of chunk maps when output_format == \"chunks\".",
|
||||
"json": "Structured parser payload when output_format == \"json\" or unset.",
|
||||
"text": "Plain-text payload when output_format == \"text\".",
|
||||
"markdown": "Markdown payload when output_format == \"markdown\".",
|
||||
"html": "HTML payload when output_format == \"html\".",
|
||||
"name": "Upstream document name (used for title_tks and the title-blended embedding).",
|
||||
}
|
||||
}
|
||||
|
||||
// Outputs returns the parameter metadata. Mirrors python set_output
|
||||
// contract for Tokenizer.
|
||||
func (c *TokenizerComponent) Outputs() map[string]string {
|
||||
return map[string]string{
|
||||
"chunks": "Tokenized chunk list (each entry gains content_ltks / content_sm_ltks / title_tks and, when embedding is requested, q_<n>_vec).",
|
||||
"embedding_token_consumption": "Non-negative token count consumed by the embedding call. Omitted when no embedding ran.",
|
||||
"output_format": "Always \"chunks\" (matches python set_output).",
|
||||
"_created_time": "RFC3339Nano creation timestamp (TrackElapsed).",
|
||||
"_elapsed_time": "Wall-clock seconds (TrackElapsed).",
|
||||
}
|
||||
}
|
||||
|
||||
// Parallelism is fixed at 1 — embedding calls are batched in one
|
||||
// round-trip (plan §2 AD-5a "Tokenizer: 1 (embedding calls batched,
|
||||
// not fanned)").
|
||||
func (c *TokenizerComponent) Parallelism() int { return 1 }
|
||||
|
||||
// Invoke computes tokens + embeddings for the upstream chunks.
|
||||
//
|
||||
// Failure modes:
|
||||
//
|
||||
// - "embedding" requested but EncodeFunc is nil → returns an
|
||||
// error (fail-loud: same contract as python when LLMBundle is
|
||||
// unconstructable).
|
||||
// - Empty chunks list → returns an empty chunks output without
|
||||
// panicking (python tokenizer.py:121 treats this as valid).
|
||||
// - Per-chunk empty cleaned text → chunk is skipped from the
|
||||
// embedding batch (python tokenizer.py:80-82 `if not cleaned_txt:
|
||||
// continue`), but the chunk still carries tokenized fields if
|
||||
// `full_text` is in `search_method`.
|
||||
func (c *TokenizerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
tenantID := getStringOr(inputs, "tenant_id", "")
|
||||
modelID := getStringOr(inputs, "model_id", "")
|
||||
upstream, err := decodeTokenizerFromUpstream(inputs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
chunks := chunksFromTokenizerUpstream(upstream)
|
||||
name := upstream.Name
|
||||
titleStem := titleExtRE.ReplaceAllString(name, "")
|
||||
|
||||
// TrackElapsed wraps the whole pipeline (tokenize + embed) so the
|
||||
// upstream caller sees consistent _created_time / _elapsed_time
|
||||
// stamps matching python `ProcessBase` (helpers.go TrackElapsed).
|
||||
return runtime.TrackElapsed("Tokenizer", func() (map[string]any, error) {
|
||||
// content_with_weight fallback — populate each chunk's
|
||||
// "text" from the python-equivalent field when empty.
|
||||
// Done before tokenizeChunks so the chunker's emitted
|
||||
// text is the authoritative input.
|
||||
normalizeChunkTextFallback(chunks)
|
||||
|
||||
// full_text pass — tokenize each chunk's text fields. Mirrors
|
||||
// python tokenizer.py:130-185.
|
||||
if contains(c.param.SearchMethod, "full_text") {
|
||||
if err := tokenizeChunks(chunks, titleStem); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
out := map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": schema.ChunkDocsToMaps(chunks),
|
||||
}
|
||||
|
||||
// embedding pass — batched single call (plan §AD-5a).
|
||||
if contains(c.param.SearchMethod, "embedding") {
|
||||
if EncodeFunc == nil {
|
||||
return nil, fmt.Errorf("Tokenizer: embedding requested but EncodeFunc is unset")
|
||||
}
|
||||
embedder := EncodeFunc(tenantID, modelID)
|
||||
if embedder == nil {
|
||||
return nil, fmt.Errorf("Tokenizer: embedding requested but encoder resolution returned nil")
|
||||
}
|
||||
|
||||
// Build the batched text list + index pairs.
|
||||
texts := make([]string, 0, len(chunks))
|
||||
pairs := make([]int, 0, len(chunks))
|
||||
for i, ck := range chunks {
|
||||
txt := concatFields(ck, c.param.Fields)
|
||||
txt = htmlTableRE.ReplaceAllString(txt, " ")
|
||||
txt = strings.TrimSpace(txt)
|
||||
if txt == "" {
|
||||
continue
|
||||
}
|
||||
texts = append(texts, txt)
|
||||
pairs = append(pairs, i)
|
||||
}
|
||||
|
||||
if len(texts) > 0 {
|
||||
var (
|
||||
vects [][]float64
|
||||
encErr error
|
||||
)
|
||||
timeoutErr := runtime.WithTimeout(ctx, tokenizerTimeout, func(timeoutCtx context.Context) error {
|
||||
vects, encErr = embedder.Encode(texts)
|
||||
return encErr
|
||||
})
|
||||
if timeoutErr != nil {
|
||||
return nil, fmt.Errorf("Tokenizer: encode: %w", timeoutErr)
|
||||
}
|
||||
if len(vects) != len(pairs) {
|
||||
return nil, fmt.Errorf("Tokenizer: encode returned %d vectors for %d chunks", len(vects), len(pairs))
|
||||
}
|
||||
for k, idx := range pairs {
|
||||
ck := &chunks[idx]
|
||||
v := vects[k]
|
||||
if err := ck.SetExtraValue(fmt.Sprintf("q_%d_vec", len(v)), v); err != nil {
|
||||
return nil, fmt.Errorf("Tokenizer: vector marshal: %w", err)
|
||||
}
|
||||
}
|
||||
// token_count: best-effort approximation matching the
|
||||
// python contract — the Go Embedder doesn't surface
|
||||
// per-call token usage, so we sum
|
||||
// `NumTokensFromString` for each chunk text.
|
||||
tokenCount := 0
|
||||
for _, t := range texts {
|
||||
tokenCount += tokenizer.NumTokensFromString(t)
|
||||
}
|
||||
out["embedding_token_consumption"] = tokenCount
|
||||
out["chunks"] = schema.ChunkDocsToMaps(chunks)
|
||||
}
|
||||
}
|
||||
|
||||
return out, nil
|
||||
})
|
||||
}
|
||||
|
||||
func decodeTokenizerFromUpstream(inputs map[string]any) (schema.TokenizerFromUpstream, error) {
|
||||
var out schema.TokenizerFromUpstream
|
||||
if inputs == nil {
|
||||
return out, fmt.Errorf("Tokenizer: inputs map is nil")
|
||||
}
|
||||
data, err := json.Marshal(stripRuntimeTimestamps(inputs))
|
||||
if err != nil {
|
||||
return out, fmt.Errorf("Tokenizer: encode inputs: %w", err)
|
||||
}
|
||||
if err := json.Unmarshal(data, &out); err != nil {
|
||||
return out, fmt.Errorf("Tokenizer: decode inputs: %w", err)
|
||||
}
|
||||
if err := out.Validate(); err != nil {
|
||||
return out, fmt.Errorf("Tokenizer: input error: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func stripRuntimeTimestamps(inputs map[string]any) map[string]any {
|
||||
out := make(map[string]any, len(inputs))
|
||||
for k, v := range inputs {
|
||||
if k == "_created_time" || k == "_elapsed_time" {
|
||||
continue
|
||||
}
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func chunksFromTokenizerUpstream(in schema.TokenizerFromUpstream) []schema.ChunkDoc {
|
||||
switch in.OutputFormat {
|
||||
case schema.PayloadFormatChunks:
|
||||
return cloneChunkDocs(in.Chunks)
|
||||
case schema.PayloadFormatMarkdown:
|
||||
return textPayloadToChunks(in.MarkdownResult)
|
||||
case schema.PayloadFormatText:
|
||||
return textPayloadToChunks(in.TextResult)
|
||||
case schema.PayloadFormatHTML:
|
||||
return textPayloadToChunks(in.HTMLResult)
|
||||
default:
|
||||
return cloneChunkDocs(in.JSONResult)
|
||||
}
|
||||
}
|
||||
|
||||
func textPayloadToChunks(payload *string) []schema.ChunkDoc {
|
||||
if payload == nil || strings.TrimSpace(*payload) == "" {
|
||||
return []schema.ChunkDoc{}
|
||||
}
|
||||
return []schema.ChunkDoc{{Text: *payload}}
|
||||
}
|
||||
|
||||
func cloneChunkDocs(in []schema.ChunkDoc) []schema.ChunkDoc {
|
||||
if len(in) == 0 {
|
||||
return []schema.ChunkDoc{}
|
||||
}
|
||||
out := make([]schema.ChunkDoc, len(in))
|
||||
for i := range in {
|
||||
out[i] = cloneTokenizerChunkDoc(in[i])
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneTokenizerChunkDoc(in schema.ChunkDoc) schema.ChunkDoc {
|
||||
out := in
|
||||
if in.TKNums != nil {
|
||||
v := *in.TKNums
|
||||
out.TKNums = &v
|
||||
}
|
||||
if in.ChunkOrderInt != nil {
|
||||
v := *in.ChunkOrderInt
|
||||
out.ChunkOrderInt = &v
|
||||
}
|
||||
if in.PageNumber != nil {
|
||||
v := *in.PageNumber
|
||||
out.PageNumber = &v
|
||||
}
|
||||
if in.Extra != nil {
|
||||
out.Extra = make(map[string]json.RawMessage, len(in.Extra))
|
||||
for k, v := range in.Extra {
|
||||
out.Extra[k] = append(json.RawMessage(nil), v...)
|
||||
}
|
||||
}
|
||||
if len(in.PDFPositions) > 0 {
|
||||
out.PDFPositions = append(json.RawMessage(nil), in.PDFPositions...)
|
||||
}
|
||||
if len(in.Positions) > 0 {
|
||||
out.Positions = append(json.RawMessage(nil), in.Positions...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// normalizeChunkTextFallback populates each chunk's "text" key
|
||||
// from "content_with_weight" when "text" is absent or empty. Mirrors
|
||||
// the python rag/flow/tokenizer.py:111 fallback so a chunk that
|
||||
// arrives from the parser path with only the structured
|
||||
// content_with_weight field still tokenizes.
|
||||
//
|
||||
// The function mutates the input slice in place; callers should
|
||||
// not retain separate copies of the chunks map. If both fields
|
||||
// are present, the existing "text" wins — preserves the python
|
||||
// contract where the chunker's emitted text is authoritative.
|
||||
func normalizeChunkTextFallback(chunks []schema.ChunkDoc) {
|
||||
for i := range chunks {
|
||||
if chunks[i].Text != "" {
|
||||
continue
|
||||
}
|
||||
if chunks[i].ContentWithWeight != "" {
|
||||
chunks[i].Text = chunks[i].ContentWithWeight
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// tokenizeChunks annotates each chunk with title_tks, content_ltks,
|
||||
// and (when applicable) question_tks / important_tks / summary fields.
|
||||
// Mirrors python tokenizer.py:130-185.
|
||||
func tokenizeChunks(chunks []schema.ChunkDoc, titleStem string) error {
|
||||
for i := range chunks {
|
||||
ck := &chunks[i]
|
||||
ck.ChunkOrderInt = intPtr(i)
|
||||
titleTk, err := tokenizer.Tokenize(titleStem)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: title tokenize: %w", err)
|
||||
}
|
||||
titleSmTk, err := tokenizer.FineGrainedTokenize(titleTk)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: title fine-grain: %w", err)
|
||||
}
|
||||
ck.TitleTks = titleTk
|
||||
ck.TitleSmTks = titleSmTk
|
||||
|
||||
// Question / keyword / summary fields are optional. The python
|
||||
// path branches on each independently.
|
||||
if q := ck.Questions; q != "" {
|
||||
if err := ck.SetExtraValue("question_kwd", strings.Split(q, "\n")); err != nil {
|
||||
return fmt.Errorf("Tokenizer: question keywords marshal: %w", err)
|
||||
}
|
||||
qt, err := tokenizer.Tokenize(q)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: question tokenize: %w", err)
|
||||
}
|
||||
if err := ck.SetExtraValue("question_tks", qt); err != nil {
|
||||
return fmt.Errorf("Tokenizer: question tokens marshal: %w", err)
|
||||
}
|
||||
}
|
||||
if kw := ck.Keywords; kw != "" {
|
||||
if err := ck.SetExtraValue("important_kwd", strings.Split(kw, ",")); err != nil {
|
||||
return fmt.Errorf("Tokenizer: keyword list marshal: %w", err)
|
||||
}
|
||||
it, err := tokenizer.Tokenize(kw)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: keyword tokenize: %w", err)
|
||||
}
|
||||
if err := ck.SetExtraValue("important_tks", it); err != nil {
|
||||
return fmt.Errorf("Tokenizer: keyword tokens marshal: %w", err)
|
||||
}
|
||||
}
|
||||
if s := ck.Summary; s != "" {
|
||||
st, err := tokenizer.Tokenize(s)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: summary tokenize: %w", err)
|
||||
}
|
||||
ck.ContentLtks = st
|
||||
smt, err := tokenizer.FineGrainedTokenize(st)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: summary fine-grain: %w", err)
|
||||
}
|
||||
ck.ContentSmLtks = smt
|
||||
} else if t := ck.Text; t != "" {
|
||||
tt, err := tokenizer.Tokenize(t)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: text tokenize: %w", err)
|
||||
}
|
||||
ck.ContentLtks = tt
|
||||
smt, err := tokenizer.FineGrainedTokenize(tt)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: text fine-grain: %w", err)
|
||||
}
|
||||
ck.ContentSmLtks = smt
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// concatFields concatenates the configured fields of a chunk into
|
||||
// a single string. Mirrors python tokenizer.py:69-79 which
|
||||
// concatenates `param.fields` (string or list-of-strings per chunk).
|
||||
func concatFields(ck schema.ChunkDoc, fields []string) string {
|
||||
var b strings.Builder
|
||||
for _, f := range fields {
|
||||
switch f {
|
||||
case "text":
|
||||
b.WriteString(ck.Text)
|
||||
case "content_with_weight":
|
||||
b.WriteString(ck.ContentWithWeight)
|
||||
case "questions":
|
||||
b.WriteString(ck.Questions)
|
||||
case "keywords":
|
||||
b.WriteString(ck.Keywords)
|
||||
case "summary":
|
||||
b.WriteString(ck.Summary)
|
||||
default:
|
||||
if s, ok := ck.GetExtraString(f); ok {
|
||||
b.WriteString(s)
|
||||
continue
|
||||
}
|
||||
if values, ok := ck.GetExtraStringSlice(f); ok {
|
||||
b.WriteString(strings.Join(values, "\n"))
|
||||
}
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func getStringOr(m map[string]any, key, def string) string {
|
||||
if v, ok := getStringLocal(m, key); ok && v != "" {
|
||||
return v
|
||||
}
|
||||
return def
|
||||
}
|
||||
|
||||
// getStringLocal mirrors file.go's getString; we keep a local copy
|
||||
// so the tokenizer package does not depend on the file package's
|
||||
// helper signature. Reads either a string or a byte slice (JSON
|
||||
// decoding yields string for string fields by default).
|
||||
func getStringLocal(m map[string]any, key string) (string, bool) {
|
||||
v, ok := m[key]
|
||||
if !ok || v == nil {
|
||||
return "", false
|
||||
}
|
||||
switch s := v.(type) {
|
||||
case string:
|
||||
return s, true
|
||||
case []byte:
|
||||
return string(s), true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func contains(s []string, v string) bool {
|
||||
for _, x := range s {
|
||||
if x == v {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func intPtr(v int) *int { return &v }
|
||||
|
||||
// init registers Tokenizer under CategoryIngestion (plan §4
|
||||
// Phase 2.4). The metadata drives Phase 4's GET /api/v1/components
|
||||
// listing.
|
||||
func init() {
|
||||
c := &TokenizerComponent{}
|
||||
runtime.MustRegister(ComponentNameTokenizer, runtime.CategoryIngestion,
|
||||
func(_ string, params map[string]any) (runtime.Component, error) {
|
||||
return NewTokenizerComponent(params)
|
||||
},
|
||||
runtime.Metadata{
|
||||
Version: "1.0.0",
|
||||
Inputs: c.Inputs(),
|
||||
Outputs: c.Outputs(),
|
||||
})
|
||||
}
|
||||
648
internal/ingestion/component/tokenizer_test.go
Normal file
648
internal/ingestion/component/tokenizer_test.go
Normal file
@@ -0,0 +1,648 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/tokenizer"
|
||||
)
|
||||
|
||||
var tokenizerPoolInitErr error
|
||||
|
||||
// TestMain initializes the tokenizer pool before any test runs.
|
||||
// The tokenizer package needs the C++ RAGAnalyzer dictionaries
|
||||
// (see internal/tokenizer.Init) for `Tokenize` /
|
||||
// `FineGrainedTokenize`; without it, `tokenizeChunks` errors out
|
||||
// with "tokenizer pool not initialized". Tests in other packages
|
||||
// initialize the pool at startup; this package must do the same
|
||||
// because the Tokenizer component touches tokenizer.Tokenize.
|
||||
//
|
||||
// If Init fails (e.g., dict path missing in some CI sandboxes),
|
||||
// we log the failure but still run the tests. Cases that exercise
|
||||
// tokenizeChunks will fail rather than skip when the pool is not
|
||||
// initialized.
|
||||
func TestMain(m *testing.M) {
|
||||
cfg := &tokenizer.PoolConfig{
|
||||
DictPath: os.Getenv("RAGFLOW_DICT_PATH"),
|
||||
MinSize: 1,
|
||||
MaxSize: 2,
|
||||
IdleTimeout: 30 * time.Second,
|
||||
AcquireTimeout: 5 * time.Second,
|
||||
}
|
||||
if cfg.DictPath == "" {
|
||||
cfg.DictPath = "/usr/share/infinity/resource"
|
||||
}
|
||||
tokenizerPoolInitErr = tokenizer.Init(cfg)
|
||||
if tokenizerPoolInitErr != nil {
|
||||
fmt.Fprintf(os.Stderr, "tokenizer pool init failed (tests will skip tokenize-dependent cases): %v\n", tokenizerPoolInitErr)
|
||||
}
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
|
||||
func requireTokenizerPool(t *testing.T) {
|
||||
t.Helper()
|
||||
if tokenizerPoolInitErr != nil {
|
||||
t.Skipf("tokenizer pool unavailable: %v", tokenizerPoolInitErr)
|
||||
}
|
||||
}
|
||||
|
||||
// stubEmbedder records every call and returns canned vectors.
|
||||
// Matches the Embedder contract: len(vectors) == len(texts).
|
||||
type stubEmbedder struct {
|
||||
calls atomic.Int32
|
||||
dim int
|
||||
delay time.Duration
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *stubEmbedder) Encode(texts []string) ([][]float64, error) {
|
||||
s.calls.Add(1)
|
||||
if s.delay > 0 {
|
||||
time.Sleep(s.delay)
|
||||
}
|
||||
if s.err != nil {
|
||||
return nil, s.err
|
||||
}
|
||||
out := make([][]float64, len(texts))
|
||||
for i := range texts {
|
||||
v := make([]float64, s.dim)
|
||||
v[0] = float64(i + 1) // mark the index so callers can verify alignment
|
||||
out[i] = v
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// withStubEmbedder installs a stub Embedder and restores the previous
|
||||
// EncodeFunc on cleanup. Returns the stub so the test can assert on
|
||||
// call count / latency.
|
||||
func withStubEmbedder(t *testing.T, dim int) *stubEmbedder {
|
||||
t.Helper()
|
||||
stub := &stubEmbedder{dim: dim}
|
||||
prev := EncodeFunc
|
||||
EncodeFunc = func(_, _ string) Embedder { return stub }
|
||||
t.Cleanup(func() { EncodeFunc = prev })
|
||||
return stub
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Registered verifies init() enrollment
|
||||
// under runtime.CategoryIngestion (Phase 4 / API endpoint depends
|
||||
// on this contract).
|
||||
func TestTokenizerComponent_Registered(t *testing.T) {
|
||||
factory, cat, md, ok := runtime.DefaultRegistry.Lookup("Tokenizer")
|
||||
if !ok {
|
||||
t.Fatal("Tokenizer not registered in runtime.DefaultRegistry")
|
||||
}
|
||||
if cat != runtime.CategoryIngestion {
|
||||
t.Errorf("category = %q, want %q", cat, runtime.CategoryIngestion)
|
||||
}
|
||||
if factory == nil {
|
||||
t.Error("factory is nil")
|
||||
}
|
||||
if len(md.Inputs) == 0 {
|
||||
t.Error("metadata.Inputs empty")
|
||||
}
|
||||
if len(md.Outputs) == 0 {
|
||||
t.Error("metadata.Outputs empty")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Invoke_HappyPath drives three chunks
|
||||
// through both full_text tokenization and embedding. Verifies that
|
||||
// every chunk gains `content_ltks`, `content_sm_ltks`, and a
|
||||
// `q_<n>_vec` vector keyed by the embedder's vector dimension.
|
||||
func TestTokenizerComponent_Invoke_HappyPath(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
const dim = 4
|
||||
withStubEmbedder(t, dim)
|
||||
|
||||
c, err := NewTokenizerComponent(map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenizerComponent: %v", err)
|
||||
}
|
||||
|
||||
chunks := []map[string]any{
|
||||
{"text": "alpha chunk text"},
|
||||
{"text": "bravo chunk text"},
|
||||
{"text": "charlie chunk text"},
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"tenant_id": "t1",
|
||||
"model_id": "embd-1",
|
||||
"name": "doc.pdf",
|
||||
"output_format": "chunks",
|
||||
"chunks": chunks,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
|
||||
gotChunks, ok := out["chunks"].([]map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("chunks type = %T, want []map[string]any", out["chunks"])
|
||||
}
|
||||
if len(gotChunks) != 3 {
|
||||
t.Fatalf("len(chunks) = %d, want 3", len(gotChunks))
|
||||
}
|
||||
for i, ck := range gotChunks {
|
||||
if ck["content_ltks"] == nil {
|
||||
t.Errorf("chunk[%d].content_ltks missing", i)
|
||||
}
|
||||
if ck["content_sm_ltks"] == nil {
|
||||
t.Errorf("chunk[%d].content_sm_ltks missing", i)
|
||||
}
|
||||
if ck["title_tks"] == nil {
|
||||
t.Errorf("chunk[%d].title_tks missing", i)
|
||||
}
|
||||
key := "q_4_vec"
|
||||
v, ok := ck[key].([]float64)
|
||||
if !ok {
|
||||
t.Errorf("chunk[%d].%s missing or wrong type: %T", i, key, ck[key])
|
||||
continue
|
||||
}
|
||||
if len(v) != dim {
|
||||
t.Errorf("chunk[%d].%s len = %d, want %d", i, key, len(v), dim)
|
||||
}
|
||||
if v[0] != float64(i+1) {
|
||||
t.Errorf("chunk[%d].%s[0] = %v, want %d (index alignment)", i, key, v[0], i+1)
|
||||
}
|
||||
}
|
||||
if out["output_format"] != "chunks" {
|
||||
t.Errorf("output_format = %v, want chunks", out["output_format"])
|
||||
}
|
||||
if out["embedding_token_consumption"] == nil {
|
||||
t.Error("embedding_token_consumption missing")
|
||||
}
|
||||
if out["_elapsed_time"] == nil {
|
||||
t.Error("_elapsed_time missing")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Invoke_EmptyChunks covers the no-op branch:
|
||||
// empty chunk list → empty output, no panic, no encoder call.
|
||||
func TestTokenizerComponent_Invoke_EmptyChunks(t *testing.T) {
|
||||
stub := withStubEmbedder(t, 4)
|
||||
c, err := NewTokenizerComponent(map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenizerComponent: %v", err)
|
||||
}
|
||||
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) != 0 {
|
||||
t.Errorf("chunks len = %d, want 0", len(chunks))
|
||||
}
|
||||
if stub.calls.Load() != 0 {
|
||||
t.Errorf("embedder called %d times on empty input, want 0", stub.calls.Load())
|
||||
}
|
||||
if out["output_format"] != "chunks" {
|
||||
t.Errorf("output_format = %v, want chunks", out["output_format"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Invoke_NilChunks covers the nil-input
|
||||
// branch: nil chunks list is treated as zero-length (matches
|
||||
// python `kwargs.get("chunks")` with None).
|
||||
func TestTokenizerComponent_Invoke_NilChunks(t *testing.T) {
|
||||
withStubEmbedder(t, 4)
|
||||
c, _ := NewTokenizerComponent(map[string]any{})
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"output_format": "chunks",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) != 0 {
|
||||
t.Errorf("chunks len = %d, want 0", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Invoke_Unicode asserts CJK input
|
||||
// produces finite, non-negative token counts (plan §8 Q2: Go
|
||||
// `NumTokensFromString` over-counts CJK on tiktoken-init failure;
|
||||
// Python returns 0 — both are valid as long as the count is
|
||||
// finite).
|
||||
func TestTokenizerComponent_Invoke_Unicode(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
withStubEmbedder(t, 4)
|
||||
c, _ := NewTokenizerComponent(map[string]any{})
|
||||
|
||||
inputs := []string{
|
||||
"中文测试文本",
|
||||
"こんにちは世界",
|
||||
"한국어 텍스트",
|
||||
"English mixed 中文 français 日本語",
|
||||
}
|
||||
chunks := make([]map[string]any, len(inputs))
|
||||
for i, txt := range inputs {
|
||||
chunks[i] = map[string]any{"text": txt}
|
||||
}
|
||||
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": chunks,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
gotChunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(gotChunks) != len(inputs) {
|
||||
t.Fatalf("chunks len = %d, want %d", len(gotChunks), len(inputs))
|
||||
}
|
||||
for i, ck := range gotChunks {
|
||||
// Direct call to verify the count contract.
|
||||
tokens := tokenizer.NumTokensFromString(inputs[i])
|
||||
if tokens < 0 {
|
||||
t.Errorf("chunk[%d] token count negative: %d", i, tokens)
|
||||
}
|
||||
if ck["content_ltks"] == nil {
|
||||
t.Errorf("chunk[%d].content_ltks missing", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenizerComponent_Invoke_TextPayload(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
withStubEmbedder(t, 4)
|
||||
c, _ := NewTokenizerComponent(map[string]any{
|
||||
"search_method": []any{"full_text"},
|
||||
})
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "note.txt",
|
||||
"output_format": "text",
|
||||
"text": "plain payload",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
got, _ := out["chunks"].([]map[string]any)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("chunks len = %d, want 1", len(got))
|
||||
}
|
||||
if got[0]["text"] != "plain payload" {
|
||||
t.Errorf("text = %v, want plain payload", got[0]["text"])
|
||||
}
|
||||
if got[0]["content_ltks"] == nil {
|
||||
t.Errorf("content_ltks missing: %v", got[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenizerComponent_Invoke_JSONPayload(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
withStubEmbedder(t, 4)
|
||||
c, _ := NewTokenizerComponent(map[string]any{
|
||||
"search_method": []any{"full_text"},
|
||||
})
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "note.pdf",
|
||||
"output_format": "json",
|
||||
"json": []map[string]any{{"text": "row one"}, {"text": "row two"}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
got, _ := out["chunks"].([]map[string]any)
|
||||
if len(got) != 2 {
|
||||
t.Fatalf("chunks len = %d, want 2", len(got))
|
||||
}
|
||||
if got[0]["content_ltks"] == nil || got[1]["content_ltks"] == nil {
|
||||
t.Errorf("content_ltks missing: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Invoke_BatchedEmbedding asserts the
|
||||
// embedding client is called ONCE with all chunks (not fanned per
|
||||
// chunk — plan §AD-5a). 3 chunks → 1 call.
|
||||
func TestTokenizerComponent_Invoke_BatchedEmbedding(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
stub := withStubEmbedder(t, 8)
|
||||
c, _ := NewTokenizerComponent(map[string]any{})
|
||||
chunks := []map[string]any{
|
||||
{"text": "one"},
|
||||
{"text": "two"},
|
||||
{"text": "three"},
|
||||
}
|
||||
if _, err := c.Invoke(context.Background(), map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": chunks,
|
||||
}); err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got := stub.calls.Load(); got != 1 {
|
||||
t.Errorf("embedder calls = %d, want 1 (single batched call)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Invoke_FullTextOnly covers the
|
||||
// search_method=["full_text"] branch: no embedding, no encoder
|
||||
// call, but tokenized fields present.
|
||||
func TestTokenizerComponent_Invoke_FullTextOnly(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
stub := withStubEmbedder(t, 4)
|
||||
c, _ := NewTokenizerComponent(map[string]any{
|
||||
"search_method": []any{"full_text"},
|
||||
})
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{{"text": "alpha bravo"}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if stub.calls.Load() != 0 {
|
||||
t.Errorf("embedder should not be called, got %d", stub.calls.Load())
|
||||
}
|
||||
if out["embedding_token_consumption"] != nil {
|
||||
t.Errorf("embedding_token_consumption should be absent, got %v", out["embedding_token_consumption"])
|
||||
}
|
||||
got, _ := out["chunks"].([]map[string]any)
|
||||
if len(got) == 0 || got[0]["content_ltks"] == nil {
|
||||
t.Errorf("content_ltks missing: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Invoke_EmbedNoEncodeFunc covers the
|
||||
// "embedding requested but EncodeFunc is nil" branch — must
|
||||
// return a clear error, not panic.
|
||||
func TestTokenizerComponent_Invoke_EmbedNoEncodeFunc(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
prev := EncodeFunc
|
||||
EncodeFunc = nil
|
||||
t.Cleanup(func() { EncodeFunc = prev })
|
||||
|
||||
c, _ := NewTokenizerComponent(map[string]any{})
|
||||
_, err := c.Invoke(context.Background(), map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{{"text": "alpha"}},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected error when embedding requested without EncodeFunc, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "EncodeFunc") {
|
||||
t.Errorf("error should mention EncodeFunc: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Invoke_EmbedderError covers the
|
||||
// propagation of an error from the embedding driver.
|
||||
func TestTokenizerComponent_Invoke_EmbedderError(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
stub := withStubEmbedder(t, 4)
|
||||
stub.err = errors.New("simulated upstream error")
|
||||
|
||||
c, _ := NewTokenizerComponent(map[string]any{})
|
||||
_, err := c.Invoke(context.Background(), map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{{"text": "alpha"}},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected error from embedder, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "simulated upstream error") {
|
||||
t.Errorf("error should chain embedder error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Invoke_EncoderCountMismatch covers the
|
||||
// "embedder returned wrong number of vectors" defensive branch.
|
||||
func TestTokenizerComponent_Invoke_EncoderCountMismatch(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
stub := withStubEmbedder(t, 4)
|
||||
// Inject an embedder that returns the wrong number of vectors
|
||||
// regardless of input.
|
||||
wrong := &countMismatchedEmbedder{want: 1}
|
||||
prev := EncodeFunc
|
||||
EncodeFunc = func(_, _ string) Embedder { return wrong }
|
||||
t.Cleanup(func() {
|
||||
EncodeFunc = prev
|
||||
_ = stub
|
||||
})
|
||||
|
||||
c, _ := NewTokenizerComponent(map[string]any{})
|
||||
_, err := c.Invoke(context.Background(), map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{{"text": "a"}, {"text": "b"}, {"text": "c"}},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected error from count mismatch, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "vectors") {
|
||||
t.Errorf("error should mention vectors: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
type countMismatchedEmbedder struct{ want int }
|
||||
|
||||
func (c *countMismatchedEmbedder) Encode(texts []string) ([][]float64, error) {
|
||||
out := make([][]float64, c.want)
|
||||
for i := range out {
|
||||
out[i] = make([]float64, 4)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Invoke_HonorsTimeout installs an
|
||||
// embedder that blocks past a (test-shrunk) tokenizerTimeout and
|
||||
// asserts the component returns context.DeadlineExceeded.
|
||||
func TestTokenizerComponent_Invoke_HonorsTimeout(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
prevTimeout := tokenizerTimeout
|
||||
tokenizerTimeout = 50 * time.Millisecond
|
||||
t.Cleanup(func() { tokenizerTimeout = prevTimeout })
|
||||
|
||||
stub := withStubEmbedder(t, 4)
|
||||
stub.delay = 500 * time.Millisecond
|
||||
|
||||
c, _ := NewTokenizerComponent(map[string]any{})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
_, err := c.Invoke(ctx, map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{{"text": "alpha"}},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected timeout error, got nil")
|
||||
}
|
||||
if !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Errorf("expected context.DeadlineExceeded, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_InputsOutputs_NonEmpty verifies Phase 4
|
||||
// API metadata shape.
|
||||
func TestTokenizerComponent_InputsOutputs_NonEmpty(t *testing.T) {
|
||||
c, _ := NewTokenizerComponent(map[string]any{})
|
||||
ins := c.(*TokenizerComponent).Inputs()
|
||||
outs := c.(*TokenizerComponent).Outputs()
|
||||
if len(ins) == 0 {
|
||||
t.Error("Inputs() empty")
|
||||
}
|
||||
if len(outs) == 0 {
|
||||
t.Error("Outputs() empty")
|
||||
}
|
||||
for _, key := range []string{"chunks", "output_format"} {
|
||||
if _, ok := outs[key]; !ok {
|
||||
t.Errorf("Outputs() missing %q", key)
|
||||
}
|
||||
}
|
||||
for _, key := range []string{"chunks", "name"} {
|
||||
if _, ok := ins[key]; !ok {
|
||||
t.Errorf("Inputs() missing %q", key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Parallelism locks the fan-out to 1 (plan
|
||||
// §AD-5a: "embedding calls batched, not fanned").
|
||||
func TestTokenizerComponent_Parallelism(t *testing.T) {
|
||||
c, _ := NewTokenizerComponent(map[string]any{})
|
||||
if got := c.(*TokenizerComponent).Parallelism(); got != 1 {
|
||||
t.Errorf("Parallelism() = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_NewTokenizerComponent_Defaults verifies
|
||||
// the Python default param values propagate.
|
||||
func TestTokenizerComponent_NewTokenizerComponent_Defaults(t *testing.T) {
|
||||
c, err := NewTokenizerComponent(nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NewTokenizerComponent(nil): %v", err)
|
||||
}
|
||||
tc := c.(*TokenizerComponent)
|
||||
if tc.param.FilenameEmbdWeight != 0.1 {
|
||||
t.Errorf("filename_embd_weight = %v, want 0.1", tc.param.FilenameEmbdWeight)
|
||||
}
|
||||
if len(tc.param.Fields) != 1 || tc.param.Fields[0] != "text" {
|
||||
t.Errorf("fields = %v, want [text]", tc.param.Fields)
|
||||
}
|
||||
if len(tc.param.SearchMethod) != 2 {
|
||||
t.Errorf("search_method len = %d, want 2", len(tc.param.SearchMethod))
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_NewTokenizerComponent_BadParam covers
|
||||
// the param-validation branch (invalid search_method value).
|
||||
func TestTokenizerComponent_NewTokenizerComponent_BadParam(t *testing.T) {
|
||||
_, err := NewTokenizerComponent(map[string]any{
|
||||
"search_method": []any{"unknown"},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected param validation error, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Smoke_EndToEnd is the BLOCKER smoke test
|
||||
// (plan §8 R3). Drives 1 chunk of ~1000 tokens through the real
|
||||
// tokenizer and a stub embedder with no artificial latency, then
|
||||
// asserts:
|
||||
//
|
||||
// - non-zero vector returned
|
||||
// - latency well under 5s (real stub returns in <1ms; we assert
|
||||
// < 5s as the §R3 ceiling)
|
||||
// - no panic
|
||||
//
|
||||
// Documented result: this stub embedding completes in well under
|
||||
// 5s (typical observed latency < 5ms on the test host). The
|
||||
// production path against a real embedding API was not exercised
|
||||
// in this CI sandbox; the helper `withStubEmbedder` deliberately
|
||||
// avoids the network round-trip while still exercising the full
|
||||
// wiring (TrackElapsed, WithTimeout, batched Encode, vector
|
||||
// stamping).
|
||||
func TestTokenizerComponent_Smoke_EndToEnd(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
const dim = 1024
|
||||
withStubEmbedder(t, dim)
|
||||
|
||||
// Build a chunk of ~1000 tokens. Each word ≈ 1 token for English
|
||||
// under cl100k_base. We pad with a recognizable sentinel so we
|
||||
// can later check tokenization fidelity if desired.
|
||||
words := make([]string, 0, 1000)
|
||||
for i := 0; i < 1000; i++ {
|
||||
words = append(words, "ragflow")
|
||||
}
|
||||
chunkText := strings.Join(words, " ")
|
||||
// Sanity-check the count is in the expected ballpark (cl100k_base
|
||||
// may over- or under-count; we only assert the order of magnitude).
|
||||
preflightTokens := tokenizer.NumTokensFromString(chunkText)
|
||||
if preflightTokens < 100 || preflightTokens > 5000 {
|
||||
t.Logf("preflight token count = %d (acceptable range 100-5000)", preflightTokens)
|
||||
}
|
||||
|
||||
c, _ := NewTokenizerComponent(map[string]any{})
|
||||
chunks := []map[string]any{
|
||||
{"text": chunkText},
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"tenant_id": "tenant-smoke",
|
||||
"model_id": "embd-smoke",
|
||||
"name": "smoke.pdf",
|
||||
"output_format": "chunks",
|
||||
"chunks": chunks,
|
||||
})
|
||||
elapsed := time.Since(start)
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if elapsed >= 5*time.Second {
|
||||
t.Errorf("elapsed %v exceeds §R3 ceiling of 5s", elapsed)
|
||||
}
|
||||
|
||||
got, ok := out["chunks"].([]map[string]any)
|
||||
if !ok || len(got) != 1 {
|
||||
t.Fatalf("chunks output malformed: %v", out["chunks"])
|
||||
}
|
||||
vec, ok := got[0]["q_1024_vec"].([]float64)
|
||||
if !ok {
|
||||
t.Fatalf("q_1024_vec missing or wrong type: %T", got[0]["q_1024_vec"])
|
||||
}
|
||||
if len(vec) != dim {
|
||||
t.Errorf("vector len = %d, want %d", len(vec), dim)
|
||||
}
|
||||
// "non-zero vector" assertion: at least one element is non-zero.
|
||||
nonZero := 0
|
||||
for _, x := range vec {
|
||||
if x != 0 {
|
||||
nonZero++
|
||||
}
|
||||
}
|
||||
if nonZero == 0 {
|
||||
t.Error("vector is all zeros (smoke contract: non-zero vector returned)")
|
||||
}
|
||||
|
||||
// No panic == pass; explicitly assert log message.
|
||||
t.Logf("smoke complete: chunks=%d elapsed=%v tokens≈%d vec_dim=%d",
|
||||
len(got), elapsed, preflightTokens, len(vec))
|
||||
}
|
||||
@@ -23,6 +23,7 @@ import (
|
||||
"ragflow/internal/dao"
|
||||
"ragflow/internal/engine"
|
||||
"ragflow/internal/entity"
|
||||
"ragflow/internal/ingestion/pipeline"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
@@ -188,11 +189,6 @@ func (e *Ingestor) Start() error {
|
||||
TaskHandle: taskHandle,
|
||||
}
|
||||
|
||||
// Register in currentTasks immediately so heartbeat sees PENDING state
|
||||
//e.tasksMu.Lock()
|
||||
//e.currentTasks[task.ID] = taskCtx
|
||||
//e.tasksMu.Unlock()
|
||||
|
||||
// Push to task channel; if full, reject the task (backpressure)
|
||||
select {
|
||||
case e.taskChan <- taskCtx:
|
||||
@@ -200,10 +196,6 @@ func (e *Ingestor) Start() error {
|
||||
default:
|
||||
common.Info(fmt.Sprintf("No available slot for task %s, failed", task.ID))
|
||||
|
||||
//e.tasksMu.Lock()
|
||||
//delete(e.currentTasks, task.ID)
|
||||
//e.tasksMu.Unlock()
|
||||
|
||||
err = taskHandle.Nack()
|
||||
if err != nil {
|
||||
common.Error(fmt.Sprintf("error nack task %s", taskMessage.TaskID), err)
|
||||
@@ -214,275 +206,6 @@ func (e *Ingestor) Start() error {
|
||||
}
|
||||
}
|
||||
|
||||
//// Connect connects to the admin and establishes a bidirectional stream
|
||||
//func (e *Ingestor) Connect(serverAddr string) error {
|
||||
// e.serverAddr = serverAddr
|
||||
// conn, err := grpc.Dial(serverAddr,
|
||||
// grpc.WithTransportCredentials(insecure.NewCredentials()),
|
||||
// grpc.WithBlock(),
|
||||
// grpc.WithTimeout(5*time.Second),
|
||||
// )
|
||||
// if err != nil {
|
||||
// return fmt.Errorf("fail to connect admin server: %s", err.Error())
|
||||
// }
|
||||
// e.conn = conn
|
||||
//
|
||||
// e.client = common.NewIngestionManagerClient(conn)
|
||||
//
|
||||
// stream, err := e.client.Action(e.ctx)
|
||||
// if err != nil {
|
||||
// conn.Close()
|
||||
// return err
|
||||
// }
|
||||
// e.stream = stream
|
||||
//
|
||||
// common.Info(fmt.Sprintf("Ingestor %s connected to admin", e.id))
|
||||
//
|
||||
// // 1. Send registration message
|
||||
// if err = e.sendRegister(); err != nil {
|
||||
// conn.Close()
|
||||
// return err
|
||||
// }
|
||||
//
|
||||
// // Ensure worker pool is started on first task
|
||||
// e.startWorkerPool()
|
||||
//
|
||||
// // 2. Start receive loop
|
||||
// go e.receiveLoop()
|
||||
//
|
||||
// // 3. Start heartbeat loop
|
||||
// go e.heartbeatLoop()
|
||||
//
|
||||
// return nil
|
||||
//}
|
||||
|
||||
//func (e *Ingestor) sendRegister() error {
|
||||
// msg := &common.IngestionMessage{
|
||||
// IngestorId: e.id,
|
||||
// MessageType: "REGISTER",
|
||||
// RegisterInfo: &common.RegisterInfo{
|
||||
// MaxConcurrency: e.maxConcurrency,
|
||||
// SupportedDocTypes: e.supportedDocTypes,
|
||||
// Version: e.version,
|
||||
// Name: e.name,
|
||||
// },
|
||||
// }
|
||||
// return e.stream.Send(msg)
|
||||
//}
|
||||
//
|
||||
//func (e *Ingestor) sendHeartbeat() error {
|
||||
// e.tasksMu.RLock()
|
||||
//
|
||||
// cutoff := time.Now().Add(-10 * time.Minute)
|
||||
// var toDeleteTask []string
|
||||
// taskStates := make([]*common.TaskState, 0, len(e.currentTasks))
|
||||
//
|
||||
// for tid, tc := range e.currentTasks {
|
||||
// // Check if task is in a terminal state and expired beyond 10 minutes
|
||||
// if (tc.Status == "CANCELED" || tc.Status == "COMPLETED" || tc.Status == "REJECTED") &&
|
||||
// !tc.EndTime.IsZero() && tc.EndTime.Before(cutoff) {
|
||||
// toDeleteTask = append(toDeleteTask, tid)
|
||||
// } else {
|
||||
// taskStates = append(taskStates, &common.TaskState{
|
||||
// TaskId: tid,
|
||||
// Status: tc.Status,
|
||||
// EstimatedRemainingTimeSeconds: int64(tc.estimatedRemainingTime),
|
||||
// ErrorMessage: tc.ErrorMessage,
|
||||
// StartTime: tc.StartTime.UnixNano(),
|
||||
// ComeFrom: tc.Task.ComeFrom,
|
||||
// })
|
||||
// }
|
||||
// }
|
||||
// e.tasksMu.RUnlock()
|
||||
//
|
||||
// // Delete expired terminal tasks from currentTasks
|
||||
// if len(toDeleteTask) > 0 {
|
||||
// e.tasksMu.Lock()
|
||||
// for _, id := range toDeleteTask {
|
||||
// delete(e.currentTasks, id)
|
||||
// }
|
||||
// e.tasksMu.Unlock()
|
||||
// }
|
||||
//
|
||||
// var pid = int64(os.Getpid())
|
||||
// p, err := process.NewProcess(int32(pid))
|
||||
// if err != nil {
|
||||
// log.Fatal(err)
|
||||
// }
|
||||
//
|
||||
// var cpuPercent float64
|
||||
// cpuPercent, err = p.Percent(100 * time.Millisecond)
|
||||
// if err != nil {
|
||||
// cpuPercent = math.NaN()
|
||||
// common.Info(fmt.Sprintf("Fail to read CPU usage: %v", err))
|
||||
// }
|
||||
//
|
||||
// RssUsage := math.NaN()
|
||||
// VmsUsage := math.NaN()
|
||||
// memInfo, err := p.MemoryInfo()
|
||||
// if err == nil {
|
||||
// RssUsage = float64(memInfo.RSS)
|
||||
// VmsUsage = float64(memInfo.VMS)
|
||||
// } else {
|
||||
// common.Info(fmt.Sprintf("Fail to read memory usage: %v", err))
|
||||
// }
|
||||
// msg := &common.IngestionMessage{
|
||||
// IngestorId: e.id,
|
||||
// MessageType: "HEARTBEAT",
|
||||
// HeartbeatInfo: &common.HeartbeatInfo{
|
||||
// TaskStates: taskStates,
|
||||
// DeleteTaskIds: toDeleteTask,
|
||||
// CpuUsage: float32(cpuPercent),
|
||||
// VmsUsage: float32(VmsUsage),
|
||||
// RssUsage: float32(RssUsage),
|
||||
// ProcessId: pid,
|
||||
// },
|
||||
// }
|
||||
// return e.stream.Send(msg)
|
||||
//}
|
||||
//
|
||||
//func (e *Ingestor) sendTaskResult(taskID, status, errorMsg string) error {
|
||||
// msg := &common.IngestionMessage{
|
||||
// IngestorId: e.id,
|
||||
// MessageType: "TASK_RESULT",
|
||||
// TaskResult: &common.TaskResult{
|
||||
// TaskId: taskID,
|
||||
// Status: status,
|
||||
// ErrorMessage: errorMsg,
|
||||
// },
|
||||
// }
|
||||
// return e.stream.Send(msg)
|
||||
//}
|
||||
//
|
||||
//func (e *Ingestor) sendTaskProgress(taskID string, progress int32, info string) error {
|
||||
// msg := &common.IngestionMessage{
|
||||
// IngestorId: e.id,
|
||||
// MessageType: "TASK_PROGRESS",
|
||||
// TaskProgress: &common.TaskProgress{
|
||||
// TaskId: taskID,
|
||||
// Progress: progress,
|
||||
// Info: info,
|
||||
// },
|
||||
// }
|
||||
// return e.stream.Send(msg)
|
||||
//}
|
||||
//
|
||||
//func (e *Ingestor) receiveLoop() {
|
||||
// for {
|
||||
// msg, err := e.stream.Recv()
|
||||
// if err != nil {
|
||||
// if e.ctx.Err() != nil {
|
||||
// common.Info(fmt.Sprintf("Ingestor %s context cancelled, receive loop exiting", e.id))
|
||||
// return
|
||||
// }
|
||||
// common.Info(fmt.Sprintf("Receive error: %v", err))
|
||||
// common.Info("Admin connection lost, attempting to reconnect")
|
||||
// e.reconnect()
|
||||
// return
|
||||
// }
|
||||
//
|
||||
// switch msg.MessageType {
|
||||
// case "TASK_ASSIGNMENT":
|
||||
// e.handleTaskAssignment(msg.TaskAssignment)
|
||||
//
|
||||
// case "ACK":
|
||||
// common.Info(fmt.Sprintf("Received ACK: task=%s, success=%v, msg=%s",
|
||||
// msg.AckInfo.TaskId, msg.AckInfo.Success, msg.AckInfo.Message))
|
||||
//
|
||||
// case "ERROR":
|
||||
// common.Info(fmt.Sprintf("Received error from admin: %s", msg.ErrorMessage))
|
||||
//
|
||||
// default:
|
||||
// common.Info(fmt.Sprintf("Unknown admin message type: %s", msg.MessageType))
|
||||
// }
|
||||
// }
|
||||
//}
|
||||
//
|
||||
//func (e *Ingestor) handleTaskAssignment(task *common.TaskAssignment) {
|
||||
// if task == nil {
|
||||
// return
|
||||
// }
|
||||
//
|
||||
// common.Info(fmt.Sprintf("Received task: %s, task_type=%s", task.TaskId, task.TaskType))
|
||||
//
|
||||
// switch task.TaskType {
|
||||
// case "shutdown_ingestor":
|
||||
// if e.id == task.AssignedTo {
|
||||
// e.handleShutdownIngestor()
|
||||
// return
|
||||
// }
|
||||
//
|
||||
// common.Error("unmatched ingestor id", fmt.Errorf("attempt to shutdown ingestor: %s, current ingestor: %s, mismatched", task.AssignedTo, e.id))
|
||||
// return
|
||||
// case "cancel_ingestion_task":
|
||||
// e.handleCancelTask(task.TaskId)
|
||||
// return
|
||||
// case "start_ingestion_task":
|
||||
// // create ingestion task log
|
||||
// err := e.ingestionTaskLogDAO.Create(&entity.IngestionTaskLog{
|
||||
// TaskID: task.TaskId,
|
||||
// Action: "CREATED",
|
||||
// })
|
||||
// if err != nil {
|
||||
// common.Fatal(fmt.Sprintf("Failed to create ingestion task log for task %s: %v", task.TaskId, err))
|
||||
// return
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// // Construct TaskContext with a cancellable context
|
||||
// ctx, cancel := context.WithCancel(e.ctx)
|
||||
// taskCtx := &TaskContext{
|
||||
// Ctx: ctx,
|
||||
// CancelFunc: cancel,
|
||||
// Task: task,
|
||||
// Status: "QUEUED",
|
||||
// }
|
||||
//
|
||||
// // Register in currentTasks immediately so heartbeat sees PENDING state
|
||||
// e.tasksMu.Lock()
|
||||
// e.currentTasks[task.TaskId] = taskCtx
|
||||
// e.tasksMu.Unlock()
|
||||
//
|
||||
// common.Info("wait for 10 seconds")
|
||||
// time.Sleep(time.Second * 10)
|
||||
// // Push to task channel; if full, reject the task (backpressure)
|
||||
// select {
|
||||
// case e.taskChan <- taskCtx:
|
||||
// common.Info(fmt.Sprintf("Task %s queued (channel: %d/%d)", task.TaskId, len(e.taskChan), cap(e.taskChan)))
|
||||
// default:
|
||||
// common.Info(fmt.Sprintf("No available slot for task %s, rejecting", task.TaskId))
|
||||
// //e.tasksMu.Lock()
|
||||
// //delete(e.currentTasks, task.TaskId)
|
||||
// //e.tasksMu.Unlock()
|
||||
// taskCtx.Status = "REJECTED"
|
||||
// taskCtx.EndTime = time.Now()
|
||||
// e.sendTaskResult(taskCtx.Task.TaskId, "REJECTED", "task rejected before execution")
|
||||
// }
|
||||
//}
|
||||
//
|
||||
//func (e *Ingestor) handleCancelTask(taskID string) {
|
||||
// e.tasksMu.Lock()
|
||||
// taskCtx, exists := e.currentTasks[taskID]
|
||||
// e.tasksMu.Unlock()
|
||||
//
|
||||
// if !exists {
|
||||
// common.Info(fmt.Sprintf("Cancel request for unknown task %s, ignoring", taskID))
|
||||
// return
|
||||
// }
|
||||
//
|
||||
// common.Info(fmt.Sprintf("Cancelling task %s (current status: %s)", taskID, taskCtx.Status))
|
||||
// taskCtx.CancelFunc()
|
||||
//}
|
||||
//
|
||||
//func (e *Ingestor) handleShutdownIngestor() {
|
||||
// common.Info(fmt.Sprintf("Shutdown task received, initiating graceful shutdown of ingestor %s", e.id))
|
||||
// select {
|
||||
// case e.ShutdownCh <- struct{}{}:
|
||||
// default:
|
||||
// }
|
||||
// return
|
||||
//}
|
||||
|
||||
func (e *Ingestor) startWorkerPool() {
|
||||
e.startOnce.Do(func() {
|
||||
for i := int32(0); i < e.maxConcurrency; i++ {
|
||||
@@ -511,77 +234,123 @@ func (e *Ingestor) workerLoop(id int32) {
|
||||
}
|
||||
|
||||
func (e *Ingestor) executeTask(taskCtx *TaskContext) {
|
||||
defer func() {
|
||||
//e.tasksMu.Lock()
|
||||
//delete(e.currentTasks, taskCtx.Task.TaskId)
|
||||
//e.tasksMu.Unlock()
|
||||
}()
|
||||
|
||||
ctx := taskCtx.Ctx
|
||||
task := taskCtx.Task
|
||||
common.Info(fmt.Sprintf("Starting task %s", task.ID))
|
||||
|
||||
latestLog, err := e.ingestionTaskLogDAO.LatestLogByTaskID(task.ID)
|
||||
if err != nil {
|
||||
latestLog = &entity.IngestionTaskLog{
|
||||
ID: 0,
|
||||
TaskID: task.ID,
|
||||
Checkpoint: entity.JSONMap{
|
||||
"current_step": 1,
|
||||
"total_step": 5,
|
||||
},
|
||||
}
|
||||
err = e.ingestionTaskLogDAO.Create(latestLog)
|
||||
if err != nil {
|
||||
common.Error(fmt.Sprintf("Failed to create task log for task %s", task.ID), err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
var checkpointMap map[string]interface{}
|
||||
checkpointMap = latestLog.Checkpoint
|
||||
currentStep, ok := common.GetInt(checkpointMap["current_step"])
|
||||
if !ok {
|
||||
common.Fatal(fmt.Sprintf("Failed to get current step from task log for task %s", task.ID))
|
||||
// Execute the canonical ingestion canvas DSL carried by the task.
|
||||
// The Go ingestion path no longer synthesizes a parallel `stages[]`
|
||||
// schema; the only accepted format is the template/canvas DSL.
|
||||
dslBytes := defaultPipelineDSL(task)
|
||||
if len(dslBytes) == 0 {
|
||||
err := fmt.Errorf("task %s missing canonical ingestion DSL in schema.pipeline or schema.dsl", task.ID)
|
||||
common.Error(fmt.Sprintf("Failed to load pipeline DSL for task %s", task.ID), err)
|
||||
e.failTask(taskCtx, err)
|
||||
return
|
||||
}
|
||||
totalStep, ok := common.GetInt(checkpointMap["total_step"])
|
||||
if !ok {
|
||||
common.Fatal(fmt.Sprintf("Failed to get current step from task log for task %s", task.ID))
|
||||
pl, err := pipeline.NewPipelineFromDSL(dslBytes, task.ID)
|
||||
if err != nil {
|
||||
common.Error(fmt.Sprintf("Failed to compile pipeline for task %s", task.ID), err)
|
||||
e.failTask(taskCtx, err)
|
||||
return
|
||||
}
|
||||
for i := currentStep; i < totalStep; i++ {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
// Task canceled
|
||||
common.Info(fmt.Sprintf("Task %s stopped", task.ID))
|
||||
return
|
||||
case <-time.After(5000 * time.Millisecond):
|
||||
common.Info(fmt.Sprintf("Task %s is running step %d", task.ID, i))
|
||||
checkpointMap["current_step"] = i + 1
|
||||
latestLog.Checkpoint = checkpointMap
|
||||
latestLog.ID++
|
||||
err = latestLog.UpdateCreateDateAndTime()
|
||||
if err != nil {
|
||||
common.Error(fmt.Sprintf("Failed to update date and time of task log for task %s", task.ID), err)
|
||||
return
|
||||
}
|
||||
|
||||
err = e.ingestionTaskLogDAO.Create(latestLog)
|
||||
if err != nil {
|
||||
common.Error(fmt.Sprintf("Failed to create task log for task %s", task.ID), err)
|
||||
return
|
||||
}
|
||||
inputs := map[string]any{
|
||||
"doc_id": task.DocumentID,
|
||||
}
|
||||
_, runErr := pl.Run(ctx, inputs)
|
||||
if runErr != nil {
|
||||
if errors.Is(runErr, context.Canceled) || errors.Is(runErr, context.DeadlineExceeded) {
|
||||
common.Info(fmt.Sprintf("Task %s cancelled: %v", task.ID, runErr))
|
||||
// STOPPED is a terminal status — the task will not be
|
||||
// re-attempted by the consumer. Ack the message so the
|
||||
// queue does not redeliver it (Nack here would race
|
||||
// with the STOPPED write and let another consumer pick
|
||||
// up a "stopped" task).
|
||||
_ = e.ingestionTaskDAO.UpdateStatus(task.ID, common.STOPPED)
|
||||
_ = e.ackOrNack(taskCtx, true)
|
||||
return
|
||||
}
|
||||
common.Error(fmt.Sprintf("Task %s pipeline failed", task.ID), runErr)
|
||||
e.failTask(taskCtx, runErr)
|
||||
return
|
||||
}
|
||||
|
||||
err = e.ingestionTaskDAO.UpdateStatus(task.ID, common.COMPLETED)
|
||||
if err != nil {
|
||||
if err = e.ingestionTaskDAO.UpdateStatus(task.ID, common.COMPLETED); err != nil {
|
||||
common.Error(fmt.Sprintf("Task %s update status failed", task.ID), err)
|
||||
_ = e.ackOrNack(taskCtx, true)
|
||||
return
|
||||
}
|
||||
|
||||
common.Info(fmt.Sprintf("Task %s completed", task.ID))
|
||||
_ = e.ackOrNack(taskCtx, true)
|
||||
}
|
||||
|
||||
// defaultPipelineDSL returns the canonical ingestion canvas DSL bytes carried
|
||||
// by the task schema. The ingestion runtime accepts only the template/canvas
|
||||
// DSL shape; it does not synthesize a separate linear stages[] schema.
|
||||
func defaultPipelineDSL(task *entity.IngestionTask) []byte {
|
||||
if task != nil && task.Schema != nil {
|
||||
if raw, ok := task.Schema["pipeline"]; ok {
|
||||
switch v := raw.(type) {
|
||||
case []byte:
|
||||
if len(v) > 0 {
|
||||
return v
|
||||
}
|
||||
case string:
|
||||
if v != "" {
|
||||
return []byte(v)
|
||||
}
|
||||
}
|
||||
}
|
||||
if raw, ok := task.Schema["dsl"]; ok {
|
||||
switch v := raw.(type) {
|
||||
case []byte:
|
||||
if len(v) > 0 {
|
||||
return v
|
||||
}
|
||||
case string:
|
||||
if v != "" {
|
||||
return []byte(v)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// failTask updates the task to FAILED and Acks the message
|
||||
// (terminal-failure path: even on error, the message must be
|
||||
// removed from the queue, otherwise the broker redelivers it
|
||||
// indefinitely). This fixes the pre-existing bug that the
|
||||
// placeholder sleep loop never called Ack at all (plan §8 Q3).
|
||||
func (e *Ingestor) failTask(taskCtx *TaskContext, runErr error) {
|
||||
if err := e.ingestionTaskDAO.UpdateStatus(taskCtx.Task.ID, common.FAILED); err != nil {
|
||||
common.Error(fmt.Sprintf("Task %s update status (failed) error", taskCtx.Task.ID), err)
|
||||
}
|
||||
_ = e.ackOrNack(taskCtx, true)
|
||||
common.Error(fmt.Sprintf("Task %s failed: %v", taskCtx.Task.ID, runErr), runErr)
|
||||
}
|
||||
|
||||
// ackOrNack centralises the post-execution NATS message
|
||||
// disposition. ack=true removes the message from the queue
|
||||
// (success OR terminal-failure); ack=false re-queues (rare;
|
||||
// we use it only on context cancellation to let another
|
||||
// worker pick it up).
|
||||
func (e *Ingestor) ackOrNack(taskCtx *TaskContext, ack bool) error {
|
||||
if taskCtx == nil || taskCtx.TaskHandle == nil {
|
||||
return nil
|
||||
}
|
||||
var err error
|
||||
if ack {
|
||||
err = taskCtx.TaskHandle.Ack()
|
||||
} else {
|
||||
err = taskCtx.TaskHandle.Nack()
|
||||
}
|
||||
if err != nil {
|
||||
common.Error(fmt.Sprintf("Task %s ack/nack error", taskCtx.Task.ID), err)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (e *Ingestor) executeTasklet(taskCtx *TaskContext) {
|
||||
@@ -636,85 +405,6 @@ func (e *Ingestor) executeTasklet(taskCtx *TaskContext) {
|
||||
common.Info(fmt.Sprintf("Tasklet %s completed", tasklet.ID))
|
||||
}
|
||||
|
||||
//
|
||||
//func (e *Ingestor) heartbeatLoop() {
|
||||
// ticker := time.NewTicker(5 * time.Second)
|
||||
// defer ticker.Stop()
|
||||
//
|
||||
// for {
|
||||
// select {
|
||||
// case <-e.ctx.Done():
|
||||
// return
|
||||
// case <-ticker.C:
|
||||
// if err := e.sendHeartbeat(); err != nil {
|
||||
// common.Info(fmt.Sprintf("Failed to send heartbeat: %v", err))
|
||||
// if e.ctx.Err() != nil {
|
||||
// common.Info(fmt.Sprintf("Ingestor %s context cancelled, heartbeat loop exiting", e.id))
|
||||
// return
|
||||
// }
|
||||
// common.Info(fmt.Sprintf("Admin connection lost, attempting to reconnect"))
|
||||
// e.reconnect()
|
||||
// return
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
//}
|
||||
//
|
||||
//// reconnect closes the old connection and establishes a new one with exponential backoff.
|
||||
//// Only one reconnection attempt runs at a time; concurrent callers return immediately.
|
||||
//func (e *Ingestor) reconnect() {
|
||||
// if e.ctx.Err() != nil {
|
||||
// common.Info(fmt.Sprintf("Ingestor %s is shutting down, skipping reconnection", e.id))
|
||||
// return
|
||||
// }
|
||||
//
|
||||
// if !e.reconnectMu.TryLock() {
|
||||
// return
|
||||
// }
|
||||
// defer e.reconnectMu.Unlock()
|
||||
//
|
||||
// common.Info(fmt.Sprintf("Ingestor %s attempting to reconnect to admin at %s", e.id, e.serverAddr))
|
||||
//
|
||||
// // Close old stream and connection
|
||||
// if e.stream != nil {
|
||||
// e.stream.CloseSend()
|
||||
// }
|
||||
// if e.conn != nil {
|
||||
// e.conn.Close()
|
||||
// }
|
||||
//
|
||||
// backoff := 1 * time.Second
|
||||
// maxBackoff := 30 * time.Second
|
||||
//
|
||||
// for {
|
||||
// conn, err := grpc.Dial(e.serverAddr,
|
||||
// grpc.WithTransportCredentials(insecure.NewCredentials()),
|
||||
// grpc.WithBlock(),
|
||||
// grpc.WithTimeout(5*time.Second),
|
||||
// )
|
||||
// if err != nil {
|
||||
// common.Info(fmt.Sprintf("Reconnect dial failed: %v, retrying in %v", err, backoff))
|
||||
// time.Sleep(backoff)
|
||||
// backoff *= 2
|
||||
// if backoff > maxBackoff {
|
||||
// backoff = maxBackoff
|
||||
// }
|
||||
// continue
|
||||
// }
|
||||
// e.conn = conn
|
||||
// e.client = common.NewIngestionManagerClient(conn)
|
||||
//
|
||||
// stream, err := e.client.Action(e.ctx)
|
||||
// if err != nil {
|
||||
// conn.Close()
|
||||
// common.Info(fmt.Sprintf("Reconnect create stream failed: %v, retrying in %v", err, backoff))
|
||||
// time.Sleep(backoff)
|
||||
// backoff *= 2
|
||||
// if backoff > maxBackoff {
|
||||
// backoff = maxBackoff
|
||||
// }
|
||||
// continue
|
||||
// }
|
||||
// e.stream = stream
|
||||
//
|
||||
// if err = e.sendRegister(); err != nil {
|
||||
|
||||
478
internal/ingestion/ingestion_service_test.go
Normal file
478
internal/ingestion/ingestion_service_test.go
Normal file
@@ -0,0 +1,478 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// e2e tests for the Ingestor.
|
||||
//
|
||||
// These tests exercise Ingestor.executeTask end-to-end against
|
||||
// an in-memory SQLite + memory storage backend, with a stub
|
||||
// ProgressSink that records per-stage invocations. The DSL
|
||||
// runs against the production pipeline package; the test
|
||||
// supplies a custom 1-stage DSL that points at a stub
|
||||
// component so no real storage / LLM backend is required.
|
||||
//
|
||||
// The pattern mirrors internal/service/agent_run_e2e_test.go
|
||||
// (in-memory SQLite + miniredis, no Docker). The Ingestor's
|
||||
// runnable surface is exerciseTask, which is package-private,
|
||||
// so the test lives in the `ingestion` package itself.
|
||||
package ingestion
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/dao"
|
||||
"ragflow/internal/entity"
|
||||
_ "ragflow/internal/ingestion/component" // blank import: registers ingestion factories
|
||||
"ragflow/internal/storage"
|
||||
|
||||
"github.com/glebarez/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Stub component for e2e testing.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const (
|
||||
stubIngestionComp = "StubIngest"
|
||||
stubTokenizerComp = "MockTokenizer"
|
||||
)
|
||||
|
||||
type stubIngest struct {
|
||||
counter *int32
|
||||
}
|
||||
|
||||
func (s *stubIngest) Parallelism() int { return 1 }
|
||||
func (s *stubIngest) Inputs() map[string]string {
|
||||
return map[string]string{"x": "any"}
|
||||
}
|
||||
func (s *stubIngest) Outputs() map[string]string {
|
||||
return map[string]string{"y": "any"}
|
||||
}
|
||||
func (s *stubIngest) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
atomic.AddInt32(s.counter, 1)
|
||||
return map[string]any{"y": "stub-output"}, nil
|
||||
}
|
||||
|
||||
var stubCounter int32
|
||||
var stubTokenizerCounter int32
|
||||
|
||||
func newStubIngest(_ string, _ map[string]any) (runtime.Component, error) {
|
||||
return &stubIngest{counter: &stubCounter}, nil
|
||||
}
|
||||
|
||||
type stubTokenizer struct {
|
||||
counter *int32
|
||||
}
|
||||
|
||||
func (s *stubTokenizer) Parallelism() int { return 1 }
|
||||
func (s *stubTokenizer) Inputs() map[string]string {
|
||||
return map[string]string{"chunks": "any"}
|
||||
}
|
||||
func (s *stubTokenizer) Outputs() map[string]string {
|
||||
return map[string]string{"tokens": "any"}
|
||||
}
|
||||
func (s *stubTokenizer) Invoke(_ context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
atomic.AddInt32(s.counter, 1)
|
||||
return map[string]any{"tokens": inputs["chunks"]}, nil
|
||||
}
|
||||
|
||||
func newStubTokenizer(_ string, _ map[string]any) (runtime.Component, error) {
|
||||
return &stubTokenizer{counter: &stubTokenizerCounter}, nil
|
||||
}
|
||||
|
||||
func init() {
|
||||
runtime.MustRegister(stubIngestionComp, runtime.CategoryIngestion, newStubIngest, runtime.Metadata{
|
||||
Inputs: map[string]string{"x": "any"},
|
||||
Outputs: map[string]string{"y": "any"},
|
||||
})
|
||||
runtime.MustRegister(stubTokenizerComp, runtime.CategoryIngestion, newStubTokenizer, runtime.Metadata{
|
||||
Inputs: map[string]string{"chunks": "any"},
|
||||
Outputs: map[string]string{"tokens": "any"},
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Stub TaskHandle — captures Ack/Nack calls.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type stubTaskHandle struct {
|
||||
acked int32
|
||||
nacked int32
|
||||
taskID string
|
||||
}
|
||||
|
||||
func (h *stubTaskHandle) GetMessage() common.TaskMessage {
|
||||
return common.TaskMessage{TaskID: h.taskID, TaskType: common.TaskTypeIngestionTask}
|
||||
}
|
||||
func (h *stubTaskHandle) Ack() error {
|
||||
atomic.AddInt32(&h.acked, 1)
|
||||
return nil
|
||||
}
|
||||
func (h *stubTaskHandle) Nack() error {
|
||||
atomic.AddInt32(&h.nacked, 1)
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// DB setup
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func setupIngestorTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
// Use shared in-memory SQLite so the Ingestor's
|
||||
// dao.NewIngestionTaskLogDAO() (which references the
|
||||
// global dao.DB) and this test's local *gorm.DB see
|
||||
// the same data. Each test gets a unique DSN so
|
||||
// parallel test runs don't collide. The DSN name is
|
||||
// sanitised to avoid SQLite parsing surprises on
|
||||
// special characters in the test name.
|
||||
dsn := sharedIngestorCacheDSN(t.Name())
|
||||
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatalf("sqlite: %v", err)
|
||||
}
|
||||
if err := db.AutoMigrate(
|
||||
&entity.IngestionTask{},
|
||||
&entity.IngestionTaskLog{},
|
||||
); err != nil {
|
||||
t.Fatalf("migrate: %v", err)
|
||||
}
|
||||
return db
|
||||
}
|
||||
|
||||
// sharedIngestorCacheDSN builds a unique per-test DSN that
|
||||
// points at a shared-cache in-memory SQLite database. The
|
||||
// test name is sanitised to alphanumerics + underscores so
|
||||
// the DSN parser doesn't choke on Go test names like
|
||||
// "TestPipeline_Resume_RoundTripViaTaskLogSink" or unicode.
|
||||
// A monotonic counter uniquifies the DSN across
|
||||
// `go test -count=N` iterations so each invocation sees a
|
||||
// fresh DB.
|
||||
var sharedIngestorCacheCounter atomic.Uint64
|
||||
|
||||
func sharedIngestorCacheDSN(testName string) string {
|
||||
var b strings.Builder
|
||||
b.Grow(len(testName) + 48)
|
||||
b.WriteString("file:test-")
|
||||
for _, r := range testName {
|
||||
switch {
|
||||
case r >= 'A' && r <= 'Z', r >= 'a' && r <= 'z', r >= '0' && r <= '9':
|
||||
b.WriteRune(r)
|
||||
default:
|
||||
b.WriteByte('_')
|
||||
}
|
||||
}
|
||||
b.WriteByte('-')
|
||||
b.WriteString(fmt.Sprintf("%d", sharedIngestorCacheCounter.Add(1)))
|
||||
b.WriteString("?mode=memory&cache=shared")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestIngestor_ExecuteTask_PipelineRuns
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// TestIngestor_ExecuteTask_PipelineRuns is the load-bearing
|
||||
// e2e test for Phase 3. It:
|
||||
//
|
||||
// 1. Spins up in-memory SQLite + memory storage.
|
||||
// 2. Creates an Ingestor + seeds an IngestionTask row with
|
||||
// a custom 1-stage DSL that points at the stub
|
||||
// component.
|
||||
// 3. Calls executeTask on a TaskContext wired with the
|
||||
// stub TaskHandle.
|
||||
// 4. Verifies:
|
||||
// a. The stub component was invoked exactly once.
|
||||
// b. The IngestionTaskLog has at least one row
|
||||
// recording the stub's stage completion.
|
||||
// c. The IngestionTask status is COMPLETED.
|
||||
// d. The stub TaskHandle was Ack()ed exactly once
|
||||
// (plan §8 Q3 — fixes the pre-existing no-Ack bug).
|
||||
func TestIngestor_ExecuteTask_PipelineRuns(t *testing.T) {
|
||||
atomic.StoreInt32(&stubCounter, 0)
|
||||
|
||||
db := setupIngestorTestDB(t)
|
||||
origDB := dao.DB
|
||||
dao.DB = db
|
||||
t.Cleanup(func() { dao.DB = origDB })
|
||||
|
||||
// Inject memory storage so the Ingestor's
|
||||
// storage.GetStorageFactory().GetStorage() call returns
|
||||
// a usable backend. Production uses MinIO; tests use the
|
||||
// in-memory mock.
|
||||
origStorage := storage.GetStorageFactory().GetStorage()
|
||||
storage.GetStorageFactory().SetStorage(storage.NewMemoryStorage())
|
||||
t.Cleanup(func() { storage.GetStorageFactory().SetStorage(origStorage) })
|
||||
|
||||
taskID := "ingestor-e2e-task"
|
||||
task := &entity.IngestionTask{
|
||||
ID: taskID,
|
||||
UserID: "user-e2e",
|
||||
DocumentID: "doc-e2e",
|
||||
DatasetID: "ds-e2e",
|
||||
Status: common.RUNNING,
|
||||
}
|
||||
if err := db.Create(task).Error; err != nil {
|
||||
t.Fatalf("create task: %v", err)
|
||||
}
|
||||
|
||||
// Canonical template wrapper carrying a tiny canvas DSL.
|
||||
dsl := map[string]any{
|
||||
"dsl": map[string]any{
|
||||
"components": map[string]any{
|
||||
"begin": map[string]any{
|
||||
"obj": map[string]any{"component_name": "Begin", "params": map[string]any{}},
|
||||
"downstream": []string{"stub"},
|
||||
},
|
||||
"stub": map[string]any{
|
||||
"obj": map[string]any{"component_name": stubIngestionComp, "params": map[string]any{}},
|
||||
"upstream": []string{"begin"},
|
||||
},
|
||||
},
|
||||
"path": []string{"begin", "stub"},
|
||||
"graph": map[string]any{"nodes": []any{}},
|
||||
},
|
||||
}
|
||||
dslBytes, _ := json.Marshal(dsl)
|
||||
task.Schema = entity.JSONMap{"pipeline": []byte(dslBytes)}
|
||||
if err := db.Save(task).Error; err != nil {
|
||||
t.Fatalf("save task: %v", err)
|
||||
}
|
||||
|
||||
handle := &stubTaskHandle{taskID: taskID}
|
||||
taskCtx := &TaskContext{
|
||||
Ctx: context.Background(),
|
||||
Task: task,
|
||||
TaskHandle: handle,
|
||||
}
|
||||
|
||||
ing := NewIngestor("test-ingestor", 1, []string{"pdf"})
|
||||
ing.executeTask(taskCtx)
|
||||
|
||||
// (a) stub invoked
|
||||
if got := atomic.LoadInt32(&stubCounter); got != 1 {
|
||||
t.Errorf("expected stub invoked 1 time, got %d", got)
|
||||
}
|
||||
|
||||
// (b) task is COMPLETED.
|
||||
var reloaded entity.IngestionTask
|
||||
if err := db.Where("id = ?", taskID).First(&reloaded).Error; err != nil {
|
||||
t.Fatalf("reload task: %v", err)
|
||||
}
|
||||
if reloaded.Status != common.COMPLETED {
|
||||
t.Errorf("expected status=COMPLETED, got %s", reloaded.Status)
|
||||
}
|
||||
|
||||
// (c) Ack() called exactly once (plan §8 Q3).
|
||||
if got := atomic.LoadInt32(&handle.acked); got != 1 {
|
||||
t.Errorf("expected Ack() called 1 time, got %d", got)
|
||||
}
|
||||
if got := atomic.LoadInt32(&handle.nacked); got != 0 {
|
||||
t.Errorf("expected Nack() called 0 times, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestIngestor_ExecuteTask_MissingDSL verifies the task fails cleanly when
|
||||
// no canonical DSL is present on the task schema.
|
||||
func TestIngestor_ExecuteTask_MissingDSL(t *testing.T) {
|
||||
db := setupIngestorTestDB(t)
|
||||
origDB := dao.DB
|
||||
dao.DB = db
|
||||
t.Cleanup(func() { dao.DB = origDB })
|
||||
|
||||
taskID := "missing-dsl-task"
|
||||
task := &entity.IngestionTask{
|
||||
ID: taskID,
|
||||
UserID: "user-1",
|
||||
DocumentID: "doc-1",
|
||||
DatasetID: "ds-1",
|
||||
Status: common.RUNNING,
|
||||
}
|
||||
if err := db.Create(task).Error; err != nil {
|
||||
t.Fatalf("create task: %v", err)
|
||||
}
|
||||
|
||||
handle := &stubTaskHandle{taskID: taskID}
|
||||
taskCtx := &TaskContext{
|
||||
Ctx: context.Background(),
|
||||
Task: task,
|
||||
TaskHandle: handle,
|
||||
}
|
||||
|
||||
ing := NewIngestor("test-missing-dsl", 1, []string{"pdf"})
|
||||
ing.executeTask(taskCtx)
|
||||
|
||||
var reloaded entity.IngestionTask
|
||||
if err := db.Where("id = ?", taskID).First(&reloaded).Error; err != nil {
|
||||
t.Fatalf("reload task: %v", err)
|
||||
}
|
||||
if reloaded.Status != common.FAILED {
|
||||
t.Errorf("expected status=FAILED, got %s", reloaded.Status)
|
||||
}
|
||||
if got := atomic.LoadInt32(&handle.acked); got != 1 {
|
||||
t.Errorf("expected Ack() called 1 time on FAILED, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestIngestor_ExecuteTask_Cancellation verifies the ctx
|
||||
// cancellation path: cancel the context, expect the task
|
||||
// to be STOPPED and the message Nack()ed.
|
||||
func TestIngestor_ExecuteTask_Cancellation(t *testing.T) {
|
||||
atomic.StoreInt32(&stubCounter, 0)
|
||||
|
||||
db := setupIngestorTestDB(t)
|
||||
origDB := dao.DB
|
||||
dao.DB = db
|
||||
t.Cleanup(func() { dao.DB = origDB })
|
||||
|
||||
origStorage := storage.GetStorageFactory().GetStorage()
|
||||
storage.GetStorageFactory().SetStorage(storage.NewMemoryStorage())
|
||||
t.Cleanup(func() { storage.GetStorageFactory().SetStorage(origStorage) })
|
||||
|
||||
taskID := "cancel-task"
|
||||
task := &entity.IngestionTask{
|
||||
ID: taskID,
|
||||
UserID: "u",
|
||||
DocumentID: "d",
|
||||
DatasetID: "ds",
|
||||
Status: common.RUNNING,
|
||||
}
|
||||
if err := db.Create(task).Error; err != nil {
|
||||
t.Fatalf("create task: %v", err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
// Cancel immediately so the run sees a cancelled ctx.
|
||||
cancel()
|
||||
|
||||
dsl := map[string]any{
|
||||
"dsl": map[string]any{
|
||||
"components": map[string]any{
|
||||
"begin": map[string]any{
|
||||
"obj": map[string]any{"component_name": "Begin", "params": map[string]any{}},
|
||||
"downstream": []string{"stub"},
|
||||
},
|
||||
"stub": map[string]any{
|
||||
"obj": map[string]any{"component_name": stubIngestionComp, "params": map[string]any{}},
|
||||
"upstream": []string{"begin"},
|
||||
},
|
||||
},
|
||||
"path": []string{"begin", "stub"},
|
||||
"graph": map[string]any{"nodes": []any{}},
|
||||
},
|
||||
}
|
||||
dslBytes, _ := json.Marshal(dsl)
|
||||
task.Schema = entity.JSONMap{"pipeline": []byte(dslBytes)}
|
||||
if err := db.Save(task).Error; err != nil {
|
||||
t.Fatalf("save task: %v", err)
|
||||
}
|
||||
|
||||
handle := &stubTaskHandle{taskID: taskID}
|
||||
taskCtx := &TaskContext{
|
||||
Ctx: ctx,
|
||||
Task: task,
|
||||
TaskHandle: handle,
|
||||
}
|
||||
|
||||
ing := NewIngestor("test-cancel", 1, []string{"pdf"})
|
||||
ing.executeTask(taskCtx)
|
||||
|
||||
// Status: STOPPED on cancellation.
|
||||
var reloaded entity.IngestionTask
|
||||
if err := db.Where("id = ?", taskID).First(&reloaded).Error; err != nil {
|
||||
t.Fatalf("reload: %v", err)
|
||||
}
|
||||
// The cancellation may either STOPPED (cancellation
|
||||
// path) or FAILED (stage Invoke failed because ctx was
|
||||
// done) — either is acceptable. The key invariant is
|
||||
// that the message was NOT silently dropped without
|
||||
// Ack/Nack.
|
||||
if reloaded.Status != common.STOPPED && reloaded.Status != common.FAILED {
|
||||
t.Errorf("expected status=STOPPED or FAILED on cancel, got %s", reloaded.Status)
|
||||
}
|
||||
// Total message disposition: Ack() OR Nack() must have
|
||||
// been called (the bug pre-Phase-3 was: neither).
|
||||
totalDisp := atomic.LoadInt32(&handle.acked) + atomic.LoadInt32(&handle.nacked)
|
||||
if totalDisp != 1 {
|
||||
t.Errorf("expected exactly one Ack/Nack on cancel, got ack=%d nack=%d", handle.acked, handle.nacked)
|
||||
}
|
||||
}
|
||||
|
||||
// TestIngestor_ExecuteTask_MalformedDSL verifies executeTask
|
||||
// fails cleanly on a bad DSL: task is FAILED, Ack() called.
|
||||
func TestIngestor_ExecuteTask_MalformedDSL(t *testing.T) {
|
||||
atomic.StoreInt32(&stubCounter, 0)
|
||||
|
||||
db := setupIngestorTestDB(t)
|
||||
origDB := dao.DB
|
||||
dao.DB = db
|
||||
t.Cleanup(func() { dao.DB = origDB })
|
||||
|
||||
origStorage := storage.GetStorageFactory().GetStorage()
|
||||
storage.GetStorageFactory().SetStorage(storage.NewMemoryStorage())
|
||||
t.Cleanup(func() { storage.GetStorageFactory().SetStorage(origStorage) })
|
||||
|
||||
taskID := "malformed-task"
|
||||
task := &entity.IngestionTask{
|
||||
ID: taskID,
|
||||
UserID: "u",
|
||||
DocumentID: "d",
|
||||
DatasetID: "ds",
|
||||
Status: common.RUNNING,
|
||||
Schema: entity.JSONMap{"pipeline": []byte("not-json")},
|
||||
}
|
||||
if err := db.Create(task).Error; err != nil {
|
||||
t.Fatalf("create: %v", err)
|
||||
}
|
||||
|
||||
handle := &stubTaskHandle{taskID: taskID}
|
||||
taskCtx := &TaskContext{
|
||||
Ctx: context.Background(),
|
||||
Task: task,
|
||||
TaskHandle: handle,
|
||||
}
|
||||
|
||||
ing := NewIngestor("test-malformed", 1, []string{"pdf"})
|
||||
ing.executeTask(taskCtx)
|
||||
|
||||
var reloaded entity.IngestionTask
|
||||
if err := db.Where("id = ?", taskID).First(&reloaded).Error; err != nil {
|
||||
t.Fatalf("reload: %v", err)
|
||||
}
|
||||
if reloaded.Status != common.FAILED {
|
||||
t.Errorf("expected status=FAILED on malformed DSL, got %s", reloaded.Status)
|
||||
}
|
||||
if got := atomic.LoadInt32(&handle.acked); got != 1 {
|
||||
t.Errorf("expected Ack() on FAILED, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// tiny helper for strings.Contains — Go 1.21 has it but the
|
||||
// codebase pins strings.Contains.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
var _ = strings.Contains
|
||||
var _ = time.Now
|
||||
var _ = fmt.Sprintf
|
||||
@@ -1,117 +0,0 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/net/html"
|
||||
)
|
||||
|
||||
const (
|
||||
Official string = "official"
|
||||
)
|
||||
|
||||
type HTMLParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewHTMLParser(libType string) (*HTMLParser, error) {
|
||||
switch libType {
|
||||
case Official:
|
||||
return &HTMLParser{
|
||||
libType: Official,
|
||||
}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported HTML library type: %s", libType)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *HTMLParser) Parse(filename string, data []byte) error {
|
||||
switch p.libType {
|
||||
case Official:
|
||||
return p.OfficialHTMLParse(data)
|
||||
default:
|
||||
return fmt.Errorf("unsupported HTML library type: %s", p.libType)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *HTMLParser) OfficialHTMLParse(data []byte) error {
|
||||
doc, _ := html.Parse(strings.NewReader(string(data)))
|
||||
p.WalkIterative(doc)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *HTMLParser) WalkIterative(root *html.Node) {
|
||||
if root == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Stack: stores node and its depth
|
||||
type item struct {
|
||||
node *html.Node
|
||||
depth int
|
||||
}
|
||||
stack := []item{{root, 0}}
|
||||
|
||||
for len(stack) > 0 {
|
||||
// Pop the top of the stack
|
||||
current := stack[len(stack)-1]
|
||||
stack = stack[:len(stack)-1]
|
||||
|
||||
indent := strings.Repeat(" ", current.depth)
|
||||
|
||||
// Handle different node types
|
||||
switch current.node.Type {
|
||||
case html.ElementNode:
|
||||
// Print opening tag
|
||||
fmt.Printf("%s<%s", indent, current.node.Data)
|
||||
// Optionally print attributes
|
||||
for _, attr := range current.node.Attr {
|
||||
fmt.Printf(" %s=%q", attr.Key, attr.Val)
|
||||
}
|
||||
fmt.Println(">")
|
||||
|
||||
case html.TextNode:
|
||||
// Print text content (trim extra whitespace)
|
||||
text := strings.TrimSpace(current.node.Data)
|
||||
if text != "" {
|
||||
fmt.Printf("%stext: %q\n", indent, text)
|
||||
}
|
||||
|
||||
case html.CommentNode:
|
||||
fmt.Printf("%scomment: %s\n", indent, current.node.Data)
|
||||
|
||||
case html.DoctypeNode:
|
||||
fmt.Printf("%sDOCTYPE: %s\n", indent, current.node.Data)
|
||||
}
|
||||
|
||||
// Push children onto stack in reverse order to maintain original sequence
|
||||
var children []*html.Node
|
||||
for child := current.node.FirstChild; child != nil; child = child.NextSibling {
|
||||
children = append([]*html.Node{child}, children...) // Reverse order
|
||||
}
|
||||
for _, child := range children {
|
||||
stack = append(stack, item{child, current.depth + 1})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *HTMLParser) String() string {
|
||||
return "HTMLParser"
|
||||
}
|
||||
@@ -1,71 +0,0 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
"github.com/gomarkdown/markdown/ast"
|
||||
"github.com/gomarkdown/markdown/parser"
|
||||
)
|
||||
|
||||
const (
|
||||
GoMarkdown = "go_markdown"
|
||||
)
|
||||
|
||||
type MarkdownParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewMarkdownParser(libType string) (*MarkdownParser, error) {
|
||||
switch libType {
|
||||
case GoMarkdown:
|
||||
return &MarkdownParser{
|
||||
libType: GoMarkdown,
|
||||
}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported Markdown library type: %s", libType)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *MarkdownParser) Parse(filename string, data []byte) error {
|
||||
switch p.libType {
|
||||
case GoMarkdown:
|
||||
return p.GoMarkdownParse(data)
|
||||
default:
|
||||
return fmt.Errorf("unsupported Markdown library type: %s", p.libType)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *MarkdownParser) GoMarkdownParse(data []byte) error {
|
||||
// create Markdown parser with extensions
|
||||
extensions := parser.CommonExtensions | parser.AutoHeadingIDs | parser.NoEmptyLineBeforeBlock
|
||||
parser.NewWithExtensions(extensions)
|
||||
markdownParser := parser.NewWithExtensions(extensions)
|
||||
doc := markdownParser.Parse(data)
|
||||
|
||||
fmt.Print("--- AST tree:\n")
|
||||
ast.Print(os.Stdout, doc)
|
||||
fmt.Print("\n")
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *MarkdownParser) String() string {
|
||||
return "MarkdownParser"
|
||||
}
|
||||
@@ -1,104 +0,0 @@
|
||||
//go:build !cgo || !office
|
||||
|
||||
package parser
|
||||
|
||||
import "fmt"
|
||||
|
||||
// OfficeOxide is the lib_type identifier for office_oxide backend.
|
||||
const OfficeOxide = "office_oxide"
|
||||
|
||||
type DOCParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewDOCParser(libType string) (*DOCParser, error) {
|
||||
return nil, fmt.Errorf("DOC parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *DOCParser) Parse(_ string, _ []byte) error {
|
||||
return fmt.Errorf("DOC parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *DOCParser) String() string {
|
||||
return "DOCParser(no-cgo)"
|
||||
}
|
||||
|
||||
type DOCXParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewDOCXParser(libType string) (*DOCXParser, error) {
|
||||
return nil, fmt.Errorf("DOCX parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *DOCXParser) Parse(_ string, _ []byte) error {
|
||||
return fmt.Errorf("DOCX parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *DOCXParser) String() string {
|
||||
return "DOCXParser(no-cgo)"
|
||||
}
|
||||
|
||||
type XLSParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewXLSParser(libType string) (*XLSParser, error) {
|
||||
return nil, fmt.Errorf("XLS parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *XLSParser) Parse(_ string, _ []byte) error {
|
||||
return fmt.Errorf("XLS parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *XLSParser) String() string {
|
||||
return "XLSParser(no-cgo)"
|
||||
}
|
||||
|
||||
type XLSXParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewXLSXParser(libType string) (*XLSXParser, error) {
|
||||
return nil, fmt.Errorf("XLSX parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *XLSXParser) Parse(_ string, _ []byte) error {
|
||||
return fmt.Errorf("XLSX parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *XLSXParser) String() string {
|
||||
return "XLSXParser(no-cgo)"
|
||||
}
|
||||
|
||||
type PPTParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewPPTParser(libType string) (*PPTParser, error) {
|
||||
return nil, fmt.Errorf("PPT parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *PPTParser) Parse(_ string, _ []byte) error {
|
||||
return fmt.Errorf("PPT parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *PPTParser) String() string {
|
||||
return "PPTParser(no-cgo)"
|
||||
}
|
||||
|
||||
type PPTXParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewPPTXParser(libType string) (*PPTXParser, error) {
|
||||
return nil, fmt.Errorf("PPTX parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *PPTXParser) Parse(_ string, _ []byte) error {
|
||||
return fmt.Errorf("PPTX parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *PPTXParser) String() string {
|
||||
return "PPTXParser(no-cgo)"
|
||||
}
|
||||
@@ -1,81 +0,0 @@
|
||||
//go:build cgo && office
|
||||
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
officeOxide "github.com/yfedoseev/office_oxide/go"
|
||||
)
|
||||
|
||||
type XLSParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewXLSParser(libType string) (*XLSParser, error) {
|
||||
switch libType {
|
||||
case OfficeOxide:
|
||||
return &XLSParser{
|
||||
libType: OfficeOxide,
|
||||
}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported XLS library type: %s", libType)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *XLSParser) Parse(filename string, data []byte) error {
|
||||
switch p.libType {
|
||||
case OfficeOxide:
|
||||
return p.OfficeOxideParse(data)
|
||||
default:
|
||||
return fmt.Errorf("unsupported XLS library type: %s", p.libType)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *XLSParser) OfficeOxideParse(data []byte) error {
|
||||
doc, err := officeOxide.OpenFromBytes(data, "xls")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer doc.Close()
|
||||
|
||||
docFormat, err := doc.Format()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
fmt.Println("Document format:", docFormat)
|
||||
|
||||
docContext, err := doc.PlainText()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Println("Document context:", docContext)
|
||||
|
||||
md, err := doc.ToMarkdown()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Println("Document Markdown:", md)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *XLSParser) String() string {
|
||||
return "XLSParser"
|
||||
}
|
||||
@@ -1,81 +0,0 @@
|
||||
//go:build cgo && office
|
||||
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
officeOxide "github.com/yfedoseev/office_oxide/go"
|
||||
)
|
||||
|
||||
type XLSXParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewXLSXParser(libType string) (*XLSXParser, error) {
|
||||
switch libType {
|
||||
case OfficeOxide:
|
||||
return &XLSXParser{
|
||||
libType: OfficeOxide,
|
||||
}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported XLSX library type: %s", libType)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *XLSXParser) Parse(filename string, data []byte) error {
|
||||
switch p.libType {
|
||||
case OfficeOxide:
|
||||
return p.OfficeOxideParse(data)
|
||||
default:
|
||||
return fmt.Errorf("unsupported XLSX library type: %s", p.libType)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *XLSXParser) OfficeOxideParse(data []byte) error {
|
||||
doc, err := officeOxide.OpenFromBytes(data, "xlsx")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer doc.Close()
|
||||
|
||||
docFormat, err := doc.Format()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
fmt.Println("Document format:", docFormat)
|
||||
|
||||
docContext, err := doc.PlainText()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Println("Document context:", docContext)
|
||||
|
||||
md, err := doc.ToMarkdown()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Println("Document Markdown:", md)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *XLSXParser) String() string {
|
||||
return "XLSXParser"
|
||||
}
|
||||
43
internal/ingestion/pipeline/errors.go
Normal file
43
internal/ingestion/pipeline/errors.go
Normal file
@@ -0,0 +1,43 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package pipeline
|
||||
|
||||
import "errors"
|
||||
|
||||
var (
|
||||
errNilDSL = errors.New("pipeline: nil DSL")
|
||||
errEmptyStages = errors.New("pipeline: DSL has no components")
|
||||
errUnknownComponent = errors.New("pipeline: unknown component")
|
||||
errUnknownStage = errors.New("pipeline: unknown stage")
|
||||
)
|
||||
|
||||
type stageError struct {
|
||||
Stage string
|
||||
Reason string
|
||||
}
|
||||
|
||||
func (e *stageError) Error() string {
|
||||
return "pipeline: stage " + e.Stage + ": " + e.Reason
|
||||
}
|
||||
|
||||
type sinkError struct {
|
||||
Reason string
|
||||
}
|
||||
|
||||
func (e *sinkError) Error() string {
|
||||
return "pipeline: sink: " + e.Reason
|
||||
}
|
||||
141
internal/ingestion/pipeline/pipeline.go
Normal file
141
internal/ingestion/pipeline/pipeline.go
Normal file
@@ -0,0 +1,141 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package pipeline
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"ragflow/internal/agent/canvas"
|
||||
_ "ragflow/internal/agent/component"
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
// Pipeline is a compiled ingestion canvas plus task-scoped metadata.
|
||||
type Pipeline struct {
|
||||
taskID string
|
||||
canvas *canvas.Canvas
|
||||
}
|
||||
|
||||
// NewPipelineFromDSL compiles the canonical ingestion canvas DSL.
|
||||
// It accepts either the inner canvas DSL or the template wrapper whose
|
||||
// top-level `dsl` field carries that canvas.
|
||||
func NewPipelineFromDSL(dsl []byte, taskID string) (*Pipeline, error) {
|
||||
var raw map[string]any
|
||||
if err := json.Unmarshal(dsl, &raw); err != nil {
|
||||
return nil, fmt.Errorf("pipeline: decode DSL: %w", err)
|
||||
}
|
||||
canvasDSL, err := unwrapCanvasDSL(raw)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cnv, err := canvas.DecodeFromDSL(canvasDSL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("pipeline: decode canvas DSL: %w", err)
|
||||
}
|
||||
return &Pipeline{
|
||||
taskID: taskID,
|
||||
canvas: cnv,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func unwrapCanvasDSL(raw map[string]any) (map[string]any, error) {
|
||||
if len(raw) == 0 {
|
||||
return nil, errNilDSL
|
||||
}
|
||||
if rawDSL, ok := raw["dsl"]; ok {
|
||||
canvasDSL, ok := rawDSL.(map[string]any)
|
||||
if !ok || len(canvasDSL) == 0 {
|
||||
return nil, errNilDSL
|
||||
}
|
||||
return canvasDSL, nil
|
||||
}
|
||||
return raw, nil
|
||||
}
|
||||
|
||||
func mergeInto(dst, src map[string]any) map[string]any {
|
||||
if src == nil {
|
||||
return dst
|
||||
}
|
||||
if dst == nil {
|
||||
dst = make(map[string]any, len(src))
|
||||
}
|
||||
for k, v := range src {
|
||||
dst[k] = v
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
func cloneMapOrEmpty(m map[string]any) map[string]any {
|
||||
if m == nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
out := make(map[string]any, len(m))
|
||||
for k, v := range m {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func stageTimeout() time.Duration {
|
||||
return defaultStageTimeout
|
||||
}
|
||||
|
||||
var defaultStageTimeout = 60 * time.Second
|
||||
|
||||
// Run executes the full ingestion graph described by the canonical DSL.
|
||||
// There is no pipeline-layer partial resume entry point: execution always
|
||||
// starts from the graph entry and component-level replay decisions belong to
|
||||
// the components themselves.
|
||||
func (p *Pipeline) Run(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
if p == nil {
|
||||
return nil, fmt.Errorf("pipeline: Run on nil pipeline")
|
||||
}
|
||||
if p.canvas == nil {
|
||||
return nil, fmt.Errorf("pipeline: canvas is nil")
|
||||
}
|
||||
if runtime.DefaultFactory() == nil {
|
||||
runtime.InstallDefaultRegistryFactory()
|
||||
}
|
||||
if runtime.DefaultFactory() == nil {
|
||||
return nil, fmt.Errorf("pipeline: Run: runtime default component factory is not installed")
|
||||
}
|
||||
|
||||
compiled, err := canvas.Compile(ctx, p.canvas)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("pipeline: Run: compile canvas: %w", err)
|
||||
}
|
||||
|
||||
runState := canvas.NewCanvasState("", p.taskID)
|
||||
runCtx := canvas.WithState(ctx, runState)
|
||||
runCtx = canvas.WithComponentTimeoutOverride(runCtx, stageTimeout())
|
||||
|
||||
current := cloneMapOrEmpty(inputs)
|
||||
out, err := compiled.Workflow.Invoke(runCtx, current)
|
||||
if err != nil {
|
||||
return current, fmt.Errorf("pipeline: run canvas workflow: %w", err)
|
||||
}
|
||||
if out == nil {
|
||||
current["state"] = runState.Snapshot()
|
||||
return current, nil
|
||||
}
|
||||
merged := mergeInto(current, out)
|
||||
merged["state"] = runState.Snapshot()
|
||||
return merged, nil
|
||||
}
|
||||
152
internal/ingestion/pipeline/pipeline_test.go
Normal file
152
internal/ingestion/pipeline/pipeline_test.go
Normal file
@@ -0,0 +1,152 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package pipeline
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
type mockCanvasStage struct {
|
||||
output map[string]any
|
||||
called bool
|
||||
}
|
||||
|
||||
func (m *mockCanvasStage) Invoke(_ context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
m.called = true
|
||||
out := cloneMapOrEmpty(inputs)
|
||||
for k, v := range m.output {
|
||||
out[k] = v
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (m *mockCanvasStage) Parallelism() int { return 1 }
|
||||
func (m *mockCanvasStage) Inputs() map[string]string { return map[string]string{"name": "string"} }
|
||||
func (m *mockCanvasStage) Outputs() map[string]string { return map[string]string{"output": "any"} }
|
||||
|
||||
func TestPipelineRunHappyPath(t *testing.T) {
|
||||
stageA := &mockCanvasStage{output: map[string]any{"a": 1}}
|
||||
stageB := &mockCanvasStage{output: map[string]any{"b": 2}}
|
||||
|
||||
const (
|
||||
nameA = "p.RunStageA"
|
||||
nameB = "p.RunStageB"
|
||||
)
|
||||
runtime.MustRegister(nameA, runtime.CategoryIngestion,
|
||||
func(_ string, _ map[string]any) (runtime.Component, error) { return stageA, nil },
|
||||
runtime.Metadata{Version: "1.0.0"})
|
||||
runtime.MustRegister(nameB, runtime.CategoryIngestion,
|
||||
func(_ string, _ map[string]any) (runtime.Component, error) { return stageB, nil },
|
||||
runtime.Metadata{Version: "1.0.0"})
|
||||
|
||||
pipe, err := NewPipelineFromDSL([]byte(`{
|
||||
"dsl": {
|
||||
"components": {
|
||||
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["a"]},
|
||||
"a": {"obj": {"component_name": "`+nameA+`", "params": {}}, "upstream": ["begin"], "downstream": ["b"]},
|
||||
"b": {"obj": {"component_name": "`+nameB+`", "params": {}}, "upstream": ["a"]}
|
||||
},
|
||||
"path": ["begin", "a", "b"],
|
||||
"graph": {"nodes": []}
|
||||
}
|
||||
}`), "task-canvas-happy")
|
||||
if err != nil {
|
||||
t.Fatalf("NewPipelineFromDSL: %v", err)
|
||||
}
|
||||
|
||||
out, err := pipe.Run(context.Background(), map[string]any{"name": "doc-canvas"})
|
||||
if err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
if !stageA.called || !stageB.called {
|
||||
t.Fatalf("expected both stages to run, got A=%v B=%v", stageA.called, stageB.called)
|
||||
}
|
||||
if got := out["name"]; got != "doc-canvas" {
|
||||
t.Fatalf("name = %v, want doc-canvas", got)
|
||||
}
|
||||
gotB, ok := out["b"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("b = %T, want map[string]any", out["b"])
|
||||
}
|
||||
if got := gotB["b"]; got != 2 {
|
||||
t.Fatalf("b.b = %v, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPipelineRunNilPipeline(t *testing.T) {
|
||||
var p *Pipeline
|
||||
if _, err := p.Run(context.Background(), nil); err == nil {
|
||||
t.Fatal("expected error for nil pipeline")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPipelineRunStageErrorBubbles(t *testing.T) {
|
||||
const name = "p.RunErrStage"
|
||||
runtime.MustRegister(name, runtime.CategoryIngestion,
|
||||
func(_ string, _ map[string]any) (runtime.Component, error) { return &errCanvasStage{}, nil },
|
||||
runtime.Metadata{Version: "1.0.0"})
|
||||
|
||||
pipe, err := NewPipelineFromDSL([]byte(`{
|
||||
"dsl": {
|
||||
"components": {
|
||||
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["err"]},
|
||||
"err": {"obj": {"component_name": "`+name+`", "params": {}}, "upstream": ["begin"]}
|
||||
},
|
||||
"path": ["begin", "err"],
|
||||
"graph": {"nodes": []}
|
||||
}
|
||||
}`), "task-canvas-err")
|
||||
if err != nil {
|
||||
t.Fatalf("NewPipelineFromDSL: %v", err)
|
||||
}
|
||||
|
||||
if _, err := pipe.Run(context.Background(), map[string]any{"name": "x"}); err == nil {
|
||||
t.Fatal("expected stage error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewPipelineFromDSLUnwrapsTemplateDSL(t *testing.T) {
|
||||
pipe, err := NewPipelineFromDSL([]byte(`{
|
||||
"id": "template-1",
|
||||
"title": "template",
|
||||
"dsl": {
|
||||
"components": {
|
||||
"begin": {"obj": {"component_name": "Begin", "params": {}}}
|
||||
},
|
||||
"path": ["begin"],
|
||||
"graph": {"nodes": []}
|
||||
}
|
||||
}`), "task-template")
|
||||
if err != nil {
|
||||
t.Fatalf("NewPipelineFromDSL: %v", err)
|
||||
}
|
||||
if pipe.canvas == nil {
|
||||
t.Fatal("expected decoded canvas")
|
||||
}
|
||||
}
|
||||
|
||||
type errCanvasStage struct{}
|
||||
|
||||
func (e *errCanvasStage) Invoke(_ context.Context, _ map[string]any) (map[string]any, error) {
|
||||
return nil, &stageError{Stage: "p.RunErrStage", Reason: "intentional"}
|
||||
}
|
||||
func (e *errCanvasStage) Parallelism() int { return 1 }
|
||||
func (e *errCanvasStage) Inputs() map[string]string { return nil }
|
||||
func (e *errCanvasStage) Outputs() map[string]string { return nil }
|
||||
1001
internal/ingestion/pipeline/template_integration_test.go
Normal file
1001
internal/ingestion/pipeline/template_integration_test.go
Normal file
File diff suppressed because it is too large
Load Diff
76
internal/ingestion/wire/wire.go
Normal file
76
internal/ingestion/wire/wire.go
Normal file
@@ -0,0 +1,76 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Single owner for the ingestion-component registration surface.
|
||||
//
|
||||
// Plan §2 AD-2 places every component under the same registry as
|
||||
// the agent canvas; in practice that means cmd entries (server,
|
||||
// ingestor, ...) must blank-import the ingestion component
|
||||
// packages so their init() runs. The historical pattern was
|
||||
// for each cmd entry to repeat the blank imports, which is
|
||||
// fragile: a new entry that forgets the blank import sees
|
||||
// "unknown component" at run time.
|
||||
//
|
||||
// This package centralises the import list. cmd entries should
|
||||
// blank-import this package and call RegisterComponents once
|
||||
// at startup. The function is a no-op at run time — its only
|
||||
// purpose is to ensure the blank imports here are evaluated,
|
||||
// which in turn triggers init() in each component package.
|
||||
//
|
||||
// LOCATION NOTE: this lives under ingestion/ (not
|
||||
// ingestion/pipeline/) because pipeline.go imports the
|
||||
// chunker subpackage, and chunker imports the parent
|
||||
// ingestion package, which imports pipeline. A wire file
|
||||
// inside pipeline would close that cycle. By sitting
|
||||
// alongside pipeline under ingestion/, this package can
|
||||
// blank-import the component packages without going through
|
||||
// pipeline.
|
||||
package wire
|
||||
|
||||
import (
|
||||
// Component registration: the blank imports trigger each
|
||||
// package's init() which calls runtime.DefaultRegistry.Register.
|
||||
_ "ragflow/internal/ingestion/component" // File / Parser / Tokenizer / Extractor
|
||||
_ "ragflow/internal/ingestion/component/chunker" // 4 chunker variants
|
||||
)
|
||||
|
||||
// RegisterComponents is the single bootstrap entry point that
|
||||
// guarantees every ingestion component is registered before the
|
||||
// pipeline runner resolves a name. It is a no-op at run time
|
||||
// (the actual registration happens via the package-level init()
|
||||
// triggered by the blank imports above); the function exists
|
||||
// only so cmd entries have a deterministic symbol to call.
|
||||
//
|
||||
// Usage from a cmd entry:
|
||||
//
|
||||
// import _ "ragflow/internal/ingestion/wire"
|
||||
//
|
||||
// The blank import alone is sufficient — the Go toolchain will
|
||||
// evaluate the blank imports at link time. The function is
|
||||
// exported for callers that prefer a function-call style over a
|
||||
// blank import:
|
||||
//
|
||||
// import "ragflow/internal/ingestion/wire"
|
||||
// ...
|
||||
// wire.RegisterComponents()
|
||||
//
|
||||
// Both styles are equivalent; the blank-import form is the
|
||||
// convention in this repo today.
|
||||
func RegisterComponents() {
|
||||
// Intentional no-op. The init() functions in the blank-imported
|
||||
// packages above have already registered the components by
|
||||
// the time this function is callable.
|
||||
}
|
||||
95
internal/parser/chunk/execute.go
Normal file
95
internal/parser/chunk/execute.go
Normal file
@@ -0,0 +1,95 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Internal chunk execution entrypoint used by production callers.
|
||||
package chunk
|
||||
|
||||
import "fmt"
|
||||
|
||||
// Run executes the internal chunk steps against `text` using typed
|
||||
// options. The sequence is preprocess -> split -> postprocess.
|
||||
func Run(text string, opts ChunkOptions) (*ChunkContext, error) {
|
||||
if err := opts.validate(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ctx := &ChunkContext{
|
||||
Origin: text,
|
||||
TextAfterPreprocess: text,
|
||||
}
|
||||
|
||||
var stages []Operator
|
||||
var stageNames []string
|
||||
|
||||
if opts.NormalizeNewlines || opts.StripWhitespace || opts.RemoveEmptyLines {
|
||||
pre, err := NewPreprocessOperator(map[string]interface{}{
|
||||
"normalize_newlines": opts.NormalizeNewlines,
|
||||
"strip_whitespace": opts.StripWhitespace,
|
||||
"remove_empty_lines": opts.RemoveEmptyLines,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("chunk: build preprocess: %w", err)
|
||||
}
|
||||
stages = append(stages, pre)
|
||||
stageNames = append(stageNames, "preprocess")
|
||||
}
|
||||
|
||||
split, err := NewSplitOperator(map[string]interface{}{
|
||||
"strategy": opts.SplitStrategy,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("chunk: build split: %w", err)
|
||||
}
|
||||
stages = append(stages, split)
|
||||
stageNames = append(stageNames, "split")
|
||||
|
||||
postCfg := map[string]interface{}{}
|
||||
if opts.MergeTargetSize > 0 {
|
||||
postCfg["merge"] = map[string]interface{}{
|
||||
"target_size": float64(opts.MergeTargetSize),
|
||||
"strategy": "greedy",
|
||||
}
|
||||
}
|
||||
if opts.FilterMinLength > 0 {
|
||||
postCfg["filter"] = map[string]interface{}{
|
||||
"min_length": float64(opts.FilterMinLength),
|
||||
}
|
||||
}
|
||||
if len(postCfg) > 0 {
|
||||
post, err := NewPostprocessOperator(postCfg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("chunk: build postprocess: %w", err)
|
||||
}
|
||||
stages = append(stages, post)
|
||||
stageNames = append(stageNames, "postprocess")
|
||||
}
|
||||
|
||||
for i, op := range stages {
|
||||
if err := op.Prepare(ctx); err != nil {
|
||||
return ctx, fmt.Errorf("%s: prepare: %w", stageNames[i], err)
|
||||
}
|
||||
if err := op.Execute(ctx); err != nil {
|
||||
return ctx, fmt.Errorf("%s: execute: %w", stageNames[i], err)
|
||||
}
|
||||
if err := op.Finish(ctx); err != nil {
|
||||
return ctx, fmt.Errorf("%s: finish: %w", stageNames[i], err)
|
||||
}
|
||||
}
|
||||
if len(ctx.ResultChunks) == 0 && len(ctx.SplitChunks) > 0 {
|
||||
ctx.ResultChunks = ctx.SplitChunks
|
||||
}
|
||||
return ctx, nil
|
||||
}
|
||||
53
internal/parser/chunk/execute_test.go
Normal file
53
internal/parser/chunk/execute_test.go
Normal file
@@ -0,0 +1,53 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package chunk
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestRun_SentenceSplitHandlesASCIIBoundaries(t *testing.T) {
|
||||
ctx, err := Run("One. Two! Three? Four;", ChunkOptions{SplitStrategy: "sentence"})
|
||||
if err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
if len(ctx.ResultChunks) != 4 {
|
||||
t.Fatalf("got %d chunks, want 4", len(ctx.ResultChunks))
|
||||
}
|
||||
want := []string{"One", "Two", "Three", "Four"}
|
||||
for i, chunk := range ctx.ResultChunks {
|
||||
if chunk.Content != want[i] {
|
||||
t.Errorf("chunk[%d] = %q, want %q", i, chunk.Content, want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_PostprocessHonorsTypedNumericOptions(t *testing.T) {
|
||||
ctx, err := Run("a\nbb\nccc", ChunkOptions{
|
||||
SplitStrategy: "paragraph",
|
||||
MergeTargetSize: 100,
|
||||
FilterMinLength: 2,
|
||||
RemoveEmptyLines: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
if len(ctx.ResultChunks) != 1 {
|
||||
t.Fatalf("got %d chunks, want 1", len(ctx.ResultChunks))
|
||||
}
|
||||
if ctx.ResultChunks[0].Content != "a bb ccc" {
|
||||
t.Errorf("merged content = %q, want %q", ctx.ResultChunks[0].Content, "a bb ccc")
|
||||
}
|
||||
}
|
||||
@@ -14,22 +14,30 @@
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package parser
|
||||
package chunk
|
||||
|
||||
type PDFParser struct {
|
||||
ParserType string // DeepDoc, PaddleOCR, MinerU
|
||||
Model string // DeepDoc@buildin@ragflow
|
||||
LibType string // pdf_oxide, used by DeepDoc
|
||||
import "unicode"
|
||||
|
||||
// DetectLanguage returns a best-effort language code based on the
|
||||
// proportion of CJK characters.
|
||||
func DetectLanguage(text string) string {
|
||||
cjk := 0
|
||||
total := 0
|
||||
for _, r := range text {
|
||||
if unicode.Is(unicode.Han, r) {
|
||||
cjk++
|
||||
}
|
||||
if unicode.IsLetter(r) {
|
||||
total++
|
||||
}
|
||||
}
|
||||
if total > 0 && float64(cjk)/float64(total) > 0.3 {
|
||||
return "zh"
|
||||
}
|
||||
return "en"
|
||||
}
|
||||
|
||||
func NewPDFParser() *PDFParser {
|
||||
return &PDFParser{}
|
||||
}
|
||||
|
||||
func (p *PDFParser) Parse(filename string, data []byte) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *PDFParser) String() string {
|
||||
return "PDFParser"
|
||||
// RuneCount returns the number of runes in text.
|
||||
func RuneCount(text string) int {
|
||||
return len([]rune(text))
|
||||
}
|
||||
59
internal/parser/chunk/options.go
Normal file
59
internal/parser/chunk/options.go
Normal file
@@ -0,0 +1,59 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Typed options for the internal chunk execution path.
|
||||
|
||||
package chunk
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// ChunkOptions is the typed configuration for production callers.
|
||||
// It intentionally models only the option subset used by the current
|
||||
// production call sites.
|
||||
type ChunkOptions struct {
|
||||
// Preprocess flags. Any combination of the three is honoured; all
|
||||
// false means "no preprocess stage" (the engine skips the stage
|
||||
// rather than running an identity preprocess).
|
||||
NormalizeNewlines bool
|
||||
StripWhitespace bool
|
||||
RemoveEmptyLines bool
|
||||
|
||||
// Split configuration. Callers must select a strategy; an unset
|
||||
// strategy degrades to "sentence" inside the operator.
|
||||
SplitStrategy string
|
||||
|
||||
// Postprocess configuration. Zero values mean "do not run that
|
||||
// step"; non-zero values enable it.
|
||||
MergeTargetSize int
|
||||
|
||||
// FilterMinLength > 0 drops chunks shorter than that (rune count).
|
||||
FilterMinLength int
|
||||
}
|
||||
|
||||
// validate ensures the typed option set is internally consistent.
|
||||
// The check is cheap; an option set that fails validation will not
|
||||
// produce meaningful results at run time.
|
||||
func (o ChunkOptions) validate() error {
|
||||
if o.MergeTargetSize < 0 {
|
||||
return fmt.Errorf("chunk: MergeTargetSize must be >= 0 (got %d)", o.MergeTargetSize)
|
||||
}
|
||||
if o.FilterMinLength < 0 {
|
||||
return fmt.Errorf("chunk: FilterMinLength must be >= 0 (got %d)", o.FilterMinLength)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
200
internal/parser/chunk/postprocess.go
Normal file
200
internal/parser/chunk/postprocess.go
Normal file
@@ -0,0 +1,200 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package chunk
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type mergeConfig struct {
|
||||
TargetSize int `json:"target_size"`
|
||||
}
|
||||
|
||||
type filterConfig struct {
|
||||
MinLength int `json:"min_length"`
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// PostprocessOperator
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type PostprocessOperator struct {
|
||||
merge *mergeConfig
|
||||
filter *filterConfig
|
||||
}
|
||||
|
||||
func NewPostprocessOperator(config map[string]interface{}) (*PostprocessOperator, error) {
|
||||
op := &PostprocessOperator{}
|
||||
|
||||
// Merge
|
||||
if m, ok := config["merge"].(map[string]interface{}); ok {
|
||||
op.merge = &mergeConfig{}
|
||||
if ts, ok := m["target_size"].(float64); ok {
|
||||
op.merge.TargetSize = int(ts)
|
||||
} else {
|
||||
op.merge.TargetSize = 500
|
||||
}
|
||||
}
|
||||
|
||||
// Filter
|
||||
if f, ok := config["filter"].(map[string]interface{}); ok {
|
||||
op.filter = &filterConfig{}
|
||||
if v, ok := f["min_length"].(float64); ok {
|
||||
op.filter.MinLength = int(v)
|
||||
}
|
||||
}
|
||||
|
||||
return op, nil
|
||||
}
|
||||
|
||||
func (o *PostprocessOperator) Prepare(chunkCtx *ChunkContext) error {
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (o *PostprocessOperator) Execute(chunkCtx *ChunkContext) error {
|
||||
chunks := chunkCtx.SplitChunks
|
||||
if len(chunks) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 1. Merge
|
||||
if o.merge != nil {
|
||||
chunks = o.mergeChunks(chunks)
|
||||
}
|
||||
|
||||
// 2. Filter
|
||||
if o.filter != nil {
|
||||
chunks = o.filterChunks(chunks)
|
||||
}
|
||||
|
||||
// Re-index
|
||||
for i := range chunks {
|
||||
chunks[i].Index = i
|
||||
chunks[i].Size = len(chunks[i].GetContent())
|
||||
}
|
||||
|
||||
chunkCtx.ResultChunks = chunks
|
||||
return nil
|
||||
}
|
||||
|
||||
func (o *PostprocessOperator) Finish(chunkCtx *ChunkContext) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (o *PostprocessOperator) String() string {
|
||||
var buf strings.Builder
|
||||
buf.WriteString("postprocess:\n")
|
||||
|
||||
if o.merge != nil {
|
||||
fmt.Fprintf(&buf, " merge:\n")
|
||||
fmt.Fprintf(&buf, " target_size: %d\n", o.merge.TargetSize)
|
||||
}
|
||||
|
||||
if o.filter != nil {
|
||||
fmt.Fprintf(&buf, " filter:\n")
|
||||
fmt.Fprintf(&buf, " min_length: %d\n", o.filter.MinLength)
|
||||
}
|
||||
|
||||
return buf.String()
|
||||
}
|
||||
|
||||
// mergeChunks greedily merges small chunks into larger ones up to target_size.
|
||||
func (o *PostprocessOperator) mergeChunks(chunks []ChunkData) []ChunkData {
|
||||
target := o.merge.TargetSize
|
||||
if target <= 0 {
|
||||
target = 500
|
||||
}
|
||||
|
||||
var merged []ChunkData
|
||||
var buf strings.Builder
|
||||
var bufMeta map[string]interface{}
|
||||
firstIndex := 0
|
||||
|
||||
for i, c := range chunks {
|
||||
// If this single chunk already exceeds target, flush first then add
|
||||
if len([]rune(c.Content)) >= target {
|
||||
if buf.Len() > 0 {
|
||||
merged = append(merged, ChunkData{
|
||||
Content: buf.String(),
|
||||
Index: firstIndex,
|
||||
Metadata: bufMeta,
|
||||
})
|
||||
buf.Reset()
|
||||
bufMeta = nil
|
||||
}
|
||||
merged = append(merged, c)
|
||||
firstIndex = i + 1
|
||||
continue
|
||||
}
|
||||
|
||||
if buf.Len() == 0 {
|
||||
buf.WriteString(c.Content)
|
||||
bufMeta = c.Metadata
|
||||
firstIndex = c.Index
|
||||
} else {
|
||||
nextLen := len([]rune(c.Content))
|
||||
// If adding this chunk would exceed target, flush current and start new
|
||||
if buf.Len()+nextLen+1 > target {
|
||||
merged = append(merged, ChunkData{
|
||||
Content: buf.String(),
|
||||
Index: firstIndex,
|
||||
Metadata: bufMeta,
|
||||
})
|
||||
buf.Reset()
|
||||
buf.WriteString(c.Content)
|
||||
bufMeta = c.Metadata
|
||||
firstIndex = c.Index
|
||||
} else {
|
||||
buf.WriteString(" ")
|
||||
buf.WriteString(c.Content)
|
||||
// Merge metadata (last wins for overlapping keys)
|
||||
if c.Metadata != nil && bufMeta == nil {
|
||||
bufMeta = make(map[string]interface{})
|
||||
}
|
||||
for k, v := range c.Metadata {
|
||||
bufMeta[k] = v
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Flush remaining
|
||||
if buf.Len() > 0 {
|
||||
merged = append(merged, ChunkData{
|
||||
Content: buf.String(),
|
||||
Index: firstIndex,
|
||||
Metadata: bufMeta,
|
||||
})
|
||||
}
|
||||
|
||||
return merged
|
||||
}
|
||||
|
||||
// filterChunks removes chunks outside the length bounds.
|
||||
func (o *PostprocessOperator) filterChunks(chunks []ChunkData) []ChunkData {
|
||||
filtered := make([]ChunkData, 0, len(chunks))
|
||||
for _, c := range chunks {
|
||||
l := len([]rune(c.Content))
|
||||
if o.filter.MinLength > 0 && l < o.filter.MinLength {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, c)
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
@@ -23,9 +23,7 @@ import (
|
||||
)
|
||||
|
||||
type SplitOperator struct {
|
||||
strategy string
|
||||
boundaries []string
|
||||
keepSeparators bool
|
||||
strategy string
|
||||
}
|
||||
|
||||
func NewSplitOperator(config map[string]interface{}) (*SplitOperator, error) {
|
||||
@@ -37,23 +35,6 @@ func NewSplitOperator(config map[string]interface{}) (*SplitOperator, error) {
|
||||
}
|
||||
}
|
||||
|
||||
if params, ok := config["params"].(map[string]interface{}); ok {
|
||||
if b, ok := params["boundaries"]; ok {
|
||||
if boundStrs, ok := b.([]interface{}); ok {
|
||||
for _, bs := range boundStrs {
|
||||
if s, ok := bs.(string); ok {
|
||||
op.boundaries = append(op.boundaries, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if ks, ok := params["keep_separators"]; ok {
|
||||
if b, ok := ks.(bool); ok {
|
||||
op.keepSeparators = b
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return op, nil
|
||||
}
|
||||
|
||||
@@ -90,20 +71,13 @@ func (o *SplitOperator) String() string {
|
||||
var buf strings.Builder
|
||||
buf.WriteString("split:\n")
|
||||
fmt.Fprintf(&buf, " strategy: %q\n", o.strategy)
|
||||
fmt.Fprintf(&buf, " boundaries:\n")
|
||||
for _, r := range o.boundaries {
|
||||
fmt.Fprintf(&buf, " - %q\n", r)
|
||||
}
|
||||
fmt.Fprintf(&buf, " keep_separators: %t\n", o.keepSeparators)
|
||||
return buf.String()
|
||||
}
|
||||
|
||||
// splitSentences splits text at multi-rune boundaries, optionally keeping separators.
|
||||
func (o *SplitOperator) splitSentences(text string) []ChunkData {
|
||||
if len(o.boundaries) == 0 {
|
||||
o.boundaries = []string{"。", "！", "？", "\n"}
|
||||
}
|
||||
var sentenceBoundaries = []string{"。", "！", "？", ".", "!", "?", ";", "\n"}
|
||||
|
||||
// splitSentences splits text at the built-in sentence boundaries.
|
||||
func (o *SplitOperator) splitSentences(text string) []ChunkData {
|
||||
var chunks []ChunkData
|
||||
var buf strings.Builder
|
||||
i := 0
|
||||
@@ -111,7 +85,7 @@ func (o *SplitOperator) splitSentences(text string) []ChunkData {
|
||||
for i < len(text) {
|
||||
// Try to match any boundary at current position (first match wins)
|
||||
matchedBound := ""
|
||||
for _, bound := range o.boundaries {
|
||||
for _, bound := range sentenceBoundaries {
|
||||
if bound != "" && i+len(bound) <= len(text) && text[i:i+len(bound)] == bound {
|
||||
matchedBound = bound
|
||||
break
|
||||
@@ -119,17 +93,17 @@ func (o *SplitOperator) splitSentences(text string) []ChunkData {
|
||||
}
|
||||
|
||||
if matchedBound != "" {
|
||||
if o.keepSeparators {
|
||||
buf.WriteString(matchedBound)
|
||||
}
|
||||
if buf.Len() > 0 {
|
||||
chunks = append(chunks, ChunkData{
|
||||
Content: buf.String(),
|
||||
Index: len(chunks),
|
||||
Metadata: map[string]interface{}{
|
||||
"language": DetectLanguage(buf.String()),
|
||||
},
|
||||
})
|
||||
content := strings.TrimSpace(buf.String())
|
||||
if content != "" {
|
||||
chunks = append(chunks, ChunkData{
|
||||
Content: content,
|
||||
Index: len(chunks),
|
||||
Metadata: map[string]interface{}{
|
||||
"language": DetectLanguage(content),
|
||||
},
|
||||
})
|
||||
}
|
||||
buf.Reset()
|
||||
}
|
||||
i += len(matchedBound)
|
||||
@@ -142,13 +116,16 @@ func (o *SplitOperator) splitSentences(text string) []ChunkData {
|
||||
|
||||
// flush remaining text
|
||||
if buf.Len() > 0 {
|
||||
chunks = append(chunks, ChunkData{
|
||||
Content: buf.String(),
|
||||
Index: len(chunks),
|
||||
Metadata: map[string]interface{}{
|
||||
"language": DetectLanguage(buf.String()),
|
||||
},
|
||||
})
|
||||
content := strings.TrimSpace(buf.String())
|
||||
if content != "" {
|
||||
chunks = append(chunks, ChunkData{
|
||||
Content: content,
|
||||
Index: len(chunks),
|
||||
Metadata: map[string]interface{}{
|
||||
"language": DetectLanguage(content),
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return chunks
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build cgo && office
|
||||
//go:build cgo
|
||||
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
@@ -39,43 +39,30 @@ func NewDOCParser(libType string) (*DOCParser, error) {
|
||||
}
|
||||
}
|
||||
|
||||
func (p *DOCParser) Parse(filename string, data []byte) error {
|
||||
switch p.libType {
|
||||
case OfficeOxide:
|
||||
return p.OfficeOxideParse(data)
|
||||
default:
|
||||
return fmt.Errorf("unsupported DOC library type: %s", p.libType)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *DOCParser) OfficeOxideParse(data []byte) error {
|
||||
doc, err := officeOxide.OpenFromBytes(data, "doc")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer doc.Close()
|
||||
|
||||
docFormat, err := doc.Format()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
fmt.Println("Document format:", docFormat)
|
||||
|
||||
docContext, err := doc.PlainText()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Println("Document context:", docContext)
|
||||
|
||||
md, err := doc.ToMarkdown()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Println("Document Markdown:", md)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *DOCParser) String() string {
|
||||
return "DOCParser"
|
||||
}
|
||||
|
||||
// ParseWithResult captures the office_oxide PlainText output for
|
||||
// the DOC family. Python parser.py routes .doc through tika;
|
||||
// the Go side uses office_oxide which supports DOC via PlainText.
|
||||
// OutputFormat="text" — the python side falls back to text for
|
||||
// legacy DOC files since structured extraction is unreliable.
|
||||
func (p *DOCParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
doc, err := officeOxide.OpenFromBytes(data, "doc")
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("doc open: %w", err)}
|
||||
}
|
||||
defer doc.Close()
|
||||
|
||||
text, err := doc.PlainText()
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("doc plain-text: %w", err)}
|
||||
}
|
||||
|
||||
return ParseResult{
|
||||
OutputFormat: "text",
|
||||
File: map[string]any{"name": filename, "format": "doc"},
|
||||
Text: text,
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build cgo && office
|
||||
//go:build cgo
|
||||
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
@@ -24,10 +24,6 @@ import (
|
||||
officeOxide "github.com/yfedoseev/office_oxide/go"
|
||||
)
|
||||
|
||||
const (
|
||||
OfficeOxide string = "office_oxide"
|
||||
)
|
||||
|
||||
type DOCXParser struct {
|
||||
libType string
|
||||
}
|
||||
@@ -43,43 +39,29 @@ func NewDOCXParser(libType string) (*DOCXParser, error) {
|
||||
}
|
||||
}
|
||||
|
||||
func (p *DOCXParser) Parse(filename string, data []byte) error {
|
||||
|
||||
switch p.libType {
|
||||
case OfficeOxide:
|
||||
return p.OfficeOxideParse(data)
|
||||
default:
|
||||
return fmt.Errorf("unsupported DOCX library type: %s", p.libType)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *DOCXParser) OfficeOxideParse(data []byte) error {
|
||||
|
||||
// ParseWithResult captures the office_oxide ToMarkdown output
|
||||
// instead of discarding it. Returns OutputFormat="markdown" with
|
||||
// the rendered Markdown on the matching key. Mirrors the python
|
||||
// parser.py:Docx branch which uses rag.app.naive.Docx to render
|
||||
// markdown; the Go side delegates to office_oxide for now and
|
||||
// will switch to a native Markdown renderer in a follow-up.
|
||||
func (p *DOCXParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
doc, err := officeOxide.OpenFromBytes(data, "docx")
|
||||
if err != nil {
|
||||
return err
|
||||
return ParseResult{Err: fmt.Errorf("docx open: %w", err)}
|
||||
}
|
||||
defer doc.Close()
|
||||
|
||||
docFormat, err := doc.Format()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
fmt.Println("Document format:", docFormat)
|
||||
|
||||
docContext, err := doc.PlainText()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Println("Document context:", docContext)
|
||||
|
||||
md, err := doc.ToMarkdown()
|
||||
if err != nil {
|
||||
return err
|
||||
return ParseResult{Err: fmt.Errorf("docx to-markdown: %w", err)}
|
||||
}
|
||||
|
||||
return ParseResult{
|
||||
OutputFormat: "markdown",
|
||||
File: map[string]any{"name": filename, "format": "docx"},
|
||||
Markdown: md,
|
||||
}
|
||||
fmt.Println("Document Markdown:", md)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *DOCXParser) String() string {
|
||||
185
internal/parser/parser/html_parser.go
Normal file
185
internal/parser/parser/html_parser.go
Normal file
@@ -0,0 +1,185 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/net/html"
|
||||
)
|
||||
|
||||
const (
|
||||
Official string = "official"
|
||||
)
|
||||
|
||||
type HTMLParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewHTMLParser(libType string) (*HTMLParser, error) {
|
||||
switch libType {
|
||||
case Official:
|
||||
return &HTMLParser{
|
||||
libType: Official,
|
||||
}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported HTML library type: %s", libType)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *HTMLParser) String() string {
|
||||
return "HTMLParser"
|
||||
}
|
||||
|
||||
// ParseWithResult emits one item per block-level HTML element
|
||||
// (headings, paragraphs, lists, pre blocks). The walker is a
|
||||
// pure-Go replacement for the previous `fmt.Printf` debug output:
|
||||
// it descends the html.Parse tree, joins the leaf text of each
|
||||
// block-level element, and emits the python-compatible
|
||||
// `{text, doc_type_kwd:"text"}` shape.
|
||||
//
|
||||
// Phase 2.5 (Slice 1) of port-rag-flow-pipeline-to-go.md makes
|
||||
// HTMLParser a ParseResultProducer so the dispatch seam routes
|
||||
// the html family through the structured path. Inline formatting
|
||||
// (bold / links / images) is intentionally NOT surfaced as a
|
||||
// separate ck_type — the python HtmlParser collapses inline
|
||||
// formatting into the parent block's text.
|
||||
func (p *HTMLParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
if p.libType != Official {
|
||||
return ParseResult{Err: fmt.Errorf("unsupported HTML library type: %s", p.libType)}
|
||||
}
|
||||
doc, err := html.Parse(bytes.NewReader(data))
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("html parse: %w", err)}
|
||||
}
|
||||
var items []map[string]any
|
||||
walkHTMLBlocks(doc, &items)
|
||||
if items == nil {
|
||||
items = []map[string]any{{"text": "", "doc_type_kwd": "text"}}
|
||||
}
|
||||
return ParseResult{
|
||||
OutputFormat: "json",
|
||||
File: map[string]any{
|
||||
"name": filename,
|
||||
"encoding": "utf-8",
|
||||
},
|
||||
JSON: items,
|
||||
}
|
||||
}
|
||||
|
||||
// walkHTMLBlocks emits one normalized item per block-level
|
||||
// descendant of root. Inline elements (b, i, a, span, …) are
|
||||
// collapsed into the parent's text via leafText. <script> and
|
||||
// <style> blocks are skipped entirely so they don't pollute the
|
||||
// downstream chunker input.
|
||||
func walkHTMLBlocks(root *html.Node, out *[]map[string]any) {
|
||||
for child := root.FirstChild; child != nil; child = child.NextSibling {
|
||||
if child.Type != html.ElementNode {
|
||||
continue
|
||||
}
|
||||
tag := child.Data
|
||||
switch tag {
|
||||
case "script", "style", "noscript":
|
||||
// Skip executable / stylistic blocks entirely.
|
||||
continue
|
||||
case "html", "head", "body":
|
||||
// Wrapper elements: descend into their children.
|
||||
walkHTMLBlocks(child, out)
|
||||
continue
|
||||
}
|
||||
text := htmlLeafText(child)
|
||||
if strings.TrimSpace(text) == "" {
|
||||
continue
|
||||
}
|
||||
*out = append(*out, map[string]any{
|
||||
"text": strings.TrimSpace(text),
|
||||
"doc_type_kwd": "text",
|
||||
"ck_type": htmlTagToCkType(tag),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// htmlTagToCkType maps HTML block tags to the python `ck_type`
|
||||
// vocabulary used downstream by TitleChunker and similar
|
||||
// components. Tags not in the map fall back to "text".
|
||||
func htmlTagToCkType(tag string) string {
|
||||
switch tag {
|
||||
case "h1", "h2", "h3", "h4", "h5", "h6":
|
||||
return "heading"
|
||||
case "p":
|
||||
return "paragraph"
|
||||
case "ul", "ol", "li":
|
||||
return "list"
|
||||
case "pre", "code":
|
||||
return "code"
|
||||
case "table", "tr", "td", "th":
|
||||
return "table"
|
||||
case "blockquote":
|
||||
return "quote"
|
||||
case "img":
|
||||
return "image"
|
||||
}
|
||||
return "text"
|
||||
}
|
||||
|
||||
// htmlLeafText joins the visible text of an HTML node and its
|
||||
// descendants. <script>/<style> subtrees are skipped (mirrors
|
||||
// the python html.parser behaviour). The output preserves
|
||||
// whitespace runs so headings like "<h1>Hello world</h1>"
|
||||
// round-trip with their spacing intact.
|
||||
func htmlLeafText(n *html.Node) string {
|
||||
var b strings.Builder
|
||||
walkHTMLLeaf(n, &b)
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func walkHTMLLeaf(n *html.Node, b *strings.Builder) {
|
||||
switch n.Type {
|
||||
case html.TextNode:
|
||||
b.WriteString(n.Data)
|
||||
case html.ElementNode:
|
||||
if n.Data == "script" || n.Data == "style" {
|
||||
return
|
||||
}
|
||||
// Add a line break between block children so headings,
|
||||
// paragraphs, and list items don't run together.
|
||||
switch n.Data {
|
||||
case "h1", "h2", "h3", "h4", "h5", "h6", "p", "li", "pre",
|
||||
"tr", "blockquote":
|
||||
if b.Len() > 0 && !strings.HasSuffix(b.String(), "\n") {
|
||||
b.WriteString("\n")
|
||||
}
|
||||
}
|
||||
for child := n.FirstChild; child != nil; child = child.NextSibling {
|
||||
walkHTMLLeaf(child, b)
|
||||
}
|
||||
if isBlockTag(n.Data) && b.Len() > 0 && !strings.HasSuffix(b.String(), "\n") {
|
||||
b.WriteString("\n")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func isBlockTag(tag string) bool {
|
||||
switch tag {
|
||||
case "h1", "h2", "h3", "h4", "h5", "h6", "p", "li", "pre",
|
||||
"tr", "blockquote", "div", "section", "article", "header", "footer":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
167
internal/parser/parser/markdown_parser.go
Normal file
167
internal/parser/parser/markdown_parser.go
Normal file
@@ -0,0 +1,167 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/gomarkdown/markdown/ast"
|
||||
"github.com/gomarkdown/markdown/parser"
|
||||
)
|
||||
|
||||
const (
|
||||
GoMarkdown = "go_markdown"
|
||||
)
|
||||
|
||||
type MarkdownParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewMarkdownParser(libType string) (*MarkdownParser, error) {
|
||||
switch libType {
|
||||
case GoMarkdown:
|
||||
return &MarkdownParser{
|
||||
libType: GoMarkdown,
|
||||
}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported Markdown library type: %s", libType)
|
||||
}
|
||||
}
|
||||
|
||||
// ParseWithResult implements ParseResultProducer (plan §6.5) and
|
||||
// returns a structured markdown payload that mirrors the Python
|
||||
// parser's `output_format == "json"` shape. Each top-level block
|
||||
// emits one item with `text` + `doc_type_kwd: "text"`. The legacy
|
||||
// debug-print path has been removed; callers consume ParseResult directly.
|
||||
func (p *MarkdownParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
if p.libType != GoMarkdown {
|
||||
return ParseResult{Err: fmt.Errorf("unsupported Markdown library type: %s", p.libType)}
|
||||
}
|
||||
doc := markdownNew().Parse(data)
|
||||
|
||||
var items []map[string]any
|
||||
walkMarkdownBlocks(doc, &items)
|
||||
if items == nil {
|
||||
// No blocks emitted — surface a single empty item so the
|
||||
// downstream chunker sees a non-nil JSON slice (mirrors the
|
||||
// Python contract of always producing at least one item).
|
||||
items = []map[string]any{{"text": "", "doc_type_kwd": "text"}}
|
||||
}
|
||||
return ParseResult{
|
||||
OutputFormat: "json",
|
||||
File: map[string]any{
|
||||
"name": filename,
|
||||
},
|
||||
JSON: items,
|
||||
}
|
||||
}
|
||||
|
||||
func (p *MarkdownParser) String() string {
|
||||
return "MarkdownParser"
|
||||
}
|
||||
|
||||
// markdownNew is a thin constructor so the extension set is owned
|
||||
// in one place (both Parse and ParseWithResult consume it).
|
||||
func markdownNew() *parser.Parser {
|
||||
extensions := parser.CommonExtensions | parser.AutoHeadingIDs | parser.NoEmptyLineBeforeBlock
|
||||
return parser.NewWithExtensions(extensions)
|
||||
}
|
||||
|
||||
// walkMarkdownBlocks emits one normalized item per top-level block.
|
||||
// Headings (LeafBlock / H*) are emitted with the heading text so a
|
||||
// downstream title chunker can find them; paragraphs (Paragraph)
|
||||
// emit the leaf-node text. The walker is intentionally a
|
||||
// best-effort pass — full TOC / outline handling lands with the
|
||||
// deepdoc/parser port — but it satisfies the per-block "emit
|
||||
// text+doc_type_kwd" contract enough for a Phase-1a migration
|
||||
// test to verify wire-shape parity.
|
||||
func walkMarkdownBlocks(doc ast.Node, out *[]map[string]any) {
|
||||
for _, child := range doc.GetChildren() {
|
||||
switch n := child.(type) {
|
||||
case *ast.Heading:
|
||||
*out = append(*out, map[string]any{
|
||||
"text": headingText(n),
|
||||
"doc_type_kwd": "text",
|
||||
"ck_type": "heading",
|
||||
})
|
||||
case *ast.Paragraph:
|
||||
*out = append(*out, map[string]any{
|
||||
"text": leafText(n),
|
||||
"doc_type_kwd": "text",
|
||||
"ck_type": "text",
|
||||
})
|
||||
case *ast.List:
|
||||
*out = append(*out, map[string]any{
|
||||
"text": leafText(n),
|
||||
"doc_type_kwd": "text",
|
||||
"ck_type": "list",
|
||||
})
|
||||
case *ast.CodeBlock:
|
||||
*out = append(*out, map[string]any{
|
||||
"text": leafText(n),
|
||||
"doc_type_kwd": "text",
|
||||
"ck_type": "code",
|
||||
})
|
||||
default:
|
||||
// Block types we don't yet normalize (HTML, tables,
|
||||
// images, definitions) are best-effort: emit the leaf
|
||||
// text without a ck_type so downstream components can
|
||||
// still treat them as text chunks.
|
||||
txt := leafText(n)
|
||||
if strings.TrimSpace(txt) != "" {
|
||||
*out = append(*out, map[string]any{
|
||||
"text": txt,
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// headingText returns the inline-text of a heading node by
|
||||
// concatenating every Leaf / Text child. Empty headings emit "".
|
||||
func headingText(h *ast.Heading) string {
|
||||
var buf bytes.Buffer
|
||||
for _, c := range h.GetChildren() {
|
||||
buf.WriteString(leafText(c))
|
||||
}
|
||||
return strings.TrimSpace(buf.String())
|
||||
}
|
||||
|
||||
// leafText mirrors gomarkdown's leaf walker: walks every descendant
|
||||
// leaf (Text or Inline content) and returns the concatenated UTF-8.
|
||||
// Non-text containers that have no leaf descendants return "".
|
||||
func leafText(n ast.Node) string {
|
||||
var buf bytes.Buffer
|
||||
walkLeaf(n, &buf)
|
||||
return strings.TrimSpace(buf.String())
|
||||
}
|
||||
|
||||
func walkLeaf(n ast.Node, buf *bytes.Buffer) {
|
||||
switch t := n.(type) {
|
||||
case *ast.Text:
|
||||
buf.Write(t.Literal)
|
||||
case *ast.Code:
|
||||
buf.Write(t.Literal)
|
||||
default:
|
||||
for _, c := range n.GetChildren() {
|
||||
walkLeaf(c, buf)
|
||||
}
|
||||
}
|
||||
}
|
||||
12
internal/parser/parser/office_parsers_cgo_shared.go
Normal file
12
internal/parser/parser/office_parsers_cgo_shared.go
Normal file
@@ -0,0 +1,12 @@
|
||||
//go:build cgo
|
||||
|
||||
package parser
|
||||
|
||||
import "html"
|
||||
|
||||
// OfficeOxide is the lib_type identifier for office_oxide backend.
|
||||
const OfficeOxide = "office_oxide"
|
||||
|
||||
func htmlEscape(s string) string {
|
||||
return html.EscapeString(s)
|
||||
}
|
||||
136
internal/parser/parser/office_parsers_no_cgo.go
Normal file
136
internal/parser/parser/office_parsers_no_cgo.go
Normal file
@@ -0,0 +1,136 @@
|
||||
//go:build !cgo
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// OfficeOxide is the lib_type identifier for office_oxide backend.
|
||||
const OfficeOxide = "office_oxide"
|
||||
|
||||
// ErrOfficeCGORequired is returned by ParseWithResult on every
|
||||
// office-parser family (DOC / DOCX / PPT / PPTX / XLS / XLSX)
|
||||
// when the build is not CGO-enabled. The CGO build's
|
||||
// implementation captures the office_oxide PlainText / ToMarkdown
|
||||
// output; this stub mirrors that surface so the package compiles
|
||||
// and existing tests pass.
|
||||
var ErrOfficeCGORequired = errors.New("parser: office family requires CGO (office_oxide)")
|
||||
|
||||
// docxParseWithResultNoCGO is the no-CGO stub for the DOCX
|
||||
// family. The CGO build's implementation lives in docx_parser.go
|
||||
// under //go:build cgo.
|
||||
func (p *DOCXParser) ParseWithResult(filename string, _ []byte) ParseResult {
|
||||
return ParseResult{
|
||||
File: map[string]any{"name": filename},
|
||||
Err: fmt.Errorf("%w: docx", ErrOfficeCGORequired),
|
||||
}
|
||||
}
|
||||
|
||||
func (p *DOCParser) ParseWithResult(filename string, _ []byte) ParseResult {
|
||||
return ParseResult{
|
||||
File: map[string]any{"name": filename},
|
||||
Err: fmt.Errorf("%w: doc", ErrOfficeCGORequired),
|
||||
}
|
||||
}
|
||||
|
||||
func (p *XLSParser) ParseWithResult(filename string, _ []byte) ParseResult {
|
||||
return ParseResult{
|
||||
File: map[string]any{"name": filename},
|
||||
Err: fmt.Errorf("%w: xls", ErrOfficeCGORequired),
|
||||
}
|
||||
}
|
||||
|
||||
func (p *XLSXParser) ParseWithResult(filename string, _ []byte) ParseResult {
|
||||
return ParseResult{
|
||||
File: map[string]any{"name": filename},
|
||||
Err: fmt.Errorf("%w: xlsx", ErrOfficeCGORequired),
|
||||
}
|
||||
}
|
||||
|
||||
func (p *PPTParser) ParseWithResult(filename string, _ []byte) ParseResult {
|
||||
return ParseResult{
|
||||
File: map[string]any{"name": filename},
|
||||
Err: fmt.Errorf("%w: ppt", ErrOfficeCGORequired),
|
||||
}
|
||||
}
|
||||
|
||||
func (p *PPTXParser) ParseWithResult(filename string, _ []byte) ParseResult {
|
||||
return ParseResult{
|
||||
File: map[string]any{"name": filename},
|
||||
Err: fmt.Errorf("%w: pptx", ErrOfficeCGORequired),
|
||||
}
|
||||
}
|
||||
|
||||
type DOCParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewDOCParser(libType string) (*DOCParser, error) {
|
||||
return nil, fmt.Errorf("DOC parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *DOCParser) String() string {
|
||||
return "DOCParser(no-cgo)"
|
||||
}
|
||||
|
||||
type DOCXParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewDOCXParser(libType string) (*DOCXParser, error) {
|
||||
return nil, fmt.Errorf("DOCX parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *DOCXParser) String() string {
|
||||
return "DOCXParser(no-cgo)"
|
||||
}
|
||||
|
||||
type XLSParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewXLSParser(libType string) (*XLSParser, error) {
|
||||
return nil, fmt.Errorf("XLS parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *XLSParser) String() string {
|
||||
return "XLSParser(no-cgo)"
|
||||
}
|
||||
|
||||
type XLSXParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewXLSXParser(libType string) (*XLSXParser, error) {
|
||||
return nil, fmt.Errorf("XLSX parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *XLSXParser) String() string {
|
||||
return "XLSXParser(no-cgo)"
|
||||
}
|
||||
|
||||
type PPTParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewPPTParser(libType string) (*PPTParser, error) {
|
||||
return nil, fmt.Errorf("PPT parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *PPTParser) String() string {
|
||||
return "PPTParser(no-cgo)"
|
||||
}
|
||||
|
||||
type PPTXParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewPPTXParser(libType string) (*PPTXParser, error) {
|
||||
return nil, fmt.Errorf("PPTX parser requires CGO (office_oxide)")
|
||||
}
|
||||
|
||||
func (p *PPTXParser) String() string {
|
||||
return "PPTXParser(no-cgo)"
|
||||
}
|
||||
88
internal/parser/parser/parse_result.go
Normal file
88
internal/parser/parser/parse_result.go
Normal file
@@ -0,0 +1,88 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// ParseResult is the structured output contract for the Go parser
|
||||
// library. Port-rag-flow-pipeline-to-go.md §6.5 mandates that
|
||||
// parsers surface enough data to reconstruct a Python-compatible
|
||||
// stage-boundary payload:
|
||||
//
|
||||
// output_format ∈ {"json","markdown","text","html"}
|
||||
// file (enriched metadata)
|
||||
// exactly one payload family populated (matching output_format)
|
||||
// err
|
||||
//
|
||||
// Go parser callers now consume only the structured ParseResult
|
||||
// contract. The legacy `Parse(filename, []byte) error` interface has
|
||||
// been removed so parser dispatch, ingestion, and service paths all
|
||||
// share the same typed payload contract.
|
||||
|
||||
package parser
|
||||
|
||||
// ParseResult is the structured return value of a successful parse.
|
||||
// Exactly one of the payload fields (JSON / Markdown / Text / HTML)
|
||||
// is populated on success, matching the Python contract — see
|
||||
// port-rag-flow-pipeline-to-go.md §4.2:
|
||||
//
|
||||
// - OutputFormat = "json" → JSON populated
|
||||
// - OutputFormat = "markdown" → Markdown populated
|
||||
// - OutputFormat = "text" → Text populated
|
||||
// - OutputFormat = "html" → HTML populated
|
||||
//
|
||||
// On failure (Err != nil), all payload fields are zero values and
|
||||
// OutputFormat is empty.
|
||||
type ParseResult struct {
|
||||
// OutputFormat is the wire-compatible format the parser
|
||||
// chose. Empty when Err is non-nil.
|
||||
OutputFormat string
|
||||
|
||||
// File is the enriched file metadata the parser emits. In
|
||||
// Python this is the dict form of the original `file`
|
||||
// descriptor, augmented with format-specific keys (e.g.
|
||||
// `outline` on the PDF path, `page_count` for paginated
|
||||
// formats). Nil when the parser did not enrich.
|
||||
File map[string]any
|
||||
|
||||
// JSON is the structured payload when OutputFormat == "json".
|
||||
// Shape depends on the parser family: PDF emits
|
||||
// `[]map[string]any` with `text` + `doc_type_kwd` keys (and
|
||||
// optional `image` / `layout` / `positions` fields);
|
||||
// markdown / html / text emit normalized
|
||||
// `{text, doc_type_kwd}` items; image emits OCR/VLM result
|
||||
// items. Exactly one payload family is populated on success.
|
||||
JSON []map[string]any
|
||||
|
||||
// Markdown is the string payload when OutputFormat ==
|
||||
// "markdown". Empty otherwise.
|
||||
Markdown string
|
||||
|
||||
// Text is the string payload when OutputFormat == "text".
|
||||
// Empty otherwise.
|
||||
Text string
|
||||
|
||||
// HTML is the string payload when OutputFormat == "html".
|
||||
// Empty otherwise.
|
||||
HTML string
|
||||
|
||||
// Err is the failure reason. On non-nil Err, all payload
|
||||
// fields are zero values.
|
||||
Err error
|
||||
}
|
||||
|
||||
// ParseResultProducer is the parser package's single structured-output
|
||||
// contract. Every parser returned by GetParser must implement it.
|
||||
type ParseResultProducer interface {
|
||||
ParseWithResult(filename string, data []byte) ParseResult
|
||||
}
|
||||
212
internal/parser/parser/parse_result_test.go
Normal file
212
internal/parser/parser/parse_result_test.go
Normal file
@@ -0,0 +1,212 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestParseResult_Contract pins the wire-shape guarantees
|
||||
// port-rag-flow-pipeline-to-go.md §6.5 requires:
|
||||
//
|
||||
// - exactly one payload family populated on success
|
||||
// - empty payload fields on failure (Err != nil)
|
||||
// - OutputFormat is "" on failure
|
||||
//
|
||||
// These are the public-key invariants a component-level caller
|
||||
// (component/parser.go) relies on. The test does not exercise any
|
||||
// specific parser implementation — it pins the contract that the
|
||||
// ParseResult type itself enforces, so future parser additions
|
||||
// inherit the same guarantees.
|
||||
func TestParseResult_Contract(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
in ParseResult
|
||||
wantErr bool
|
||||
wantFmt string // expected OutputFormat on success
|
||||
}{
|
||||
{
|
||||
name: "json family only",
|
||||
in: ParseResult{
|
||||
OutputFormat: "json",
|
||||
JSON: []map[string]any{
|
||||
{"text": "hello", "doc_type_kwd": "text"},
|
||||
},
|
||||
},
|
||||
wantFmt: "json",
|
||||
},
|
||||
{
|
||||
name: "markdown family only",
|
||||
in: ParseResult{
|
||||
OutputFormat: "markdown",
|
||||
Markdown: "# Title\n\nbody",
|
||||
},
|
||||
wantFmt: "markdown",
|
||||
},
|
||||
{
|
||||
name: "text family only",
|
||||
in: ParseResult{
|
||||
OutputFormat: "text",
|
||||
Text: "raw text",
|
||||
},
|
||||
wantFmt: "text",
|
||||
},
|
||||
{
|
||||
name: "html family only",
|
||||
in: ParseResult{
|
||||
OutputFormat: "html",
|
||||
HTML: "<p>x</p>",
|
||||
},
|
||||
wantFmt: "html",
|
||||
},
|
||||
{
|
||||
name: "failure clears payload",
|
||||
in: ParseResult{Err: errSentinel},
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if tc.wantErr {
|
||||
if tc.in.Err == nil {
|
||||
t.Fatal("Err: want non-nil, got nil")
|
||||
}
|
||||
if tc.in.OutputFormat != "" {
|
||||
t.Errorf("OutputFormat on failure = %q, want empty", tc.in.OutputFormat)
|
||||
}
|
||||
if tc.in.JSON != nil || tc.in.Markdown != "" || tc.in.Text != "" || tc.in.HTML != "" {
|
||||
t.Errorf("payload fields must be zero on failure; got %+v", tc.in)
|
||||
}
|
||||
return
|
||||
}
|
||||
if tc.in.Err != nil {
|
||||
t.Errorf("Err on success: want nil, got %v", tc.in.Err)
|
||||
}
|
||||
if tc.in.OutputFormat != tc.wantFmt {
|
||||
t.Errorf("OutputFormat = %q, want %q", tc.in.OutputFormat, tc.wantFmt)
|
||||
}
|
||||
// Exactly-one-payload-family: pick the family declared
|
||||
// by OutputFormat; every other field must be zero.
|
||||
active := payloadFamily(tc.in)
|
||||
for _, other := range allPayloadFamilies {
|
||||
if other == active {
|
||||
continue
|
||||
}
|
||||
if !isPayloadZero(tc.in, other) {
|
||||
t.Errorf("OutputFormat=%s: secondary family %s should be zero", tc.in.OutputFormat, other)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// errSentinel is a tiny stand-in so the test does not need to
|
||||
// import errors just to declare one.
|
||||
var errSentinel = sentinelErr("sentinel")
|
||||
|
||||
type sentinelErr string
|
||||
|
||||
func (s sentinelErr) Error() string { return string(s) }
|
||||
|
||||
// TestMarkdownParser_ParseWithResult pins the migration exemplar
|
||||
// — port-rag-flow-pipeline-to-go.md §6.5 marks MarkdownParser as
|
||||
// the first ported format whose ParseWithResult surface emulates
|
||||
// the Python "json" output_format. The fixture is intentionally
|
||||
// tiny: 1 heading, 1 paragraph, 1 unordered list item, no nested
|
||||
// formatting.
|
||||
func TestMarkdownParser_ParseWithResult(t *testing.T) {
|
||||
p, err := NewMarkdownParser(GoMarkdown)
|
||||
if err != nil {
|
||||
t.Fatalf("NewMarkdownParser: %v", err)
|
||||
}
|
||||
src := []byte("# Title\n\nFirst paragraph.\n\n- Item one\n")
|
||||
res := p.ParseWithResult("doc.md", src)
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if res.OutputFormat != "json" {
|
||||
t.Errorf("OutputFormat = %q, want json", res.OutputFormat)
|
||||
}
|
||||
if res.File == nil || res.File["name"] != "doc.md" {
|
||||
t.Errorf("File.name = %v, want doc.md", res.File)
|
||||
}
|
||||
if res.JSON == nil {
|
||||
t.Fatal("JSON: want non-nil slice, got nil")
|
||||
}
|
||||
// 1 heading + 1 paragraph + 1 list = 3 items.
|
||||
if len(res.JSON) != 3 {
|
||||
t.Fatalf("JSON len = %d, want 3 (heading, paragraph, list)", len(res.JSON))
|
||||
}
|
||||
heading := res.JSON[0]
|
||||
if txt, _ := heading["text"].(string); !strings.Contains(txt, "Title") {
|
||||
t.Errorf("JSON[0].text = %q, want contains 'Title'", txt)
|
||||
}
|
||||
if got, _ := heading["ck_type"].(string); got != "heading" {
|
||||
t.Errorf("JSON[0].ck_type = %q, want heading", got)
|
||||
}
|
||||
for i, it := range res.JSON {
|
||||
if _, ok := it["doc_type_kwd"]; !ok {
|
||||
t.Errorf("JSON[%d] missing doc_type_kwd: %+v", i, it)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseResultProducer_PDFIsProducer pins the explicit
|
||||
// PDFParser ParseResultProducer wiring. The dispatch seam routes
|
||||
// PDFs through ParseWithResult regardless of whether the current
|
||||
// build has the cgo-backed DeepDOC engine enabled.
|
||||
func TestParseResultProducer_PDFIsProducer(t *testing.T) {
|
||||
pdf := &PDFParser{}
|
||||
if _, ok := any(pdf).(ParseResultProducer); !ok {
|
||||
t.Error("PDFParser must implement ParseResultProducer so the " +
|
||||
"dispatch seam routes PDFs through ParseWithResult")
|
||||
}
|
||||
}
|
||||
|
||||
// payloadFamily returns the field name that should be populated
|
||||
// given the declared OutputFormat. "" when OutputFormat is
|
||||
// unrecognized.
|
||||
func payloadFamily(r ParseResult) string {
|
||||
switch r.OutputFormat {
|
||||
case "json":
|
||||
return "json"
|
||||
case "markdown":
|
||||
return "markdown"
|
||||
case "text":
|
||||
return "text"
|
||||
case "html":
|
||||
return "html"
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
var allPayloadFamilies = []string{"json", "markdown", "text", "html"}
|
||||
|
||||
func isPayloadZero(r ParseResult, family string) bool {
|
||||
switch family {
|
||||
case "json":
|
||||
return r.JSON == nil
|
||||
case "markdown":
|
||||
return r.Markdown == ""
|
||||
case "text":
|
||||
return r.Text == ""
|
||||
case "html":
|
||||
return r.HTML == ""
|
||||
}
|
||||
return false
|
||||
}
|
||||
199
internal/parser/parser/parse_with_result_test.go
Normal file
199
internal/parser/parser/parse_with_result_test.go
Normal file
@@ -0,0 +1,199 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Slice 1 tests for port-rag-flow-pipeline-to-go.md Phase 2.5.
|
||||
// These pin the new ParseWithResult contracts for the parsers
|
||||
// that did not previously satisfy ParseResultProducer:
|
||||
//
|
||||
// - HTMLParser — block-level walker that emits the python-compatible
|
||||
// {text, doc_type_kwd, ck_type} shape.
|
||||
// - TextParser — paragraph-splitting for the text&code family
|
||||
// (.txt / .py / .js / .java / .c / .cpp / .h / .php / .go / .ts
|
||||
// / .sh / .cs / .kt / .sql).
|
||||
//
|
||||
// MarkdownParser's ParseWithResult is already pinned by
|
||||
// parse_result_test.go (prior slice). PDFParser and the office
|
||||
// variants remain deferred to a follow-up slice that wires
|
||||
// them to the existing internal/deepdoc/parser/pdf pipeline and
|
||||
// office_oxide libraries.
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/utility"
|
||||
)
|
||||
|
||||
// TestTextParser_ParseWithResult_ParaSplit pins the paragraph-split
|
||||
// rule. A blank-line-separated input yields one item per
|
||||
// paragraph; the python TxtParser does the same.
|
||||
func TestTextParser_ParseWithResult_ParaSplit(t *testing.T) {
|
||||
p, err := NewTextParser("")
|
||||
if err != nil {
|
||||
t.Fatalf("NewTextParser: %v", err)
|
||||
}
|
||||
src := []byte("First paragraph.\n\nSecond paragraph.\n\nThird.")
|
||||
res := p.ParseWithResult("doc.txt", src)
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if res.OutputFormat != "json" {
|
||||
t.Errorf("OutputFormat = %q, want json", res.OutputFormat)
|
||||
}
|
||||
if got, want := res.File["name"], "doc.txt"; got != want {
|
||||
t.Errorf("File.name = %v, want %v", got, want)
|
||||
}
|
||||
if len(res.JSON) != 3 {
|
||||
t.Fatalf("JSON len = %d, want 3 (one per paragraph)", len(res.JSON))
|
||||
}
|
||||
if got, want := res.JSON[0]["text"], "First paragraph."; got != want {
|
||||
t.Errorf("JSON[0].text = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := res.JSON[2]["text"], "Third."; got != want {
|
||||
t.Errorf("JSON[2].text = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTextParser_ParseWithResult_Empty pins the empty-input
|
||||
// fallback (one empty item, not nil) so the downstream chunker
|
||||
// sees a non-nil JSON slice. Mirrors the MarkdownParser convention
|
||||
// at markdown_parser.go:71-76.
|
||||
func TestTextParser_ParseWithResult_Empty(t *testing.T) {
|
||||
p, _ := NewTextParser("")
|
||||
res := p.ParseWithResult("empty.txt", []byte{})
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if len(res.JSON) != 1 {
|
||||
t.Errorf("JSON len = %d, want 1 (empty-input fallback)", len(res.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
// TestTextParser_ParseWithResult_LongParagraphSlicing pins the
|
||||
// maxItemBytes boundary behaviour. A single paragraph longer
|
||||
// than 8192 bytes is sliced at the nearest line boundary.
|
||||
func TestTextParser_ParseWithResult_LongParagraphSlicing(t *testing.T) {
|
||||
p, _ := NewTextParser("")
|
||||
long := strings.Repeat("a", 9000)
|
||||
res := p.ParseWithResult("long.txt", []byte(long))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if len(res.JSON) < 2 {
|
||||
t.Errorf("JSON len = %d, want >=2 (sliced at maxItemBytes)", len(res.JSON))
|
||||
}
|
||||
for i, it := range res.JSON {
|
||||
if txt, _ := it["text"].(string); len(txt) > 8192 {
|
||||
t.Errorf("JSON[%d].text len = %d, exceeds maxItemBytes=8192", i, len(txt))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestTextParser_ParseWithResult_InvalidUTF8 pins the UTF-8
|
||||
// validation rule. Invalid bytes produce an error in the result
|
||||
// (matching the python TxtParser's behaviour).
|
||||
func TestTextParser_ParseWithResult_InvalidUTF8(t *testing.T) {
|
||||
p, _ := NewTextParser("")
|
||||
bad := []byte{0xff, 0xfe, 0xfd}
|
||||
res := p.ParseWithResult("bad.txt", bad)
|
||||
if res.Err == nil {
|
||||
t.Fatal("want error for invalid UTF-8, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestHTMLParser_ParseWithResult_BlockSplit pins the HTML walker.
|
||||
// Three block elements (heading, paragraph, list) yield three
|
||||
// items with the python-compatible ck_type vocabulary.
|
||||
func TestHTMLParser_ParseWithResult_BlockSplit(t *testing.T) {
|
||||
p, err := NewHTMLParser(Official)
|
||||
if err != nil {
|
||||
t.Fatalf("NewHTMLParser: %v", err)
|
||||
}
|
||||
src := []byte(`<!DOCTYPE html><html><body>
|
||||
<h1>Title</h1>
|
||||
<p>First paragraph.</p>
|
||||
<ul><li>Item one</li></ul>
|
||||
</body></html>`)
|
||||
res := p.ParseWithResult("doc.html", src)
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if res.OutputFormat != "json" {
|
||||
t.Errorf("OutputFormat = %q, want json", res.OutputFormat)
|
||||
}
|
||||
if len(res.JSON) != 3 {
|
||||
t.Fatalf("JSON len = %d, want 3 (h1, p, ul)", len(res.JSON))
|
||||
}
|
||||
if got, want := res.JSON[0]["ck_type"], "heading"; got != want {
|
||||
t.Errorf("JSON[0].ck_type = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := res.JSON[0]["text"], "Title"; got != want {
|
||||
t.Errorf("JSON[0].text = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := res.JSON[1]["ck_type"], "paragraph"; got != want {
|
||||
t.Errorf("JSON[1].ck_type = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := res.JSON[1]["text"], "First paragraph."; got != want {
|
||||
t.Errorf("JSON[1].text = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := res.JSON[2]["ck_type"], "list"; got != want {
|
||||
t.Errorf("JSON[2].ck_type = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := res.JSON[2]["text"], "Item one"; got != want {
|
||||
t.Errorf("JSON[2].text = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHTMLParser_ParseWithResult_SkipsScriptAndStyle pins the
|
||||
// rule that <script> / <style> subtrees are skipped entirely so
|
||||
// they don't pollute the downstream chunker input.
|
||||
func TestHTMLParser_ParseWithResult_SkipsScriptAndStyle(t *testing.T) {
|
||||
p, _ := NewHTMLParser(Official)
|
||||
src := []byte(`<html><body>
|
||||
<p>Visible.</p>
|
||||
<script>alert("x")</script>
|
||||
<style>body { color: red; }</style>
|
||||
<p>Also visible.</p>
|
||||
</body></html>`)
|
||||
res := p.ParseWithResult("doc.html", src)
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if len(res.JSON) != 2 {
|
||||
t.Errorf("JSON len = %d, want 2 (script+style skipped)", len(res.JSON))
|
||||
}
|
||||
for _, it := range res.JSON {
|
||||
if txt, _ := it["text"].(string); strings.Contains(txt, "alert") || strings.Contains(txt, "color") {
|
||||
t.Errorf("item text leaks script/style content: %q", txt)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetParser_RoutesTextAndCode pins the parser-type switch
|
||||
// routing for the text&code family. After the Slice 1 additions
|
||||
// `utility.FileTypeTXT` resolves to a TextParser that satisfies
|
||||
// ParseResultProducer.
|
||||
func TestGetParser_RoutesTextAndCode(t *testing.T) {
|
||||
p, err := GetParser(utility.FileTypeTXT, map[string]string{"lib_type": ""})
|
||||
if err != nil {
|
||||
t.Fatalf("GetParser(FileTypeTXT): %v", err)
|
||||
}
|
||||
if _, ok := p.(ParseResultProducer); !ok {
|
||||
t.Fatal("TextParser does not implement ParseResultProducer")
|
||||
}
|
||||
}
|
||||
@@ -21,7 +21,7 @@ import (
|
||||
"ragflow/internal/utility"
|
||||
)
|
||||
|
||||
func GetParser(fileType utility.FileType, config map[string]string) (FileParser, error) {
|
||||
func GetParser(fileType utility.FileType, config map[string]string) (ParseResultProducer, error) {
|
||||
libType, ok := config["lib_type"]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("missing lib_type config")
|
||||
@@ -45,15 +45,9 @@ func GetParser(fileType utility.FileType, config map[string]string) (FileParser,
|
||||
return NewHTMLParser(Official)
|
||||
case utility.FileTypeMarkdown:
|
||||
return NewMarkdownParser(GoMarkdown)
|
||||
case utility.FileTypeTXT:
|
||||
return NewTextParser(libType)
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported file type: %s", fileType)
|
||||
}
|
||||
}
|
||||
|
||||
// FileParser defines the interface for all file parsers.
|
||||
type FileParser interface {
|
||||
// Parse parses the input text.
|
||||
Parse(filename string, data []byte) error
|
||||
|
||||
String() string
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user