feat(sensitivity): add sensitivity tagging with pattern-based auto-classification (#85)

- New pkg/sensitivity with Classifier, Level enum (None/PII/InternalIP/Credentials)
- Built-in patterns: email, phone, credit card, SSN, AWS keys, OpenAI keys,
  GitHub tokens, Slack tokens, generic secrets
- Configurable internal domain detection (.internal, .corp, .local)
- StoreEntry accepts Sensitivity (explicit) and AutoClassify (pattern-based)
- RecallResult includes MaxSensitivity and SensitiveChunks metadata
- Sensitivity stored in SQLite, does not affect dedup or ranking
- 17 classifier tests + 4 benchmarks (all <1ms)
- 7 memory integration tests for sensitivity propagation

Closes #82

Co-authored-by: Ona <no-reply@ona.com>

Author: Siddhant-K-code

pkg/memory/sensitivity_test.go

@@ -0,0 +1,184 @@
+package memory
+
+import (
+	"context"
+	"testing"
+
+	"github.com/Siddhant-K-code/distill/pkg/sensitivity"
+)
+
+func TestStore_ExplicitSensitivity(t *testing.T) {
+	s := newTestStore(t)
+	ctx := context.Background()
+
+	_, err := s.Store(ctx, StoreRequest{
+		Entries: []StoreEntry{
+			{Text: "Q3 pricing: customer A at $120k", Sensitivity: sensitivity.InternalIP, Embedding: makeEmbedding(0, 8)},
+		},
+	})
+	if err != nil {
+		t.Fatalf("Store: %v", err)
+	}
+
+	recall, err := s.Recall(ctx, RecallRequest{
+		Query: "pricing", QueryEmbedding: makeEmbedding(0, 8), MaxResults: 10,
+	})
+	if err != nil {
+		t.Fatalf("Recall: %v", err)
+	}
+	if recall.MaxSensitivity != sensitivity.InternalIP {
+		t.Errorf("expected MaxSensitivity=InternalIP, got %s", recall.MaxSensitivity)
+	}
+	if len(recall.SensitiveChunks) != 1 {
+		t.Fatalf("expected 1 sensitive chunk, got %d", len(recall.SensitiveChunks))
+	}
+	if recall.SensitiveChunks[0].Sensitivity != sensitivity.InternalIP {
+		t.Errorf("expected chunk sensitivity InternalIP, got %s", recall.SensitiveChunks[0].Sensitivity)
+	}
+}
+
+func TestStore_AutoClassify_Credentials(t *testing.T) {
+	s := newTestStore(t)
+	ctx := context.Background()
+
+	_, err := s.Store(ctx, StoreRequest{
+		Entries: []StoreEntry{
+			{Text: "API key: sk-proj-abc123def456ghi789jkl012", AutoClassify: true, Embedding: makeEmbedding(0, 8)},
+		},
+	})
+	if err != nil {
+		t.Fatalf("Store: %v", err)
+	}
+
+	recall, _ := s.Recall(ctx, RecallRequest{
+		Query: "key", QueryEmbedding: makeEmbedding(0, 8), MaxResults: 10,
+	})
+	if recall.MaxSensitivity != sensitivity.Credentials {
+		t.Errorf("expected MaxSensitivity=Credentials, got %s", recall.MaxSensitivity)
+	}
+}
+
+func TestStore_AutoClassify_PII(t *testing.T) {
+	s := newTestStore(t)
+	ctx := context.Background()
+
+	_, err := s.Store(ctx, StoreRequest{
+		Entries: []StoreEntry{
+			{Text: "Contact alice@example.com for the report", AutoClassify: true, Embedding: makeEmbedding(0, 8)},
+		},
+	})
+	if err != nil {
+		t.Fatalf("Store: %v", err)
+	}
+
+	recall, _ := s.Recall(ctx, RecallRequest{
+		Query: "contact", QueryEmbedding: makeEmbedding(0, 8), MaxResults: 10,
+	})
+	if recall.MaxSensitivity != sensitivity.PII {
+		t.Errorf("expected MaxSensitivity=PII, got %s", recall.MaxSensitivity)
+	}
+}
+
+func TestStore_AutoClassify_NoMatch(t *testing.T) {
+	s := newTestStore(t)
+	ctx := context.Background()
+
+	_, err := s.Store(ctx, StoreRequest{
+		Entries: []StoreEntry{
+			{Text: "The service uses REST APIs", AutoClassify: true, Embedding: makeEmbedding(0, 8)},
+		},
+	})
+	if err != nil {
+		t.Fatalf("Store: %v", err)
+	}
+
+	recall, _ := s.Recall(ctx, RecallRequest{
+		Query: "service", QueryEmbedding: makeEmbedding(0, 8), MaxResults: 10,
+	})
+	if recall.MaxSensitivity != sensitivity.None {
+		t.Errorf("expected MaxSensitivity=None, got %s", recall.MaxSensitivity)
+	}
+	if len(recall.SensitiveChunks) != 0 {
+		t.Errorf("expected 0 sensitive chunks, got %d", len(recall.SensitiveChunks))
+	}
+}
+
+func TestStore_AutoClassify_ExplicitOverride(t *testing.T) {
+	s := newTestStore(t)
+	ctx := context.Background()
+
+	// Explicit sensitivity is higher than what auto-classify would find
+	_, err := s.Store(ctx, StoreRequest{
+		Entries: []StoreEntry{
+			{
+				Text:         "Normal text with no patterns",
+				Sensitivity:  sensitivity.Credentials,
+				AutoClassify: true,
+				Embedding:    makeEmbedding(0, 8),
+			},
+		},
+	})
+	if err != nil {
+		t.Fatalf("Store: %v", err)
+	}
+
+	recall, _ := s.Recall(ctx, RecallRequest{
+		Query: "text", QueryEmbedding: makeEmbedding(0, 8), MaxResults: 10,
+	})
+	// Explicit Credentials should be preserved even though auto-classify finds None
+	if recall.MaxSensitivity != sensitivity.Credentials {
+		t.Errorf("expected MaxSensitivity=Credentials (explicit), got %s", recall.MaxSensitivity)
+	}
+}
+
+func TestRecall_MaxSensitivity_MultipleEntries(t *testing.T) {
+	s := newTestStore(t)
+	ctx := context.Background()
+
+	_, err := s.Store(ctx, StoreRequest{
+		Entries: []StoreEntry{
+			{Text: "Normal architecture notes", Embedding: makeEmbedding(0, 8)},
+			{Text: "Contact bob@company.com", Sensitivity: sensitivity.PII, Embedding: makeEmbedding(1.5, 8)},
+			{Text: "AWS key AKIAIOSFODNN7EXAMPLE", Sensitivity: sensitivity.Credentials, Embedding: makeEmbedding(3.0, 8)},
+		},
+	})
+	if err != nil {
+		t.Fatalf("Store: %v", err)
+	}
+
+	recall, _ := s.Recall(ctx, RecallRequest{
+		Query: "all", QueryEmbedding: makeEmbedding(0.5, 8), MaxResults: 10,
+	})
+	if recall.MaxSensitivity != sensitivity.Credentials {
+		t.Errorf("expected MaxSensitivity=Credentials, got %s", recall.MaxSensitivity)
+	}
+	if len(recall.SensitiveChunks) != 2 {
+		t.Errorf("expected 2 sensitive chunks, got %d", len(recall.SensitiveChunks))
+	}
+}
+
+func TestRecall_SensitivityDoesNotAffectRanking(t *testing.T) {
+	s := newTestStore(t)
+	ctx := context.Background()
+
+	_, err := s.Store(ctx, StoreRequest{
+		Entries: []StoreEntry{
+			{Text: "Highly relevant auth info", Embedding: makeEmbedding(0, 8), Sensitivity: sensitivity.Credentials},
+			{Text: "Less relevant payment info", Embedding: makeEmbedding(1.5, 8)},
+		},
+	})
+	if err != nil {
+		t.Fatalf("Store: %v", err)
+	}
+
+	recall, _ := s.Recall(ctx, RecallRequest{
+		Query: "auth", QueryEmbedding: makeEmbedding(0, 8), MaxResults: 10,
+	})
+	// The sensitive entry should still be ranked first (closest embedding)
+	if len(recall.Memories) < 1 {
+		t.Fatal("expected at least 1 memory")
+	}
+	if recall.Memories[0].Sensitivity != sensitivity.Credentials {
+		t.Errorf("expected first result to have Credentials sensitivity, got %s", recall.Memories[0].Sensitivity)
+	}
+}

pkg/memory/sqlite.go

@@ -9,16 +9,18 @@ import (
 	"time"
 
 	distillmath "github.com/Siddhant-K-code/distill/pkg/math"
+	"github.com/Siddhant-K-code/distill/pkg/sensitivity"
 	_ "modernc.org/sqlite"
 )
 
 // SQLiteStore implements Store using SQLite for local persistent storage.
 // Uses a single connection (SetMaxOpenConns(1)) so SQLite's internal
 // serialization handles concurrency. No application-level mutex needed.
 type SQLiteStore struct {
-	db       *sql.DB
-	cfg      Config
-	handlers []MemoryEventHandler
+	db         *sql.DB
+	cfg        Config
+	handlers   []MemoryEventHandler
+	classifier *sensitivity.Classifier
 }
 
 // NewSQLiteStore creates a new SQLite-backed memory store.
@@ -49,7 +51,11 @@ func NewSQLiteStore(dsn string, cfg Config) (*SQLiteStore, error) {
 		return nil, fmt.Errorf("enable foreign keys: %w", err)
 	}
 
-	s := &SQLiteStore{db: db, cfg: cfg}
+	s := &SQLiteStore{
+		db:         db,
+		cfg:        cfg,
+		classifier: sensitivity.New(sensitivity.DefaultConfig()),
+	}
 	if err := s.migrate(); err != nil {
 		_ = db.Close()
 		return nil, fmt.Errorf("migrate: %w", err)
@@ -68,6 +74,7 @@ func (s *SQLiteStore) migrate() error {
 		session_id      TEXT DEFAULT '',
 		metadata        TEXT DEFAULT '{}',
 		decay_level     INTEGER DEFAULT 0,
+		sensitivity     INTEGER DEFAULT 0,
 		created_at      TEXT NOT NULL,
 		last_referenced TEXT NOT NULL,
 		access_count    INTEGER DEFAULT 0,
@@ -98,6 +105,7 @@ func (s *SQLiteStore) migrate() error {
 		{"expired_at", "TEXT DEFAULT ''"},
 		{"superseded_by", "TEXT DEFAULT ''"},
 		{"expires_at", "TEXT DEFAULT ''"},
+		{"sensitivity", "INTEGER DEFAULT 0"},
 	} {
 		_, _ = s.db.Exec("ALTER TABLE memories ADD COLUMN " + col.name + " " + col.def)
 	}
@@ -149,10 +157,19 @@ func (s *SQLiteStore) Store(ctx context.Context, req StoreRequest) (*StoreResult
 			expiresAt = entry.ExpiresAt.UTC().Format(time.RFC3339Nano)
 		}
 
+		// Determine sensitivity level
+		sens := entry.Sensitivity
+		if entry.AutoClassify {
+			classified := s.classifier.Classify(entry.Text)
+			if classified.Level > sens {
+				sens = classified.Level
+			}
+		}
+
 		_, err := s.db.ExecContext(ctx,
-			`INSERT INTO memories (id, text, embedding, source, session_id, metadata, decay_level, created_at, last_referenced, access_count, expires_at)
-			 VALUES (?, ?, ?, ?, ?, ?, 0, ?, ?, 0, ?)`,
-			id, entry.Text, embBlob, entry.Source, sessionID, string(metaJSON), now, now, expiresAt,
+			`INSERT INTO memories (id, text, embedding, source, session_id, metadata, decay_level, sensitivity, created_at, last_referenced, access_count, expires_at)
+			 VALUES (?, ?, ?, ?, ?, ?, 0, ?, ?, ?, 0, ?)`,
+			id, entry.Text, embBlob, entry.Source, sessionID, string(metaJSON), int(sens), now, now, expiresAt,
 		)
 		if err != nil {
 			return nil, fmt.Errorf("insert memory: %w", err)
@@ -237,7 +254,7 @@ func (s *SQLiteStore) Recall(ctx context.Context, req RecallRequest) (*RecallRes
 	}
 
 	// Build query with optional tag filter and expiry exclusion
-	query := "SELECT m.id, m.text, m.embedding, m.source, m.decay_level, m.last_referenced FROM memories m"
+	query := "SELECT m.id, m.text, m.embedding, m.source, m.decay_level, m.sensitivity, m.last_referenced FROM memories m"
 	var args []interface{}
 	var conditions []string
 
@@ -273,11 +290,12 @@ func (s *SQLiteStore) Recall(ctx context.Context, req RecallRequest) (*RecallRes
 		id, text, source, refStr string
 		embBlob                  []byte
 		decayLevel               int
+		sensitivity              int
 	}
 	var rawRows []rawRow
 	for rows.Next() {
 		var r rawRow
-		if err := rows.Scan(&r.id, &r.text, &r.embBlob, &r.source, &r.decayLevel, &r.refStr); err != nil {
+		if err := rows.Scan(&r.id, &r.text, &r.embBlob, &r.source, &r.decayLevel, &r.sensitivity, &r.refStr); err != nil {
 			_ = rows.Close()
 			return nil, err
 		}
@@ -323,6 +341,7 @@ func (s *SQLiteStore) Recall(ctx context.Context, req RecallRequest) (*RecallRes
 				Tags:           tags,
 				Relevance:      relevance,
 				DecayLevel:     DecayLevel(r.decayLevel),
+				Sensitivity:    sensitivity.Level(r.sensitivity),
 				LastReferenced: lastRef,
 			},
 			relevance: relevance,
@@ -360,6 +379,9 @@ func (s *SQLiteStore) Recall(ctx context.Context, req RecallRequest) (*RecallRes
 	// Entries with relevance >= 0.7 are considered stable candidates.
 	hint := buildCacheBoundaryHint(results)
 
+	// Build sensitivity metadata from returned memories.
+	maxSens, sensitiveChunks := buildSensitivityMetadata(results)
+
 	return &RecallResult{
 		Memories: results,
 		Stats: RecallStats{
@@ -368,7 +390,9 @@ func (s *SQLiteStore) Recall(ctx context.Context, req RecallRequest) (*RecallRes
 			Returned:     len(results),
 			TokenCount:   tokenCount,
 		},
-		CacheHint: hint,
+		CacheHint:       hint,
+		MaxSensitivity:  maxSens,
+		SensitiveChunks: sensitiveChunks,
 	}, nil
 }
 
@@ -395,6 +419,25 @@ func buildCacheBoundaryHint(memories []RecalledMemory) *CacheBoundaryHint {
 	}
 }
 
+// buildSensitivityMetadata derives MaxSensitivity and SensitiveChunks from
+// the recalled memories. Only entries with non-zero sensitivity are included.
+func buildSensitivityMetadata(memories []RecalledMemory) (sensitivity.Level, []SensitiveChunk) {
+	var maxSens sensitivity.Level
+	var chunks []SensitiveChunk
+	for _, m := range memories {
+		if m.Sensitivity > maxSens {
+			maxSens = m.Sensitivity
+		}
+		if m.Sensitivity > sensitivity.None {
+			chunks = append(chunks, SensitiveChunk{
+				ChunkID:     m.ID,
+				Sensitivity: m.Sensitivity,
+			})
+		}
+	}
+	return maxSens, chunks
+}
+
 // Forget removes memories matching the given criteria.
 func (s *SQLiteStore) Forget(ctx context.Context, req ForgetRequest) (*ForgetResult, error) {