diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d6fe6162..d5c57299 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -167,6 +167,33 @@ jobs: version: "~> v2" args: release --snapshot --clean --skip=publish + kiwi-helper: + strategy: + fail-fast: false + matrix: + os: + - ubuntu-latest + - macos-latest + runs-on: ${{ matrix.os }} + timeout-minutes: 20 + steps: + - name: Checkout + uses: actions/checkout@v7.0.1 + + - name: Setup Go + uses: actions/setup-go@v7 + with: + go-version-file: tools/discrawl-kiwi/go.mod + cache-dependency-path: tools/discrawl-kiwi/go.sum + + - name: Install Kiwi 0.23.2 + working-directory: tools/discrawl-kiwi + run: bash install-kiwi.sh + + - name: Test Go binding helper + working-directory: tools/discrawl-kiwi + run: go test -count=1 ./... + secrets: runs-on: ubuntu-latest timeout-minutes: 15 @@ -192,3 +219,33 @@ jobs: - name: Scan working tree run: | "$(go env GOPATH)/bin/gitleaks" dir . --no-banner --redact + + ja-helper: + runs-on: ubuntu-latest + timeout-minutes: 20 + steps: + - name: Checkout + uses: actions/checkout@v7.0.1 + - name: Setup Go + uses: actions/setup-go@v7 + with: + go-version-file: tools/discrawl-ja/go.mod + cache-dependency-path: tools/discrawl-ja/go.sum + - name: Test Japanese helper + working-directory: tools/discrawl-ja + run: go test -count=1 ./... + + zh-helper: + runs-on: ubuntu-latest + timeout-minutes: 20 + steps: + - name: Checkout + uses: actions/checkout@v7.0.1 + - name: Setup Go + uses: actions/setup-go@v7 + with: + go-version-file: tools/discrawl-zh/go.mod + cache-dependency-path: tools/discrawl-zh/go.sum + - name: Test Chinese helper + working-directory: tools/discrawl-zh + run: go test -count=1 ./... diff --git a/README.md b/README.md index 1837d5cf..756f8f33 100644 --- a/README.md +++ b/README.md @@ -124,6 +124,7 @@ The full documentation lives at **[discrawl.sh](https://discrawl.sh/)**: - [Command reference](docs/README.md) - [Sync sources](docs/guides/sync-sources.md) - [Search modes](docs/guides/search-modes.md) +- [Multilingual lexical benchmark](docs/benchmarks/multilingual-lexical.md) - [Git snapshot workflows](docs/guides/git-snapshots.md) - [Configuration](docs/configuration.md) diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 00000000..88962579 --- /dev/null +++ b/THIRD_PARTY_NOTICES.md @@ -0,0 +1,39 @@ +# Third-party notices + +## Optional Korean lexical helper + +The separately built `tools/discrawl-kiwi` helper depends on: + +- [Kiwi](https://github.com/bab2min/Kiwi), copyright Minchul Lee, + licensed under GNU LGPL 2.1 or later. +- [github.com/codingpot/kiwigo](https://github.com/codingpot/kiwigo), a Go + binding for Kiwi, licensed under GNU LGPL 2.1. + +These dependencies are optional and are not linked into the default Discrawl +binary. Distributors who provide the helper or Kiwi native binaries must +satisfy their applicable LGPL notice, source-access, and relinking +requirements. Kiwi's full license text is available from its source +repository and the GNU project: + +https://www.gnu.org/licenses/old-licenses/lgpl-2.1.html + +## Optional Japanese lexical helper + +`tools/discrawl-ja` depends on: + +- [Kagome](https://github.com/ikawaha/kagome), MIT +- [kagome-dict IPA](https://github.com/ikawaha/kagome-dict), MIT wrapper around + mecab-ipadic-2.7.0-20070801 / ICOT Free Software + +These dependencies are optional and are not linked into the default Discrawl +binary. Preserve the IPADIC/ICOT notice when distributing the helper. + +## Optional Chinese lexical helper + +`tools/discrawl-zh` depends on [GSE](https://github.com/go-ego/gse), Apache-2.0. +It is optional and is not linked into the default Discrawl binary. + +## In-process Arabic analyzer + +The Arabic light stemmer follows the Lucene/Bleve prefix-and-suffix contract. +It is implemented in Discrawl itself and does not depend on Python. diff --git a/docs/benchmarks/multilingual-lexical.md b/docs/benchmarks/multilingual-lexical.md new file mode 100644 index 00000000..db6cf949 --- /dev/null +++ b/docs/benchmarks/multilingual-lexical.md @@ -0,0 +1,45 @@ +# Multilingual lexical benchmark + +This targeted microbenchmark checks the failure mode the optional multilingual +fields are designed to fix: a query matching a useful subword inside an +unsegmented Korean, Japanese, Chinese, or Arabic surface form. + +It is not a general relevance benchmark. The fixtures deliberately contain +compound forms and attached Arabic proclitics that SQLite FTS5 `unicode61` +cannot retrieve as independent terms. + +## Reproduce + +```bash +# Korean helper: official Kiwi 0.23.2 + discrawl-kiwi +# Japanese helper: go build ./tools/discrawl-ja +# Chinese helper: go build ./tools/discrawl-zh +# Arabic: in-process, no helper + +DISCRAWL_TOKENIZER_E2E=1 \ +DISCRAWL_KIWI_HELPER=/tmp/discrawl-kiwi \ +DISCRAWL_KIWI_MODEL=/tmp/kiwi-model/models/cong/base \ +DISCRAWL_JA_HELPER=/tmp/discrawl-ja \ +DISCRAWL_ZH_HELPER=/tmp/discrawl-zh \ +go test ./internal/store \ + -run TestMultilingualLexicalQualityBenchmark \ + -count=1 -v +``` + +## Result + +Measured on macOS arm64 with native Kiwi 0.23.2, Kagome Search, GSE CutSearch, +and the in-process Arabic analyzer: + +| Language | `unicode61` recall@5 | Multilingual recall@5 | +| --- | ---: | ---: | +| Korean / Kiwi via kiwigo | 0/5 | 5/5 | +| Japanese / Kagome Search | 0/5 | 5/5 | +| Chinese / GSE search mode | 0/5 | 5/5 | +| Arabic / in-process light stem | 0/5 | 5/5 | +| **Total** | **0/20** | **20/20** | + +The benchmark therefore supports a narrow claim: configured language fields +substantially improve recall for these segmentation cases. It does not claim a +universal 100-point gain on natural Discord query distributions. No Python +tokenizer is used. diff --git a/docs/commands/search.md b/docs/commands/search.md index 10466eef..afcd849a 100644 --- a/docs/commands/search.md +++ b/docs/commands/search.md @@ -18,7 +18,7 @@ discrawl --json search "websocket closed" ## Modes -- `fts` (default) - SQLite FTS5 with `unicode61` tokenizer; newest matches first +- `fts` (default) - SQLite FTS5 with `unicode61`; optional multilingual lexical fields use language-specific tokenizers and RRF - `semantic` - embeds the query, scores against locally stored vectors; errors out if embeddings are disabled or no compatible vectors exist - `hybrid` - runs both, deduplicates by message id, falls back to FTS when semantic is unavailable @@ -36,6 +36,12 @@ discrawl --json search "websocket closed" User query terms are parameterized and quoted before `MATCH`, so tokens like `AND`, `OR`, `NOT`, `NEAR`, and `*` are searched as input terms instead of FTS operators. Punctuation still follows FTS5 tokenization rules. +When `[search.lexical].languages` is non-empty, Discrawl searches the default +field plus every configured language field, ranks each field independently, +and merges the lists with reciprocal rank fusion. See +[Search modes](../guides/search-modes.html#optional-multilingual-lexical-fields) +for setup and dependency details. + Ambiguous channel names fail with candidate guild/channel ids instead of silently searching multiple channels. Use `discrawl channels resolve --json`, then keep the numeric id for repeatable workflows. diff --git a/docs/configuration.md b/docs/configuration.md index c894186b..4e3712f1 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -77,6 +77,13 @@ full_cache = false [search] default_mode = "fts" +[search.lexical] +languages = [] # optional: "ko", "ja", "zh", "ar" +kiwi_command = "discrawl-kiwi" +kiwi_model = "" +ja_command = "discrawl-ja" +zh_command = "discrawl-zh" + [search.embeddings] enabled = false provider = "openai" @@ -140,6 +147,12 @@ Set `discord.token_source = "keyring"` if you want to require keyring lookup and - `sync.exclude_channel_ids` and `sync.exclude_channel_kinds` apply to historical sync, live tail events, and repair syncs; exclusions always win over category inclusion - `sync.exclude_channel_kinds` accepts Discrawl kinds such as `text`, `announcement`, `forum`, `thread_public`, `thread_private`, and `thread_announcement` - a non-zero `sync.repair_offset` aligns periodic repairs to local wall-clock boundaries; for example, `repair_every = "6h"` with `repair_offset = "2h"` targets 02:00, 08:00, 14:00, and 20:00 local time +- `[search.lexical].languages` enables opt-in multilingual FTS fields. Supported presets are Korean (`ko`, Kiwi through `github.com/codingpot/kiwigo`), Japanese (`ja`, Kagome Search through `discrawl-ja`), Chinese (`zh`, GSE CutSearch through `discrawl-zh`), and Arabic (`ar`, in-process light stemming). +- Korean, Japanese, and Chinese use separately built Go helpers so the default Discrawl binary stays small and CGO-free. Arabic is implemented in-process. +- Optional lexical helpers are installed and built separately; Discrawl never downloads packages. +- Tokenizer helpers load lazily on first indexing or search use. Commands that only inspect metadata do not start helpers. +- Each enabled language adds an independent FTS5 table. Index and query text pass through the same tokenizer, and results from the default plus language-specific tables are merged with reciprocal rank fusion. +- After adding or changing `search.lexical.languages`, run a writer command such as `discrawl sync` once so the configured lexical tables are built. Read-only commands never mutate the archive; new and edited messages update the tables automatically during later syncs. - changing `[search.embeddings]` provider/model/input version retargets pending jobs and resets prior attempts; existing vectors for another identity remain in SQLite but are not used for semantic search - `[search.embeddings].dimensions` is an optional positive OpenAI projection size. Changing it requires `embed --rebuild` so stored message vectors and query vectors use the same dimensions. - `[search.embeddings].vector_backend` accepts `exact` or optional `turbovec`; turbovec requires Python plus the `turbovec` package and embedding dimensions divisible by 8. diff --git a/docs/guides/search-modes.md b/docs/guides/search-modes.md index a251bea5..c338f984 100644 --- a/docs/guides/search-modes.md +++ b/docs/guides/search-modes.md @@ -11,10 +11,60 @@ ## FTS details - backed by SQLite FTS5 with the default `unicode61` tokenizer +- optional `[search.lexical]` languages add independent tokenizer-specific FTS tables and merge their ranked results with reciprocal rank fusion +- supported presets are Korean with native Kiwi through `kiwigo`, Japanese with Kagome Search, Chinese with GSE search mode, and Arabic with in-process light stemming - user query terms are parameterized and quoted before `MATCH`, so tokens like `AND`, `OR`, `NOT`, `NEAR`, and `*` are searched as input terms instead of FTS operators - punctuation still follows FTS5 tokenization rules - by default, `search` skips rows with no searchable content (attachment text, attachment filenames, embeds, and replies still count as content); use `--include-empty` to opt back in +### Optional multilingual lexical fields + +Install Kiwi 0.23.2's dynamic library and base model, then build the Go helper: + +```bash +git clone https://github.com/openclaw/discrawl +cd discrawl/tools/discrawl-kiwi +go build -o ~/.local/share/discrawl/bin/discrawl-kiwi . +``` + +`github.com/codingpot/kiwigo` links to the system Kiwi C API. Its upstream +installation expects Kiwi headers and dynamic libraries under `/usr/local`; +the model is the `kiwi_model_v0.23.2_base.tgz` release asset. + +Build the optional Japanese and Chinese helpers: + +```bash +cd tools/discrawl-ja && go build -o ~/.local/share/discrawl/bin/discrawl-ja . +cd ../discrawl-zh && go build -o ~/.local/share/discrawl/bin/discrawl-zh . +``` + +Configure the fields: + +```toml +[search.lexical] +languages = ["ko", "ja", "zh", "ar"] +kiwi_command = "~/.local/share/discrawl/bin/discrawl-kiwi" +kiwi_model = "~/.local/share/discrawl/models/kiwi/base" +ja_command = "~/.local/share/discrawl/bin/discrawl-ja" +zh_command = "~/.local/share/discrawl/bin/discrawl-zh" +``` + +Discrawl does not download or install helper packages; build the optional +helpers separately and configure their absolute paths. + +Every message is analyzed into each configured field. This deliberately avoids +language detection, so mixed-language Discord messages remain searchable +through every enabled analyzer. Disk usage and indexing work increase with the +number of fields; query-time RRF deduplicates message ids without mixing the +different BM25 term statistics into one field. + +None of the lexical analyzers use Python. Korean, Japanese, and Chinese run as +persistent Go helpers; Arabic is in-process. Helpers load lazily on first use +and are not linked into the default Discrawl binary. + +See [Multilingual lexical benchmark](../benchmarks/multilingual-lexical.html) +for the reproducible targeted quality check and its storage tradeoff. + ## Semantic and hybrid prerequisites - `[search.embeddings]` configured in the Discrawl config file diff --git a/internal/cli/cli.go b/internal/cli/cli.go index 2c30dcc9..5eb4d0d5 100644 --- a/internal/cli/cli.go +++ b/internal/cli/cli.go @@ -573,10 +573,7 @@ func (r *runtime) shouldAutoUpdateShare(mode shareUpdateMode) bool { func (r *runtime) autoUpdateShareIfLockAvailable(dbPath string, updateMode shareUpdateMode) error { locked, err := r.tryWithSyncLock(func() error { - storeFactory := r.openStore - if storeFactory == nil { - storeFactory = store.Open - } + storeFactory := r.localStoreFactory() var openErr error r.store, openErr = storeFactory(r.ctx, dbPath) if openErr != nil { @@ -598,10 +595,7 @@ func (r *runtime) autoUpdateShareIfLockAvailable(dbPath string, updateMode share } func (r *runtime) openLocalStore(dbPath string, updateMode shareUpdateMode, fn func() error) error { - storeFactory := r.openStore - if storeFactory == nil { - storeFactory = store.Open - } + storeFactory := r.localStoreFactory() var err error r.store, err = storeFactory(r.ctx, dbPath) if err != nil { @@ -656,7 +650,7 @@ func (r *runtime) withExistingLocalStoreReadOnly(fn func() error) error { func (r *runtime) openLocalStoreReadOnly(dbPath string, fn func() error) error { r.store = nil - s, err := store.OpenReadOnly(r.ctx, dbPath) + s, err := r.openConfiguredReadOnlyStore(dbPath) if err != nil { if errors.Is(err, os.ErrNotExist) { return fn() @@ -673,7 +667,7 @@ func (r *runtime) openLocalStoreReadOnly(dbPath string, fn func() error) error { func (r *runtime) openExistingLocalStoreReadOnly(dbPath string, fn func() error) error { r.store = nil - s, err := store.OpenReadOnly(r.ctx, dbPath) + s, err := r.openConfiguredReadOnlyStore(dbPath) if err != nil { if errors.Is(err, os.ErrNotExist) { return fn() @@ -688,6 +682,19 @@ func (r *runtime) openExistingLocalStoreReadOnly(dbPath string, fn func() error) return fn() } +func (r *runtime) openConfiguredReadOnlyStore(path string) (*store.Store, error) { + if len(r.cfg.Search.Lexical.Languages) == 0 { + return store.OpenReadOnly(r.ctx, path) + } + return store.OpenReadOnlyWithOptions(r.ctx, path, store.OpenOptions{ + LexicalLanguages: r.cfg.Search.Lexical.Languages, + LexicalKiwiCommand: r.cfg.Search.Lexical.KiwiCommand, + LexicalKiwiModel: r.cfg.Search.Lexical.KiwiModel, + LexicalJaCommand: r.cfg.Search.Lexical.JaCommand, + LexicalZhCommand: r.cfg.Search.Lexical.ZhCommand, + }) +} + func (r *runtime) withServicesAuto(withDiscord, autoShareUpdate bool, fn func() error) error { return r.withServicesAutoLocked(withDiscord, autoShareUpdate, false, fn) } @@ -731,10 +738,7 @@ func (r *runtime) withServicesUpdateLockedOperation(withDiscord bool, updateMode } func (r *runtime) openServices(dbPath string, withDiscord bool, updateMode shareUpdateMode, fn func() error) error { - storeFactory := r.openStore - if storeFactory == nil { - storeFactory = store.Open - } + storeFactory := r.localStoreFactory() var err error r.store, err = storeFactory(r.ctx, dbPath) if err != nil { @@ -757,6 +761,21 @@ func (r *runtime) openServices(dbPath string, withDiscord bool, updateMode share return fn() } +func (r *runtime) localStoreFactory() func(context.Context, string) (*store.Store, error) { + if r.openStore != nil { + return r.openStore + } + return func(ctx context.Context, path string) (*store.Store, error) { + return store.OpenWithOptions(ctx, path, store.OpenOptions{ + LexicalLanguages: r.cfg.Search.Lexical.Languages, + LexicalKiwiCommand: r.cfg.Search.Lexical.KiwiCommand, + LexicalKiwiModel: r.cfg.Search.Lexical.KiwiModel, + LexicalJaCommand: r.cfg.Search.Lexical.JaCommand, + LexicalZhCommand: r.cfg.Search.Lexical.ZhCommand, + }) + } +} + func (r *runtime) ensureDiscordServices() error { discordFactory := r.newDiscord if discordFactory == nil { diff --git a/internal/config/config.go b/internal/config/config.go index d9caa565..a7135bf9 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -67,8 +67,17 @@ type SyncConfig struct { } type SearchConfig struct { - DefaultMode string `toml:"default_mode"` - Embeddings EmbeddingsConfig `toml:"embeddings"` + DefaultMode string `toml:"default_mode"` + Lexical LexicalSearchConfig `toml:"lexical"` + Embeddings EmbeddingsConfig `toml:"embeddings"` +} + +type LexicalSearchConfig struct { + Languages []string `toml:"languages,omitempty"` + KiwiCommand string `toml:"kiwi_command"` + KiwiModel string `toml:"kiwi_model"` + JaCommand string `toml:"ja_command"` + ZhCommand string `toml:"zh_command"` } type ShareConfig struct { @@ -154,6 +163,11 @@ func Default() Config { }, Search: SearchConfig{ DefaultMode: "fts", + Lexical: LexicalSearchConfig{ + KiwiCommand: "discrawl-kiwi", + JaCommand: "discrawl-ja", + ZhCommand: "discrawl-zh", + }, Embeddings: EmbeddingsConfig{ Enabled: false, Provider: "openai", @@ -300,6 +314,37 @@ func (c *Config) Normalize() error { if c.Search.DefaultMode == "" { c.Search.DefaultMode = "fts" } + c.Search.Lexical.KiwiCommand = strings.TrimSpace(c.Search.Lexical.KiwiCommand) + if c.Search.Lexical.KiwiCommand == "" { + c.Search.Lexical.KiwiCommand = "discrawl-kiwi" + } + c.Search.Lexical.KiwiModel = strings.TrimSpace(c.Search.Lexical.KiwiModel) + c.Search.Lexical.JaCommand = strings.TrimSpace(c.Search.Lexical.JaCommand) + if c.Search.Lexical.JaCommand == "" { + c.Search.Lexical.JaCommand = "discrawl-ja" + } + c.Search.Lexical.ZhCommand = strings.TrimSpace(c.Search.Lexical.ZhCommand) + if c.Search.Lexical.ZhCommand == "" { + c.Search.Lexical.ZhCommand = "discrawl-zh" + } + seenLexicalLanguages := make(map[string]struct{}, len(c.Search.Lexical.Languages)) + normalizedLexicalLanguages := make([]string, 0, len(c.Search.Lexical.Languages)) + for _, language := range c.Search.Lexical.Languages { + language = strings.ToLower(strings.TrimSpace(language)) + switch language { + case "ko", "ja", "zh", "ar": + case "": + continue + default: + return fmt.Errorf("unsupported search.lexical language %q; use ko, ja, zh, or ar", language) + } + if _, ok := seenLexicalLanguages[language]; ok { + return fmt.Errorf("duplicate search.lexical language %q", language) + } + seenLexicalLanguages[language] = struct{}{} + normalizedLexicalLanguages = append(normalizedLexicalLanguages, language) + } + c.Search.Lexical.Languages = normalizedLexicalLanguages c.Search.Embeddings.Provider = strings.ToLower(strings.TrimSpace(c.Search.Embeddings.Provider)) c.Search.Embeddings.Model = strings.TrimSpace(c.Search.Embeddings.Model) c.Search.Embeddings.BaseURL = strings.TrimRight(strings.TrimSpace(c.Search.Embeddings.BaseURL), "/") diff --git a/internal/config/config_test.go b/internal/config/config_test.go index d6bbffff..a3c5f327 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -249,6 +249,49 @@ func TestDefaultSyncConcurrencyBounds(t *testing.T) { require.Equal(t, 32, defaultSyncConcurrency()) } +func TestLoadConfigMultilingualLexicalSearch(t *testing.T) { + t.Parallel() + + path := filepath.Join(t.TempDir(), "config.toml") + require.NoError(t, os.WriteFile(path, []byte(` +version = 1 + +[discord] +token_source = "env" + +[search.lexical] +languages = ["ko", "ja", "zh", "ar"] +kiwi_command = "/opt/discrawl/bin/discrawl-kiwi" +kiwi_model = "/opt/discrawl/models/kiwi/base" +ja_command = "/opt/discrawl/bin/discrawl-ja" +zh_command = "/opt/discrawl/bin/discrawl-zh" +`), 0o600)) + + cfg, err := Load(path) + require.NoError(t, err) + require.Equal(t, []string{"ko", "ja", "zh", "ar"}, cfg.Search.Lexical.Languages) + require.Equal(t, "/opt/discrawl/bin/discrawl-kiwi", cfg.Search.Lexical.KiwiCommand) + require.Equal(t, "/opt/discrawl/models/kiwi/base", cfg.Search.Lexical.KiwiModel) + require.Equal(t, "/opt/discrawl/bin/discrawl-ja", cfg.Search.Lexical.JaCommand) + require.Equal(t, "/opt/discrawl/bin/discrawl-zh", cfg.Search.Lexical.ZhCommand) +} + +func TestNormalizeRejectsUnsupportedLexicalLanguage(t *testing.T) { + t.Parallel() + + cfg := Default() + cfg.Search.Lexical.Languages = []string{"ko", "klingon"} + require.ErrorContains(t, cfg.Normalize(), `unsupported search.lexical language "klingon"`) +} + +func TestNormalizeRejectsDuplicateLexicalLanguages(t *testing.T) { + t.Parallel() + + cfg := Default() + cfg.Search.Lexical.Languages = []string{"ko", "ko"} + require.ErrorContains(t, cfg.Normalize(), `duplicate search.lexical language "ko"`) +} + func TestResolveDiscordTokenFromEnv(t *testing.T) { cfg := Default() t.Setenv(DefaultTokenEnv, "Bot env-token") diff --git a/internal/store/lexical.go b/internal/store/lexical.go new file mode 100644 index 00000000..500145bf --- /dev/null +++ b/internal/store/lexical.go @@ -0,0 +1,339 @@ +package store + +import ( + "context" + "database/sql" + "fmt" + "sort" + "strings" + "time" +) + +const lexicalFTSVersion = "2" + +type LexicalTokenizer interface { + Tokenize(context.Context, string) (string, error) + Close() error +} + +func openWithLexicalTokenizers( + ctx context.Context, + path string, + tokenizers map[string]LexicalTokenizer, +) (*Store, error) { + base, err := openBaseStore(ctx, path) + if err != nil { + closeLexicalTokenizers(tokenizers) + return nil, err + } + store := &Store{ + db: base.DB(), + q: newStoreQueries(base.DB()), + path: path, + baseClose: base.Close, + lexicalTokenizers: tokenizers, + } + if err := store.migrate(ctx); err != nil { + _ = store.Close() + return nil, err + } + if err := store.ensureLexicalFTS(ctx); err != nil { + _ = store.Close() + return nil, err + } + if err := store.invalidateDisabledLexicalVersions(ctx); err != nil { + _ = store.Close() + return nil, err + } + return store, nil +} + +func (s *Store) lexicalLanguages() []string { + languages := make([]string, 0, len(s.lexicalTokenizers)) + for language := range s.lexicalTokenizers { + languages = append(languages, language) + } + sort.Strings(languages) + return languages +} + +func (s *Store) tokenizeLexical(ctx context.Context, text string) (map[string]string, error) { + if len(s.lexicalTokenizers) == 0 || strings.TrimSpace(text) == "" { + return nil, nil + } + tokenized := make(map[string]string, len(s.lexicalTokenizers)) + for _, language := range s.lexicalLanguages() { + content, err := s.lexicalTokenizers[language].Tokenize(ctx, text) + if err != nil { + return nil, fmt.Errorf("tokenize %s text: %w", language, err) + } + tokenized[language] = content + } + return tokenized, nil +} + +func (s *Store) ensureLexicalFTS(ctx context.Context) error { + for _, language := range s.lexicalLanguages() { + var version sql.NullString + err := s.db.QueryRowContext(ctx, ` + select cursor from sync_state where scope = ? + `, lexicalFTSScope(language)).Scan(&version) + if err == nil && version.String == lexicalFTSVersion { + continue + } + if err != nil && err != sql.ErrNoRows { + return fmt.Errorf("check %s lexical index version: %w", language, err) + } + if err := s.rebuildLexicalFTS(ctx, language); err != nil { + return err + } + if _, err := s.db.ExecContext(ctx, ` + insert into sync_state(scope, cursor, updated_at) + values(?, ?, ?) + on conflict(scope) do update set + cursor = excluded.cursor, + updated_at = excluded.updated_at + `, lexicalFTSScope(language), lexicalFTSVersion, time.Now().UTC().Format(timeLayout)); err != nil { + return fmt.Errorf("stamp %s lexical index version: %w", language, err) + } + } + return nil +} + +func (s *Store) invalidateDisabledLexicalVersions(ctx context.Context) error { + enabled := make(map[string]struct{}, len(s.lexicalTokenizers)) + for language := range s.lexicalTokenizers { + enabled[language] = struct{}{} + } + knownScopes := map[string]string{ + lexicalFTSScope("ko"): "ko", + lexicalFTSScope("ja"): "ja", + lexicalFTSScope("zh"): "zh", + lexicalFTSScope("ar"): "ar", + } + disabledScopes, err := func() ([]string, error) { + rows, err := s.db.QueryContext(ctx, ` + select scope + from sync_state + where scope in (?, ?, ?, ?) + `, + lexicalFTSScope("ko"), + lexicalFTSScope("ja"), + lexicalFTSScope("zh"), + lexicalFTSScope("ar"), + ) + if err != nil { + return nil, fmt.Errorf("query lexical index versions: %w", err) + } + defer func() { _ = rows.Close() }() + var scopes []string + for rows.Next() { + var scope string + if err := rows.Scan(&scope); err != nil { + return nil, fmt.Errorf("scan lexical index version: %w", err) + } + language, ok := knownScopes[scope] + if !ok { + continue + } + if _, ok := enabled[language]; !ok { + scopes = append(scopes, scope) + } + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate lexical index versions: %w", err) + } + return scopes, nil + }() + if err != nil { + return err + } + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer rollback(tx) + for _, scope := range disabledScopes { + language := knownScopes[scope] + if _, err := tx.ExecContext(ctx, "drop table if exists "+lexicalFTSTable(language)); err != nil { + return fmt.Errorf("drop disabled %s lexical index: %w", language, err) + } + if _, err := tx.ExecContext( + ctx, + `delete from sync_state where scope = ?`, + scope, + ); err != nil { + return fmt.Errorf("invalidate disabled lexical index %s: %w", scope, err) + } + } + return tx.Commit() +} + +func (s *Store) rebuildLexicalIndexes(ctx context.Context) error { + for _, language := range s.lexicalLanguages() { + if err := s.rebuildLexicalFTS(ctx, language); err != nil { + return err + } + } + return nil +} + +func (s *Store) rebuildLexicalFTS(ctx context.Context, language string) error { + table := lexicalFTSTable(language) + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer rollback(tx) + if _, err := tx.ExecContext(ctx, "drop table if exists "+table); err != nil { + return fmt.Errorf("drop %s: %w", table, err) + } + if _, err := tx.ExecContext(ctx, createLexicalFTSSQL(table)); err != nil { + return fmt.Errorf("create %s: %w", table, err) + } + if err := configureFTSBulkLoad(ctx, tx, table); err != nil { + return err + } + rows, err := tx.QueryContext(ctx, lexicalRebuildRowsSQL) + if err != nil { + return fmt.Errorf("query %s rebuild rows: %w", table, err) + } + defer func() { _ = rows.Close() }() + for rows.Next() { + var message MessageRecord + if err := rows.Scan( + &message.ID, + &message.GuildID, + &message.ChannelID, + &message.AuthorID, + &message.AuthorName, + &message.ChannelName, + &message.NormalizedContent, + ); err != nil { + return fmt.Errorf("scan %s rebuild row: %w", table, err) + } + content, err := s.lexicalTokenizers[language].Tokenize(ctx, message.NormalizedContent) + if err != nil { + return fmt.Errorf("tokenize %s rebuild row %s: %w", language, message.ID, err) + } + if err := insertLexicalMessageTx(ctx, tx, table, message, content); err != nil { + return err + } + } + if err := rows.Err(); err != nil { + return fmt.Errorf("iterate %s rebuild rows: %w", table, err) + } + if err := optimizeFTS(ctx, tx, table); err != nil { + return err + } + return tx.Commit() +} + +func (s *Store) upsertLexicalMessageTx( + ctx context.Context, + tx *sql.Tx, + message MessageRecord, + tokenized map[string]string, +) error { + for _, language := range s.lexicalLanguages() { + table := lexicalFTSTable(language) + rowID, ok := messageFTSRowID(message.ID) + if !ok { + continue + } + if _, err := tx.ExecContext(ctx, "delete from "+table+" where rowid = ?", rowID); err != nil { + return err + } + if message.DeletedAt == "" { + if err := insertLexicalMessageTx(ctx, tx, table, message, tokenized[language]); err != nil { + return err + } + } + } + return nil +} + +func (s *Store) deleteLexicalMessagesTx(ctx context.Context, tx *sql.Tx, column string, value any) error { + if column != "rowid" && column != "guild_id" { + return fmt.Errorf("unsupported lexical delete column %q", column) + } + for _, language := range s.lexicalLanguages() { + if _, err := tx.ExecContext(ctx, "delete from "+lexicalFTSTable(language)+" where "+column+" = ?", value); err != nil { + return err + } + } + return nil +} + +func insertLexicalMessageTx(ctx context.Context, tx *sql.Tx, table string, message MessageRecord, content string) error { + rowID, ok := messageFTSRowID(message.ID) + if !ok { + return nil + } + _, err := tx.ExecContext(ctx, ` + insert into `+table+`( + rowid, message_id, guild_id, channel_id, author_id, author_name, channel_name, content + ) values(?, ?, ?, ?, ?, ?, ?, ?) + `, rowID, message.ID, message.GuildID, message.ChannelID, nullable(message.AuthorID), message.AuthorName, message.ChannelName, content) + return err +} + +func lexicalFTSTable(language string) string { + switch language { + case "ko", "ja", "zh", "ar": + return "message_fts_" + language + default: + panic("unsupported lexical language: " + language) + } +} + +func lexicalFTSScope(language string) string { + return "schema:" + lexicalFTSTable(language) + "_version" +} + +func isMessageFTSTable(table string) bool { + return table == "message_fts" || + table == "message_fts_ko" || + table == "message_fts_ja" || + table == "message_fts_zh" || + table == "message_fts_ar" +} + +func createLexicalFTSSQL(table string) string { + return `create virtual table ` + table + ` using fts5( + message_id unindexed, + guild_id unindexed, + channel_id unindexed, + author_id unindexed, + author_name, + channel_name, + content, + tokenize = 'unicode61 remove_diacritics 0' + )` +} + +func closeLexicalTokenizers(tokenizers map[string]LexicalTokenizer) { + for _, tokenizer := range tokenizers { + _ = tokenizer.Close() + } +} + +const lexicalRebuildRowsSQL = ` + select + m.id, + m.guild_id, + m.channel_id, + coalesce(m.author_id, ''), + coalesce( + json_extract(m.raw_json, '$.member.nick'), + json_extract(m.raw_json, '$.author.global_name'), + json_extract(m.raw_json, '$.author.username'), + '' + ), + coalesce(c.name, ''), + m.normalized_content + from messages m + left join channels c on c.id = m.channel_id + where m.deleted_at is null + order by cast(m.id as integer) +` diff --git a/internal/store/lexical_arabic.go b/internal/store/lexical_arabic.go new file mode 100644 index 00000000..40be3fff --- /dev/null +++ b/internal/store/lexical_arabic.go @@ -0,0 +1,119 @@ +package store + +import ( + "context" + "strings" + "unicode" + "unicode/utf8" +) + +// Arabic light stemming follows the Lucene/Bleve prefix-and-suffix +// contract so attached proclitics remain searchable as independent terms. +var arabicPrefixes = []string{"وال", "فال", "بال", "كال", "لال", "ال", "لل", "و", "ف", "ب", "ك", "ل"} + +var arabicSuffixes = []string{"ها", "ان", "ات", "ون", "ين", "يه", "ية", "ه", "ة", "ي"} + +type arabicLexicalTokenizer struct{} + +func newArabicLexicalTokenizer() LexicalTokenizer { + return arabicLexicalTokenizer{} +} + +func (arabicLexicalTokenizer) Tokenize(_ context.Context, text string) (string, error) { + return strings.Join(tokenizeArabic(text), " "), nil +} + +func (arabicLexicalTokenizer) Close() error { + return nil +} + +func tokenizeArabic(text string) []string { + seen := make(map[string]struct{}) + var tokens []string + add := func(token string) { + token = strings.ToLower(strings.TrimSpace(token)) + if token == "" { + return + } + if _, ok := seen[token]; ok { + return + } + seen[token] = struct{}{} + tokens = append(tokens, token) + } + for _, word := range splitArabicWords(normalizeArabic(text)) { + add(word) + stripped := stripArabicPrefix(word) + add(stripped) + add(stemArabic(word)) + add(stemArabic(stripped)) + } + return tokens +} + +func splitArabicWords(text string) []string { + var words []string + var current []rune + flush := func() { + if len(current) == 0 { + return + } + words = append(words, string(current)) + current = current[:0] + } + for _, r := range text { + if unicode.IsLetter(r) { + current = append(current, r) + continue + } + flush() + } + flush() + return words +} + +func normalizeArabic(text string) string { + var b strings.Builder + b.Grow(len(text)) + for _, r := range strings.ToValidUTF8(text, "") { + switch r { + case '\u064B', '\u064C', '\u064D', '\u064E', '\u064F', '\u0650', '\u0651', '\u0652', '\u0670', '\u0640': + continue + default: + b.WriteRune(unicode.ToLower(r)) + } + } + return b.String() +} + +func stripArabicPrefix(word string) string { + for _, prefix := range arabicPrefixes { + if !strings.HasPrefix(word, prefix) { + continue + } + rest := strings.TrimPrefix(word, prefix) + if utf8.RuneCountInString(rest) < 2 { + continue + } + if prefix == "و" && utf8.RuneCountInString(word) < 4 { + continue + } + return rest + } + return word +} + +func stemArabic(word string) string { + runes := []rune(stripArabicPrefix(word)) + for _, suffix := range arabicSuffixes { + suffixRunes := []rune(suffix) + if len(runes)-len(suffixRunes) < 2 { + continue + } + if string(runes[len(runes)-len(suffixRunes):]) != suffix { + continue + } + runes = runes[:len(runes)-len(suffixRunes)] + } + return string(runes) +} diff --git a/internal/store/lexical_arabic_test.go b/internal/store/lexical_arabic_test.go new file mode 100644 index 00000000..8ab19482 --- /dev/null +++ b/internal/store/lexical_arabic_test.go @@ -0,0 +1,32 @@ +package store + +import ( + "context" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestTokenizeArabicSplitsProcliticsForSearch(t *testing.T) { + cases := map[string]string{ + "والكتاب": "كتاب", + "والمدرسة": "مدرسة", + "فالاجتماع": "اجتماع", + "بالسجل": "سجل", + "كالبرنامج": "برنامج", + } + for input, want := range cases { + tokens := tokenizeArabic(input) + require.Containsf(t, tokens, want, "input %q tokens=%v", input, tokens) + } +} + +func TestArabicTokenizerIsInProcessAndIdempotent(t *testing.T) { + tokenizer := newArabicLexicalTokenizer() + tokens, err := tokenizer.Tokenize(context.Background(), "والكتاب في المدرسة") + require.NoError(t, err) + require.Contains(t, strings.Split(tokens, " "), "كتاب") + require.NoError(t, tokenizer.Close()) + require.NoError(t, tokenizer.Close()) +} diff --git a/internal/store/lexical_helper.go b/internal/store/lexical_helper.go new file mode 100644 index 00000000..98a5a0b3 --- /dev/null +++ b/internal/store/lexical_helper.go @@ -0,0 +1,253 @@ +package store + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "time" +) + +type helperLexicalResponse struct { + Ready bool `json:"ready,omitempty"` + Tokens string `json:"tokens,omitempty"` + Error string `json:"error,omitempty"` +} + +type lockedBuffer struct { + mutex sync.Mutex + buffer bytes.Buffer +} + +func (b *lockedBuffer) Write(data []byte) (int, error) { + b.mutex.Lock() + defer b.mutex.Unlock() + return b.buffer.Write(data) +} + +func (b *lockedBuffer) String() string { + b.mutex.Lock() + defer b.mutex.Unlock() + return b.buffer.String() +} + +func startHelperLexicalTokenizer(language, command, defaultName string, extraArgs []string) (LexicalTokenizer, error) { + cmd, err := newHelperLexicalCommand(command, defaultName, extraArgs) + if err != nil { + return nil, err + } + return startHelperLexicalTokenizerCommand(language, cmd) +} + +func startHelperLexicalTokenizerCommand(language string, cmd *exec.Cmd) (LexicalTokenizer, error) { + tokenizer := &externalLexicalTokenizer{ + language: language, + command: cmd, + } + stdin, err := cmd.StdinPipe() + if err != nil { + return nil, err + } + stdout, err := cmd.StdoutPipe() + if err != nil { + _ = stdin.Close() + return nil, err + } + cmd.Stderr = &tokenizer.stderr + tokenizer.stdin = stdin + tokenizer.stdout = bufio.NewScanner(stdout) + tokenizer.stdout.Buffer(make([]byte, 4096), 8*1024*1024) + if err := cmd.Start(); err != nil { + return nil, fmt.Errorf("start %s lexical tokenizer: %w", language, err) + } + response, err := tokenizer.readStartup() + if err != nil { + _ = tokenizer.Close() + return nil, fmt.Errorf("initialize %s lexical tokenizer: %w", language, err) + } + if !response.Ready { + _ = tokenizer.Close() + if response.Error == "" { + return nil, fmt.Errorf("initialize %s lexical tokenizer", language) + } + return nil, errors.New(response.Error) + } + return tokenizer, nil +} + +func newHelperLexicalCommand(command, defaultName string, extraArgs []string) (*exec.Cmd, error) { + command, err := expandLexicalPath(command, defaultName) + if err != nil { + return nil, fmt.Errorf("resolve %s helper: %w", defaultName, err) + } + if !filepath.IsAbs(command) { + if command != defaultName { + return nil, fmt.Errorf( + "unsupported %s helper %q; use an absolute path or %s", + defaultName, + command, + defaultName, + ) + } + command, err = exec.LookPath(command) + if err != nil { + return nil, fmt.Errorf("find %s helper: %w", defaultName, err) + } + } + args := append([]string{command}, extraArgs...) + return &exec.Cmd{ + Path: command, + Args: args, + Env: lexicalWorkerEnvironment(os.Environ()), + }, nil +} + +func expandLexicalPath(path, fallback string) (string, error) { + path = strings.TrimSpace(path) + if path == "" { + return fallback, nil + } + if !strings.HasPrefix(path, "~/") { + return path, nil + } + home, err := os.UserHomeDir() + if err != nil { + return "", err + } + return filepath.Join(home, strings.TrimPrefix(path, "~/")), nil +} + +type externalLexicalTokenizer struct { + language string + command *exec.Cmd + stdin io.WriteCloser + stdout *bufio.Scanner + stderr lockedBuffer + mutex sync.Mutex +} + +func (t *externalLexicalTokenizer) readStartup() (helperLexicalResponse, error) { + result := make(chan struct { + response helperLexicalResponse + err error + }, 1) + go func() { + response, err := t.readResponse() + result <- struct { + response helperLexicalResponse + err error + }{response, err} + }() + select { + case output := <-result: + return output.response, output.err + case <-time.After(30 * time.Second): + _ = t.command.Process.Kill() + return helperLexicalResponse{}, errors.New("lexical tokenizer startup timed out after 30s") + } +} + +func (t *externalLexicalTokenizer) Tokenize(ctx context.Context, text string) (string, error) { + t.mutex.Lock() + defer t.mutex.Unlock() + if err := ctx.Err(); err != nil { + return "", err + } + request, err := json.Marshal(map[string]string{"text": text}) + if err != nil { + return "", err + } + if _, err := t.stdin.Write(append(request, '\n')); err != nil { + return "", fmt.Errorf("write %s tokenizer request: %w", t.language, err) + } + result := make(chan struct { + response helperLexicalResponse + err error + }, 1) + go func() { + response, err := t.readResponse() + result <- struct { + response helperLexicalResponse + err error + }{response, err} + }() + select { + case output := <-result: + if output.err != nil { + return "", output.err + } + if output.response.Error != "" { + return "", errors.New(output.response.Error) + } + return output.response.Tokens, nil + case <-ctx.Done(): + _ = t.command.Process.Kill() + return "", ctx.Err() + case <-time.After(30 * time.Second): + _ = t.command.Process.Kill() + return "", errors.New("lexical tokenizer response timed out after 30s") + } +} + +func (t *externalLexicalTokenizer) readResponse() (helperLexicalResponse, error) { + if !t.stdout.Scan() { + if err := t.stdout.Err(); err != nil { + return helperLexicalResponse{}, err + } + if detail := strings.TrimSpace(t.stderr.String()); detail != "" { + return helperLexicalResponse{}, errors.New(detail) + } + return helperLexicalResponse{}, io.EOF + } + var response helperLexicalResponse + if err := json.Unmarshal(t.stdout.Bytes(), &response); err != nil { + return helperLexicalResponse{}, fmt.Errorf("decode lexical tokenizer response: %w", err) + } + return response, nil +} + +func (t *externalLexicalTokenizer) Close() error { + if t == nil || t.command == nil || t.command.Process == nil { + return nil + } + _ = t.stdin.Close() + err := t.command.Wait() + if errors.Is(err, os.ErrProcessDone) { + return nil + } + return err +} + +func lexicalWorkerEnvironment(parent []string) []string { + allowed := map[string]struct{}{ + "HOME": {}, + "LANG": {}, + "LC_ALL": {}, + "PATH": {}, + "PATHEXT": {}, + "SYSTEMROOT": {}, + "TEMP": {}, + "TMP": {}, + "TMPDIR": {}, + "WINDIR": {}, + } + environment := make([]string, 0, len(allowed)) + for _, entry := range parent { + key, _, ok := strings.Cut(entry, "=") + if !ok { + continue + } + if _, ok := allowed[strings.ToUpper(key)]; ok { + environment = append(environment, entry) + } + } + return environment +} diff --git a/internal/store/lexical_helper_test.go b/internal/store/lexical_helper_test.go new file mode 100644 index 00000000..aeb36cbb --- /dev/null +++ b/internal/store/lexical_helper_test.go @@ -0,0 +1,152 @@ +package store + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestNewLexicalTokenizersDisabled(t *testing.T) { + tokenizers, err := newLexicalTokenizers(OpenOptions{}) + require.NoError(t, err) + require.Nil(t, tokenizers) +} + +func TestNewLexicalTokenizersArabicIsInProcess(t *testing.T) { + tokenizers, err := newLexicalTokenizers(OpenOptions{LexicalLanguages: []string{"ar"}}) + require.NoError(t, err) + require.Contains(t, tokenizers, "ar") + tokens, err := tokenizers["ar"].Tokenize(context.Background(), "والكتاب") + require.NoError(t, err) + require.Contains(t, tokens, "كتاب") +} + +func TestNewLexicalTokenizersRejectsUnknownLanguage(t *testing.T) { + _, err := newLexicalTokenizers(OpenOptions{LexicalLanguages: []string{"default"}}) + require.ErrorContains(t, err, "unsupported lexical language") +} + +func TestOpenWithOptionsLoadsKiwiHelperLazily(t *testing.T) { + ctx := context.Background() + s, err := OpenWithOptions(ctx, filepath.Join(t.TempDir(), "discrawl.db"), OpenOptions{ + LexicalLanguages: []string{"ko"}, + LexicalKiwiCommand: "/definitely/missing/discrawl-kiwi", + LexicalKiwiModel: "/definitely/missing/kiwi-model", + }) + require.NoError(t, err) + defer func() { _ = s.Close() }() + + err = s.UpsertMessage(ctx, MessageRecord{ + ID: "ko", GuildID: "g1", ChannelID: "c1", + CreatedAt: time.Now().UTC().Format(time.RFC3339Nano), + Content: "저녁먹음", NormalizedContent: "저녁먹음", RawJSON: `{}`, + }) + require.ErrorContains(t, err, "start ko lexical tokenizer") + require.NotContains(t, err.Error(), "Python") +} + +func lexicalHelperCommand(mode string) *exec.Cmd { + return &exec.Cmd{ + Path: os.Args[0], + Args: []string{os.Args[0], "-test.run=TestLexicalTokenizerHelperProcess", "--", mode}, + Env: append(os.Environ(), "DISCRAWL_LEXICAL_HELPER=1"), + } +} + +func TestLexicalTokenizerHelperProcess(t *testing.T) { + if os.Getenv("DISCRAWL_LEXICAL_HELPER") != "1" { + return + } + mode := os.Args[len(os.Args)-1] + switch mode { + case "startup-error": + fmt.Println(`{"error":"missing tokenizer package"}`) + os.Exit(2) + case "malformed-startup": + fmt.Println("not-json") + os.Exit(2) + case "stderr-startup": + fmt.Fprintln(os.Stderr, "tokenizer stderr") + os.Exit(2) + case "ready", "response-error", "malformed-response": + fmt.Println(`{"ready":true}`) + default: + fmt.Println(`{"ready":true}`) + } + scanner := bufio.NewScanner(os.Stdin) + for scanner.Scan() { + var request map[string]string + if err := json.Unmarshal(scanner.Bytes(), &request); err != nil { + fmt.Println(`{"error":"bad request"}`) + continue + } + if mode == "response-error" { + fmt.Println(`{"error":"tokenization failed"}`) + continue + } + if mode == "malformed-response" { + fmt.Println("not-json") + continue + } + tokens := request["text"] + " tokenized" + if request["text"] == "오늘 저녁먹음 기록" { + tokens = "오늘 저녁 먹 음 기록" + } + response, err := json.Marshal(map[string]string{"tokens": tokens}) + if err != nil { + fmt.Println(`{"error":"encode response"}`) + continue + } + fmt.Println(string(response)) + } +} + +func TestNewLexicalTokenizersCreatesLazyHelpers(t *testing.T) { + tokenizers, err := newLexicalTokenizers(OpenOptions{ + LexicalLanguages: []string{"ko", "ja", "zh", "ar"}, + LexicalKiwiCommand: "/definitely/missing/discrawl-kiwi", + LexicalJaCommand: "/definitely/missing/discrawl-ja", + LexicalZhCommand: "/definitely/missing/discrawl-zh", + }) + require.NoError(t, err) + require.ElementsMatch(t, []string{"ko", "ja", "zh", "ar"}, keys(tokenizers)) +} + +func keys(tokenizers map[string]LexicalTokenizer) []string { + out := make([]string, 0, len(tokenizers)) + for language := range tokenizers { + out = append(out, language) + } + return out +} + +func TestStartHelperLexicalTokenizerUsesReadyScript(t *testing.T) { + dir := t.TempDir() + script := filepath.Join(dir, "discrawl-ja") + require.NoError(t, os.WriteFile(script, []byte(`#!/bin/sh +echo '{"ready":true}' +read line +echo '{"tokens":"tokyo"}' +`), 0o700)) + tokenizer, err := startHelperLexicalTokenizer("ja", script, "discrawl-ja", nil) + require.NoError(t, err) + defer func() { _ = tokenizer.Close() }() + tokens, err := tokenizer.Tokenize(context.Background(), "東京") + require.NoError(t, err) + require.Equal(t, "tokyo", tokens) +} + +func TestOpenReadOnlyWithOptionsRejectsUnknownLanguage(t *testing.T) { + _, err := OpenReadOnlyWithOptions(context.Background(), filepath.Join(t.TempDir(), "discrawl.db"), OpenOptions{ + LexicalLanguages: []string{"nope"}, + }) + require.ErrorContains(t, err, "unsupported lexical language") +} diff --git a/internal/store/lexical_kiwi.go b/internal/store/lexical_kiwi.go new file mode 100644 index 00000000..c0dcf147 --- /dev/null +++ b/internal/store/lexical_kiwi.go @@ -0,0 +1,30 @@ +package store + +import ( + "fmt" + "os/exec" +) + +func startKiwiLexicalTokenizer(command, model string) (LexicalTokenizer, error) { + cmd, err := newKiwiLexicalCommand(command, model) + if err != nil { + return nil, err + } + return startKiwiLexicalTokenizerCommand(cmd) +} + +func startKiwiLexicalTokenizerCommand(cmd *exec.Cmd) (LexicalTokenizer, error) { + return startHelperLexicalTokenizerCommand("ko", cmd) +} + +func newKiwiLexicalCommand(command, model string) (*exec.Cmd, error) { + model, err := expandLexicalPath(model, "") + if err != nil { + return nil, fmt.Errorf("resolve Kiwi model: %w", err) + } + var extra []string + if model != "" { + extra = []string{"--model", model} + } + return newHelperLexicalCommand(command, "discrawl-kiwi", extra) +} diff --git a/internal/store/lexical_lazy.go b/internal/store/lexical_lazy.go new file mode 100644 index 00000000..04dcbb91 --- /dev/null +++ b/internal/store/lexical_lazy.go @@ -0,0 +1,34 @@ +package store + +import ( + "context" + "sync" +) + +type lazyLexicalTokenizer struct { + start func() (LexicalTokenizer, error) + once sync.Once + tokenizer LexicalTokenizer + err error +} + +func newLazyLexicalTokenizer(start func() (LexicalTokenizer, error)) LexicalTokenizer { + return &lazyLexicalTokenizer{start: start} +} + +func (l *lazyLexicalTokenizer) Tokenize(ctx context.Context, text string) (string, error) { + l.once.Do(func() { + l.tokenizer, l.err = l.start() + }) + if l.err != nil { + return "", l.err + } + return l.tokenizer.Tokenize(ctx, text) +} + +func (l *lazyLexicalTokenizer) Close() error { + if l == nil || l.tokenizer == nil { + return nil + } + return l.tokenizer.Close() +} diff --git a/internal/store/lexical_lazy_test.go b/internal/store/lexical_lazy_test.go new file mode 100644 index 00000000..729bf26f --- /dev/null +++ b/internal/store/lexical_lazy_test.go @@ -0,0 +1,50 @@ +package store + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestLazyLexicalTokenizerStartsOnce(t *testing.T) { + starts := 0 + tokenizer := newLazyLexicalTokenizer(func() (LexicalTokenizer, error) { + starts++ + return stubLexicalTokenizer{tokenize: func(text string) string { + return text + " tokenized" + }}, nil + }) + + first, err := tokenizer.Tokenize(context.Background(), "first") + require.NoError(t, err) + require.Equal(t, "first tokenized", first) + second, err := tokenizer.Tokenize(context.Background(), "second") + require.NoError(t, err) + require.Equal(t, "second tokenized", second) + require.Equal(t, 1, starts) + require.NoError(t, tokenizer.Close()) +} + +func TestLazyLexicalTokenizerCachesStartupFailure(t *testing.T) { + starts := 0 + tokenizer := newLazyLexicalTokenizer(func() (LexicalTokenizer, error) { + starts++ + return nil, errors.New("startup failed") + }) + + _, err := tokenizer.Tokenize(context.Background(), "first") + require.ErrorContains(t, err, "startup failed") + _, err = tokenizer.Tokenize(context.Background(), "second") + require.ErrorContains(t, err, "startup failed") + require.Equal(t, 1, starts) + require.NoError(t, tokenizer.Close()) +} + +func TestLazyLexicalTokenizerCloseBeforeStart(t *testing.T) { + tokenizer := newLazyLexicalTokenizer(func() (LexicalTokenizer, error) { + return nil, errors.New("must not start") + }) + require.NoError(t, tokenizer.Close()) +} diff --git a/internal/store/lexical_review_test.go b/internal/store/lexical_review_test.go new file mode 100644 index 00000000..4761c814 --- /dev/null +++ b/internal/store/lexical_review_test.go @@ -0,0 +1,121 @@ +package store + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestLexicalWorkerEnvironmentDropsParentSecrets(t *testing.T) { + environment := lexicalWorkerEnvironment([]string{ + "PATH=/usr/bin", + "HOME=/tmp/home", + "LANG=en_US.UTF-8", + "VIRTUAL_ENV=/tmp/venv", + "DISCORD_BOT_TOKEN=secret", + "OPENAI_API_KEY=secret", + "PIP_INDEX_URL=https://user:password@example.invalid/simple", + }) + joined := strings.Join(environment, "\n") + require.Contains(t, joined, "PATH=/usr/bin") + require.Contains(t, joined, "HOME=/tmp/home") + require.Contains(t, joined, "LANG=en_US.UTF-8") + require.NotContains(t, joined, "VIRTUAL_ENV") + require.NotContains(t, joined, "PYTHONNOUSERSITE") + require.NotContains(t, joined, "DISCORD_BOT_TOKEN") + require.NotContains(t, joined, "OPENAI_API_KEY") + require.NotContains(t, joined, "PIP_INDEX_URL") + require.NotContains(t, joined, "password") +} + +func TestKiwiCommandUsesGoHelperAndConfiguredModel(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + command, err := newKiwiLexicalCommand( + "/opt/discrawl/bin/discrawl-kiwi", + "~/models/kiwi/base", + ) + require.NoError(t, err) + require.Equal(t, "/opt/discrawl/bin/discrawl-kiwi", command.Path) + require.Equal(t, []string{ + "/opt/discrawl/bin/discrawl-kiwi", + "--model", + filepath.Join(home, "models/kiwi/base"), + }, command.Args) +} + +func TestHelperCommandsRejectArbitraryRelativeCommands(t *testing.T) { + _, err := newKiwiLexicalCommand("sh", "/tmp/model") + require.ErrorContains(t, err, "unsupported discrawl-kiwi helper") + _, err = newHelperLexicalCommand("sh", "discrawl-ja", nil) + require.ErrorContains(t, err, "unsupported discrawl-ja helper") + _, err = newHelperLexicalCommand("sh", "discrawl-zh", nil) + require.ErrorContains(t, err, "unsupported discrawl-zh helper") +} + +func TestKiwiTokenizerCommandProtocol(t *testing.T) { + tokenizer, err := startKiwiLexicalTokenizerCommand(lexicalHelperCommand("ready")) + require.NoError(t, err) + tokens, err := tokenizer.Tokenize(context.Background(), "오늘 저녁먹음 기록") + require.NoError(t, err) + require.Equal(t, "오늘 저녁 먹 음 기록", tokens) + require.NoError(t, tokenizer.Close()) + _, err = tokenizer.Tokenize(context.Background(), "text") + require.ErrorContains(t, err, "write ko tokenizer request") +} + +func TestKiwiTokenizerCommandStartupFailures(t *testing.T) { + tokenizer, err := startKiwiLexicalTokenizerCommand(lexicalHelperCommand("startup-error")) + require.Nil(t, tokenizer) + require.ErrorContains(t, err, "missing tokenizer package") + + tokenizer, err = startKiwiLexicalTokenizerCommand(lexicalHelperCommand("malformed-startup")) + require.Nil(t, tokenizer) + require.ErrorContains(t, err, "decode lexical tokenizer response") + + tokenizer, err = startKiwiLexicalTokenizerCommand(lexicalHelperCommand("stderr-startup")) + require.Nil(t, tokenizer) + require.ErrorContains(t, err, "tokenizer stderr") +} + +func TestKiwiTokenizerCommandResponseErrorAndCancellation(t *testing.T) { + tokenizer, err := startKiwiLexicalTokenizerCommand(lexicalHelperCommand("response-error")) + require.NoError(t, err) + defer func() { _ = tokenizer.Close() }() + + _, err = tokenizer.Tokenize(context.Background(), "text") + require.ErrorContains(t, err, "tokenization failed") + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err = tokenizer.Tokenize(ctx, "text") + require.ErrorIs(t, err, context.Canceled) +} + +func TestKiwiTokenizerCommandMalformedResponse(t *testing.T) { + tokenizer, err := startKiwiLexicalTokenizerCommand(lexicalHelperCommand("malformed-response")) + require.NoError(t, err) + defer func() { _ = tokenizer.Close() }() + + _, err = tokenizer.Tokenize(context.Background(), "text") + require.ErrorContains(t, err, "decode lexical tokenizer response") + + var nilTokenizer *externalLexicalTokenizer + require.NoError(t, nilTokenizer.Close()) +} + +func TestKiwiCommandDefaultHelperAndOptionalModel(t *testing.T) { + bin := t.TempDir() + helper := filepath.Join(bin, "discrawl-kiwi") + require.NoError(t, os.WriteFile(helper, []byte("#!/bin/sh\n"), 0o700)) + t.Setenv("PATH", bin) + + command, err := newKiwiLexicalCommand("", "") + require.NoError(t, err) + require.Equal(t, helper, command.Path) + require.Equal(t, []string{helper}, command.Args) +} diff --git a/internal/store/lexical_search.go b/internal/store/lexical_search.go new file mode 100644 index 00000000..fd4ffaed --- /dev/null +++ b/internal/store/lexical_search.go @@ -0,0 +1,150 @@ +package store + +import ( + "context" + "fmt" + "strings" + + "github.com/openclaw/crawlkit/vector" +) + +func (s *Store) searchMessagesMultilingual( + ctx context.Context, + opts SearchOptions, +) ([]SearchResult, error) { + if strings.TrimSpace(opts.Query) == "" { + return nil, nil + } + if opts.Limit <= 0 { + opts.Limit = 20 + } + candidateLimit := searchCandidateLimit(opts.Limit) + rankings := make([][]SearchResult, 0, len(s.lexicalTokenizers)+1) + + defaultResults, err := s.searchMessagesFTSTable( + ctx, + "message_fts", + normalizeFTSQuery(opts.Query), + opts, + candidateLimit, + ) + if err != nil { + if !shouldSearchFallback(err) { + return nil, err + } + return s.searchFallback(ctx, opts) + } + rankings = append(rankings, defaultResults) + + for _, language := range s.lexicalLanguages() { + query, err := s.lexicalTokenizers[language].Tokenize(ctx, opts.Query) + if err != nil { + return nil, fmt.Errorf("tokenize %s query: %w", language, err) + } + query = normalizeFTSQuery(query) + if query == "" { + continue + } + results, err := s.searchMessagesFTSTable( + ctx, + lexicalFTSTable(language), + query, + opts, + candidateLimit, + ) + if err != nil { + return nil, err + } + rankings = append(rankings, results) + } + return fuseLexicalSearchResults(rankings, opts.Limit), nil +} + +func (s *Store) searchMessagesFTSTable( + ctx context.Context, + table string, + queryText string, + opts SearchOptions, + limit int, +) ([]SearchResult, error) { + args := []any{queryText} + clauses := []string{table + " match ?"} + if len(opts.GuildIDs) > 0 { + clauses = append(clauses, table+".guild_id in ("+placeholders(len(opts.GuildIDs))+")") + for _, guildID := range opts.GuildIDs { + args = append(args, guildID) + } + } + if strings.TrimSpace(opts.Channel) != "" { + clauses = append(clauses, "("+table+".channel_id = ? or "+table+".channel_name like ?)") + args = append(args, opts.Channel, "%"+opts.Channel+"%") + } + if strings.TrimSpace(opts.Author) != "" { + clauses = append(clauses, "("+table+".author_id = ? or "+table+".author_name like ?)") + args = append(args, opts.Author, "%"+opts.Author+"%") + } + if !opts.IncludeEmpty { + clauses = append(clauses, "trim(coalesce(m.normalized_content, '')) <> ''") + } + args = append(args, limit) + queryCtx, cancel := withQueryTimeout(ctx) + defer cancel() + rows, err := s.db.QueryContext(queryCtx, ` + select `+table+`.message_id + from `+table+` + join messages m on m.id = `+table+`.message_id + where m.deleted_at is null + and `+strings.Join(clauses, " and ")+` + order by bm25(`+table+`) asc, `+table+`.rowid desc + limit ? + `, args...) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + ids := make([]string, 0, limit) + for rows.Next() { + var messageID string + if err := rows.Scan(&messageID); err != nil { + return nil, err + } + ids = append(ids, messageID) + } + if err := rows.Err(); err != nil { + return nil, err + } + details, err := s.searchResultDetails(queryCtx, ids) + if err != nil { + return nil, err + } + results := make([]SearchResult, 0, len(ids)) + for _, messageID := range ids { + if result, ok := details[messageID]; ok { + results = append(results, result) + } + } + return results, nil +} + +func fuseLexicalSearchResults(rankings [][]SearchResult, limit int) []SearchResult { + if limit <= 0 { + limit = 20 + } + ids := make([]func(SearchResult) string, len(rankings)) + weights := make([]float64, len(rankings)) + for i := range rankings { + ids[i] = func(result SearchResult) string { + return result.MessageID + } + weights[i] = 1 + } + fused := vector.ReciprocalRankFusion(rankings, ids, weights, rrfK) + if len(fused) > limit { + fused = fused[:limit] + } + results := make([]SearchResult, 0, len(fused)) + for _, entry := range fused { + results = append(results, entry.Item) + } + return results +} diff --git a/internal/store/lexical_tokenizers.go b/internal/store/lexical_tokenizers.go new file mode 100644 index 00000000..d209560b --- /dev/null +++ b/internal/store/lexical_tokenizers.go @@ -0,0 +1,74 @@ +package store + +import ( + "context" + "fmt" +) + +type OpenOptions struct { + LexicalLanguages []string + LexicalKiwiCommand string + LexicalKiwiModel string + LexicalJaCommand string + LexicalZhCommand string +} + +func OpenWithOptions(ctx context.Context, path string, opts OpenOptions) (*Store, error) { + tokenizers, err := newLexicalTokenizers(opts) + if err != nil { + return nil, err + } + return openWithLexicalTokenizers(ctx, path, tokenizers) +} + +func OpenReadOnlyWithOptions(ctx context.Context, path string, opts OpenOptions) (*Store, error) { + tokenizers, err := newLexicalTokenizers(opts) + if err != nil { + return nil, err + } + return openReadOnlyWithLexicalTokenizers(ctx, path, tokenizers) +} + +func newLexicalTokenizers(opts OpenOptions) (map[string]LexicalTokenizer, error) { + if len(opts.LexicalLanguages) == 0 { + return nil, nil + } + tokenizers := make(map[string]LexicalTokenizer, len(opts.LexicalLanguages)) + for _, language := range opts.LexicalLanguages { + switch language { + case "ko": + command := opts.LexicalKiwiCommand + model := opts.LexicalKiwiModel + tokenizers[language] = newLazyLexicalTokenizer(func() (LexicalTokenizer, error) { + tokenizer, err := startKiwiLexicalTokenizer(command, model) + if err != nil { + return nil, fmt.Errorf("start ko lexical tokenizer: %w", err) + } + return tokenizer, nil + }) + case "ja": + command := opts.LexicalJaCommand + tokenizers[language] = newLazyLexicalTokenizer(func() (LexicalTokenizer, error) { + tokenizer, err := startHelperLexicalTokenizer("ja", command, "discrawl-ja", nil) + if err != nil { + return nil, fmt.Errorf("start ja lexical tokenizer: %w", err) + } + return tokenizer, nil + }) + case "zh": + command := opts.LexicalZhCommand + tokenizers[language] = newLazyLexicalTokenizer(func() (LexicalTokenizer, error) { + tokenizer, err := startHelperLexicalTokenizer("zh", command, "discrawl-zh", nil) + if err != nil { + return nil, fmt.Errorf("start zh lexical tokenizer: %w", err) + } + return tokenizer, nil + }) + case "ar": + tokenizers[language] = newArabicLexicalTokenizer() + default: + return nil, fmt.Errorf("unsupported lexical language %q", language) + } + } + return tokenizers, nil +} diff --git a/internal/store/multilingual_benchmark_test.go b/internal/store/multilingual_benchmark_test.go new file mode 100644 index 00000000..82b85bb2 --- /dev/null +++ b/internal/store/multilingual_benchmark_test.go @@ -0,0 +1,134 @@ +package store + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type multilingualBenchmarkCase struct { + language string + id string + content string + query string +} + +func TestMultilingualLexicalQualityBenchmark(t *testing.T) { + if os.Getenv("DISCRAWL_TOKENIZER_E2E") != "1" { + t.Skip("set DISCRAWL_TOKENIZER_E2E=1 with optional tokenizer packages installed") + } + ctx := context.Background() + root := t.TempDir() + baseline, err := Open(ctx, filepath.Join(root, "baseline.db")) + require.NoError(t, err) + defer func() { _ = baseline.Close() }() + multilingual, err := OpenWithOptions(ctx, filepath.Join(root, "multilingual.db"), OpenOptions{ + LexicalLanguages: []string{"ko", "ja", "zh", "ar"}, + LexicalKiwiCommand: os.Getenv("DISCRAWL_KIWI_HELPER"), + LexicalKiwiModel: os.Getenv("DISCRAWL_KIWI_MODEL"), + LexicalJaCommand: os.Getenv("DISCRAWL_JA_HELPER"), + LexicalZhCommand: os.Getenv("DISCRAWL_ZH_HELPER"), + }) + require.NoError(t, err) + defer func() { _ = multilingual.Close() }() + + cases := multilingualBenchmarkCases() + base := time.Date(2026, 8, 18, 12, 0, 0, 0, time.UTC) + for i, item := range cases { + message := MessageRecord{ + ID: item.id, + GuildID: "g1", + ChannelID: "c1", + CreatedAt: base.Add(time.Duration(i) * time.Minute).Format(time.RFC3339Nano), + Content: item.content, + NormalizedContent: item.content, + RawJSON: `{}`, + } + require.NoError(t, baseline.UpsertMessage(ctx, message)) + require.NoError(t, multilingual.UpsertMessage(ctx, message)) + } + baselineBytes := sqliteDatabaseBytes(t, baseline) + multilingualBytes := sqliteDatabaseBytes(t, multilingual) + t.Logf( + "database pages: unicode61=%d bytes multilingual=%d bytes (%.2fx)", + baselineBytes, + multilingualBytes, + float64(multilingualBytes)/float64(baselineBytes), + ) + + baselineHits := make(map[string]int) + multilingualHits := make(map[string]int) + totals := make(map[string]int) + for _, item := range cases { + totals[item.language]++ + baselineResults, err := baseline.SearchMessages(ctx, SearchOptions{Query: item.query, Limit: 5}) + require.NoError(t, err) + if containsSearchResult(baselineResults, item.id) { + baselineHits[item.language]++ + } + multilingualResults, err := multilingual.SearchMessages(ctx, SearchOptions{Query: item.query, Limit: 5}) + require.NoError(t, err) + if containsSearchResult(multilingualResults, item.id) { + multilingualHits[item.language]++ + } + } + for _, language := range []string{"ko", "ja", "zh", "ar"} { + t.Logf( + "%s recall@5: unicode61=%d/%d multilingual=%d/%d", + language, + baselineHits[language], + totals[language], + multilingualHits[language], + totals[language], + ) + require.Greater(t, multilingualHits[language], baselineHits[language], language) + require.Equal(t, totals[language], multilingualHits[language], language) + } +} + +func sqliteDatabaseBytes(t *testing.T, s *Store) int64 { + t.Helper() + var pageCount int64 + var pageSize int64 + require.NoError(t, s.DB().QueryRowContext(t.Context(), `pragma page_count`).Scan(&pageCount)) + require.NoError(t, s.DB().QueryRowContext(t.Context(), `pragma page_size`).Scan(&pageSize)) + return pageCount * pageSize +} + +func containsSearchResult(results []SearchResult, messageID string) bool { + for _, result := range results { + if result.MessageID == messageID { + return true + } + } + return false +} + +func multilingualBenchmarkCases() []multilingualBenchmarkCase { + return []multilingualBenchmarkCase{ + {language: "ko", id: "ko-1", content: "오늘저녁먹음", query: "저녁"}, + {language: "ko", id: "ko-2", content: "서울맛집추천", query: "맛집"}, + {language: "ko", id: "ko-3", content: "프로젝트검색기능", query: "검색"}, + {language: "ko", id: "ko-4", content: "회의기록정리", query: "기록"}, + {language: "ko", id: "ko-5", content: "운동계획세움", query: "계획"}, + {language: "ja", id: "ja-1", content: "東京都庁に行きます", query: "東京"}, + {language: "ja", id: "ja-2", content: "自然言語処理を学ぶ", query: "言語"}, + {language: "ja", id: "ja-3", content: "検索機能を改善する", query: "検索"}, + {language: "ja", id: "ja-4", content: "会議記録を整理する", query: "記録"}, + {language: "ja", id: "ja-5", content: "機械学習モデル", query: "学習"}, + {language: "zh", id: "zh-1", content: "自然语言处理很有趣", query: "语言"}, + {language: "zh", id: "zh-2", content: "北京大学校园很美", query: "大学"}, + {language: "zh", id: "zh-3", content: "搜索功能需要改进", query: "搜索"}, + {language: "zh", id: "zh-4", content: "会议记录已经完成", query: "记录"}, + {language: "zh", id: "zh-5", content: "机器学习模型上线", query: "学习"}, + {language: "ar", id: "ar-1", content: "والكتاب مفيد للطلاب", query: "كتاب"}, + {language: "ar", id: "ar-2", content: "والمدرسة تفتح صباحا", query: "مدرسة"}, + {language: "ar", id: "ar-3", content: "فالاجتماع مهم اليوم", query: "اجتماع"}, + {language: "ar", id: "ar-4", content: "بالسجل تفاصيل كاملة", query: "سجل"}, + {language: "ar", id: "ar-5", content: "كالبرنامج سريع جدا", query: "برنامج"}, + } +} diff --git a/internal/store/multilingual_e2e_test.go b/internal/store/multilingual_e2e_test.go new file mode 100644 index 00000000..d478a33f --- /dev/null +++ b/internal/store/multilingual_e2e_test.go @@ -0,0 +1,55 @@ +package store + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestMultilingualLexicalSearchE2E(t *testing.T) { + if os.Getenv("DISCRAWL_TOKENIZER_E2E") != "1" { + t.Skip("set DISCRAWL_TOKENIZER_E2E=1 with optional tokenizer packages installed") + } + ctx := context.Background() + s, err := OpenWithOptions(ctx, filepath.Join(t.TempDir(), "discrawl.db"), OpenOptions{ + LexicalLanguages: []string{"ko", "ja", "zh", "ar"}, + LexicalKiwiCommand: os.Getenv("DISCRAWL_KIWI_HELPER"), + LexicalKiwiModel: os.Getenv("DISCRAWL_KIWI_MODEL"), + LexicalJaCommand: os.Getenv("DISCRAWL_JA_HELPER"), + LexicalZhCommand: os.Getenv("DISCRAWL_ZH_HELPER"), + }) + require.NoError(t, err) + defer func() { _ = s.Close() }() + + base := time.Date(2026, 8, 18, 12, 0, 0, 0, time.UTC) + fixtures := []struct { + id string + content string + query string + }{ + {id: "ko", content: "오늘 저녁먹음 기록", query: "저녁"}, + {id: "ja", content: "東京都庁に行きます", query: "東京"}, + {id: "zh", content: "自然语言处理很有趣", query: "语言"}, + {id: "ar", content: "والكتاب مفيد للطلاب", query: "كتاب"}, + } + for i, fixture := range fixtures { + require.NoError(t, s.UpsertMessage(ctx, MessageRecord{ + ID: fixture.id, + GuildID: "g1", + ChannelID: "c1", + CreatedAt: base.Add(time.Duration(i) * time.Minute).Format(time.RFC3339Nano), + Content: fixture.content, + NormalizedContent: fixture.content, + RawJSON: `{}`, + })) + } + for _, fixture := range fixtures { + results, err := s.SearchMessages(ctx, SearchOptions{Query: fixture.query, Limit: 10}) + require.NoError(t, err, fixture.id) + require.Contains(t, searchResultIDs(results), fixture.id, fixture.query) + } +} diff --git a/internal/store/multilingual_edges_test.go b/internal/store/multilingual_edges_test.go new file mode 100644 index 00000000..bf271f3d --- /dev/null +++ b/internal/store/multilingual_edges_test.go @@ -0,0 +1,101 @@ +package store + +import ( + "context" + "errors" + "fmt" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestMultilingualSearchPreservesMetadataFilters(t *testing.T) { + ctx := context.Background() + s, err := openWithLexicalTokenizers(ctx, filepath.Join(t.TempDir(), "discrawl.db"), map[string]LexicalTokenizer{ + "ko": stubLexicalTokenizer{tokenize: func(text string) string { return text }}, + }) + require.NoError(t, err) + defer func() { _ = s.Close() }() + + now := time.Now().UTC() + require.NoError(t, s.UpsertMessage(ctx, MessageRecord{ + ID: "wanted", GuildID: "g1", ChannelID: "c1", ChannelName: "alpha", + AuthorID: "u1", AuthorName: "alice", CreatedAt: now.Format(time.RFC3339Nano), + Content: "needle", NormalizedContent: "needle", RawJSON: `{}`, + })) + require.NoError(t, s.UpsertMessage(ctx, MessageRecord{ + ID: "other", GuildID: "g2", ChannelID: "c2", ChannelName: "beta", + AuthorID: "u2", AuthorName: "bob", CreatedAt: now.Add(time.Minute).Format(time.RFC3339Nano), + Content: "needle", NormalizedContent: "needle", RawJSON: `{}`, + })) + + results, err := s.SearchMessages(ctx, SearchOptions{ + Query: "needle", GuildIDs: []string{"g1"}, Channel: "alpha", Author: "alice", + }) + require.NoError(t, err) + require.Equal(t, []string{"wanted"}, searchResultIDs(results)) + + results, err = s.SearchMessages(ctx, SearchOptions{ + Query: "needle", GuildIDs: []string{"g1"}, Channel: "missing", + }) + require.NoError(t, err) + require.Empty(t, results) +} + +func TestMultilingualSearchReportsQueryTokenizerFailure(t *testing.T) { + ctx := context.Background() + s, err := openWithLexicalTokenizers(ctx, filepath.Join(t.TempDir(), "discrawl.db"), map[string]LexicalTokenizer{ + "ko": stubLexicalTokenizer{tokenize: func(text string) string { return text }}, + }) + require.NoError(t, err) + defer func() { _ = s.Close() }() + s.lexicalTokenizers["ko"] = failingLexicalTokenizer{err: errors.New("query tokenizer failed")} + + _, err = s.SearchMessages(ctx, SearchOptions{Query: "needle"}) + require.ErrorContains(t, err, "tokenize ko query") +} + +func TestMultilingualSearchFallsBackWhenDefaultFTSIsMissing(t *testing.T) { + ctx := context.Background() + s, err := openWithLexicalTokenizers(ctx, filepath.Join(t.TempDir(), "discrawl.db"), map[string]LexicalTokenizer{ + "ko": stubLexicalTokenizer{tokenize: func(text string) string { return text }}, + }) + require.NoError(t, err) + defer func() { _ = s.Close() }() + require.NoError(t, s.UpsertMessage(ctx, MessageRecord{ + ID: "fallback", GuildID: "g1", ChannelID: "c1", + CreatedAt: time.Now().UTC().Format(time.RFC3339Nano), + Content: "needle", NormalizedContent: "needle", RawJSON: `{}`, + })) + _, err = s.DB().ExecContext(ctx, `drop table message_fts`) + require.NoError(t, err) + + results, err := s.SearchMessages(ctx, SearchOptions{Query: "needle", Limit: 10}) + require.NoError(t, err) + require.Equal(t, []string{"fallback"}, searchResultIDs(results)) +} + +func TestMultilingualDeleteRejectsUnsafeColumn(t *testing.T) { + ctx := context.Background() + s, err := openWithLexicalTokenizers(ctx, filepath.Join(t.TempDir(), "discrawl.db"), nil) + require.NoError(t, err) + defer func() { _ = s.Close() }() + tx, err := s.DB().BeginTx(ctx, nil) + require.NoError(t, err) + defer rollback(tx) + + err = s.deleteLexicalMessagesTx(ctx, tx, "message_id", "unsafe") + require.ErrorContains(t, err, "unsupported lexical delete column") +} + +func TestFuseLexicalSearchResultsUsesDefaultLimit(t *testing.T) { + ranking := make([]SearchResult, 25) + for i := range ranking { + ranking[i].MessageID = fmt.Sprintf("message-%02d", i) + } + results := fuseLexicalSearchResults([][]SearchResult{ranking}, 0) + require.Len(t, results, 20) + require.Equal(t, "message-00", results[0].MessageID) +} diff --git a/internal/store/multilingual_lifecycle_test.go b/internal/store/multilingual_lifecycle_test.go new file mode 100644 index 00000000..d02f0248 --- /dev/null +++ b/internal/store/multilingual_lifecycle_test.go @@ -0,0 +1,272 @@ +package store + +import ( + "context" + "errors" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type failingLexicalTokenizer struct { + err error +} + +func (f failingLexicalTokenizer) Tokenize(context.Context, string) (string, error) { + return "", f.err +} + +func (failingLexicalTokenizer) Close() error { + return nil +} + +func TestMultilingualIndexesTrackBatchDeletesAndGuildPurge(t *testing.T) { + ctx := context.Background() + s, err := openWithLexicalTokenizers(ctx, filepath.Join(t.TempDir(), "discrawl.db"), map[string]LexicalTokenizer{ + "ko": stubLexicalTokenizer{tokenize: replaceLexicalTerms(map[string]string{ + "저녁먹음": "저녁 먹 음", + "회의기록": "회의 기록", + })}, + }) + require.NoError(t, err) + defer func() { _ = s.Close() }() + + now := time.Now().UTC() + require.NoError(t, s.UpsertMessages(ctx, []MessageMutation{ + {Record: MessageRecord{ + ID: "first", GuildID: "g1", ChannelID: "c1", + CreatedAt: now.Format(time.RFC3339Nano), + Content: "저녁먹음", NormalizedContent: "저녁먹음", RawJSON: `{}`, + }}, + {Record: MessageRecord{ + ID: "second", GuildID: "g2", ChannelID: "c2", + CreatedAt: now.Add(time.Minute).Format(time.RFC3339Nano), + Content: "회의기록", NormalizedContent: "회의기록", RawJSON: `{}`, + }}, + })) + + results, err := s.SearchMessages(ctx, SearchOptions{Query: "저녁", Limit: 10}) + require.NoError(t, err) + require.Equal(t, []string{"first"}, searchResultIDs(results)) + require.NoError(t, s.MarkMessageDeleted( + ctx, + "g1", + "c1", + "first", + map[string]string{"deleted_at": now.Add(time.Hour).Format(time.RFC3339Nano)}, + )) + results, err = s.SearchMessages(ctx, SearchOptions{Query: "저녁", Limit: 10}) + require.NoError(t, err) + require.Empty(t, results) + + results, err = s.SearchMessages(ctx, SearchOptions{Query: "기록", Limit: 10}) + require.NoError(t, err) + require.Equal(t, []string{"second"}, searchResultIDs(results)) + require.NoError(t, s.DeleteGuildData(ctx, "g2")) + results, err = s.SearchMessages(ctx, SearchOptions{Query: "기록", Limit: 10}) + require.NoError(t, err) + require.Empty(t, results) +} + +func TestDisabledLexicalIndexDoesNotRetainPurgedGuildData(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "discrawl.db") + enabled, err := openWithLexicalTokenizers(ctx, path, map[string]LexicalTokenizer{ + "ko": stubLexicalTokenizer{tokenize: replaceLexicalTerms(map[string]string{ + "회의기록": "회의 기록", + })}, + }) + require.NoError(t, err) + require.NoError(t, enabled.UpsertMessage(ctx, MessageRecord{ + ID: "message", GuildID: "guild", ChannelID: "channel", + CreatedAt: time.Now().UTC().Format(time.RFC3339Nano), + Content: "회의기록", NormalizedContent: "회의기록", RawJSON: `{}`, + })) + require.NoError(t, enabled.Close()) + + disabled, err := openWithLexicalTokenizers(ctx, path, nil) + require.NoError(t, err) + defer func() { _ = disabled.Close() }() + require.NoError(t, disabled.DeleteGuildData(ctx, "guild")) + + var tables int + require.NoError(t, disabled.DB().QueryRowContext( + ctx, + `select count(*) from sqlite_schema where type = 'table' and name = 'message_fts_ko'`, + ).Scan(&tables)) + require.Zero(t, tables) +} + +func TestMultilingualIndexVersionSurvivesReopen(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "discrawl.db") + tokenizers := func() map[string]LexicalTokenizer { + return map[string]LexicalTokenizer{ + "ko": stubLexicalTokenizer{tokenize: replaceLexicalTerms(map[string]string{ + "저녁먹음": "저녁 먹 음", + })}, + } + } + s, err := openWithLexicalTokenizers(ctx, path, tokenizers()) + require.NoError(t, err) + require.NoError(t, s.UpsertMessage(ctx, MessageRecord{ + ID: "ko", GuildID: "g1", ChannelID: "c1", + CreatedAt: time.Now().UTC().Format(time.RFC3339Nano), + Content: "저녁먹음", NormalizedContent: "저녁먹음", RawJSON: `{}`, + })) + require.NoError(t, s.Close()) + + s, err = openWithLexicalTokenizers(ctx, path, tokenizers()) + require.NoError(t, err) + defer func() { _ = s.Close() }() + results, err := s.SearchMessages(ctx, SearchOptions{Query: "저녁", Limit: 10}) + require.NoError(t, err) + require.Equal(t, []string{"ko"}, searchResultIDs(results)) +} + +func TestMultilingualIndexRebuildsAfterDisabledWrites(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "discrawl.db") + tokenizers := func() map[string]LexicalTokenizer { + return map[string]LexicalTokenizer{ + "ko": stubLexicalTokenizer{tokenize: replaceLexicalTerms(map[string]string{ + "저녁먹음": "저녁 먹 음", + "회의기록": "회의 기록", + })}, + } + } + + enabled, err := openWithLexicalTokenizers(ctx, path, tokenizers()) + require.NoError(t, err) + require.NoError(t, enabled.UpsertMessage(ctx, MessageRecord{ + ID: "before", GuildID: "g1", ChannelID: "c1", + CreatedAt: time.Now().UTC().Format(time.RFC3339Nano), + Content: "저녁먹음", NormalizedContent: "저녁먹음", RawJSON: `{}`, + })) + require.NoError(t, enabled.Close()) + + disabled, err := openWithLexicalTokenizers(ctx, path, nil) + require.NoError(t, err) + require.NoError(t, disabled.UpsertMessage(ctx, MessageRecord{ + ID: "during", GuildID: "g1", ChannelID: "c1", + CreatedAt: time.Now().UTC().Add(time.Minute).Format(time.RFC3339Nano), + Content: "회의기록", NormalizedContent: "회의기록", RawJSON: `{}`, + })) + require.NoError(t, disabled.Close()) + + reenabled, err := openWithLexicalTokenizers(ctx, path, tokenizers()) + require.NoError(t, err) + defer func() { _ = reenabled.Close() }() + results, err := reenabled.SearchMessages(ctx, SearchOptions{Query: "기록", Limit: 10}) + require.NoError(t, err) + require.Equal(t, []string{"during"}, searchResultIDs(results)) +} + +func TestMultilingualIndexesSearchThroughReadOnlyStore(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "discrawl.db") + tokenizers := func() map[string]LexicalTokenizer { + return map[string]LexicalTokenizer{ + "ko": stubLexicalTokenizer{tokenize: replaceLexicalTerms(map[string]string{ + "저녁먹음": "저녁 먹 음", + })}, + } + } + writer, err := openWithLexicalTokenizers(ctx, path, tokenizers()) + require.NoError(t, err) + require.NoError(t, writer.UpsertMessage(ctx, MessageRecord{ + ID: "ko", GuildID: "g1", ChannelID: "c1", + CreatedAt: time.Now().UTC().Format(time.RFC3339Nano), + Content: "저녁먹음", NormalizedContent: "저녁먹음", RawJSON: `{}`, + })) + require.NoError(t, writer.Close()) + + reader, err := openReadOnlyWithLexicalTokenizers(ctx, path, tokenizers()) + require.NoError(t, err) + defer func() { _ = reader.Close() }() + results, err := reader.SearchMessages(ctx, SearchOptions{Query: "저녁", Limit: 10}) + require.NoError(t, err) + require.Equal(t, []string{"ko"}, searchResultIDs(results)) +} + +func TestOpenReadOnlyWithOptionsWithoutLexicalLanguages(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "discrawl.db") + writer, err := Open(ctx, path) + require.NoError(t, err) + require.NoError(t, writer.Close()) + + reader, err := OpenReadOnlyWithOptions(ctx, path, OpenOptions{}) + require.NoError(t, err) + require.NoError(t, reader.Close()) +} + +func TestOpenReadOnlyWithOptionsKeepsMissingTokenizerLazy(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "discrawl.db") + writer, err := Open(ctx, path) + require.NoError(t, err) + require.NoError(t, writer.Close()) + + reader, err := OpenReadOnlyWithOptions(ctx, path, OpenOptions{ + LexicalLanguages: []string{"ko"}, + LexicalKiwiCommand: "/definitely/missing/discrawl-kiwi", + LexicalKiwiModel: "/definitely/missing/kiwi-model", + }) + require.NoError(t, err) + require.NoError(t, reader.Close()) +} + +func TestMultilingualTokenizerFailureAbortsWrite(t *testing.T) { + ctx := context.Background() + s, err := openWithLexicalTokenizers(ctx, filepath.Join(t.TempDir(), "discrawl.db"), nil) + require.NoError(t, err) + defer func() { _ = s.Close() }() + s.lexicalTokenizers = map[string]LexicalTokenizer{ + "ko": failingLexicalTokenizer{err: errors.New("tokenizer unavailable")}, + } + + err = s.UpsertMessage(ctx, MessageRecord{ + ID: "ko", GuildID: "g1", ChannelID: "c1", + CreatedAt: time.Now().UTC().Format(time.RFC3339Nano), + Content: "저녁먹음", NormalizedContent: "저녁먹음", RawJSON: `{}`, + }) + require.ErrorContains(t, err, "tokenize ko text") + var count int + require.NoError(t, s.DB().QueryRowContext(ctx, `select count(*) from messages`).Scan(&count)) + require.Zero(t, count) +} + +func TestMultilingualDeleteUpsertDoesNotRequireTokenizer(t *testing.T) { + ctx := context.Background() + s, err := openWithLexicalTokenizers(ctx, filepath.Join(t.TempDir(), "discrawl.db"), map[string]LexicalTokenizer{ + "ko": stubLexicalTokenizer{tokenize: replaceLexicalTerms(map[string]string{ + "저녁먹음": "저녁 먹 음", + })}, + }) + require.NoError(t, err) + defer func() { _ = s.Close() }() + + message := MessageRecord{ + ID: "ko", GuildID: "g1", ChannelID: "c1", + CreatedAt: time.Now().UTC().Format(time.RFC3339Nano), + Content: "저녁먹음", NormalizedContent: "저녁먹음", RawJSON: `{}`, + } + require.NoError(t, s.UpsertMessage(ctx, message)) + s.lexicalTokenizers["ko"] = failingLexicalTokenizer{err: errors.New("tokenizer unavailable")} + + message.DeletedAt = time.Now().UTC().Add(time.Minute).Format(time.RFC3339Nano) + require.NoError(t, s.UpsertMessage(ctx, message)) + var lexicalRows int + require.NoError(t, s.DB().QueryRowContext(ctx, `select count(*) from message_fts_ko`).Scan(&lexicalRows)) + require.Zero(t, lexicalRows) + s.lexicalTokenizers["ko"] = stubLexicalTokenizer{tokenize: replaceLexicalTerms(map[string]string{ + "저녁먹음": "저녁 먹 음", + })} + + results, err := s.SearchMessages(ctx, SearchOptions{Query: "저녁", Limit: 10}) + require.NoError(t, err) + require.Empty(t, results) +} diff --git a/internal/store/multilingual_search_test.go b/internal/store/multilingual_search_test.go new file mode 100644 index 00000000..1b96ddc1 --- /dev/null +++ b/internal/store/multilingual_search_test.go @@ -0,0 +1,148 @@ +package store + +import ( + "context" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type stubLexicalTokenizer struct { + tokenize func(string) string +} + +func (s stubLexicalTokenizer) Tokenize(_ context.Context, text string) (string, error) { + return s.tokenize(text), nil +} + +func (stubLexicalTokenizer) Close() error { + return nil +} + +func TestSearchMessagesMultilingualIndexesEachConfiguredAnalyzer(t *testing.T) { + ctx := context.Background() + tokenizers := map[string]LexicalTokenizer{ + "ko": stubLexicalTokenizer{tokenize: replaceLexicalTerms(map[string]string{ + "저녁먹음": "저녁 먹 음", + })}, + "ja": stubLexicalTokenizer{tokenize: replaceLexicalTerms(map[string]string{ + "東京都庁": "東京 都庁", + })}, + "zh": stubLexicalTokenizer{tokenize: replaceLexicalTerms(map[string]string{ + "自然语言处理": "自然 语言 处理", + })}, + "ar": stubLexicalTokenizer{tokenize: replaceLexicalTerms(map[string]string{ + "والكتاب": "و ال كتاب", + })}, + } + s, err := openWithLexicalTokenizers(ctx, filepath.Join(t.TempDir(), "discrawl.db"), tokenizers) + require.NoError(t, err) + defer func() { _ = s.Close() }() + + base := time.Date(2026, 8, 18, 12, 0, 0, 0, time.UTC) + messages := []MessageRecord{ + {ID: "ko", GuildID: "g1", ChannelID: "c1", CreatedAt: base.Format(time.RFC3339Nano), Content: "저녁먹음", NormalizedContent: "저녁먹음", RawJSON: `{}`}, + {ID: "ja", GuildID: "g1", ChannelID: "c1", CreatedAt: base.Add(time.Minute).Format(time.RFC3339Nano), Content: "東京都庁", NormalizedContent: "東京都庁", RawJSON: `{}`}, + {ID: "zh", GuildID: "g1", ChannelID: "c1", CreatedAt: base.Add(2 * time.Minute).Format(time.RFC3339Nano), Content: "自然语言处理", NormalizedContent: "自然语言处理", RawJSON: `{}`}, + {ID: "ar", GuildID: "g1", ChannelID: "c1", CreatedAt: base.Add(3 * time.Minute).Format(time.RFC3339Nano), Content: "والكتاب", NormalizedContent: "والكتاب", RawJSON: `{}`}, + } + for _, message := range messages { + require.NoError(t, s.UpsertMessage(ctx, message)) + } + + for query, wantID := range map[string]string{ + "저녁": "ko", + "東京": "ja", + "语言": "zh", + "كتاب": "ar", + } { + results, err := s.SearchMessages(ctx, SearchOptions{Query: query, Limit: 10}) + require.NoError(t, err, query) + require.Equal(t, []string{wantID}, searchResultIDs(results), query) + } +} + +func TestSearchMessagesMultilingualRRFCombinesAndDeduplicates(t *testing.T) { + ctx := context.Background() + tokenizers := map[string]LexicalTokenizer{ + "ko": stubLexicalTokenizer{tokenize: func(text string) string { + text = strings.ReplaceAll(text, "검색", "shared") + return strings.ReplaceAll(text, "検索", "") + }}, + "ja": stubLexicalTokenizer{tokenize: func(text string) string { + text = strings.ReplaceAll(text, "検索", "shared") + return strings.ReplaceAll(text, "검색", "") + }}, + } + s, err := openWithLexicalTokenizers(ctx, filepath.Join(t.TempDir(), "discrawl.db"), tokenizers) + require.NoError(t, err) + defer func() { _ = s.Close() }() + + base := time.Date(2026, 8, 18, 12, 0, 0, 0, time.UTC) + require.NoError(t, s.UpsertMessage(ctx, MessageRecord{ + ID: "mixed", GuildID: "g1", ChannelID: "c1", CreatedAt: base.Format(time.RFC3339Nano), + Content: "검색 検索", NormalizedContent: "검색 検索", RawJSON: `{}`, + })) + require.NoError(t, s.UpsertMessage(ctx, MessageRecord{ + ID: "single", GuildID: "g1", ChannelID: "c1", CreatedAt: base.Add(time.Minute).Format(time.RFC3339Nano), + Content: "검색 only", NormalizedContent: "검색 only", RawJSON: `{}`, + })) + + results, err := s.SearchMessages(ctx, SearchOptions{Query: "검색 検索", Limit: 10}) + require.NoError(t, err) + require.Equal(t, []string{"mixed", "single"}, searchResultIDs(results)) +} + +func TestRebuildSearchIndexesRebuildsMultilingualTables(t *testing.T) { + ctx := context.Background() + s, err := openWithLexicalTokenizers(ctx, filepath.Join(t.TempDir(), "discrawl.db"), map[string]LexicalTokenizer{ + "ko": stubLexicalTokenizer{tokenize: replaceLexicalTerms(map[string]string{"저녁먹음": "저녁 먹 음"})}, + }) + require.NoError(t, err) + defer func() { _ = s.Close() }() + + require.NoError(t, s.UpsertMessage(ctx, MessageRecord{ + ID: "ko", GuildID: "g1", ChannelID: "c1", CreatedAt: time.Now().UTC().Format(time.RFC3339Nano), + Content: "저녁먹음", NormalizedContent: "저녁먹음", RawJSON: `{}`, + })) + _, err = s.DB().ExecContext(ctx, `delete from message_fts_ko`) + require.NoError(t, err) + require.NoError(t, s.RebuildSearchIndexes(ctx)) + + results, err := s.SearchMessages(ctx, SearchOptions{Query: "저녁", Limit: 10}) + require.NoError(t, err) + require.Equal(t, []string{"ko"}, searchResultIDs(results)) +} + +func TestSearchMessagesMultilingualHonorsIncludeEmpty(t *testing.T) { + ctx := context.Background() + s, err := openWithLexicalTokenizers(ctx, filepath.Join(t.TempDir(), "discrawl.db"), map[string]LexicalTokenizer{ + "ko": stubLexicalTokenizer{tokenize: func(text string) string { return text }}, + }) + require.NoError(t, err) + defer func() { _ = s.Close() }() + + require.NoError(t, s.UpsertMessage(ctx, MessageRecord{ + ID: "empty", GuildID: "g1", ChannelID: "c1", AuthorName: "needle", + CreatedAt: time.Now().UTC().Format(time.RFC3339Nano), RawJSON: `{}`, + })) + results, err := s.SearchMessages(ctx, SearchOptions{Query: "needle", Limit: 10}) + require.NoError(t, err) + require.Empty(t, results) + + results, err = s.SearchMessages(ctx, SearchOptions{Query: "needle", Limit: 10, IncludeEmpty: true}) + require.NoError(t, err) + require.Equal(t, []string{"empty"}, searchResultIDs(results)) +} + +func replaceLexicalTerms(replacements map[string]string) func(string) string { + return func(text string) string { + for from, to := range replacements { + text = strings.ReplaceAll(text, from, to) + } + return text + } +} diff --git a/internal/store/query.go b/internal/store/query.go index 94013dea..84fd0eda 100644 --- a/internal/store/query.go +++ b/internal/store/query.go @@ -105,6 +105,9 @@ func (s *Store) SearchMessages(ctx context.Context, opts SearchOptions) ([]Searc if opts.Limit <= 0 { opts.Limit = 20 } + if len(s.lexicalTokenizers) > 0 { + return s.searchMessagesMultilingual(ctx, opts) + } args := []any{normalizeFTSQuery(opts.Query)} clauses := []string{"message_fts match ?"} if len(opts.GuildIDs) > 0 { diff --git a/internal/store/store.go b/internal/store/store.go index a12586cf..903b0d86 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -23,9 +23,11 @@ const ( var ErrSchemaVersionMismatch = errors.New("database schema version mismatch") type Store struct { - db *sql.DB - q *storedb.Queries - path string + db *sql.DB + q *storedb.Queries + path string + baseClose func() error + lexicalTokenizers map[string]LexicalTokenizer } type CatalogIntegrity struct { @@ -132,31 +134,48 @@ type ChannelRow struct { } func Open(ctx context.Context, path string) (*Store, error) { + return openWithLexicalTokenizers(ctx, path, nil) +} + +func openBaseStore(ctx context.Context, path string) (*crawlstore.Store, error) { base, err := crawlstore.Open(ctx, crawlstore.Options{Path: path}) if err != nil { return nil, err } - db := base.DB() - store := &Store{db: db, q: storedb.New(db), path: path} - if err := store.migrate(ctx); err != nil { - _ = base.Close() - return nil, err - } - return store, nil + return base, nil +} + +func newStoreQueries(db *sql.DB) *storedb.Queries { + return storedb.New(db) } func OpenReadOnly(ctx context.Context, path string) (*Store, error) { + return openReadOnlyWithLexicalTokenizers(ctx, path, nil) +} + +func openReadOnlyWithLexicalTokenizers( + ctx context.Context, + path string, + tokenizers map[string]LexicalTokenizer, +) (*Store, error) { base, err := crawlstore.OpenReadOnly(ctx, path) if err != nil { + closeLexicalTokenizers(tokenizers) return nil, err } db := base.DB() - store := &Store{db: db, q: storedb.New(db), path: path} + store := &Store{ + db: db, + q: storedb.New(db), + path: path, + baseClose: base.Close, + lexicalTokenizers: tokenizers, + } if version, err := store.schemaVersion(ctx); err != nil { - _ = base.Close() + _ = store.Close() return nil, err } else if version != storeSchemaVersion { - _ = base.Close() + _ = store.Close() return nil, fmt.Errorf("%w: got %d want %d", ErrSchemaVersionMismatch, version, storeSchemaVersion) } return store, nil @@ -166,6 +185,10 @@ func (s *Store) Close() error { if s == nil || s.db == nil { return nil } + closeLexicalTokenizers(s.lexicalTokenizers) + if s.baseClose != nil { + return s.baseClose() + } return s.db.Close() } @@ -258,6 +281,9 @@ func (s *Store) RebuildSearchIndexes(ctx context.Context) error { if err := s.rebuildFTS(ctx); err != nil { return err } + if err := s.rebuildLexicalIndexes(ctx); err != nil { + return err + } if err := s.rebuildMemberFTS(ctx); err != nil { return err } @@ -268,6 +294,9 @@ func (s *Store) RebuildMessageSearchIndex(ctx context.Context) error { if err := s.rebuildFTS(ctx); err != nil { return err } + if err := s.rebuildLexicalIndexes(ctx); err != nil { + return err + } return s.stampSearchIndexVersions(ctx, true, false) } @@ -845,7 +874,7 @@ func (s *Store) rebuildFTS(ctx context.Context) error { } func configureFTSBulkLoad(ctx context.Context, tx *sql.Tx, table string) error { - if table != "message_fts" && table != "member_fts" { + if !isMessageFTSTable(table) && table != "member_fts" { return fmt.Errorf("unsupported fts table %q", table) } stmts := []string{ @@ -862,7 +891,7 @@ func configureFTSBulkLoad(ctx context.Context, tx *sql.Tx, table string) error { } func optimizeFTS(ctx context.Context, tx *sql.Tx, table string) error { - if table != "message_fts" && table != "member_fts" { + if !isMessageFTSTable(table) && table != "member_fts" { return fmt.Errorf("unsupported fts table %q", table) } if _, err := tx.ExecContext(ctx, fmt.Sprintf("insert into %s(%s) values('optimize')", table, table)); err != nil { diff --git a/internal/store/write.go b/internal/store/write.go index 70174fec..64e52569 100644 --- a/internal/store/write.go +++ b/internal/store/write.go @@ -227,6 +227,9 @@ func (s *Store) DeleteGuildData(ctx context.Context, guildID string) error { if _, err := tx.ExecContext(ctx, `delete from message_fts where guild_id = ?`, guildID); err != nil { return err } + if err := s.deleteLexicalMessagesTx(ctx, tx, "guild_id", guildID); err != nil { + return err + } if err := qtx.DeleteMessageEventsByGuild(ctx, guildID); err != nil { return err } @@ -263,12 +266,20 @@ func (s *Store) UpsertMessage(ctx context.Context, message MessageRecord) error } func (s *Store) UpsertMessageWithOptions(ctx context.Context, message MessageRecord, opts WriteOptions) error { + var tokenized map[string]string + if message.DeletedAt == "" { + var err error + tokenized, err = s.tokenizeLexical(ctx, message.NormalizedContent) + if err != nil { + return err + } + } tx, err := s.db.BeginTx(ctx, nil) if err != nil { return err } defer rollback(tx) - if err := upsertMessageTx(ctx, tx, s.q.WithTx(tx), message, opts); err != nil { + if err := s.upsertMessageTx(ctx, tx, s.q.WithTx(tx), message, opts, tokenized); err != nil { return err } return tx.Commit() @@ -278,17 +289,27 @@ func (s *Store) UpsertMessages(ctx context.Context, messages []MessageMutation) if len(messages) == 0 { return nil } + tokenized := make([]map[string]string, len(messages)) + for i, message := range messages { + if message.Record.DeletedAt == "" { + var err error + tokenized[i], err = s.tokenizeLexical(ctx, message.Record.NormalizedContent) + if err != nil { + return err + } + } + } tx, err := s.db.BeginTx(ctx, nil) if err != nil { return err } defer rollback(tx) qtx := s.q.WithTx(tx) - for _, message := range messages { + for i, message := range messages { if err := ctx.Err(); err != nil { return err } - if err := upsertMessageTx(ctx, tx, qtx, message.Record, message.Options); err != nil { + if err := s.upsertMessageTx(ctx, tx, qtx, message.Record, message.Options, tokenized[i]); err != nil { return err } if err := replaceAttachmentsTx(ctx, qtx, message.Record.ID, message.Attachments); err != nil { @@ -314,7 +335,14 @@ func (s *Store) UpsertMessages(ctx context.Context, messages []MessageMutation) return tx.Commit() } -func upsertMessageTx(ctx context.Context, tx *sql.Tx, qtx *storedb.Queries, message MessageRecord, opts WriteOptions) error { +func (s *Store) upsertMessageTx( + ctx context.Context, + tx *sql.Tx, + qtx *storedb.Queries, + message MessageRecord, + opts WriteOptions, + tokenized map[string]string, +) error { now := time.Now().UTC().Format(timeLayout) var previousNormalized sql.NullString previousErr := sql.ErrNoRows @@ -350,6 +378,9 @@ func upsertMessageTx(ctx context.Context, tx *sql.Tx, qtx *storedb.Queries, mess if _, err := tx.ExecContext(ctx, deleteMessageFTSByRowIDSQL, rowID); err != nil { return err } + if err := s.upsertLexicalMessageTx(ctx, tx, message, tokenized); err != nil { + return err + } if message.DeletedAt != "" { if err := qtx.DeleteMessageEmbeddingsByMessage(ctx, message.ID); err != nil { return err @@ -450,6 +481,9 @@ func (s *Store) markMessageDeleted( if _, err := tx.ExecContext(ctx, deleteMessageFTSByRowIDSQL, rowID); err != nil { return err } + if err := s.deleteLexicalMessagesTx(ctx, tx, "rowid", rowID); err != nil { + return err + } } if err := qtx.DeleteMessageEmbeddingsByMessage(ctx, messageID); err != nil { return err diff --git a/tools/discrawl-ja/README.md b/tools/discrawl-ja/README.md new file mode 100644 index 00000000..a281d082 --- /dev/null +++ b/tools/discrawl-ja/README.md @@ -0,0 +1,14 @@ +# discrawl-ja + +Optional Japanese lexical helper. It uses the pure-Go +[`github.com/ikawaha/kagome/v2`](https://pkg.go.dev/github.com/ikawaha/kagome/v2) +tokenizer in Search mode with the embedded MeCab-IPADIC dictionary. + +This binary is not linked into the default Discrawl release. Build it only when +Japanese lexical fields are enabled. + +```bash +go build -o discrawl-ja . +``` + +The helper speaks the same newline-delimited JSON protocol as `discrawl-kiwi`. diff --git a/tools/discrawl-ja/go.mod b/tools/discrawl-ja/go.mod new file mode 100644 index 00000000..4766c3e6 --- /dev/null +++ b/tools/discrawl-ja/go.mod @@ -0,0 +1,10 @@ +module github.com/openclaw/discrawl/tools/discrawl-ja + +go 1.26 + +require ( + github.com/ikawaha/kagome-dict/ipa v1.2.6 + github.com/ikawaha/kagome/v2 v2.11.0 +) + +require github.com/ikawaha/kagome-dict v1.1.7 // indirect diff --git a/tools/discrawl-ja/go.sum b/tools/discrawl-ja/go.sum new file mode 100644 index 00000000..709824ae --- /dev/null +++ b/tools/discrawl-ja/go.sum @@ -0,0 +1,6 @@ +github.com/ikawaha/kagome-dict v1.1.7 h1:O/uAL+WCGhp6kT0+szxBSPaSM4i+vdArSefFvJE4Nug= +github.com/ikawaha/kagome-dict v1.1.7/go.mod h1:9tvk7/jZkvYt40foxkB9CqSAAknoQrIPfzqQd05UkFw= +github.com/ikawaha/kagome-dict/ipa v1.2.6 h1:Bcvm4jgxAAnTIKb6ckqUKBiFDN0wuanFfycMuYt7xGQ= +github.com/ikawaha/kagome-dict/ipa v1.2.6/go.mod h1:ONdTMUAKMCq9yx4s69QRtPcJLEMVM0BNNYQrMCJLWb0= +github.com/ikawaha/kagome/v2 v2.11.0 h1:R914EkRzay9qtUbsFzEbcdZ3wHwwSPvbPkuBI1oIf78= +github.com/ikawaha/kagome/v2 v2.11.0/go.mod h1:6mYPezBou+iNVnX9uNa00Sfu6S6t2zcM8Nv1EW9Y9so= diff --git a/tools/discrawl-ja/main.go b/tools/discrawl-ja/main.go new file mode 100644 index 00000000..fc37738f --- /dev/null +++ b/tools/discrawl-ja/main.go @@ -0,0 +1,81 @@ +package main + +import ( + "bufio" + "encoding/json" + "fmt" + "os" + "strings" + + "github.com/ikawaha/kagome-dict/ipa" + "github.com/ikawaha/kagome/v2/tokenizer" +) + +type request struct { + Text string `json:"text"` +} + +type response struct { + Ready bool `json:"ready,omitempty"` + Tokens string `json:"tokens,omitempty"` + Error string `json:"error,omitempty"` + Version string `json:"version,omitempty"` +} + +func main() { + analyzer, err := tokenizer.New(ipa.Dict(), tokenizer.OmitBosEos()) + if err != nil { + writeResponse(response{Error: err.Error()}) + os.Exit(2) + } + writeResponse(response{Ready: true, Version: "kagome-ipa-search"}) + scanner := bufio.NewScanner(os.Stdin) + scanner.Buffer(make([]byte, 4096), 8*1024*1024) + for scanner.Scan() { + var input request + if err := json.Unmarshal(scanner.Bytes(), &input); err != nil { + writeResponse(response{Error: fmt.Sprintf("decode request: %v", err)}) + continue + } + writeResponse(response{Tokens: strings.Join(tokenize(analyzer, input.Text), " ")}) + } + if err := scanner.Err(); err != nil { + _, _ = fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } +} + +func tokenize(analyzer *tokenizer.Tokenizer, text string) []string { + seen := make(map[string]struct{}) + var tokens []string + add := func(form string) { + form = strings.ToLower(strings.TrimSpace(form)) + if form == "" || form == "*" { + return + } + if _, ok := seen[form]; ok { + return + } + seen[form] = struct{}{} + tokens = append(tokens, form) + } + for _, token := range analyzer.Analyze(text, tokenizer.Search) { + add(token.Surface) + if base := tokenBaseForm(token); base != "" { + add(base) + } + } + return tokens +} + +func tokenBaseForm(token tokenizer.Token) string { + features := token.Features() + if len(features) > 6 { + return features[6] + } + return "" +} + +func writeResponse(output response) { + _ = json.NewEncoder(os.Stdout).Encode(output) +} diff --git a/tools/discrawl-ja/main_test.go b/tools/discrawl-ja/main_test.go new file mode 100644 index 00000000..fd7dc724 --- /dev/null +++ b/tools/discrawl-ja/main_test.go @@ -0,0 +1,26 @@ +package main + +import ( + "testing" + + "github.com/ikawaha/kagome-dict/ipa" + "github.com/ikawaha/kagome/v2/tokenizer" +) + +func TestTokenizeSearchSplitsCompounds(t *testing.T) { + analyzer, err := tokenizer.New(ipa.Dict(), tokenizer.OmitBosEos()) + if err != nil { + t.Fatal(err) + } + tokens := tokenize(analyzer, "東京都庁に行きます") + found := false + for _, token := range tokens { + if token == "東京" { + found = true + break + } + } + if !found { + t.Fatalf("expected 東京 in %v", tokens) + } +} diff --git a/tools/discrawl-kiwi/README.md b/tools/discrawl-kiwi/README.md new file mode 100644 index 00000000..465d1a28 --- /dev/null +++ b/tools/discrawl-kiwi/README.md @@ -0,0 +1,47 @@ +# discrawl-kiwi + +`discrawl-kiwi` is Discrawl's optional Korean lexical analyzer helper. It uses +the existing [`github.com/codingpot/kiwigo`](https://pkg.go.dev/github.com/codingpot/kiwigo) +Go binding and Kiwi's public C API. It does not use Python. + +## Native prerequisites + +- Kiwi 0.23.2 headers and dynamic library +- Kiwi 0.23.2 base model +- a C/C++ toolchain supported by CGO + +The `kiwigo` build currently looks for headers and libraries under +`/usr/local/include` and `/usr/local/lib`. The official Kiwi release assets are: + +- `kiwi___v0.23.2.tgz` +- `kiwi_model_v0.23.2_base.tgz` + +Build: + +```bash +bash install-kiwi.sh +go build -o discrawl-kiwi . +``` + +Run: + +```bash +discrawl-kiwi --model /path/to/models/cong/base +``` + +The helper speaks newline-delimited JSON over stdin/stdout and stays alive so +the model is loaded once: + +```text +{"ready":true,"version":"0.23.2"} +{"text":"오늘 저녁먹음 기록"} +{"tokens":"오늘 저녁 먹 음 기록"} +``` + +## License boundary + +Kiwi and `github.com/codingpot/kiwigo` are licensed under +LGPL-2.1-or-later. Discrawl invokes this separately distributed helper as an +optional process, and the helper dynamically links to the replaceable Kiwi +library. Distributors of the helper or Kiwi binary assets must include the +applicable LGPL notices and corresponding Kiwi source access. diff --git a/tools/discrawl-kiwi/go.mod b/tools/discrawl-kiwi/go.mod new file mode 100644 index 00000000..54853644 --- /dev/null +++ b/tools/discrawl-kiwi/go.mod @@ -0,0 +1,14 @@ +module github.com/openclaw/discrawl/tools/discrawl-kiwi + +go 1.26 + +require ( + github.com/codingpot/kiwigo v0.0.0-20260812004023-ba165a3d4d4e + github.com/stretchr/testify v1.7.0 +) + +require ( + github.com/davecgh/go-spew v1.1.0 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect +) diff --git a/tools/discrawl-kiwi/go.sum b/tools/discrawl-kiwi/go.sum new file mode 100644 index 00000000..145ac1a7 --- /dev/null +++ b/tools/discrawl-kiwi/go.sum @@ -0,0 +1,15 @@ +github.com/codingpot/kiwigo v0.0.0-20260812004023-ba165a3d4d4e h1:nAYTYAwQxCBScKgc5uWTTI2MJym/5i+zImK57kL0e8I= +github.com/codingpot/kiwigo v0.0.0-20260812004023-ba165a3d4d4e/go.mod h1:dBnWtj6rEmAUXB20taKR3tYwyMkJiT1ekNCTuO824WU= +github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= +github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/tools/discrawl-kiwi/install-kiwi.sh b/tools/discrawl-kiwi/install-kiwi.sh new file mode 100644 index 00000000..d19010aa --- /dev/null +++ b/tools/discrawl-kiwi/install-kiwi.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +set -euo pipefail + +version=v0.23.2 +case "$(uname -s)" in + Darwin) platform=mac ;; + Linux) platform=lnx ;; + *) echo "unsupported Kiwi build platform: $(uname -s)" >&2; exit 1 ;; +esac +case "$(uname -m)" in + arm64|aarch64) + if [[ "$(uname -s)" == Linux ]]; then + architecture=aarch64 + else + architecture=arm64 + fi + ;; + x86_64|amd64) architecture=x86_64 ;; + *) echo "unsupported Kiwi architecture: $(uname -m)" >&2; exit 1 ;; +esac + +archive="kiwi_${platform}_${architecture}_${version}.tgz" +url="https://github.com/bab2min/Kiwi/releases/download/${version}/${archive}" +work="$(mktemp -d)" +trap 'rm -rf "$work"' EXIT + +curl --fail --location "$url" --output "$work/kiwi.tgz" +case "$archive" in + kiwi_lnx_aarch64_v0.23.2.tgz) expected_sha256=7e093121a367087d21e7c696bcc69a505935b07798d1e95c87f3b66a646c124e ;; + kiwi_lnx_x86_64_v0.23.2.tgz) expected_sha256=0b6694a795891de22fb14ae46825403af02063450126282c18448d6562b97174 ;; + kiwi_mac_arm64_v0.23.2.tgz) expected_sha256=ac124e32e013e2089cb4d842e2b735a1e6b4f3b126cdf692d78fda1130b8a382 ;; + kiwi_mac_x86_64_v0.23.2.tgz) expected_sha256=422c4284cc73a7499a714090e4d2f1c039dbc565aa2b425e5a0c0656d7b483a5 ;; + *) echo "missing pinned Kiwi checksum for $archive" >&2; exit 1 ;; +esac +actual_sha256="$(shasum -a 256 "$work/kiwi.tgz" | awk '{print $1}')" +if [[ "$actual_sha256" != "$expected_sha256" ]]; then + echo "Kiwi archive checksum mismatch for $archive" >&2 + exit 1 +fi +tar -xzf "$work/kiwi.tgz" -C "$work" +sudo cp -R "$work/include/kiwi" /usr/local/include/ +sudo cp -P "$work"/lib/libkiwi* /usr/local/lib/ +if [[ "$(uname -s)" == Linux ]]; then + sudo ldconfig +fi diff --git a/tools/discrawl-kiwi/main.go b/tools/discrawl-kiwi/main.go new file mode 100644 index 00000000..d99a49cc --- /dev/null +++ b/tools/discrawl-kiwi/main.go @@ -0,0 +1,95 @@ +package main + +import ( + "bufio" + "encoding/json" + "errors" + "flag" + "fmt" + "os" + "strings" + + kiwi "github.com/codingpot/kiwigo" +) + +type request struct { + Text string `json:"text"` +} + +type response struct { + Ready bool `json:"ready,omitempty"` + Tokens string `json:"tokens,omitempty"` + Error string `json:"error,omitempty"` + Version string `json:"version,omitempty"` +} + +func main() { + model := flag.String("model", "", "path to the Kiwi base model directory") + flag.Parse() + if strings.TrimSpace(*model) == "" { + writeResponse(response{Error: "Kiwi model path is required; pass --model"}) + os.Exit(2) + } + analyzer, err := kiwi.New(*model, kiwi.WithNumThread(0)) + if err != nil { + writeResponse(response{Error: err.Error()}) + os.Exit(2) + } + defer analyzer.Close() + + writeResponse(response{Ready: true, Version: kiwi.KiwiVersion()}) + scanner := bufio.NewScanner(os.Stdin) + scanner.Buffer(make([]byte, 4096), 8*1024*1024) + for scanner.Scan() { + var input request + if err := json.Unmarshal(scanner.Bytes(), &input); err != nil { + writeResponse(response{Error: fmt.Sprintf("decode request: %v", err)}) + continue + } + tokens, err := tokenize(analyzer, input.Text) + if err != nil { + writeResponse(response{Error: err.Error()}) + continue + } + writeResponse(response{Tokens: strings.Join(tokens, " ")}) + } + if err := scanner.Err(); err != nil { + _, _ = fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } +} + +type kiwiAnalyzer interface { + Analyze(string, ...kiwi.AnalyzeOptionFunc) ([]kiwi.TokenResult, error) +} + +func tokenize(analyzer kiwiAnalyzer, text string) ([]string, error) { + results, err := analyzer.Analyze(text, kiwi.WithTopN(1)) + if err != nil { + return nil, fmt.Errorf("analyze Korean text: %w", err) + } + if len(results) == 0 { + return nil, errors.New("Kiwi returned no analysis") + } + tokens := make([]string, 0, len(results[0].Tokens)) + seen := make(map[string]struct{}, len(results[0].Tokens)) + for _, token := range results[0].Tokens { + if strings.HasPrefix(string(token.Tag), "S") { + continue + } + form := strings.ToLower(strings.TrimSpace(token.Form)) + if form == "" { + continue + } + if _, ok := seen[form]; ok { + continue + } + seen[form] = struct{}{} + tokens = append(tokens, form) + } + return tokens, nil +} + +func writeResponse(output response) { + _ = json.NewEncoder(os.Stdout).Encode(output) +} diff --git a/tools/discrawl-kiwi/main_test.go b/tools/discrawl-kiwi/main_test.go new file mode 100644 index 00000000..aa64a212 --- /dev/null +++ b/tools/discrawl-kiwi/main_test.go @@ -0,0 +1,40 @@ +package main + +import ( + "errors" + "strings" + "testing" + + kiwi "github.com/codingpot/kiwigo" + "github.com/stretchr/testify/require" +) + +type fakeAnalyzer struct { + results []kiwi.TokenResult + err error +} + +func (f fakeAnalyzer) Analyze(string, ...kiwi.AnalyzeOptionFunc) ([]kiwi.TokenResult, error) { + return f.results, f.err +} + +func TestTokenizeFiltersPunctuationAndDuplicates(t *testing.T) { + tokens, err := tokenize(fakeAnalyzer{results: []kiwi.TokenResult{{Tokens: []kiwi.TokenInfo{ + {Form: "오늘", Tag: kiwi.POS_NNG}, + {Form: ".", Tag: kiwi.POS_SF}, + {Form: "오늘", Tag: kiwi.POS_NNG}, + {Form: "먹", Tag: kiwi.POS_VV}, + }}}}, "오늘먹음") + require.NoError(t, err) + require.Equal(t, []string{"오늘", "먹"}, tokens) +} + +func TestTokenizeReportsAnalyzerFailures(t *testing.T) { + _, err := tokenize(fakeAnalyzer{err: errors.New("failed")}, "text") + require.Error(t, err) + require.True(t, strings.Contains(err.Error(), "failed")) + + _, err = tokenize(fakeAnalyzer{}, "text") + require.Error(t, err) + require.True(t, strings.Contains(err.Error(), "no analysis")) +} diff --git a/tools/discrawl-zh/README.md b/tools/discrawl-zh/README.md new file mode 100644 index 00000000..8cd8b019 --- /dev/null +++ b/tools/discrawl-zh/README.md @@ -0,0 +1,14 @@ +# discrawl-zh + +Optional Chinese lexical helper. It uses the pure-Go +[`github.com/go-ego/gse`](https://pkg.go.dev/github.com/go-ego/gse) +search-mode segmenter with the embedded default dictionary. + +This binary is not linked into the default Discrawl release. Build it only when +Chinese lexical fields are enabled. + +```bash +go build -o discrawl-zh . +``` + +The helper speaks the same newline-delimited JSON protocol as `discrawl-kiwi`. diff --git a/tools/discrawl-zh/go.mod b/tools/discrawl-zh/go.mod new file mode 100644 index 00000000..5aab131a --- /dev/null +++ b/tools/discrawl-zh/go.mod @@ -0,0 +1,7 @@ +module github.com/openclaw/discrawl/tools/discrawl-zh + +go 1.26 + +require github.com/go-ego/gse v1.0.2 + +require github.com/vcaesar/cedar v0.30.0 // indirect diff --git a/tools/discrawl-zh/go.sum b/tools/discrawl-zh/go.sum new file mode 100644 index 00000000..b28ed36b --- /dev/null +++ b/tools/discrawl-zh/go.sum @@ -0,0 +1,6 @@ +github.com/go-ego/gse v1.0.2 h1:+27lYFPhQEhA9igtdOsJPRKYL/k3TwYsxBF5jr6KFv4= +github.com/go-ego/gse v1.0.2/go.mod h1:Fy35G+q7VV7Et1zIKO8o/sW1kkugV3znXap/lF/11zc= +github.com/vcaesar/cedar v0.30.0 h1:9fSDpM7FTjjUdPiBUUa0MWYMRGSEcqgFXvppZcZ4d7Y= +github.com/vcaesar/cedar v0.30.0/go.mod h1:lyuGvALuZZDPNXwpzv/9LyxW+8Y6faN7zauFezNsnik= +github.com/vcaesar/tt v0.20.1 h1:D/jUeeVCNbq3ad8M7hhtB3J9x5RZ6I1n1eZ0BJp7M+4= +github.com/vcaesar/tt v0.20.1/go.mod h1:cH2+AwGAJm19Wa6xvEa+0r+sXDJBT0QgNQey6mwqLeU= diff --git a/tools/discrawl-zh/main.go b/tools/discrawl-zh/main.go new file mode 100644 index 00000000..392e56b3 --- /dev/null +++ b/tools/discrawl-zh/main.go @@ -0,0 +1,66 @@ +package main + +import ( + "bufio" + "encoding/json" + "fmt" + "os" + "strings" + + "github.com/go-ego/gse" +) + +type request struct { + Text string `json:"text"` +} + +type response struct { + Ready bool `json:"ready,omitempty"` + Tokens string `json:"tokens,omitempty"` + Error string `json:"error,omitempty"` + Version string `json:"version,omitempty"` +} + +func main() { + var seg gse.Segmenter + if err := seg.LoadDictEmbed(); err != nil { + writeResponse(response{Error: err.Error()}) + os.Exit(2) + } + writeResponse(response{Ready: true, Version: "gse-search"}) + scanner := bufio.NewScanner(os.Stdin) + scanner.Buffer(make([]byte, 4096), 8*1024*1024) + for scanner.Scan() { + var input request + if err := json.Unmarshal(scanner.Bytes(), &input); err != nil { + writeResponse(response{Error: fmt.Sprintf("decode request: %v", err)}) + continue + } + writeResponse(response{Tokens: strings.Join(unique(seg.CutSearch(input.Text, true)), " ")}) + } + if err := scanner.Err(); err != nil { + _, _ = fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } +} + +func unique(forms []string) []string { + seen := make(map[string]struct{}, len(forms)) + tokens := make([]string, 0, len(forms)) + for _, form := range forms { + form = strings.ToLower(strings.TrimSpace(form)) + if form == "" { + continue + } + if _, ok := seen[form]; ok { + continue + } + seen[form] = struct{}{} + tokens = append(tokens, form) + } + return tokens +} + +func writeResponse(output response) { + _ = json.NewEncoder(os.Stdout).Encode(output) +} diff --git a/tools/discrawl-zh/main_test.go b/tools/discrawl-zh/main_test.go new file mode 100644 index 00000000..c0099eb5 --- /dev/null +++ b/tools/discrawl-zh/main_test.go @@ -0,0 +1,25 @@ +package main + +import ( + "testing" + + "github.com/go-ego/gse" +) + +func TestCutSearchKeepsCompoundParts(t *testing.T) { + var seg gse.Segmenter + if err := seg.LoadDictEmbed(); err != nil { + t.Fatal(err) + } + tokens := unique(seg.CutSearch("自然语言处理很有趣", true)) + found := false + for _, token := range tokens { + if token == "语言" { + found = true + break + } + } + if !found { + t.Fatalf("expected 语言 in %v", tokens) + } +}