Port PR14140 and PR16881 to GO (#17102)

### Summary

Port
https://github.com/infiniflow/ragflow/pull/14140/
https://github.com/infiniflow/ragflow/pull/16881
This commit is contained in:
qinling0210
2026-07-21 17:46:23 +08:00
committed by GitHub
parent 20b760f266
commit 0189ca3700
19 changed files with 570 additions and 152 deletions

View File

@@ -73,6 +73,12 @@ type analyzerPool struct {
wg sync.WaitGroup
}
// defaultLanguage is applied to every analyzer instance on pool acquisition
// to clear any sticky language state left by a previous task. Ingestion
// callers should use the public API variants that accept an explicit
// language override.
const defaultLanguage = "English"
var (
globalPool *analyzerPool
poolOnce sync.Once
@@ -395,8 +401,9 @@ func Close() {
}
}
// withAnalyzer executes the given function with an exclusive analyzer instance
func withAnalyzer(fn func(*rag.Analyzer) error) error {
// withAnalyzer acquires an analyzer instance, applies the given language
// (with "" mapping to defaultLanguage), and executes fn.
func withAnalyzer(lang string, fn func(*rag.Analyzer) error) error {
if globalPool == nil {
return fmt.Errorf("tokenizer pool not initialized")
}
@@ -407,11 +414,16 @@ func withAnalyzer(fn func(*rag.Analyzer) error) error {
}
defer globalPool.release(instance)
if lang == "" {
lang = defaultLanguage
}
instance.analyzer.SetLanguage(lang)
return fn(instance.analyzer)
}
// withAnalyzerResult executes the given function with an exclusive analyzer instance and returns a result
func withAnalyzerResult[T any](fn func(*rag.Analyzer) (T, error)) (T, error) {
// withAnalyzerResult is the result-returning variant of withAnalyzer.
func withAnalyzerResult[T any](lang string, fn func(*rag.Analyzer) (T, error)) (T, error) {
var result T
if globalPool == nil {
return result, fmt.Errorf("tokenizer pool not initialized")
@@ -423,32 +435,63 @@ func withAnalyzerResult[T any](fn func(*rag.Analyzer) (T, error)) (T, error) {
}
defer globalPool.release(instance)
if lang == "" {
lang = defaultLanguage
}
instance.analyzer.SetLanguage(lang)
return fn(instance.analyzer)
}
// Tokenize tokenizes the text and returns a space-separated string of tokens
type Tokenizer struct {
lang string
}
// New returns a request-scoped tokenizer. Empty language falls back to English.
func New(lang string) Tokenizer {
return Tokenizer{lang: lang}
}
var defaultTokenizer = New("")
// Tokenize tokenizes the text and returns a space-separated string of tokens.
// Example: "hello world" -> "hello world"
//
// NOTE: For Infinity engine, returns input unchanged to match python's behavior
// NOTE: For Infinity engine, returns input unchanged to match python's behavior.
func Tokenize(text string) (string, error) {
return defaultTokenizer.Tokenize(text)
}
// Tokenize tokenizes the text using the tokenizer's request-scoped language.
func (t Tokenizer) Tokenize(text string) (string, error) {
if engineTypeProvider() == "infinity" {
return text, nil
}
return withAnalyzerResult(func(a *rag.Analyzer) (string, error) {
return withAnalyzerResult(t.lang, func(a *rag.Analyzer) (string, error) {
return a.Tokenize(text)
})
}
// TokenizeWithPosition tokenizes the text and returns a list of tokens with position information
// TokenizeWithPosition tokenizes the text and returns a list of tokens with position information.
func TokenizeWithPosition(text string) ([]rag.TokenWithPosition, error) {
return withAnalyzerResult(func(a *rag.Analyzer) ([]rag.TokenWithPosition, error) {
return defaultTokenizer.TokenizeWithPosition(text)
}
// TokenizeWithPosition tokenizes the text using the tokenizer's request-scoped language.
func (t Tokenizer) TokenizeWithPosition(text string) ([]rag.TokenWithPosition, error) {
return withAnalyzerResult(t.lang, func(a *rag.Analyzer) ([]rag.TokenWithPosition, error) {
return a.TokenizeWithPosition(text)
})
}
// Analyze analyzes the text and returns all tokens
// Analyze analyzes the text and returns all tokens.
func Analyze(text string) ([]rag.Token, error) {
return withAnalyzerResult(func(a *rag.Analyzer) ([]rag.Token, error) {
return defaultTokenizer.Analyze(text)
}
// Analyze analyzes the text using the tokenizer's request-scoped language.
func (t Tokenizer) Analyze(text string) ([]rag.Token, error) {
return withAnalyzerResult(t.lang, func(a *rag.Analyzer) ([]rag.Token, error) {
return a.Analyze(text)
})
}
@@ -462,16 +505,23 @@ func SetFineGrained(fineGrained bool) {
common.Debug("SetFineGrained is no-op in pool mode", zap.Bool("fine_grained", fineGrained))
}
// FineGrainedTokenize performs fine-grained tokenization on space-separated tokens
// FineGrainedTokenize performs fine-grained tokenization on space-separated
// tokens.
// Input: space-separated tokens (e.g., "hello world 测试")
// Output: space-separated fine-grained tokens (e.g., "hello world 测 试")
//
// NOTE: For Infinity engine, returns input unchanged to match python's behavior
// NOTE: For Infinity engine, returns input unchanged to match python's behavior.
func FineGrainedTokenize(tokens string) (string, error) {
return defaultTokenizer.FineGrainedTokenize(tokens)
}
// FineGrainedTokenize performs fine-grained tokenization using the tokenizer's
// request-scoped language.
func (t Tokenizer) FineGrainedTokenize(tokens string) (string, error) {
if engineTypeProvider() == "infinity" {
return tokens, nil
}
return withAnalyzerResult(func(a *rag.Analyzer) (string, error) {
return withAnalyzerResult(t.lang, func(a *rag.Analyzer) (string, error) {
return a.FineGrainedTokenize(tokens)
})
}
@@ -490,7 +540,7 @@ func IsInitialized() bool {
// GetTermFreq returns the frequency of a term (matching Python rag_tokenizer.freq)
// Returns: frequency value, or 0 if term not found
func GetTermFreq(term string) int32 {
result, _ := withAnalyzerResult(func(a *rag.Analyzer) (int32, error) {
result, _ := withAnalyzerResult("", func(a *rag.Analyzer) (int32, error) {
return a.GetTermFreq(term), nil
})
return result
@@ -499,7 +549,7 @@ func GetTermFreq(term string) int32 {
// GetTermTag returns the POS tag of a term (matching Python rag_tokenizer.tag)
// Returns: POS tag string (e.g., "n", "v", "ns"), or empty string if term not found or no tag
func GetTermTag(term string) string {
result, _ := withAnalyzerResult(func(a *rag.Analyzer) (string, error) {
result, _ := withAnalyzerResult("", func(a *rag.Analyzer) (string, error) {
return a.GetTermTag(term), nil
})
return result

View File

@@ -172,6 +172,67 @@ func TestConcurrentTokenize(t *testing.T) {
t.Log("=== Test completed successfully ===")
}
func TestConcurrentTokenizeLanguageIsolation(t *testing.T) {
restore := saveEngineType()
defer restore()
RegisterEngineType(func() string { return "" })
cfg := &PoolConfig{
DictPath: "",
MinSize: 2,
MaxSize: 8,
IdleTimeout: 3 * time.Second,
AcquireTimeout: 5 * time.Second,
}
if err := Init(cfg); err != nil {
t.Fatalf("Failed to initialize pool: %v", err)
}
defer Close()
sample := findEnglishDutchDifferentiator(t)
const goroutinesPerLang = 8
const requestsPerGoroutine = 20
var wg sync.WaitGroup
start := make(chan struct{})
errors := make(chan string, goroutinesPerLang*requestsPerGoroutine*2)
run := func(tok Tokenizer, lang, want string) {
defer wg.Done()
<-start
for i := 0; i < requestsPerGoroutine; i++ {
got, err := tok.Tokenize(sample.input)
if err != nil {
errors <- fmt.Sprintf("lang=%s req=%d unexpected error: %v", lang, i, err)
return
}
if got != want {
errors <- fmt.Sprintf("lang=%s req=%d got %q want %q", lang, i, got, want)
return
}
}
}
for i := 0; i < goroutinesPerLang; i++ {
wg.Add(2)
go run(New("English"), "English", sample.english)
go run(New("Dutch"), "Dutch", sample.dutch)
}
close(start)
wg.Wait()
close(errors)
for err := range errors {
t.Error(err)
}
if t.Failed() {
t.Fatalf("concurrent language isolation failed for input %q (English=%q Dutch=%q)", sample.input, sample.english, sample.dutch)
}
}
// TestConcurrentTokenizeWithPosition tests concurrent tokenization with position info
func TestConcurrentTokenizeWithPosition(t *testing.T) {
cfg := &PoolConfig{

View File

@@ -19,8 +19,15 @@ package tokenizer
import (
"strings"
"testing"
"time"
)
type languageDifferentiator struct {
input string
english string
dutch string
}
// saveEngineType saves the current engineTypeProvider and returns a function
// to restore it. Use this when a test modifies the engine type to avoid
// leaking global state between tests.
@@ -223,6 +230,79 @@ func TestGetTermTag_PoolNotInitialized(t *testing.T) {
}
}
func TestTokenize_DefaultLanguageResetsAnalyzerState(t *testing.T) {
restore := saveEngineType()
defer restore()
RegisterEngineType(func() string { return "" })
if err := Init(&PoolConfig{
DictPath: "",
MinSize: 1,
MaxSize: 1,
IdleTimeout: 5 * time.Second,
AcquireTimeout: 5 * time.Second,
}); err != nil {
t.Fatalf("Failed to initialize pool: %v", err)
}
defer Close()
sample := findEnglishDutchDifferentiator(t)
dutchGot, err := New("Dutch").Tokenize(sample.input)
if err != nil {
t.Fatalf("Tokenize(Dutch, %q) unexpected error: %v", sample.input, err)
}
if dutchGot != sample.dutch {
t.Fatalf("Tokenize(Dutch, %q) = %q, want %q", sample.input, dutchGot, sample.dutch)
}
defaultGot, err := Tokenize(sample.input)
if err != nil {
t.Fatalf("Tokenize(default, %q) unexpected error: %v", sample.input, err)
}
if defaultGot != sample.english {
t.Fatalf("Tokenize(default, %q) = %q, want explicit English result %q", sample.input, defaultGot, sample.english)
}
if defaultGot == dutchGot {
t.Fatalf("Tokenize(default, %q) unexpectedly inherited Dutch analyzer state: %q", sample.input, defaultGot)
}
}
func findEnglishDutchDifferentiator(t *testing.T) languageDifferentiator {
t.Helper()
candidates := []string{
"running",
"jumps",
"ponies",
"studies",
"wolves",
"relational",
"conditionally",
}
for _, input := range candidates {
english, err := New("English").Tokenize(input)
if err != nil {
t.Fatalf("Tokenize(English, %q) unexpected error: %v", input, err)
}
dutch, err := New("Dutch").Tokenize(input)
if err != nil {
t.Fatalf("Tokenize(Dutch, %q) unexpected error: %v", input, err)
}
if english != dutch {
return languageDifferentiator{
input: input,
english: english,
dutch: dutch,
}
}
}
t.Skip("no differentiating tokenizer sample found for English vs Dutch")
return languageDifferentiator{}
}
// ---------------------------------------------------------------------------
// Global state tests
// ---------------------------------------------------------------------------