From 2fe8a13c8e05e01b666e3a555f39e0ef32c92576 Mon Sep 17 00:00:00 2001 From: Audrey Tang <20723+audreyt@users.noreply.github.com> Date: Tue, 20 Oct 2026 08:43:16 +0800 Subject: [PATCH 1/2] fix(memory): keep numbers and emphasis words out of extracted entities The file-path pattern also matched decimals, versions, and addresses with no letter (0.52, 127.0.0.1) and dotted abbreviations (e.g, i.e), and the acronym pattern admitted ordinary words written in capitals for emphasis or as labels (NEVER, ONLY, NOTE, FIXED). On a real store these were among the most frequent stored entities, and entity edges linked memories that shared nothing but a decimal or an emphasis word. File-path matches now need a letter and a dot-separated segment longer than one character, and the acronym stoplist gains the emphasis words while leaving out words that double as technical names (HEAD, POST, PATH). Validated with make test. --- CHANGELOG.md | 5 +++ internal/memory/graph/entity.go | 50 ++++++++++++++++++++++++++-- internal/memory/graph/entity_test.go | 28 ++++++++++++++++ 3 files changed, 80 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a620d7a0f..14d803570 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -59,6 +59,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/). stanza. `brew` commands that load the cask printed "Calling `postflight` is deprecated! Use `postflight_steps` instead." on every invocation. The quarantine strip now runs through the declarative `postflight_steps` DSL. +- Entity extraction no longer treats numbers, versions, and addresses with no + letter (`0.52`, `127.0.0.1`) or dotted abbreviations (`e.g`, `i.e`) as file + paths, and no longer admits ordinary words written in capitals (`NEVER`, + `ONLY`, `NOTE`, `FIXED`) as acronyms. Such tokens linked unrelated memories + through entity edges. ## [0.2.7] - 2026-09-01 diff --git a/internal/memory/graph/entity.go b/internal/memory/graph/entity.go index 1f6572a0e..f59acd0ff 100644 --- a/internal/memory/graph/entity.go +++ b/internal/memory/graph/entity.go @@ -37,13 +37,16 @@ func ValidEntityMode(mode EntityMode) bool { } } +// filePathPattern matches file paths and dotted names (e.g., ./cmd/root.go, +// /etc/config.yml, example.org). Its matches must also pass looksLikePath. +var filePathPattern = regexp.MustCompile(`(?:^|[\s"'(])([.\w/-]+\.\w{1,10})(?:[\s"'),.]|$)`) + var entityPatterns = []*regexp.Regexp{ // CamelCase identifiers (e.g., MyClass, HttpServer) regexp.MustCompile(`\b([A-Z][a-z]+(?:[A-Z][a-z]+)+)\b`), // ALLCAPS acronyms 2-6 chars (e.g., API, HTTP, gRPC, SQL) regexp.MustCompile(`\b([A-Z]{2,6})\b`), - // File paths (e.g., ./cmd/root.go, /etc/config.yml) - regexp.MustCompile(`(?:^|[\s"'(])([.\w/-]+\.\w{1,10})(?:[\s"'),.]|$)`), + filePathPattern, // URLs regexp.MustCompile(`https?://[^\s"'<>)]+`), // @mentions @@ -102,7 +105,12 @@ var techDictionary = map[string]bool{ "MAGMA": true, "MCP": true, "RLM": true, } -// acronymStopwords filters out common English acronyms that aren't tech entities. +// acronymStopwords filters out all-caps tokens that aren't tech entities: +// short common English words, and ordinary words written in capitals for +// emphasis or as labels (NEVER, ONLY, NOTE, FIXED). In a real store the latter +// were among the most frequent "entities" and linked unrelated memories. +// Words that double as technical names (HEAD, POST, PATH, HOME, REST, DELETE, +// LINE) are deliberately absent. var acronymStopwords = map[string]bool{ "IN": true, "ON": true, "AT": true, "TO": true, "BY": true, "OR": true, "AN": true, "IF": true, "IS": true, "IT": true, @@ -117,6 +125,23 @@ var acronymStopwords = map[string]bool{ "SHE": true, "TWO": true, "USE": true, "BOY": true, "DID": true, "GET": true, "HIM": true, "HIS": true, "LET": true, "PUT": true, "TOP": true, "TOO": true, "ANY": true, + // Ordinary words written in capitals for emphasis or as labels. + "OK": true, "YES": true, "TRUE": true, "FALSE": true, + "ONLY": true, "BOTH": true, "NEVER": true, "ALWAYS": true, "MUST": true, + "STILL": true, "SAME": true, "ALSO": true, "EVERY": true, "OTHER": true, + "OWN": true, "REAL": true, "WORTH": true, "HARD": true, "FULL": true, + "EMPTY": true, "ZERO": true, "KEEP": true, "SKIP": true, "LEFT": true, + "WITH": true, "FROM": true, "THAT": true, "THIS": true, "THAN": true, + "WHAT": true, "WHY": true, "DOES": true, "CANNOT": true, "NEEDS": true, + "BEFORE": true, "AFTER": true, "FIRST": true, "SECOND": true, "THIRD": true, + "FINAL": true, "THREE": true, "FOUR": true, "FIVE": true, "HIGH": true, + "PASS": true, "FAIL": true, "FAILED": true, "FIX": true, "FIXED": true, + "DONE": true, "WRONG": true, "STALE": true, "ARMED": true, "OPEN": true, + "CLOSED": true, "FOUND": true, "SENT": true, "SAFE": true, "CLEAN": true, + "VERIFY": true, "CHECK": true, "REVIEW": true, "DRAFT": true, "PROOF": true, + "NOTE": true, "NOTES": true, "LESSON": true, "RULE": true, "FACT": true, + "FACTS": true, "DEFECT": true, "CAUSE": true, "GOTCHA": true, "TRAP": true, + "BUG": true, "ISSUE": true, "ERROR": true, } // ExtractEntities extracts named entities from text using regex patterns @@ -137,6 +162,9 @@ func ExtractEntities(text string) []string { if acronymStopwords[entity] { continue } + if pat == filePathPattern && !looksLikePath(entity) { + continue + } seen[entity] = true entities = append(entities, entity) } @@ -154,6 +182,22 @@ func ExtractEntities(text string) []string { return entities } +// looksLikePath reports whether a file-path match could name a file or +// domain. The path pattern also matches numbers, versions, and addresses with +// no letter (0.52, 127.0.0.1) and dotted single letters (e.g, i.e, J.Y, p.4); +// in a real store those were thousands of entities that named nothing. +func looksLikePath(s string) bool { + if !strings.ContainsFunc(s, func(r rune) bool { return (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') }) { + return false + } + for _, segment := range strings.Split(s, ".") { + if len(segment) > 1 { + return true + } + } + return false +} + // ResolveEntities returns the final entity list for a content/provided pair. func ResolveEntities(content string, provided []string, mode EntityMode) []string { return ResolveEntitiesIndexed(content, provided, mode, nil) diff --git a/internal/memory/graph/entity_test.go b/internal/memory/graph/entity_test.go index 2a8087505..c84764cb3 100644 --- a/internal/memory/graph/entity_test.go +++ b/internal/memory/graph/entity_test.go @@ -399,6 +399,34 @@ func TestResolveEntitiesIndexed_NilIndexEqualsResolveEntities(t *testing.T) { } } +func TestExtractEntities_FilePathsSkipNumbersAndAbbreviations(t *testing.T) { + entities := ExtractEntities("Scores 0.52 and 127.0.0.1, e.g. the J.Y. case on p.4, live in cmd/root.go, llama.cpp and example.org") + has := toSet(entities) + for _, want := range []string{"cmd/root.go", "llama.cpp", "example.org"} { + if !has[want] { + t.Errorf("want path %q, got %v", want, entities) + } + } + for _, noise := range []string{"0.52", "127.0.0.1", "e.g", "J.Y", "p.4"} { + if has[noise] { + t.Errorf("number or abbreviation %q must not be an entity, got %v", noise, entities) + } + } +} + +func TestExtractEntities_EmphasisWordsAreNotAcronyms(t *testing.T) { + entities := ExtractEntities("NEVER skip the API check; ONLY use HTTP. Status: OK, FIXED") + has := toSet(entities) + if !has["API"] || !has["HTTP"] { + t.Errorf("want acronyms API and HTTP, got %v", entities) + } + for _, word := range []string{"NEVER", "ONLY", "OK", "FIXED"} { + if has[word] { + t.Errorf("emphasis word %q must not be an entity, got %v", word, entities) + } + } +} + func equalAsSets(a, b []string) bool { if len(a) != len(b) { return false From 47a2204f7de65f2f03b02c8c98cac65dc2562c3c Mon Sep 17 00:00:00 2001 From: Audrey Tang <20723+audreyt@users.noreply.github.com> Date: Tue, 20 Oct 2026 08:44:37 +0800 Subject: [PATCH 2/2] feat(memory): add an opt-in entity extraction model With MNEMON_EXTRACT_MODEL and MNEMON_EXTRACT_ENDPOINT set, remember and import ask a System One extraction model such as Cloudflare Clairic for the entities of a memory that arrives without any. The request runs before the write transaction, and deterministic code keeps only names that appear in the memory as whole tokens, so the model can add the multi-word names the regex misses but cannot invent or truncate one. Model entities replace only the all-caps acronym guess; the exact patterns and the known-entity index still apply, because the model splits paths and versions that those patterns match exactly. A missing or failing model prints one warning and falls back to heuristic extraction, and import stops asking after the first failure. Memories stored with entities and recall queries never call the model. Validated with make test and a live remember against Clairic on Workers AI. --- CHANGELOG.md | 9 + Makefile | 1 + README.md | 23 ++ SECURITY.md | 1 + cmd/memory/extract.go | 42 +++ cmd/memory/extract_test.go | 190 +++++++++++ cmd/memory/import.go | 8 +- cmd/memory/remember.go | 8 +- docs/USAGE.md | 12 + docs/zh/README.md | 4 + docs/zh/USAGE.md | 4 + internal/memory/extract/extract.go | 297 ++++++++++++++++++ internal/memory/extract/extract_test.go | 164 ++++++++++ internal/memory/graph/engine.go | 10 +- internal/memory/graph/entity.go | 61 +++- internal/memory/graph/entity_test.go | 60 ++++ .../architecture/mnemond_structure_test.go | 3 +- 17 files changed, 885 insertions(+), 12 deletions(-) create mode 100644 cmd/memory/extract.go create mode 100644 cmd/memory/extract_test.go create mode 100644 internal/memory/extract/extract.go create mode 100644 internal/memory/extract/extract_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 14d803570..e090647e6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/). directory and the legacy `~/.mavis/` directory migrated by current releases. - `mnemon setup --eject --target minimax-code` removes only the Mnemon skill while preserving other MiniMax Code skills. +- Optional entity extraction model. With `MNEMON_EXTRACT_MODEL` and + `MNEMON_EXTRACT_ENDPOINT` set (for example Cloudflare Clairic on Workers AI), + `remember` and `import` ask the model for the entities of memories stored + without any, keep only names that appear in the memory as whole tokens, and + use them in place of the all-caps acronym guess. A missing or failing model + falls back to heuristic extraction with a one-line warning. ### Tests @@ -46,6 +52,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/). cleanup. - Added MiniMax Code coverage for skill installation, current and legacy data directory detection, and scoped eject cleanup. +- Added extraction client coverage for the Workers AI and bare response + shapes, whole-token grounding, limits, and error redaction, plus CLI coverage + for `remember` and `import` with a working, skipped, and failing model. ### Note diff --git a/Makefile b/Makefile index 3274c5c6f..f4517d43d 100644 --- a/Makefile +++ b/Makefile @@ -24,6 +24,7 @@ DETERMINISTIC_PKGS := \ ./internal/agency/authority \ ./internal/agency/artifact \ ./internal/memory/embed \ + ./internal/memory/extract \ ./internal/memory/graph \ ./internal/memory/importdraft \ ./internal/memory/internal/topk \ diff --git a/README.md b/README.md index d82b7719b..751ff8f17 100644 --- a/README.md +++ b/README.md @@ -546,6 +546,29 @@ export MNEMON_EMBED_API_KEY=pa-... mnemon embed --status ``` +**Entity extraction model** (optional, opt-in by environment): + +| Environment Variable | Default | Description | +|---|---|---| +| `MNEMON_EXTRACT_MODEL` | *(empty = disabled)* | System One extraction model name, e.g. `clairic` | +| `MNEMON_EXTRACT_ENDPOINT` | *(none, required)* | Full URL that requests are POSTed to, e.g. `https://api.cloudflare.com/client/v4/accounts//ai/run/@cf/cloudflare/clairic` | +| `MNEMON_EXTRACT_API_KEY` | *(none)* | Bearer token, e.g. a Cloudflare API token with Workers AI access | +| `MNEMON_EXTRACT_TIMEOUT` | `15s` | Time budget for one extraction (Go duration) | + +When `MNEMON_EXTRACT_MODEL` is set, `mnemon remember` and `mnemon import` ask +the model for the entities of every memory that arrives without any (no +`--entities` in `merge` mode, or `--entity-mode auto`). Mnemon keeps only names +that appear in the memory verbatim as whole tokens. They replace the all-caps +acronym guess; CamelCase, file path, URL, @mention, title-mark, dictionary, and +known-entity matches still apply. Memories that arrive with entities, and +`recall` queries, are unchanged. If the model is missing or fails, mnemon +prints a one-line warning to stderr and falls back to heuristic extraction; +`import` stops asking after the first failure. + +> **Privacy:** extraction sends the full text of each such memory to +> `MNEMON_EXTRACT_ENDPOINT`. Leave `MNEMON_EXTRACT_MODEL` unset if memories must +> not leave the machine. + ## Development ```bash diff --git a/SECURITY.md b/SECURITY.md index 0c572cc6d..f124d0c86 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -17,6 +17,7 @@ Mnemon runs locally and stores data in `~/.mnemon/`. Key security considerations - **SQLite database** — contains all stored insights; protected by filesystem permissions (`0644`). - **Hook scripts** — shell scripts executed by the LLM CLI at lifecycle events; written with `0755` permissions. - **Embedding provider connection** — optional requests send insight or query text to the configured Ollama or OpenAI-compatible server. The default local Ollama endpoint does not use TLS. If `MNEMON_EMBED_ENDPOINT` points outside a trusted local network, use HTTPS to protect content and any `MNEMON_EMBED_API_KEY` bearer token in transit. +- **Entity extraction model connection** — only when `MNEMON_EXTRACT_MODEL` is set, `remember` and `import` send the full text of each memory stored without entities to `MNEMON_EXTRACT_ENDPOINT`. A cloud endpoint therefore receives that text. Use HTTPS to protect content and the `MNEMON_EXTRACT_API_KEY` bearer token. ## Supported Versions diff --git a/cmd/memory/extract.go b/cmd/memory/extract.go new file mode 100644 index 000000000..b8fc95543 --- /dev/null +++ b/cmd/memory/extract.go @@ -0,0 +1,42 @@ +package memory + +import ( + "context" + "fmt" + "os" + + "github.com/mnemon-dev/mnemon/internal/memory/extract" + "github.com/mnemon-dev/mnemon/internal/memory/graph" +) + +// entityExtractor asks the opt-in extraction model (MNEMON_EXTRACT_MODEL) for +// a memory's entities before the write transaction. It fails open: after the +// first failure it warns once and stops asking for the rest of the command, so +// a batch import does not wait out the timeout on every insight. +type entityExtractor struct { + client *extract.Client +} + +func newEntityExtractor() *entityExtractor { + client, err := extract.NewClient() + if err != nil { + fmt.Fprintf(os.Stderr, "warning: entity extraction disabled: %v\n", err) + } + return &entityExtractor{client: client} +} + +// entities returns the model's entities for content, or nil when no model ran +// because extraction is not configured, the provided entities make it +// unnecessary, or it failed. +func (x *entityExtractor) entities(content string, provided []string, mode graph.EntityMode) []string { + if x.client == nil || !graph.ReliesOnExtraction(mode, provided) { + return nil + } + names, err := x.client.Entities(context.Background(), content) + if err != nil { + fmt.Fprintf(os.Stderr, "warning: entity extraction unavailable, using heuristic extraction: %v\n", err) + x.client = nil + return nil + } + return names +} diff --git a/cmd/memory/extract_test.go b/cmd/memory/extract_test.go new file mode 100644 index 000000000..e7bbd4c9d --- /dev/null +++ b/cmd/memory/extract_test.go @@ -0,0 +1,190 @@ +package memory + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "sync/atomic" + "testing" + + "github.com/mnemon-dev/mnemon/internal/memory/importdraft" +) + +const extractContent = "Jane Rivera briefed the CEO of Acme Foundation; notes in docs/plan.md" + +// startExtractServer fakes a Workers AI extraction endpoint that answers every +// request with items, or fails with status. It returns the request counter. +func startExtractServer(t *testing.T, status int, items []string) *atomic.Int32 { + t.Helper() + var hits atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + if status != http.StatusOK { + http.Error(w, `{"success":false,"errors":[{"message":"model unavailable"}]}`, status) + return + } + _ = json.NewEncoder(w).Encode(map[string]any{"success": true, "result": map[string]any{ + "answers": map[string]any{"entities": map[string]any{"type": "list", "value": items}}, + }}) + })) + t.Cleanup(srv.Close) + t.Setenv("MNEMON_EXTRACT_MODEL", "test-model") + t.Setenv("MNEMON_EXTRACT_ENDPOINT", srv.URL) + t.Setenv("MNEMON_EXTRACT_API_KEY", "") + t.Setenv("MNEMON_EXTRACT_TIMEOUT", "") + return &hits +} + +// stderrOf runs fn and returns what it wrote to stderr. +func stderrOf(t *testing.T, fn func()) string { + t.Helper() + old := os.Stderr + r, w, err := os.Pipe() + if err != nil { + t.Fatalf("pipe stderr: %v", err) + } + os.Stderr = w + defer func() { os.Stderr = old }() + fn() + if err := w.Close(); err != nil { + t.Fatalf("close stderr writer: %v", err) + } + data, err := io.ReadAll(r) + if err != nil { + t.Fatalf("read stderr: %v", err) + } + return string(data) +} + +func entitySet(entities []string) map[string]bool { + set := make(map[string]bool, len(entities)) + for _, e := range entities { + set[e] = true + } + return set +} + +func rememberEntities(t *testing.T) (map[string]bool, string) { + t.Helper() + var out string + var runErr error + stderr := stderrOf(t, func() { + out = captureStdout(t, func() { runErr = rememberCmd.RunE(rememberCmd, []string{extractContent}) }) + }) + if runErr != nil { + t.Fatalf("remember: %v", runErr) + } + var result struct { + Entities []string `json:"entities"` + } + if err := json.Unmarshal([]byte(out), &result); err != nil { + t.Fatalf("decode remember output: %v\n%s", err, out) + } + return entitySet(result.Entities), stderr +} + +func TestRememberUsesModelEntitiesWhenNoneProvided(t *testing.T) { + // Given an extraction model that finds two real names and invents one + configureRememberTest(t) + hits := startExtractServer(t, http.StatusOK, []string{"Jane Rivera", "Acme Foundation", "Invented Inc"}) + + // When a memory is remembered without entities + entities, stderr := rememberEntities(t) + + // Then the real names replace the acronym guess, exact paths stay, and the invented name is dropped + for _, want := range []string{"Jane Rivera", "Acme Foundation", "docs/plan.md"} { + if !entities[want] { + t.Errorf("want entity %q, got %v", want, entities) + } + } + if entities["CEO"] || entities["Invented Inc"] { + t.Errorf("entities = %v; want no acronym guess and no invented name", entities) + } + if hits.Load() != 1 || stderr != "" { + t.Fatalf("requests = %d, stderr = %q", hits.Load(), stderr) + } +} + +func TestRememberSkipsModelWhenEntitiesProvided(t *testing.T) { + // Given a configured model and caller-provided entities + configureRememberTest(t) + hits := startExtractServer(t, http.StatusOK, []string{"Jane Rivera"}) + remEntities = "zebra" + + // When the memory is remembered + entities, _ := rememberEntities(t) + + // Then the model is never asked and heuristic extraction is unchanged + if hits.Load() != 0 || !entities["zebra"] || !entities["CEO"] || entities["Jane Rivera"] { + t.Fatalf("requests = %d, entities = %v", hits.Load(), entities) + } +} + +func TestRememberFallsBackToHeuristicsWhenModelFails(t *testing.T) { + // Given an extraction model that errors + configureRememberTest(t) + startExtractServer(t, http.StatusInternalServerError, nil) + + // When a memory is remembered without entities + entities, stderr := rememberEntities(t) + + // Then heuristic extraction applies and one warning is printed + if !entities["CEO"] || entities["Jane Rivera"] { + t.Fatalf("entities = %v; want heuristic extraction", entities) + } + if !strings.HasPrefix(stderr, "warning: entity extraction unavailable") || strings.Count(stderr, "\n") != 1 { + t.Fatalf("stderr = %q, want a single warning line", stderr) + } +} + +func TestImportUsesModelOnlyForDraftsWithoutEntities(t *testing.T) { + // Given a model and two drafts, one already carrying entities + db := setupImportDiffTest(t, true) + hits := startExtractServer(t, http.StatusOK, []string{"Acme Foundation"}) + draft := importdraft.MemoryDraft{SchemaVersion: "1", Insights: []importdraft.DraftInsight{ + {Content: "Jane Rivera joined the Acme Foundation"}, + {Content: "the Acme Foundation funds Project Lark", Entities: []string{"Project Lark"}}, + }} + + // When the draft is imported + runImportDiffDraft(t, draft) + + // Then only the first draft asked the model, and its entities came from it + if hits.Load() != 1 { + t.Fatalf("requests = %d, want 1", hits.Load()) + } + insights, err := db.GetAllActiveInsights() + if err != nil { + t.Fatal(err) + } + for _, ins := range insights { + got := entitySet(ins.Entities) + if strings.HasPrefix(ins.Content, "Jane") && !got["Acme Foundation"] { + t.Errorf("model entities missing: %v", ins.Entities) + } + if strings.HasPrefix(ins.Content, "the Acme") && !got["Project Lark"] { + t.Errorf("provided entities missing: %v", ins.Entities) + } + } +} + +func TestImportStopsAskingAfterModelFailure(t *testing.T) { + // Given a failing model and two drafts without entities + setupImportDiffTest(t, true) + hits := startExtractServer(t, http.StatusInternalServerError, nil) + draft := importdraft.MemoryDraft{SchemaVersion: "1", Insights: []importdraft.DraftInsight{ + {Content: "Jane Rivera joined the Acme Foundation"}, + {Content: "Zebra deploys to staging-east"}, + }} + + // When the draft is imported + stderr := stderrOf(t, func() { runImportDiffDraft(t, draft) }) + + // Then the model was asked once and one warning was printed + if hits.Load() != 1 || strings.Count(stderr, "warning: entity extraction unavailable") != 1 { + t.Fatalf("requests = %d, stderr = %q", hits.Load(), stderr) + } +} diff --git a/cmd/memory/import.go b/cmd/memory/import.go index e684b0060..9269827a5 100644 --- a/cmd/memory/import.go +++ b/cmd/memory/import.go @@ -76,6 +76,7 @@ exports are documented in docs/IMPORT.md.`, importedSources := make(map[string]bool) refreshIDs := make(map[string]bool) results := make([]importResult, 0, len(draft.Insights)) + extractor := newEntityExtractor() for idx, di := range draft.Insights { cat := model.Category(di.Category) @@ -149,6 +150,8 @@ exports are documented in docs/IMPORT.md.`, continue } + modelEntities := extractor.entities(insight.Content, insight.Entities, graph.EntityModeMerge) + var writeErr error err = db.InTransaction(func() error { if err := db.InsertInsight(insight); err != nil { @@ -163,8 +166,9 @@ exports are documented in docs/IMPORT.md.`, } } engine := graph.NewEngineWithOptions(db, embedCache, graph.EngineOptions{ - EntityMode: graph.EntityModeMerge, - TemporalMode: graph.TemporalDisabled, + EntityMode: graph.EntityModeMerge, + TemporalMode: graph.TemporalDisabled, + ModelEntities: modelEntities, }) engine.OnInsightCreated(insight) if len(insight.Entities) > 0 { diff --git a/cmd/memory/remember.go b/cmd/memory/remember.go index 0768f38a2..ffef9a939 100644 --- a/cmd/memory/remember.go +++ b/cmd/memory/remember.go @@ -181,6 +181,9 @@ var rememberCmd = &cobra.Command{ return enc.Encode(output) } + // Like the embedding, ask the extraction model before the transaction. + modelEntities := newEntityExtractor().entities(content, entities, entityMode) + // 3. All DB writes in a single atomic transaction var ( edgeStats graph.EdgeStats @@ -205,7 +208,10 @@ var rememberCmd = &cobra.Command{ } // Run graph edge engine (includes auto semantic edges when embedded) - engine := graph.NewEngineWithEntityMode(db, embedCache, entityMode) + engine := graph.NewEngineWithOptions(db, embedCache, graph.EngineOptions{ + EntityMode: entityMode, + ModelEntities: modelEntities, + }) edgeStats = engine.OnInsightCreated(insight) // Update entities extracted by the engine diff --git a/docs/USAGE.md b/docs/USAGE.md index 03c168a59..316824044 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -411,9 +411,21 @@ Nodes are colored by category (decision, fact, insight, preference, context); ed | `MNEMON_EMBED_API_KEY` | (none) | Bearer token for OpenAI-compatible servers | | `MNEMON_EMBED_KEEP_ALIVE` | 30m | How long Ollama keeps the embedding model loaded after each request | | `MNEMON_EMBED_DIMENSIONS` | (native) | Embedding dimensions; set to truncate (e.g., `256` for Matryoshka models) | +| `MNEMON_EXTRACT_MODEL` | (empty = disabled) | System One extraction model (`clairic`) that `remember` and `import` ask for the entities of memories stored without any | +| `MNEMON_EXTRACT_ENDPOINT` | (none, required) | Full URL that extraction requests are POSTed to | +| `MNEMON_EXTRACT_API_KEY` | (none) | Bearer token for the extraction endpoint | +| `MNEMON_EXTRACT_TIMEOUT` | `15s` | Time budget for one extraction (Go duration) | | `MNEMON_MAX_INSIGHTS` | `1000` | Active-insight ceiling before auto-pruning starts; `0` disables auto-pruning | | `MNEMON_AUTO_PRUNE_MIN_AGE` | `24h` | Minimum age before automatic pruning; accepts durations such as `24h`, integer days such as `7d`, or `0` to disable the grace period | +`MNEMON_EXTRACT_MODEL` sends one request per memory stored without entities, +with the full memory text as `state` and one `list` field named `entities`. For +Clairic on Cloudflare Workers AI the endpoint is +`https://api.cloudflare.com/client/v4/accounts//ai/run/@cf/cloudflare/clairic`. +Only names found in the memory as whole tokens are kept, and they replace the +all-caps acronym guess. On failure mnemon warns once on stderr and falls back to +heuristic extraction. + Auto-prune may temporarily leave the active count above the ceiling when every eligible insight is still inside the grace period. Age is measured from local store insertion, so newly imported historical memories are protected too. Each deletion is soft, diff --git a/docs/zh/README.md b/docs/zh/README.md index 581c589e4..6b1ea1b8f 100644 --- a/docs/zh/README.md +++ b/docs/zh/README.md @@ -434,6 +434,10 @@ Sub-agent 委派是可选执行策略。当 runtime 支持时，主 agent 可以 | `MNEMON_EMBED_API_KEY` | *（无）* | OpenAI 兼容服务器（oMLX、vLLM 等）的 Bearer 令牌 | | `MNEMON_EMBED_KEEP_ALIVE` | `30m` | 每次请求后 Ollama 保持嵌入模型加载的时长（OpenAI 兼容服务器忽略此项） | | `MNEMON_EMBED_DIMENSIONS` | *（原生维度）* | 可选的 Matryoshka 维度截断 | +| `MNEMON_EXTRACT_MODEL` | *（空 = 禁用）* | 抽取模型（如 `clairic`）；`remember` 与 `import` 对未带实体的记忆调用它抽取实体 | +| `MNEMON_EXTRACT_ENDPOINT` | *（无，必填）* | 抽取请求 POST 的完整 URL；云端端点会收到这些记忆的全文 | +| `MNEMON_EXTRACT_API_KEY` | *（无）* | 抽取端点的 Bearer 令牌 | +| `MNEMON_EXTRACT_TIMEOUT` | `15s` | 单次抽取的时间预算（Go duration） | 每次自动删除均为软删除，以 `prune` 操作记录到 oplog，并通过触发命令的 `auto_pruned_ids` 字段返回具体 ID。 diff --git a/docs/zh/USAGE.md b/docs/zh/USAGE.md index fff9d05f7..41e7722c1 100644 --- a/docs/zh/USAGE.md +++ b/docs/zh/USAGE.md @@ -375,6 +375,10 @@ open graph.html | `MNEMON_EMBED_API_KEY` | （无） | OpenAI 兼容服务器的 Bearer 令牌 | | `MNEMON_EMBED_KEEP_ALIVE` | 30m | 每次请求后 Ollama 保持嵌入模型加载的时长 | | `MNEMON_EMBED_DIMENSIONS` | (原生维度) | 嵌入向量维度；可设置截断值（例如 Matryoshka 模型使用 `256`） | +| `MNEMON_EXTRACT_MODEL` | （空 = 禁用） | 抽取模型（如 `clairic`）；`remember` 与 `import` 对未带实体的记忆调用它抽取实体，只保留在记忆中完整出现的名称，并以之取代全大写缩写猜测 | +| `MNEMON_EXTRACT_ENDPOINT` | （无，必填） | 抽取请求 POST 的完整 URL；云端端点会收到这些记忆的全文 | +| `MNEMON_EXTRACT_API_KEY` | （无） | 抽取端点的 Bearer 令牌 | +| `MNEMON_EXTRACT_TIMEOUT` | `15s` | 单次抽取的时间预算（Go duration） | | `MNEMON_MAX_INSIGHTS` | `1000` | 触发自动清理的活跃洞察数量上限；设为 `0` 可关闭自动清理 | | `MNEMON_AUTO_PRUNE_MIN_AGE` | `24h` | 可被自动清理前的最短存活时间；支持 `24h`、整数天 `7d`，设为 `0` 可关闭保护期 | diff --git a/internal/memory/extract/extract.go b/internal/memory/extract/extract.go new file mode 100644 index 000000000..6776ed620 --- /dev/null +++ b/internal/memory/extract/extract.go @@ -0,0 +1,297 @@ +// Package extract is a small client for a System One extraction model such as +// Cloudflare Clairic. It asks the model for the named entities in a memory's +// content and keeps only names that appear in that content verbatim as whole +// tokens, so the model can find entities but never invent or truncate one. +// +// Extraction is opt-in: NewClient returns nil unless MNEMON_EXTRACT_MODEL is +// set, and callers treat a failed extraction as advisory, falling back to +// heuristic extraction. +package extract + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/url" + "os" + "strings" + "time" + "unicode" + "unicode/utf8" +) + +const ( + // DefaultTimeout bounds one extraction request. + DefaultTimeout = 15 * time.Second + // MaxEntities and MaxEntityBytes match the limits remember applies to + // --entities. + MaxEntities = 50 + MaxEntityBytes = 200 + + maxResponseBytes = 1 << 20 + maxErrorBytes = 4 << 10 + maxMessageRunes = 300 + + entitiesField = "entities" + entitiesInstructions = "Every named thing this memory is about: people, organizations, projects, " + + "products, tools, models, services, files, paths and domains. Copy each name exactly as written. " + + "Skip dates, numbers, version numbers, quantities and ordinary words." +) + +// ErrNoEndpoint reports a model configured without an endpoint. +var ErrNoEndpoint = errors.New("MNEMON_EXTRACT_MODEL is set but MNEMON_EXTRACT_ENDPOINT is not") + +// Client talks to one extraction endpoint. +type Client struct { + endpoint string + model string + apiKey string + timeout time.Duration + http *http.Client +} + +// NewClient builds a client from MNEMON_EXTRACT_MODEL, MNEMON_EXTRACT_ENDPOINT, +// MNEMON_EXTRACT_API_KEY, and MNEMON_EXTRACT_TIMEOUT. It returns nil, nil when +// MNEMON_EXTRACT_MODEL is empty, which means extraction is not configured. +// MNEMON_EXTRACT_ENDPOINT is the full URL requests are POSTed to; there is no +// default because no local server hosts an extraction model yet. +func NewClient() (*Client, error) { + model := strings.TrimSpace(os.Getenv("MNEMON_EXTRACT_MODEL")) + if model == "" { + return nil, nil + } + endpoint := strings.TrimSpace(os.Getenv("MNEMON_EXTRACT_ENDPOINT")) + if endpoint == "" { + return nil, ErrNoEndpoint + } + timeout := DefaultTimeout + if v := os.Getenv("MNEMON_EXTRACT_TIMEOUT"); v != "" { + if d, err := time.ParseDuration(v); err == nil && d > 0 { + timeout = d + } else { + fmt.Fprintf(os.Stderr, "warning: invalid MNEMON_EXTRACT_TIMEOUT %q, using %s\n", v, DefaultTimeout) + } + } + return &Client{ + endpoint: endpoint, + model: model, + apiKey: os.Getenv("MNEMON_EXTRACT_API_KEY"), + timeout: timeout, + http: &http.Client{ + Transport: &http.Transport{ + // Loopback servers never use a proxy, so a stray HTTPS_PROXY + // cannot capture local model traffic. + Proxy: proxyFunc(endpoint), + DialContext: (&net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}).DialContext, + }, + }, + }, nil +} + +func proxyFunc(endpoint string) func(*http.Request) (*url.URL, error) { + if u, err := url.Parse(endpoint); err == nil { + if host := u.Hostname(); host == "localhost" || net.ParseIP(host).IsLoopback() { + return func(*http.Request) (*url.URL, error) { return nil, nil } + } + } + return http.ProxyFromEnvironment +} + +// Entities asks the model for the entities in content. On success the result +// is non-nil, possibly empty, deduplicated, at most MaxEntities long, and holds +// only names found in content as whole tokens. +func (c *Client) Entities(ctx context.Context, content string) ([]string, error) { + ctx, cancel := context.WithTimeout(ctx, c.timeout) + defer cancel() + items, err := c.post(ctx, content) + if err != nil { + return nil, err + } + return groundedEntities(content, items), nil +} + +type wireField struct { + Type string `json:"type"` + Items string `json:"items"` + Instructions string `json:"instructions"` +} + +type wireRequest struct { + Model string `json:"model"` + State string `json:"state"` + Fields map[string]wireField `json:"fields"` +} + +type wireAnswer struct { + Type string `json:"type"` + Value json.RawMessage `json:"value"` +} + +type wireAnswers struct { + Answers map[string]wireAnswer `json:"answers"` +} + +// wireResponse accepts both a bare System One response ({"answers": ...}) and +// the Workers AI envelope ({"success": ..., "errors": [...], "result": {...}}). +type wireResponse struct { + wireAnswers + Result *wireAnswers `json:"result"` + Success *bool `json:"success"` + Errors []struct { + Message string `json:"message"` + } `json:"errors"` + Error string `json:"error"` +} + +func (r *wireResponse) message() string { + msg := r.Error + for _, e := range r.Errors { + if msg == "" { + msg = e.Message + } + } + if runes := []rune(msg); len(runes) > maxMessageRunes { + msg = string(runes[:maxMessageRunes]) + } + return msg +} + +func (c *Client) post(ctx context.Context, content string) ([]string, error) { + body, err := json.Marshal(wireRequest{ + Model: c.model, + State: content, + Fields: map[string]wireField{ + entitiesField: {Type: "list", Items: "text", Instructions: entitiesInstructions}, + }, + }) + if err != nil { + return nil, fmt.Errorf("marshal extraction request: %w", err) + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint, bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("build extraction request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + if c.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+c.apiKey) + } + resp, err := c.http.Do(req) + if err != nil { + return nil, fmt.Errorf("extraction request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + data, _ := io.ReadAll(io.LimitReader(resp.Body, maxErrorBytes)) + var payload wireResponse + if json.Unmarshal(data, &payload) == nil && payload.message() != "" { + return nil, fmt.Errorf("extraction server returned status %d: %s", resp.StatusCode, payload.message()) + } + return nil, fmt.Errorf("extraction server returned status %d", resp.StatusCode) + } + data, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes+1)) + if err != nil { + return nil, fmt.Errorf("read extraction response: %w", err) + } + if len(data) > maxResponseBytes { + return nil, fmt.Errorf("extraction response exceeds %d bytes", maxResponseBytes) + } + var decoded wireResponse + if err := json.Unmarshal(data, &decoded); err != nil { + return nil, fmt.Errorf("decode extraction response: %w", err) + } + if decoded.Success != nil && !*decoded.Success { + return nil, fmt.Errorf("extraction server reported failure: %s", decoded.message()) + } + answers := decoded.Answers + if decoded.Result != nil { + answers = decoded.Result.Answers + } + answer, ok := answers[entitiesField] + if !ok { + return nil, fmt.Errorf("extraction response has no %q answer", entitiesField) + } + if answer.Type != "list" { + return nil, fmt.Errorf("extraction answer %q has type %q, want list", entitiesField, answer.Type) + } + var items []string + if len(answer.Value) > 0 && string(answer.Value) != "null" { + if err := json.Unmarshal(answer.Value, &items); err != nil { + return nil, fmt.Errorf("decode extraction list: %w", err) + } + } + return items, nil +} + +// groundedEntities keeps the model's items that name something in content. +// Each item is trimmed of the spaces, quote marks, and clause punctuation a +// span can drag in at its edges, and is dropped unless it has at least two +// runes, a letter, at most MaxEntityBytes, and an occurrence in content that is +// a whole token rather than a fragment of a longer word, path, or version. +func groundedEntities(content string, items []string) []string { + out := []string{} + seen := make(map[string]bool, len(items)) + for _, item := range items { + name := strings.TrimFunc(item, isEdgeNoise) + if seen[name] || len(name) > MaxEntityBytes || utf8.RuneCountInString(name) < 2 || + !strings.ContainsFunc(name, unicode.IsLetter) || !appearsAsToken(content, name) { + continue + } + seen[name] = true + out = append(out, name) + if len(out) == MaxEntities { + break + } + } + return out +} + +func isEdgeNoise(r rune) bool { + return unicode.IsSpace(r) || strings.ContainsRune(`"'“”‘’「」『』,;:，；：、。`, r) +} + +// appearsAsToken reports whether name occurs in content without the token +// continuing on either side, directly ("TW" inside "zh-TW") or across a joiner +// ("5.1" inside "5.1.0", "w/demo" inside "~/w/demo"). +func appearsAsToken(content, name string) bool { + for start := 0; start <= len(content); { + i := strings.Index(content[start:], name) + if i < 0 { + return false + } + i += start + before, after := content[:i], content[i+len(name):] + r1, s1 := utf8.DecodeLastRuneInString(before) + r2, _ := utf8.DecodeLastRuneInString(before[:len(before)-s1]) + f1, s2 := utf8.DecodeRuneInString(after) + f2, _ := utf8.DecodeRuneInString(after[s2:]) + if !continues(r1, r2) && !continues(f1, f2) { + return true + } + start = i + 1 + } + return false +} + +// continues reports whether next, then beyond, moving away from a match, +// carry its token on. +func continues(next, beyond rune) bool { + if isWordRune(next) { + return true + } + return strings.ContainsRune(`.-/\:@`, next) && isWordRune(beyond) +} + +// isWordRune reports runes that extend a token. CJK scripts are excluded +// because they do not separate words, so any neighbour would look like one. +func isWordRune(r rune) bool { + if unicode.In(r, unicode.Han, unicode.Hiragana, unicode.Katakana, unicode.Hangul) { + return false + } + return unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' || r == '~' +} diff --git a/internal/memory/extract/extract_test.go b/internal/memory/extract/extract_test.go new file mode 100644 index 000000000..661aec16d --- /dev/null +++ b/internal/memory/extract/extract_test.go @@ -0,0 +1,164 @@ +package extract + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "reflect" + "strings" + "testing" +) + +type captured struct { + auth string + body wireRequest +} + +func serve(t *testing.T, status int, reply string) (*Client, *[]captured) { + t.Helper() + var got []captured + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req wireRequest + _ = json.NewDecoder(r.Body).Decode(&req) + got = append(got, captured{auth: r.Header.Get("Authorization"), body: req}) + w.WriteHeader(status) + _, _ = io.WriteString(w, reply) + })) + t.Cleanup(srv.Close) + t.Setenv("MNEMON_EXTRACT_MODEL", "clairic") + t.Setenv("MNEMON_EXTRACT_ENDPOINT", srv.URL+"/ai/run/@cf/cloudflare/clairic") + t.Setenv("MNEMON_EXTRACT_API_KEY", "k-secret") + t.Setenv("MNEMON_EXTRACT_TIMEOUT", "") + c, err := NewClient() + if err != nil || c == nil { + t.Fatalf("NewClient = %v, %v", c, err) + } + return c, &got +} + +func envelope(items any) string { + data, _ := json.Marshal(map[string]any{ + "success": true, + "errors": []any{}, + "result": map[string]any{"answers": map[string]any{"entities": map[string]any{"type": "list", "value": items}}}, + }) + return string(data) +} + +func TestNewClientIsNilWithoutModel(t *testing.T) { + t.Setenv("MNEMON_EXTRACT_MODEL", "") + t.Setenv("MNEMON_EXTRACT_ENDPOINT", "http://127.0.0.1:1") + if c, err := NewClient(); c != nil || err != nil { + t.Fatalf("NewClient = %v, %v; want nil, nil", c, err) + } +} + +func TestNewClientRequiresEndpoint(t *testing.T) { + t.Setenv("MNEMON_EXTRACT_MODEL", "clairic") + t.Setenv("MNEMON_EXTRACT_ENDPOINT", "") + if c, err := NewClient(); c != nil || !errors.Is(err, ErrNoEndpoint) { + t.Fatalf("NewClient = %v, %v; want nil, ErrNoEndpoint", c, err) + } +} + +func TestEntitiesSendsTypedListRequest(t *testing.T) { + // Given a Workers AI server that finds one name + c, got := serve(t, http.StatusOK, envelope([]string{"Jane Rivera"})) + content := "Jane Rivera joined the board" + + // When entities are requested + names, err := c.Entities(context.Background(), content) + + // Then the name comes back, and the request carried the model, content, and a typed list field + if err != nil || !reflect.DeepEqual(names, []string{"Jane Rivera"}) { + t.Fatalf("Entities = %v, %v", names, err) + } + if len(*got) != 1 { + t.Fatalf("requests = %d", len(*got)) + } + req := (*got)[0] + field := req.body.Fields[entitiesField] + if req.auth != "Bearer k-secret" || req.body.Model != "clairic" || req.body.State != content || + field.Type != "list" || field.Items != "text" || field.Instructions == "" { + t.Fatalf("request = %+v", req) + } +} + +func TestEntitiesAcceptsBareSystemOneResponse(t *testing.T) { + c, _ := serve(t, http.StatusOK, `{"answers":{"entities":{"type":"list","value":["Acme Foundation"]}}}`) + names, err := c.Entities(context.Background(), "funded by the Acme Foundation") + if err != nil || !reflect.DeepEqual(names, []string{"Acme Foundation"}) { + t.Fatalf("Entities = %v, %v", names, err) + } +} + +func TestEntitiesKeepsOnlyWholeTokensFromContent(t *testing.T) { + // Given model output with edge noise, fragments, a duplicate, and an invented name + content := "Added ~/w/demo for tool 5.1.0; the 「示範專案」 seal and a centre 愛; zh-TW cards; Jane Rivera joined." + items := []string{ + "~/", "w/demo", "tool 5.1", "示範專案」", "TW", " Jane Rivera ", + "Ghost Corp", "2026", "愛", "Jane Rivera", "~/w/demo", + } + c, _ := serve(t, http.StatusOK, envelope(items)) + + // When entities are requested + names, err := c.Entities(context.Background(), content) + + // Then only trimmed names that stand alone in the content survive, once each + want := []string{"示範專案", "Jane Rivera", "~/w/demo"} + if err != nil || !reflect.DeepEqual(names, want) { + t.Fatalf("Entities = %v, %v; want %v", names, err, want) + } +} + +func TestEntitiesNullListIsEmptyNotNil(t *testing.T) { + c, _ := serve(t, http.StatusOK, envelope(nil)) + names, err := c.Entities(context.Background(), "nothing named here") + if err != nil || names == nil || len(names) != 0 { + t.Fatalf("Entities = %#v, %v; want empty non-nil", names, err) + } +} + +func TestEntitiesCapsAtMaxEntities(t *testing.T) { + items := make([]string, MaxEntities+10) + for i := range items { + items[i] = fmt.Sprintf("Item%02d", i) + } + c, _ := serve(t, http.StatusOK, envelope(items)) + names, err := c.Entities(context.Background(), strings.Join(items, " ")) + if err != nil || len(names) != MaxEntities { + t.Fatalf("Entities returned %d names, %v; want %d", len(names), err, MaxEntities) + } +} + +func TestEntitiesErrorsNameTheFailureWithoutContentOrKey(t *testing.T) { + const content = "secret plan for project Zebra" + cases := []struct { + name string + status int + reply string + want string + }{ + {"http status", http.StatusUnauthorized, + `{"success":false,"errors":[{"code":10000,"message":"Authentication error"}]}`, "status 401: Authentication error"}, + {"envelope failure", http.StatusOK, `{"success":false,"errors":[{"message":"model overloaded"}]}`, "reported failure: model overloaded"}, + {"wrong answer type", http.StatusOK, `{"answers":{"entities":{"type":"text","value":"Zebra"}}}`, `type "text"`}, + {"missing answer", http.StatusOK, `{"result":{"answers":{}}}`, `no "entities" answer`}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + c, _ := serve(t, tc.status, tc.reply) + names, err := c.Entities(context.Background(), content) + if err == nil || names != nil || !strings.Contains(err.Error(), tc.want) { + t.Fatalf("Entities = %v, %v; want error containing %q", names, err, tc.want) + } + if strings.Contains(err.Error(), "k-secret") || strings.Contains(err.Error(), "secret plan") { + t.Fatalf("error leaks the key or content: %v", err) + } + }) + } +} diff --git a/internal/memory/graph/engine.go b/internal/memory/graph/engine.go index 3bfc47193..b94944527 100644 --- a/internal/memory/graph/engine.go +++ b/internal/memory/graph/engine.go @@ -36,6 +36,11 @@ const ( type EngineOptions struct { EntityMode EntityMode TemporalMode TemporalMode + // ModelEntities are entities an extraction model found in the insight's + // content. Callers fetch them before the write transaction the engine runs + // in, because network I/O must not hold the lock. Nil means no model ran; + // see ResolveEntitiesWithModel. + ModelEntities []string } // NewEngine creates a new graph edge engine. @@ -71,12 +76,13 @@ func (e *Engine) OnInsightCreated(insight *model.Insight) EdgeStats { // the existing entity index, so user vocabulary already represented in // the graph propagates into new insights without dictionary edits. On // error the index lookup falls through to nil and behavior matches the - // non-indexed extractor. + // non-indexed extractor. Entities from an extraction model, when given, + // replace the acronym guess. var knownEntities map[string]bool if e.entityMode != EntityModeProvided { knownEntities, _ = e.db.LoadKnownEntities() } - insight.Entities = ResolveEntitiesIndexed(insight.Content, insight.Entities, e.entityMode, knownEntities) + insight.Entities = ResolveEntitiesWithModel(insight.Content, insight.Entities, e.entityMode, knownEntities, e.options.ModelEntities) // 2. Temporal backbone + proximity edges if e.options.TemporalMode != TemporalDisabled { diff --git a/internal/memory/graph/entity.go b/internal/memory/graph/entity.go index f59acd0ff..7b024190c 100644 --- a/internal/memory/graph/entity.go +++ b/internal/memory/graph/entity.go @@ -37,15 +37,20 @@ func ValidEntityMode(mode EntityMode) bool { } } -// filePathPattern matches file paths and dotted names (e.g., ./cmd/root.go, -// /etc/config.yml, example.org). Its matches must also pass looksLikePath. -var filePathPattern = regexp.MustCompile(`(?:^|[\s"'(])([.\w/-]+\.\w{1,10})(?:[\s"'),.]|$)`) +var ( + // acronymPattern matches ALLCAPS acronyms 2-6 chars (e.g., API, HTTP, SQL). + // It is the one pattern that guesses whether a token names something, so + // model-extracted entities replace it (see ResolveEntitiesWithModel). + acronymPattern = regexp.MustCompile(`\b([A-Z]{2,6})\b`) + // filePathPattern matches file paths and dotted names (e.g., ./cmd/root.go, + // /etc/config.yml, example.org). Its matches must also pass looksLikePath. + filePathPattern = regexp.MustCompile(`(?:^|[\s"'(])([.\w/-]+\.\w{1,10})(?:[\s"'),.]|$)`) +) var entityPatterns = []*regexp.Regexp{ // CamelCase identifiers (e.g., MyClass, HttpServer) regexp.MustCompile(`\b([A-Z][a-z]+(?:[A-Z][a-z]+)+)\b`), - // ALLCAPS acronyms 2-6 chars (e.g., API, HTTP, gRPC, SQL) - regexp.MustCompile(`\b([A-Z]{2,6})\b`), + acronymPattern, filePathPattern, // URLs regexp.MustCompile(`https?://[^\s"'<>)]+`), @@ -147,11 +152,20 @@ var acronymStopwords = map[string]bool{ // ExtractEntities extracts named entities from text using regex patterns // and a technology dictionary for common terms regex would miss. func ExtractEntities(text string) []string { + return extractEntities(text, true) +} + +// extractEntities runs the regex and dictionary paths. guessAcronyms selects +// acronymPattern; callers holding model-extracted entities leave it out. +func extractEntities(text string, guessAcronyms bool) []string { seen := make(map[string]bool) var entities []string // 1. Regex-based extraction for _, pat := range entityPatterns { + if pat == acronymPattern && !guessAcronyms { + continue + } matches := pat.FindAllStringSubmatch(text, -1) for _, m := range matches { entity := m[len(m)-1] @@ -222,7 +236,13 @@ func ResolveEntities(content string, provided []string, mode EntityMode) []strin // ExtractEntities. Callers without DB access can pass nil to get the // pre-existing semantics. func ExtractEntitiesIndexed(text string, knownEntities map[string]bool) []string { - entities := ExtractEntities(text) + return extractEntitiesIndexed(text, knownEntities, true) +} + +// extractEntitiesIndexed is ExtractEntitiesIndexed with the acronym guess +// optional (see extractEntities). +func extractEntitiesIndexed(text string, knownEntities map[string]bool, guessAcronyms bool) []string { + entities := extractEntities(text, guessAcronyms) if len(knownEntities) == 0 { return entities } @@ -272,6 +292,35 @@ func ResolveEntitiesIndexed(content string, provided []string, mode EntityMode, } } +// ReliesOnExtraction reports whether resolving entities for mode and provided +// uses automatic extraction alone: auto mode, or merge mode with nothing +// provided. Only then is an extraction model worth asking, because entities a +// caller provides already name what the memory is about. +func ReliesOnExtraction(mode EntityMode, provided []string) bool { + switch mode { + case EntityModeAuto: + return true + case EntityModeMerge: + return len(provided) == 0 + default: + return false + } +} + +// ResolveEntitiesWithModel is ResolveEntitiesIndexed for content an extraction +// model has read. When resolution relies on extraction alone (see +// ReliesOnExtraction), the model's entities replace the acronym guess, while +// the exact paths (CamelCase, file paths, URLs, @mentions, title marks, +// dictionary) and the known-entity index still apply. A nil modelEntities +// means no model ran and is identical to ResolveEntitiesIndexed; an empty, +// non-nil one means the model found nothing. +func ResolveEntitiesWithModel(content string, provided []string, mode EntityMode, knownEntities map[string]bool, modelEntities []string) []string { + if modelEntities == nil || !ReliesOnExtraction(mode, provided) { + return ResolveEntitiesIndexed(content, provided, mode, knownEntities) + } + return mergeEntities(modelEntities, extractEntitiesIndexed(content, knownEntities, false)) +} + // splitWords splits text into ASCII-alphanumeric words preserving original casing. func splitWords(text string) []string { return strings.FieldsFunc(text, func(r rune) bool { diff --git a/internal/memory/graph/entity_test.go b/internal/memory/graph/entity_test.go index c84764cb3..01bd0f03f 100644 --- a/internal/memory/graph/entity_test.go +++ b/internal/memory/graph/entity_test.go @@ -427,6 +427,66 @@ func TestExtractEntities_EmphasisWordsAreNotAcronyms(t *testing.T) { } } +// --- Model-extracted entities --- + +const modelContent = "Jane Rivera briefed the CEO on HttpServer and Athena; notes in docs/plan.md, built with Go" + +var modelKnown = map[string]bool{"Athena": true} + +func TestResolveEntitiesWithModel_ReplacesAcronymGuessWhenNothingProvided(t *testing.T) { + if !toSet(ResolveEntitiesIndexed(modelContent, nil, EntityModeMerge, modelKnown))["CEO"] { + t.Fatal("baseline should guess the acronym CEO") + } + for _, mode := range []EntityMode{EntityModeMerge, EntityModeAuto} { + entities := ResolveEntitiesWithModel(modelContent, nil, mode, modelKnown, []string{"Jane Rivera"}) + has := toSet(entities) + for _, want := range []string{"Jane Rivera", "HttpServer", "Athena", "docs/plan.md", "Go"} { + if !has[want] { + t.Errorf("mode=%s: want %q, got %v", mode, want, entities) + } + } + if has["CEO"] { + t.Errorf("mode=%s: model entities replace the acronym guess, got %v", mode, entities) + } + } +} + +func TestResolveEntitiesWithModel_EmptyModelStillDropsAcronymGuess(t *testing.T) { + entities := ResolveEntitiesWithModel(modelContent, nil, EntityModeMerge, modelKnown, []string{}) + has := toSet(entities) + if has["CEO"] || !has["HttpServer"] { + t.Errorf("a model that found nothing keeps exact paths only, got %v", entities) + } +} + +func TestResolveEntitiesWithModel_AutoModeIgnoresProvided(t *testing.T) { + entities := ResolveEntitiesWithModel(modelContent, []string{"zebra"}, EntityModeAuto, modelKnown, []string{"Jane Rivera"}) + has := toSet(entities) + if has["zebra"] || !has["Jane Rivera"] || has["CEO"] { + t.Errorf("auto mode should use the model and ignore provided entities, got %v", entities) + } +} + +func TestResolveEntitiesWithModel_ProvidedEntitiesBypassTheModel(t *testing.T) { + for _, mode := range []EntityMode{EntityModeMerge, EntityModeProvided} { + got := ResolveEntitiesWithModel(modelContent, []string{"zebra"}, mode, modelKnown, []string{"Jane Rivera"}) + want := ResolveEntitiesIndexed(modelContent, []string{"zebra"}, mode, modelKnown) + if !equalAsSets(got, want) { + t.Errorf("mode=%s: provided entities should bypass the model; got %v want %v", mode, got, want) + } + } +} + +func TestResolveEntitiesWithModel_NilModelEqualsIndexed(t *testing.T) { + for _, mode := range []EntityMode{EntityModeMerge, EntityModeProvided, EntityModeAuto} { + got := ResolveEntitiesWithModel(modelContent, nil, mode, modelKnown, nil) + want := ResolveEntitiesIndexed(modelContent, nil, mode, modelKnown) + if !equalAsSets(got, want) { + t.Errorf("mode=%s: nil model should equal ResolveEntitiesIndexed; got %v want %v", mode, got, want) + } + } +} + func equalAsSets(a, b []string) bool { if len(a) != len(b) { return false diff --git a/test/mnemond/architecture/mnemond_structure_test.go b/test/mnemond/architecture/mnemond_structure_test.go index e853c7061..6250b4fbf 100644 --- a/test/mnemond/architecture/mnemond_structure_test.go +++ b/test/mnemond/architecture/mnemond_structure_test.go @@ -69,6 +69,7 @@ func assertMnemondPackageGraph(t *testing.T, root string) { t.Helper() want := map[string][]string{ "internal/memory/embed": {}, + "internal/memory/extract": {}, "internal/memory/internal/topk": {}, "internal/memory/model": {}, "internal/memory/importdraft": {"internal/memory/model"}, @@ -94,7 +95,7 @@ func assertMnemondPackageGraph(t *testing.T, root string) { }, "cmd/agency": {"internal/agency/attach", "internal/agency/client", "internal/daemon"}, "cmd/memory": { - "internal/memory/embed", "internal/memory/graph", "internal/memory/importdraft", + "internal/memory/embed", "internal/memory/extract", "internal/memory/graph", "internal/memory/importdraft", "internal/memory/model", "internal/memory/search", "internal/memory/setup", "internal/memory/setup/assets", "internal/memory/store", "internal/memory/web", "web", },