internal/ingestion/pipeline/image_ocr_test.go
Rama main · réplica de lectura del repositorio original.
// 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 (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"image"
"image/png"
"net/http"
"net/http/httptest"
"reflect"
"strings"
"testing"
"github.com/glebarez/sqlite"
"gorm.io/gorm"
"gorm.io/gorm/logger"
"ragflow/internal/agent/runtime"
"ragflow/internal/entity"
"ragflow/internal/ingestion/component"
"ragflow/internal/ingestion/component/chunker"
"ragflow/internal/ingestion/component/schema"
)
type imageOCREmbedder struct{ texts []string }
// MaxTokens gives the fixture enough room for the invoice's OCR text.
func (*imageOCREmbedder) MaxTokens() int { return 2048 }
// BatchSize bounds the embedding fixture's request size.
func (*imageOCREmbedder) BatchSize() int { return 16 }
// Encode captures the exact indexed text and supplies deterministic vectors.
func (e *imageOCREmbedder) Encode(_ context.Context, texts []string) ([]component.EmbeddingResult, error) {
e.texts = append(e.texts, texts...)
out := make([]component.EmbeddingResult, len(texts))
for i := range out {
out[i] = component.EmbeddingResult{Vector: []float64{0.2, 0.8}, TokenCount: 8}
}
return out, nil
}
// TestImageOCRSavedConfigToSearchableChunk runs the real config cleaner, SQLite
// persistence, resolver, OCR driver, Parser, decorated OneChunker and Tokenizer.
// Only OCR inference, object storage and embedding inference are test doubles.
func TestImageOCRSavedConfigToSearchableChunk(t *testing.T) {
var imageBytes bytes.Buffer
if err := png.Encode(&imageBytes, image.NewRGBA(image.Rect(0, 0, 12, 12))); err != nil {
t.Fatal(err)
}
const text = "Invoice INV-21030 total 42."
seen := make(chan []byte, 1)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/layout-parsing" || r.Header.Get("Authorization") != "Bearer fixture-key" {
t.Errorf("wrong selected instance request: %s %s", r.URL.Path, r.Header.Get("Authorization"))
}
var request struct {
File string `json:"file"`
FileType int `json:"fileType"`
Algorithm string `json:"algorithm"`
}
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
t.Error(err)
}
if request.FileType != 1 || request.Algorithm != "PP-OCRv6" {
t.Errorf("not the selected image model: %+v", request)
}
data, _ := base64.StdEncoding.DecodeString(request.File)
seen <- data
w.Header().Set("Content-Type", "application/json")
_, _ = fmt.Fprintf(w, `{"errorCode":0,"result":{"layoutParsingResults":[{"markdown":{"text":%q}}]}}`, text)
}))
defer server.Close()
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
if err != nil {
t.Fatal(err)
}
sqlDB, err := db.DB()
if err != nil {
t.Fatal(err)
}
sqlDB.SetMaxOpenConns(1)
defer sqlDB.Close()
if err := db.AutoMigrate(&entity.Knowledgebase{}, &entity.UserTenant{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil {
t.Fatal(err)
}
registry, err := DefaultRegistry()
if err != nil {
t.Fatal(err)
}
template, ok := registry.Get("picture")
if !ok {
t.Fatal("picture template missing")
}
dsl, err := json.Marshal(template.DSL)
if err != nil {
t.Fatal(err)
}
const parserID = "Parser:ViewsCaptureLight"
config := buildTestParserConfig(t, dsl, map[string]any{parserID: map[string]any{
"image": map[string]any{"ocr_enabled": true, "ocr_model_id": "ocr-model"},
"enable_vision_enhancement": false,
}})
for _, row := range []any{
&entity.Knowledgebase{ID: "kb-1", TenantID: "tenant-1", Name: "Image OCR fixture", ParserConfig: config},
&entity.TenantModelProvider{ID: "ocr-provider", UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "PaddleOCR.local"},
&entity.TenantModelInstance{ID: "ocr-instance", ProviderID: "ocr-provider", InstanceName: "selected",
Status: "active", APIKey: "fixture-key", Extra: fmt.Sprintf(`{"base_url":%q}`, server.URL)},
&entity.TenantModel{ID: "ocr-model", ProviderID: "ocr-provider", InstanceID: "ocr-instance",
ModelName: "PP-OCRv6", Status: "active", ModelType: int(entity.ModelTypeOCR)},
} {
if err := db.Create(row).Error; err != nil {
t.Fatal(err)
}
}
var saved entity.Knowledgebase
if err := db.Where("id = ?", "kb-1").Take(&saved).Error; err != nil {
t.Fatal(err)
}
params := saved.ParserConfig[parserID].(map[string]any)
if params["image"].(map[string]any)["ocr_model_id"] != "ocr-model" {
t.Fatal("image model selection did not survive save/readback")
}
parserComponent, err := component.NewParserComponent(params)
if err != nil {
t.Fatal(err)
}
inputs := map[string]any{"name": "invoice.png", "file_type": "image", "binary": imageBytes.Bytes(),
"tenant_id": "tenant-1", "kb_id": "kb-1", "doc_id": "doc-1"}
parsed, err := parserComponent.Invoke(t.Context(), db, inputs)
if err != nil {
t.Fatal(err)
}
select {
case data := <-seen:
if !bytes.Equal(data, imageBytes.Bytes()) {
t.Fatal("OCR received different image bytes")
}
default:
t.Fatal("OCR model not invoked")
}
factory, _, _, ok := runtime.DefaultRegistry.Lookup("OneChunker")
if !ok {
t.Fatal("OneChunker not registered")
}
chunkComponent, err := factory("OneChunker", nil)
if err != nil {
t.Fatal(err)
}
originalUploader := chunker.ChunkImageUploader
var uploaded []byte
chunker.ChunkImageUploader = func(_ context.Context, kbID, chunkID string, data []byte) (string, error) {
uploaded = append([]byte(nil), data...)
return kbID + "-" + chunkID, nil
}
t.Cleanup(func() { chunker.ChunkImageUploader = originalUploader })
for key, value := range parsed {
inputs[key] = value
}
chunked, err := chunkComponent.Invoke(t.Context(), db, inputs)
if err != nil {
t.Fatal(err)
}
embedder := &imageOCREmbedder{}
tokenizer, err := component.NewTokenizerComponentWithResolver(map[string]any{
"search_method": []string{"embedding"}, "filename_embd_weight": 0,
}, func(context.Context, string, string) (component.Embedder, string, error) {
return embedder, "fixture-embedding", nil
})
if err != nil {
t.Fatal(err)
}
for key, value := range chunked {
inputs[key] = value
}
indexable, err := tokenizer.Invoke(t.Context(), db, inputs)
if err != nil {
t.Fatal(err)
}
chunks, _ := schema.ChunkCarrierToMaps(indexable["chunks"])
if len(chunks) != 1 || chunks[0]["text"] != text {
t.Fatalf("searchable OCR text was lost: %+v", chunks)
}
if !reflect.DeepEqual(embedder.texts, []string{"invoice.png", text}) || !bytes.Equal(uploaded, imageBytes.Bytes()) {
t.Fatalf("embedding input=%q image preserved=%v", embedder.texts, bytes.Equal(uploaded, imageBytes.Bytes()))
}
if id, _ := chunks[0]["img_id"].(string); !strings.HasPrefix(id, "kb-1-") {
t.Fatalf("indexed image reference lost: %+v", chunks[0])
}
if _, ok := chunks[0]["q_2_vec"]; !ok {
t.Fatalf("embedding vector missing: %+v", chunks[0])
}
}