diff --git a/.env.example b/.env.example index 96e8af6..caa5eab 100644 --- a/.env.example +++ b/.env.example @@ -3,4 +3,6 @@ EMBEDDING_DIM=768 SERVER_ADDR=:8080 OLLAMA_ADDR=11434 MAX_MEMORY_DISTANCE=0.5 -TOP_K_MEMORIES=10 \ No newline at end of file +TOP_K_MEMORIES=10 +MIN_EVENT_WEIGHT=0.3 +DUPLICATE_DISTANCE=0.05 \ No newline at end of file diff --git a/README.md b/README.md index 653236b..43d9cbb 100644 --- a/README.md +++ b/README.md @@ -33,11 +33,12 @@ Memcortex is a Proof of Concept (PoC) designed to equip conversational agents an Memcortex stores user/application memories as both text and vectors in Weaviate and exposes a memory manager + middleware that: -1. Embeds incoming text using `nomic-embed-text` embeddings. -2. Stores memories in a `Memory_idx` class on Weaviate. -3. Runs vector searches to retrieve top‑K relevant memories for a user. -4. Injects retrieved memories into the prompt before it reaches the LLM. -5. Optionally persists new memories asynchronously. +1. **Classifies** incoming messages through an event-driven pipeline to determine if they are worth persisting. +2. **Embeds** meaningful text using `nomic-embed-text` embeddings. +3. **Deduplicates** against existing memories using vector similarity before saving. +4. **Stores** memories in a `Memory_idx` class on Weaviate with metadata (event type, weight, TTL). +5. **Retrieves** top-K relevant memories via vector search, scoped per user. +6. **Injects** retrieved memories into the prompt before it reaches the LLM. This pattern is ideal for building chatbots, agents, and personalization layers that must "remember" details across sessions. @@ -55,13 +56,15 @@ flowchart LR C -->|retrieve| D[Weaviate Vector Store] D -->|top-K| C C -->|inject memories| E[LLM Handler] - E -->|call LLM API| F[Ollama / Custom LLM] - F -->|response| E - E -->|save message| G(Background Save Worker) - G --> D + E -->|response| A + E -->|save message| F{Event Classifier} + F -->|trivial / low-value| G[Discard] + F -->|meaningful| H{Duplicate Check} + H -->|duplicate| G + H -->|new| I[Embed + Store] + I --> D subgraph Infra D - F end ``` @@ -72,20 +75,62 @@ User -> API Server (/chat) -> Memory Middleware: - Embed user query via Ollama - Query Weaviate vector index (top-K) - - Re-rank / filter / format - - Inject into prompt - -> LLM Handler -> Local or remote LLM - -> Return response -Background worker: saves new user messages into Weaviate (embedding -> object) + - Format and inject into request context + -> Chat Handler: + - Return response with related memories + - Classify message (preference / goal / correction / profile / achievement) + - If low-value or trivial -> discard + - If meaningful -> embed via Ollama + - Check for near-duplicate memories -> skip if found + - Save to Weaviate with metadata (event_type, weight, ttl_days) ``` --- +## Event-Driven Memory + +Not every message is worth remembering. Memcortex includes an event classifier that scores each incoming message and only persists meaningful events. + +### How it works + +1. Each message is scored by a multi-signal heuristic classifier. +2. Positive signals (pattern matches) boost the score. Negative signals (questions, negations, third-person references) apply penalties. +3. If the final weight is below `MIN_EVENT_WEIGHT`, the message is discarded. +4. If the message passes classification, it is embedded and checked against existing memories for near-duplicates. +5. New, unique memories are stored with metadata: `event_type`, `weight`, and `ttl_days`. + +### Event types + +| Type | Weight | TTL | Examples | +|------|--------|-----|----------| +| `correction` | 0.95 | never | "Actually my role is...", "I meant..." | +| `profile` | 0.90 | never | "My name is...", "I work at...", "Call me..." | +| `preference` | 0.90 | never | "I prefer...", "My favorite...", "I'd rather..." | +| `goal` | 0.85 | 90 days | "I plan to...", "I'm working on...", "Remind me to..." | +| `achievement` | 0.80 | 365 days | "I shipped...", "I got promoted...", "I launched..." | +| `general` | 0.15-0.45 | 30 days | Unmatched messages, scored by word count | +| `trivial` | 0.00 | -- | "ok", "thanks", "lol", short filler | + +### Penalties + +| Signal | Penalty | Rationale | +|--------|---------|-----------| +| Question (`?`, "what", "how do you"...) | -0.25 | Queries are not facts to remember | +| Negation ("don't", "not", "never"...) | -0.15 | Weakens the signal of the statement | +| Third-person ("he", "she", "they" without "I") | -0.35 | Not about the user | + +### Configuration + +- `MIN_EVENT_WEIGHT=0.3` -- Threshold for saving. Set to `0` to save everything (disables filtering). +- `DUPLICATE_DISTANCE=0.05` -- Vector distance below which a memory is considered a duplicate. Set to `0` to disable dedup. + +--- + ## Quickstart (developer) Prereqs: -* Go 1.20+ +* Go 1.24+ * Docker & Docker Compose 1. Copy repo and set module path (or `go mod init github.com/yourname/memcortex`). @@ -128,19 +173,22 @@ The first request will save the memory asynchronously. Later requests will retri ``` memcortex/ -├─ cmd/server/main.go # App entry point +├─ cmd/server/main.go # App entry point ├─ internal/ +│ ├─ classifier/ +│ │ ├─ classifier.go # Event classifier (heuristic scoring + penalties) +│ │ └─ classifier_test.go # Classifier unit tests │ ├─ embedder/ -│ │ └─ ollama.go # Contains OpenAI Embedder Logic (for text-embedding-3-small) +│ │ └─ ollama.go # Ollama embedding client │ ├─ handlers/ -│ │ └─ chat.go # Chat endpoint handler +│ │ └─ chat.go # Chat endpoint handler │ ├─ memory/ -│ │ ├─ manager.go # High-level RAG orchestration -│ │ ├─ queue.go -│ │ └─ store.go # Weaviate storage wrapper +│ │ ├─ manager.go # RAG orchestration, classification gate, duplicate detection +│ │ ├─ queue.go # Embedding queue with worker pool +│ │ └─ store.go # Weaviate storage wrapper │ └─ middleware/ -│ └─ memory_middleware.go # Context injection middleware -├─ .env.example # Environment file +│ └─ memory_middleware.go # Context injection middleware +├─ .env.example ├─ docker-compose.yml ├─ Dockerfile ├─ Dockerfile.ollama @@ -157,8 +205,10 @@ EMBEDDING_MODEL=nomic-embed-text EMBEDDING_DIM=768 SERVER_ADDR=:8080 OLLAMA_ADDR=11434 -MAX_MEMORY_DISTANCE=0.5 // This describes the vector search distance +MAX_MEMORY_DISTANCE=0.5 # Vector search distance threshold for retrieval TOP_K_MEMORIES=10 +MIN_EVENT_WEIGHT=0.3 # Minimum classifier weight to persist a memory (0 = save everything) +DUPLICATE_DISTANCE=0.05 # Max vector distance to consider a memory a duplicate (0 = disable) ``` --- diff --git a/cmd/server/main.go b/cmd/server/main.go index e04604d..48729b8 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -8,6 +8,7 @@ import ( "strings" "github.com/joho/godotenv" + "github.com/sobowalebukola/memcortex/internal/classifier" ollama "github.com/sobowalebukola/memcortex/internal/embedder" "github.com/sobowalebukola/memcortex/internal/handlers" "github.com/sobowalebukola/memcortex/internal/memory" @@ -41,8 +42,11 @@ func main() { log.Fatalf("failed to create memory store: %v", err) } emb := ollama.NewEmbeddingClient(os.Getenv("EMBEDDING_MODEL")) + cls := classifier.New() - m := memory.NewManager(store, emb) + m := memory.NewManager(store, emb, cls) + + log.Printf("Event classifier enabled (min_weight=%.2f)", cls.MinWeight) log.Println("MemCortex initialized successfully!") diff --git a/internal/classifier/classifier.go b/internal/classifier/classifier.go new file mode 100644 index 0000000..8b35c27 --- /dev/null +++ b/internal/classifier/classifier.go @@ -0,0 +1,299 @@ +package classifier + +import ( + "os" + "slices" + "strconv" + "strings" +) + +// EventType categorizes the kind of meaningful event a message represents. +type EventType string + +const ( + Preference EventType = "preference" + Goal EventType = "goal" + Correction EventType = "correction" + Profile EventType = "profile" + Achievement EventType = "achievement" + General EventType = "general" + Trivial EventType = "trivial" +) + +// Result is the output of classifying a single message. +type Result struct { + ShouldSave bool + Type EventType + Weight float64 + TTLDays int // 0 = no expiry +} + +// Classifier scores incoming messages to decide whether they are worth persisting. +// It uses a multi-signal approach: positive pattern rules boost weight, negative +// signals (questions, negations, third-person references) apply penalties. +type Classifier struct { + MinWeight float64 +} + +type rule struct { + eventType EventType + baseWeight float64 + ttlDays int + patterns []string +} + +// rules are checked in order; the highest-baseWeight match wins. +var rules = []rule{ + { + eventType: Correction, + baseWeight: 0.95, + ttlDays: 0, // corrections never expire + patterns: []string{ + "actually my ", "actually i ", "i meant ", "correction:", + "i'm actually ", "my role is actually", "i was wrong", + "that's not right", "let me correct", "to clarify,", + "i should clarify", "to be precise", + }, + }, + { + eventType: Profile, + baseWeight: 0.90, + ttlDays: 0, + patterns: []string{ + "my name is", "i work at", "i'm a ", "i am a ", + "i live in", "my email is", "my job is", "i work as", + "call me ", "i'm based in", "i'm from ", + "my phone", "my address", "my birthday", + "i specialize in", "my title is", "my team is", + }, + }, + { + eventType: Preference, + baseWeight: 0.90, + ttlDays: 0, + patterns: []string{ + "i prefer", "i like ", "my favorite", "i enjoy ", + "i love ", "i hate ", "i dislike", "i don't like", + "i always use", "i never use", "my go-to ", + "i'm a fan of", "i tend to use", "i usually ", + "i'd rather ", "my preferred", "i opt for", + "i'm not a fan of", + }, + }, + { + eventType: Goal, + baseWeight: 0.85, + ttlDays: 90, + patterns: []string{ + "remind me to", "i want to learn", "i need to finish", + "i plan to", "my goal is", "i'm going to start", + "i'm working on", "i'm trying to", "i aim to", + "i intend to", "by next ", "my deadline", + "i'm studying", "i'm building", "i'm learning", + "i want to build", "i need to complete", + }, + }, + { + eventType: Achievement, + baseWeight: 0.80, + ttlDays: 365, + patterns: []string{ + "i completed", "i finished", "i achieved", + "i got promoted", "i graduated", "i passed ", + "i built ", "i shipped ", "i launched", + "i earned", "i published", "i released", + "i solved", "i won ", "i was promoted", + "i got certified", "i was hired", + }, + }, +} + +func New() *Classifier { + minWeight := 0.3 + if v, err := strconv.ParseFloat(os.Getenv("MIN_EVENT_WEIGHT"), 64); err == nil && v > 0 { + minWeight = v + } + return &Classifier{MinWeight: minWeight} +} + +func (c *Classifier) Classify(text string) Result { + lower := strings.ToLower(strings.TrimSpace(text)) + + // Trivial messages are never worth saving. + if len(lower) < 10 || isTrivial(lower) { + return Result{Type: Trivial, Weight: 0.0} + } + + // Find the best matching rule. When multiple rules match with equal + // weight, prefer the longest pattern (most specific match). + var bestRule *rule + var bestLen int + for i := range rules { + for _, p := range rules[i].patterns { + if strings.Contains(lower, p) { + if bestRule == nil || + rules[i].baseWeight > bestRule.baseWeight || + (rules[i].baseWeight == bestRule.baseWeight && len(p) > bestLen) { + bestRule = &rules[i] + bestLen = len(p) + } + } + } + } + best := bestRule + + // Compute penalty from negative signals. + pen := penalty(lower) + + if best != nil { + w := clamp(best.baseWeight-pen, 0, 1) + return Result{ + ShouldSave: w >= c.MinWeight, + Type: best.eventType, + Weight: w, + TTLDays: best.ttlDays, + } + } + + // No pattern matched — score as general content. + w := clamp(generalWeight(lower)-pen, 0, 1) + return Result{ + ShouldSave: w >= c.MinWeight, + Type: General, + Weight: w, + TTLDays: 30, + } +} + +// penalty reduces weight based on signals that the message is not a +// meaningful first-person statement worth persisting. +func penalty(text string) float64 { + p := 0.0 + + if isQuestion(text) { + p += 0.25 + } + if hasNegation(text) { + p += 0.15 + } + if isThirdPerson(text) { + p += 0.35 + } + + return p +} + +// generalWeight scores unmatched messages by word count as a rough proxy +// for substantive content. +func generalWeight(text string) float64 { + words := len(strings.Fields(text)) + switch { + case words > 25: + return 0.45 + case words > 15: + return 0.35 + case words > 8: + return 0.25 + default: + return 0.15 + } +} + +func isQuestion(text string) bool { + if strings.Contains(text, "?") { + return true + } + starters := []string{ + "what ", "who ", "where ", "when ", "why ", "how ", + "do you", "can you", "could you", "would you", + "is there", "are there", "does ", "did ", + } + for _, s := range starters { + if strings.HasPrefix(text, s) { + return true + } + } + return false +} + +func hasNegation(text string) bool { + // Pad with spaces so we match whole words only. + // Without this, "nothing" would falsely match "not ". + padded := " " + text + " " + negations := []string{ + " don't ", " doesn't ", " didn't ", " not ", " never ", + " no longer ", " can't ", " won't ", " wouldn't ", + } + for _, n := range negations { + if strings.Contains(padded, n) { + return true + } + } + return false +} + +// isThirdPerson returns true when the message talks about someone else +// rather than the user themselves. Uses space-padded matching to avoid +// substring false positives (e.g. "the" matching "he"). +func isThirdPerson(text string) bool { + padded := " " + text + " " + + firstPerson := []string{" i ", " i'm ", " i've ", " i'd ", " my ", " me "} + thirdPerson := []string{" he ", " she ", " they ", " his ", " her ", " their "} + + hasFirst := false + for _, fp := range firstPerson { + if strings.Contains(padded, fp) { + hasFirst = true + break + } + } + + hasThird := false + for _, tp := range thirdPerson { + if strings.Contains(padded, tp) { + hasThird = true + break + } + } + + // Only penalize when message references others but not the user. + return hasThird && !hasFirst +} + +var trivialPhrases = []string{ + // Acknowledgements + "ok", "okay", "sure", "yes", "no", "yeah", "nah", "yep", "nope", + "k", "kk", "yup", "ya", "ye", + // Gratitude + "thanks", "thank you", "thx", "ty", "np", "no problem", + // Greetings + "hi", "hello", "hey", "yo", "sup", "howdy", + // Farewells + "bye", "goodbye", "see you", "later", "cya", + // Reactions + "lol", "haha", "heh", "lmao", "rofl", + // Agreement / filler + "got it", "cool", "nice", "great", "good", "fine", "alright", + "sounds good", "makes sense", "fair enough", "right", "true", + "same", "agreed", "exactly", "indeed", "absolutely", + // Hesitation + "hmm", "hm", "um", "uh", "idk", "idc", + // Bare question words + "what", "why", "how", "when", "where", +} + +func isTrivial(text string) bool { + cleaned := strings.TrimRight(text, ".!?,") + return slices.Contains(trivialPhrases, cleaned) +} + +func clamp(v, min, max float64) float64 { + if v < min { + return min + } + if v > max { + return max + } + return v +} diff --git a/internal/classifier/classifier_test.go b/internal/classifier/classifier_test.go new file mode 100644 index 0000000..87fd565 --- /dev/null +++ b/internal/classifier/classifier_test.go @@ -0,0 +1,291 @@ +package classifier + +import ( + "testing" +) + +func newTestClassifier(t *testing.T) *Classifier { + t.Setenv("MIN_EVENT_WEIGHT", "0.3") + return New() +} + +// --- Trivial messages should always be skipped --- + +func TestClassify_Trivial(t *testing.T) { + c := newTestClassifier(t) + cases := []string{ + "ok", "OK", "Ok.", "okay", "sure", "yes", "no", + "yeah", "nah", "yep", "nope", + "thanks", "thanks!", "thank you", "thx", "ty", + "hi", "hello", "hey", "yo", + "bye", "goodbye", "see you", "later", + "lol", "haha", "heh", "lmao", + "got it", "cool", "nice", "great", "good", "fine", + "sounds good", "makes sense", "fair enough", + "hmm", "um", "idk", + "what", "why", "how", + // Under 10 chars, always trivial + "short", "nah ok", "k", + } + for _, input := range cases { + r := c.Classify(input) + if r.ShouldSave { + t.Errorf("Classify(%q): expected skip, got save (type=%s, weight=%.2f)", input, r.Type, r.Weight) + } + if r.Type != Trivial { + t.Errorf("Classify(%q): expected type=trivial, got %s", input, r.Type) + } + } +} + +// --- Each event type should be correctly identified --- + +func TestClassify_Preference(t *testing.T) { + c := newTestClassifier(t) + cases := []string{ + "I prefer dark mode for all my editors", + "I like using Vim for Go development", + "My favorite language is Rust", + "I always use tabs instead of spaces", + "I'd rather use PostgreSQL than MySQL", + "My preferred framework is Next.js", + "I'm a fan of functional programming", + "I tend to use Linux for development", + "I'm not a fan of microservices", + } + for _, input := range cases { + r := c.Classify(input) + if !r.ShouldSave { + t.Errorf("Classify(%q): expected save, got skip (weight=%.2f)", input, r.Weight) + } + if r.Type != Preference { + t.Errorf("Classify(%q): expected type=preference, got %s", input, r.Type) + } + if r.TTLDays != 0 { + t.Errorf("Classify(%q): preferences should not expire, got TTL=%d", input, r.TTLDays) + } + } +} + +func TestClassify_Profile(t *testing.T) { + c := newTestClassifier(t) + cases := []string{ + "My name is Alex and I work at Google", + "I'm a senior backend engineer", + "I live in San Francisco", + "I work as a data scientist at Meta", + "Call me Alex", + "I'm based in Berlin", + "I specialize in distributed systems", + "My title is Staff Engineer", + } + for _, input := range cases { + r := c.Classify(input) + if !r.ShouldSave { + t.Errorf("Classify(%q): expected save, got skip (weight=%.2f)", input, r.Weight) + } + if r.Type != Profile { + t.Errorf("Classify(%q): expected type=profile, got %s", input, r.Type) + } + if r.TTLDays != 0 { + t.Errorf("Classify(%q): profile info should not expire, got TTL=%d", input, r.TTLDays) + } + } +} + +func TestClassify_Correction(t *testing.T) { + c := newTestClassifier(t) + cases := []string{ + "Actually my role is backend engineer, not frontend", + "I meant Python, not JavaScript", + "I'm actually a contractor, not full-time", + "That's not right, my team is platform engineering", + "To clarify, I work on the payments service", + "I was wrong about the deadline, it's next Friday", + } + for _, input := range cases { + r := c.Classify(input) + if !r.ShouldSave { + t.Errorf("Classify(%q): expected save, got skip (weight=%.2f)", input, r.Weight) + } + if r.Type != Correction { + t.Errorf("Classify(%q): expected type=correction, got %s", input, r.Type) + } + } +} + +func TestClassify_Goal(t *testing.T) { + c := newTestClassifier(t) + cases := []string{ + "Remind me to study Rust this month", + "I want to learn Kubernetes by next quarter", + "I plan to migrate our services to gRPC", + "I'm working on a side project in Go", + "I'm trying to get my AWS certification", + "I'm building a memory layer for LLMs", + "I'm learning Japanese on the side", + } + for _, input := range cases { + r := c.Classify(input) + if !r.ShouldSave { + t.Errorf("Classify(%q): expected save, got skip (weight=%.2f)", input, r.Weight) + } + if r.Type != Goal { + t.Errorf("Classify(%q): expected type=goal, got %s", input, r.Type) + } + if r.TTLDays != 90 { + t.Errorf("Classify(%q): goals should have TTL=90, got %d", input, r.TTLDays) + } + } +} + +func TestClassify_Achievement(t *testing.T) { + c := newTestClassifier(t) + cases := []string{ + "I shipped the new payment service last week", + "I got promoted to Staff Engineer", + "I graduated from MIT with a CS degree", + "I launched my first open source project", + "I earned my AWS Solutions Architect certification", + "I was hired at Stripe last month", + } + for _, input := range cases { + r := c.Classify(input) + if !r.ShouldSave { + t.Errorf("Classify(%q): expected save, got skip (weight=%.2f)", input, r.Weight) + } + if r.Type != Achievement { + t.Errorf("Classify(%q): expected type=achievement, got %s", input, r.Type) + } + if r.TTLDays != 365 { + t.Errorf("Classify(%q): achievements should have TTL=365, got %d", input, r.TTLDays) + } + } +} + +// --- Negative signals should reduce weight --- + +func TestClassify_QuestionPenalty(t *testing.T) { + c := newTestClassifier(t) + + // A direct preference statement should have full weight. + statement := c.Classify("I prefer dark mode for all my code editors") + // The same content as a question should be penalized. + question := c.Classify("Do you know what I prefer for my editors?") + + if question.Weight >= statement.Weight { + t.Errorf("Question (weight=%.2f) should be penalized below statement (weight=%.2f)", + question.Weight, statement.Weight) + } +} + +func TestClassify_ThirdPersonPenalty(t *testing.T) { + c := newTestClassifier(t) + + // Third-person without first-person should be heavily penalized. + r := c.Classify("He works at Google as a senior engineer") + if r.ShouldSave { + t.Errorf("Third-person message should be skipped, got save (type=%s, weight=%.2f)", r.Type, r.Weight) + } + + // Mixed first+third person should NOT be penalized (user talking about themselves and others). + r = c.Classify("I work at Google and he works at Meta") + if !r.ShouldSave { + t.Errorf("Mixed first+third person should be saved, got skip (weight=%.2f)", r.Weight) + } +} + +func TestClassify_NegationPenalty(t *testing.T) { + c := newTestClassifier(t) + + positive := c.Classify("I prefer using dark mode in all editors") + negated := c.Classify("I don't prefer using dark mode in any editors") + + if negated.Weight >= positive.Weight { + t.Errorf("Negated (weight=%.2f) should be penalized below positive (weight=%.2f)", + negated.Weight, positive.Weight) + } +} + +// --- General messages: length-based scoring --- + +func TestClassify_GeneralShort_Skipped(t *testing.T) { + c := newTestClassifier(t) + r := c.Classify("the weather is nice today") + if r.ShouldSave { + t.Errorf("Short general message should be skipped, got save (weight=%.2f)", r.Weight) + } + if r.Type != General { + t.Errorf("Expected type=general, got %s", r.Type) + } +} + +func TestClassify_GeneralLong_Saved(t *testing.T) { + c := newTestClassifier(t) + long := "We had a detailed architecture review today and decided to move the payment " + + "processing pipeline from a monolith to a set of microservices using gRPC for " + + "inter-service communication and Kafka for event streaming" + r := c.Classify(long) + if !r.ShouldSave { + t.Errorf("Long substantive message should be saved, got skip (weight=%.2f)", r.Weight) + } + if r.TTLDays != 30 { + t.Errorf("General messages should have TTL=30, got %d", r.TTLDays) + } +} + +// --- Boundary: highest-weight rule wins when multiple match --- + +func TestClassify_HighestWeightWins(t *testing.T) { + c := newTestClassifier(t) + // Contains both correction pattern ("actually my") and profile pattern ("my name is"). + // Correction has higher baseWeight (0.95) so it should win. + r := c.Classify("Actually my name is not Bob, it's Alex") + if r.Type != Correction { + t.Errorf("When correction and profile both match, correction (higher weight) should win, got %s", r.Type) + } +} + +// --- MinWeight threshold --- + +func TestClassify_CustomMinWeight(t *testing.T) { + t.Setenv("MIN_EVENT_WEIGHT", "0.99") + defer t.Setenv("MIN_EVENT_WEIGHT", "0.3") + + c := New() + // Even a strong signal shouldn't pass a very high threshold. + r := c.Classify("Actually my role is backend engineer") + // Correction base weight is 0.95, which is below 0.99. + if r.ShouldSave { + t.Errorf("With MinWeight=0.99, correction (weight=%.2f) should be skipped", r.Weight) + } +} + +// --- Edge cases --- + +func TestClassify_EmptyString(t *testing.T) { + c := newTestClassifier(t) + r := c.Classify("") + if r.ShouldSave { + t.Errorf("Empty string should be skipped") + } +} + +func TestClassify_WhitespaceOnly(t *testing.T) { + c := newTestClassifier(t) + r := c.Classify(" \t \n ") + if r.ShouldSave { + t.Errorf("Whitespace-only should be skipped") + } +} + +func TestClassify_CaseInsensitive(t *testing.T) { + c := newTestClassifier(t) + r := c.Classify("MY NAME IS ALEX AND I WORK AT GOOGLE") + if !r.ShouldSave { + t.Errorf("Classification should be case-insensitive, got skip (weight=%.2f)", r.Weight) + } + if r.Type != Profile { + t.Errorf("Expected type=profile, got %s", r.Type) + } +} diff --git a/internal/handlers/chat.go b/internal/handlers/chat.go index 4707889..4504119 100644 --- a/internal/handlers/chat.go +++ b/internal/handlers/chat.go @@ -43,13 +43,15 @@ func (h *ChatHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { response := req.Message - if err := h.Manager.Save(r.Context(), userID, req.Message); err != nil { + saved, err := h.Manager.Save(r.Context(), userID, req.Message) + if err != nil { log.Printf("Failed to save message for user %s: %v", userID, err) http.Error(w, "failed to save message", http.StatusInternalServerError) return } - - log.Printf("Saved message for user %s: %s", userID, req.Message) + if saved { + log.Printf("Saved message for user %s: %s", userID, req.Message) + } out := ChatResp{ Response: response, Memories: memories, diff --git a/internal/memory/manager.go b/internal/memory/manager.go index 269bcbe..dc8cf23 100644 --- a/internal/memory/manager.go +++ b/internal/memory/manager.go @@ -2,18 +2,22 @@ package memory import ( "context" + "log" "os" "strconv" "strings" "time" + "github.com/sobowalebukola/memcortex/internal/classifier" ollama "github.com/sobowalebukola/memcortex/internal/embedder" ) type Manager struct { - Store *Store - Embedder *ollama.EmbeddingClient - TopK int + Store *Store + Embedder *ollama.EmbeddingClient + Classifier *classifier.Classifier + TopK int + DuplicateThreshold float64 } type MemoryPrompt struct { @@ -21,7 +25,7 @@ type MemoryPrompt struct { Added string `json:"added"` } -func NewManager(store *Store, emb *ollama.EmbeddingClient) *Manager { +func NewManager(store *Store, emb *ollama.EmbeddingClient, cls *classifier.Classifier) *Manager { topKStr := os.Getenv("TOP_K_MEMORIES") if topKStr == "" { topKStr = "10" @@ -30,7 +34,18 @@ func NewManager(store *Store, emb *ollama.EmbeddingClient) *Manager { if err != nil { topK = 10 } - return &Manager{Store: store, Embedder: emb, TopK: topK} + dupThresh := 0.05 + if v, err := strconv.ParseFloat(os.Getenv("DUPLICATE_DISTANCE"), 64); err == nil && v > 0 { + dupThresh = v + } + + return &Manager{ + Store: store, + Embedder: emb, + Classifier: cls, + TopK: topK, + DuplicateThreshold: dupThresh, + } } func (m *Manager) Retrieve(ctx context.Context, userID, query string) ([]Memory, error) { @@ -44,22 +59,54 @@ func (m *Manager) Retrieve(ctx context.Context, userID, query string) ([]Memory, func (m *Manager) SaveAsync(ctx context.Context, userID, text string) { go func() { - _ = m.Save(ctx, userID, text) + _, _ = m.Save(ctx, userID, text) }() } -func (m *Manager) Save(ctx context.Context, userID, text string) error { - emb, err := m.Embedder.Embed(ctx, text) +// Save classifies the message and only persists it if it's meaningful. +// Returns (true, nil) if saved, (false, nil) if skipped, or (false, err) on failure. +func (m *Manager) Save(ctx context.Context, userID, text string) (bool, error) { + result := m.Classifier.Classify(text) + if !result.ShouldSave { + log.Printf("Skipped low-value message for user %s (type=%s, weight=%.2f): %s", + userID, result.Type, result.Weight, truncate(text, 60)) + return false, nil + } + + log.Printf("Classified message for user %s as %s (weight=%.2f, ttl=%dd)", + userID, result.Type, result.Weight, result.TTLDays) + emb, err := m.Embedder.Embed(ctx, text) if err != nil { - return err + return false, err + } + + // Check for near-duplicate memories before saving. + if m.DuplicateThreshold > 0 { + existing, searchErr := m.Store.Search(ctx, emb, userID, 1) + if searchErr == nil && len(existing) > 0 && existing[0].Distance < m.DuplicateThreshold { + log.Printf("Skipped duplicate memory for user %s (distance=%.4f): %s", + userID, existing[0].Distance, truncate(text, 60)) + return false, nil + } } + emb32 := make([]float32, len(emb)) for i, v := range emb { emb32[i] = float32(v) } - _, err = m.Store.Save(ctx, userID, text, emb32) - return err + _, err = m.Store.Save(ctx, userID, text, emb32, string(result.Type), result.Weight, result.TTLDays) + if err != nil { + return false, err + } + return true, nil +} + +func truncate(s string, max int) string { + if len(s) <= max { + return s + } + return s[:max] + "..." } func FormatMemoryPrompt(memories []Memory) []MemoryPrompt { diff --git a/internal/memory/store.go b/internal/memory/store.go index a641d3a..b74b608 100644 --- a/internal/memory/store.go +++ b/internal/memory/store.go @@ -18,6 +18,7 @@ type Memory struct { ID string `json:"id"` Text string `json:"text"` Timestamp time.Time `json:"timestamp"` + Distance float64 `json:"distance,omitempty"` } type Store struct { @@ -54,7 +55,7 @@ func NewStore(class string, dim int, host string) (*Store, error) { // --------------------------- // Save Memory // --------------------------- -func (s *Store) Save(ctx context.Context, userID, text string, embedding []float32) (string, error) { +func (s *Store) Save(ctx context.Context, userID, text string, embedding []float32, eventType string, weight float64, ttlDays int) (string, error) { if len(embedding) != s.Dim { return "", fmt.Errorf("embedding dimension mismatch") } @@ -62,10 +63,13 @@ func (s *Store) Save(ctx context.Context, userID, text string, embedding []float id := uuid.New().String() data := map[string]any{ - "user_id": userID, - "text": text, - "timestamp": time.Now().Format(time.RFC3339), - "embedding": embedding, + "user_id": userID, + "text": text, + "timestamp": time.Now().Format(time.RFC3339), + "embedding": embedding, + "event_type": eventType, + "weight": weight, + "ttl_days": ttlDays, } _, err := s.Client.Data(). @@ -164,6 +168,9 @@ func (s *Store) Search(ctx context.Context, queryEmbedding []float64, userID str if id, ok := add["id"].(string); ok { mem.ID = id } + if dist, ok := add["distance"].(float64); ok { + mem.Distance = dist + } } results = append(results, mem)