lukaszraczylo/claude-mnemonic
1
0
Fork 0
mirror of https://github.com/lukaszraczylo/claude-mnemonic.git synced 2026-10-15 05:37:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(leann-phase2): implement hybrid vector storage and graph-based search

Browse Source 
- [x] Add AST-aware code chunking for Go, Python, and TypeScript using tree-sitter
- [x] Implement LEANN-inspired hybrid vector storage with hub detection and selective embedding storage (60-80% savings)
- [x] Add observation relationship graph with CSR format and edge detection (file overlap, semantic similarity, temporal, concept)
- [x] Implement graph-aware search with two-level traversal and relationship-based ranking
- [x] Add auto-tuning system for dynamic hub threshold adjustment based on query performance
- [x] Add comprehensive metrics tracking for vector storage, queries, latency, and graph traversals
- [x] Update configuration system with graph and hybrid storage settings
- [x] Add graph stats and vector metrics endpoints to worker service
- [x] Enhance UI sidebar with advanced metrics display and graph visualization
- [x] Optimize struct field alignment throughout codebase for memory efficiency
- [x] Update documentation with LEANN Phase 2 features and performance benefits
- [x] Add tree-sitter dependency for AST parsing
This commit is contained in:
lukaszraczylo
2026-01-07 20:43:10 +00:00
parent 7ab4b07cf2
commit 74ae8ed4c1
83 changed files with 5190 additions and 603 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/worker/handlers.go
+83 -1
View File
@@ -158,10 +158,10 @@ type SessionInitRequest struct {
// SessionInitResponse is the response for session initialization.
type SessionInitResponse struct {
Reason string `json:"reason,omitempty"`
SessionDBID int64 `json:"sessionDbId"`
PromptNumber int `json:"promptNumber"`
Skipped bool `json:"skipped,omitempty"`
Reason string `json:"reason,omitempty"`
}
// DuplicatePromptWindowSeconds is the time window for detecting duplicate prompt submissions.
@@ -1312,3 +1312,85 @@ func (s *Service) handleRestart(w http.ResponseWriter, r *http.Request) {
}
}()
}
// handleGetGraphStats returns observation graph statistics.
func (s *Service) handleGetGraphStats(w http.ResponseWriter, r *http.Request) {
if s.graphSearchClient == nil {
writeJSON(w, map[string]interface{}{
"enabled": false,
"message": "Graph search not enabled",
})
return
}
stats := s.graphSearchClient.GetGraphStats()
response := map[string]interface{}{
"enabled": s.config.GraphEnabled,
"nodeCount": stats.NodeCount,
"edgeCount": stats.EdgeCount,
"avgDegree": stats.AvgDegree,
"maxDegree": stats.MaxDegree,
"minDegree": stats.MinDegree,
"medianDegree": stats.MedianDegree,
"edgeTypes": stats.EdgeTypes,
"config": map[string]interface{}{
"maxHops": s.config.GraphMaxHops,
"branchFactor": s.config.GraphBranchFactor,
"edgeWeight": s.config.GraphEdgeWeight,
"rebuildIntervalMin": s.config.GraphRebuildIntervalMin,
},
}
writeJSON(w, response)
}
// handleGetVectorMetrics returns hybrid vector storage metrics.
func (s *Service) handleGetVectorMetrics(w http.ResponseWriter, r *http.Request) {
if s.hybridMetrics == nil {
writeJSON(w, map[string]interface{}{
"enabled": false,
"message": "Vector metrics not available",
})
return
}
snapshot := s.hybridMetrics.GetSnapshot()
response := map[string]interface{}{
"queries": map[string]interface{}{
"total": snapshot.TotalQueries,
"hubOnly": snapshot.HubOnlyQueries,
"hybrid": snapshot.HybridQueries,
"onDemand": snapshot.OnDemandQueries,
"graph": snapshot.GraphQueries,
},
"latency": map[string]interface{}{
"avg": snapshot.AvgLatency.String(),
"p50": snapshot.P50Latency.String(),
"p95": snapshot.P95Latency.String(),
"p99": snapshot.P99Latency.String(),
"avgHub": snapshot.AvgHubLatency.String(),
"avgRecompute": snapshot.AvgRecomputeLatency.String(),
},
"storage": map[string]interface{}{
"totalDocuments": snapshot.TotalDocuments,
"hubDocuments": snapshot.HubDocuments,
"storedEmbeddings": snapshot.StoredEmbeddings,
"savingsPercent": snapshot.StorageSavingsPercent,
"recomputedTotal": snapshot.RecomputedTotal,
},
"cache": map[string]interface{}{
"hits": snapshot.CacheHits,
"misses": snapshot.CacheMisses,
"hitRate": snapshot.CacheHitRate,
},
"graph": map[string]interface{}{
"traversals": snapshot.GraphTraversals,
"avgDepth": snapshot.AvgTraversalDepth,
},
"uptime": snapshot.Uptime.String(),
}
writeJSON(w, response)
}
internal/worker/sdk/parser_test.go
+2 -2
View File
@@ -77,10 +77,10 @@ func TestParseObservations_TableDriven(t *testing.T) {
tests := []struct {
name string
input string
expectedCount int
expectedType models.ObservationType
expectedTitle string
checkConcepts []string
expectedCount int
}{
{
name: "valid_bugfix_observation",
@@ -300,9 +300,9 @@ func TestParseSummary_TableDriven(t *testing.T) {
tests := []struct {
name string
input string
expectedRequest string
sessionID int64
expectNil bool
expectedRequest string
}{
{
name: "empty_input",
internal/worker/sdk/processor.go
+3 -4
View File
@@ -31,15 +31,14 @@ type SyncSummaryFunc func(summary *models.SessionSummary)
// Processor handles SDK agent processing of observations and summaries using Claude Code CLI.
type Processor struct {
claudePath string
model string
observationStore *gorm.ObservationStore
summaryStore *gorm.SummaryStore
broadcastFunc BroadcastFunc
syncObservationFunc SyncObservationFunc
syncSummaryFunc SyncSummaryFunc
// Semaphore to limit concurrent Claude CLI calls (prevents API overload)
sem chan struct{}
sem chan struct{}
claudePath string
model string
}
// SetBroadcastFunc sets the broadcast callback for SSE events.
internal/worker/sdk/processor_test.go
+3 -3
View File
@@ -11,8 +11,8 @@ import (
func TestIsSelfReferentialSummary(t *testing.T) {
tests := []struct {
name string
summary *models.ParsedSummary
name string
expected bool
}{
{
@@ -281,8 +281,8 @@ func TestTruncateForLog(t *testing.T) {
tests := []struct {
name string
input string
maxLen int
expected string
maxLen int
}{
{
name: "shorter_than_max",
@@ -719,8 +719,8 @@ func TestShouldSkipTrivialOperation_EdgeCases(t *testing.T) {
// TestIsSelfReferentialSummary_MoreCases tests additional self-referential detection cases.
func TestIsSelfReferentialSummary_MoreCases(t *testing.T) {
tests := []struct {
name string
summary *models.ParsedSummary
name string
expected bool
}{
{
internal/worker/sdk/prompts.go
+3 -3
View File
@@ -24,12 +24,12 @@ var ObservationConcepts = []string{
// ToolExecution represents a tool execution for observation.
type ToolExecution struct {
ID int64
ToolName string
ToolInput string
ToolOutput string
CreatedAtEpoch int64
CWD string
ID int64
CreatedAtEpoch int64
}
// BuildObservationPrompt builds a prompt for processing a tool observation.
@@ -67,12 +67,12 @@ func BuildObservationPrompt(exec ToolExecution) string {
// SummaryRequest contains data for building a summary prompt.
type SummaryRequest struct {
SessionDBID int64
SDKSessionID string
Project string
UserPrompt string
LastUserMessage string
LastAssistantMessage string
SessionDBID int64
}
// BuildSummaryPrompt builds a prompt requesting a session summary.
internal/worker/sdk/prompts_test.go
+2 -2
View File
@@ -12,8 +12,8 @@ func TestTruncate(t *testing.T) {
tests := []struct {
name string
input string
maxLen int
expected string
maxLen int
}{
{
name: "shorter_than_max",
@@ -60,8 +60,8 @@ func TestBuildObservationPrompt(t *testing.T) {
tests := []struct {
name string
exec ToolExecution
contains []string
exec ToolExecution
}{
{
name: "basic_read_tool",
internal/worker/service.go
+188 -74
View File
@@ -12,6 +12,10 @@ import (
"github.com/go-chi/chi/v5"
"github.com/go-chi/chi/v5/middleware"
"github.com/lukaszraczylo/claude-mnemonic/internal/chunking"
"github.com/lukaszraczylo/claude-mnemonic/internal/chunking/golang"
"github.com/lukaszraczylo/claude-mnemonic/internal/chunking/python"
"github.com/lukaszraczylo/claude-mnemonic/internal/chunking/typescript"
"github.com/lukaszraczylo/claude-mnemonic/internal/config"
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
"github.com/lukaszraczylo/claude-mnemonic/internal/embedding"
@@ -20,6 +24,7 @@ import (
"github.com/lukaszraczylo/claude-mnemonic/internal/scoring"
"github.com/lukaszraczylo/claude-mnemonic/internal/search/expansion"
"github.com/lukaszraczylo/claude-mnemonic/internal/update"
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/hybrid"
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
"github.com/lukaszraczylo/claude-mnemonic/internal/watcher"
"github.com/lukaszraczylo/claude-mnemonic/internal/worker/sdk"
@@ -56,80 +61,53 @@ type RetrievalStats struct {
// Service is the main worker service orchestrator.
type Service struct {
// Version of the worker binary
version string
// Configuration
config *config.Config
// Database
store *gorm.Store
sessionStore *gorm.SessionStore
observationStore *gorm.ObservationStore
summaryStore *gorm.SummaryStore
promptStore *gorm.PromptStore
conflictStore *gorm.ConflictStore
patternStore *gorm.PatternStore
relationStore *gorm.RelationStore
// Pattern detection
patternDetector *pattern.Detector
// Domain services
sessionManager *session.Manager
sseBroadcaster *sse.Broadcaster
processor *sdk.Processor
// Vector database (sqlite-vec with local embeddings)
embedSvc *embedding.Service
vectorClient *sqlitevec.Client
vectorSync *sqlitevec.Sync
// Cross-encoder reranking (for improved search relevance)
reranker *reranking.Service
// Query expansion (for improved search recall)
queryExpander *expansion.Expander
// Importance scoring
scoreCalculator *scoring.Calculator
recalculator *scoring.Recalculator
// HTTP server
router *chi.Mux
server *http.Server
startTime time.Time
// Retrieval statistics (per-project)
retrievalStats map[string]*RetrievalStats
retrievalStatsMu sync.RWMutex
// Lifecycle
ctx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
// Initialization state (for deferred init)
ready atomic.Bool
initError error
initMu sync.RWMutex
// Background verification queue for stale observations
staleQueue chan staleVerifyRequest
staleQueueOnce sync.Once
// File watchers for auto-recreation on deletion
dbWatcher *watcher.Watcher
configWatcher *watcher.Watcher
// Self-updater
updater *update.Updater
startTime time.Time
initError error
ctx context.Context
patternDetector *pattern.Detector
queryExpander *expansion.Expander
summaryStore *gorm.SummaryStore
promptStore *gorm.PromptStore
conflictStore *gorm.ConflictStore
patternStore *gorm.PatternStore
relationStore *gorm.RelationStore
updater *update.Updater
sessionManager *session.Manager
scoreCalculator *scoring.Calculator
processor *sdk.Processor
embedSvc *embedding.Service
vectorClient *sqlitevec.Client
vectorSync *sqlitevec.Sync
graphSearchClient *hybrid.GraphSearchClient
hybridMetrics *hybrid.Metrics
graphRebuildTicker *time.Ticker
chunkingManager *chunking.Manager
observationStore *gorm.ObservationStore
reranker *reranking.Service
sseBroadcaster *sse.Broadcaster
recalculator *scoring.Recalculator
router *chi.Mux
server *http.Server
sessionStore *gorm.SessionStore
retrievalStats map[string]*RetrievalStats
configWatcher *watcher.Watcher
store *gorm.Store
cancel context.CancelFunc
dbWatcher *watcher.Watcher
staleQueue chan staleVerifyRequest
config *config.Config
version string
wg sync.WaitGroup
initMu sync.RWMutex
retrievalStatsMu sync.RWMutex
staleQueueOnce sync.Once
ready atomic.Bool
}
// staleVerifyRequest represents a request to verify a stale observation in background
type staleVerifyRequest struct {
observationID int64
cwd string
observationID int64
}
// NewService creates a new worker service with deferred initialization.
@@ -210,6 +188,9 @@ func (s *Service) initializeAsync() {
var embedSvc *embedding.Service
var vectorClient *sqlitevec.Client
var vectorSync *sqlitevec.Sync
var graphSearchClient *hybrid.GraphSearchClient
var hybridMetrics *hybrid.Metrics
var chunkingManager *chunking.Manager
var reranker *reranking.Service
@@ -218,18 +199,51 @@ func (s *Service) initializeAsync() {
log.Warn().Err(err).Msg("Embedding service creation failed - vector search disabled")
} else {
embedSvc = emb
// Create sqlite-vec client using the same DB connection
client, err := sqlitevec.NewClient(sqlitevec.Config{
// Create base sqlite-vec client using the same DB connection
baseClient, err := sqlitevec.NewClient(sqlitevec.Config{
DB: store.GetRawDB(),
}, embedSvc)
if err != nil {
log.Warn().Err(err).Msg("sqlite-vec client creation failed - vector search disabled")
} else {
vectorClient = client
vectorSync = sqlitevec.NewSync(client)
vectorClient = baseClient
// Wrap with LEANN hybrid storage client
strategy := hybrid.ParseStrategy(s.config.VectorStorageStrategy)
hybridClient := hybrid.NewClient(hybrid.Config{
BaseClient: baseClient,
DB: store.GetRawDB(),
EmbedSvc: embedSvc,
Strategy: strategy,
HubThreshold: s.config.HubThreshold,
})
// Wrap with graph-aware search client
graphConfig := hybrid.GraphConfig{
Enabled: s.config.GraphEnabled,
MaxHops: s.config.GraphMaxHops,
BranchFactor: s.config.GraphBranchFactor,
EdgeWeight: s.config.GraphEdgeWeight,
}
graphSearchClient = hybrid.NewGraphSearchClient(hybridClient, nil, graphConfig)
hybridMetrics = hybrid.NewMetrics()
vectorSync = sqlitevec.NewSync(baseClient)
// Initialize AST-aware code chunking
chunkOpts := chunking.DefaultChunkOptions()
chunkers := []chunking.Chunker{
golang.NewChunker(chunkOpts),
python.NewChunker(chunkOpts),
typescript.NewChunker(chunkOpts),
}
chunkingManager = chunking.NewManager(chunkers, chunkOpts)
log.Info().
Str("model", embedSvc.Version()).
Msg("sqlite-vec vector search enabled")
Str("storage_strategy", s.config.VectorStorageStrategy).
Bool("graph_enabled", s.config.GraphEnabled).
Msg("LEANN hybrid vector storage and graph search enabled")
}
// Create cross-encoder reranking service if enabled
@@ -284,6 +298,9 @@ func (s *Service) initializeAsync() {
s.embedSvc = embedSvc
s.vectorClient = vectorClient
s.vectorSync = vectorSync
s.graphSearchClient = graphSearchClient
s.hybridMetrics = hybridMetrics
s.chunkingManager = chunkingManager
s.reranker = reranker
s.initMu.Unlock()
@@ -411,6 +428,18 @@ func (s *Service) initializeAsync() {
s.ready.Store(true)
log.Info().Msg("Async initialization complete - service ready")
// Build initial observation graph if graph search is enabled
if graphSearchClient != nil && s.config.GraphEnabled {
s.wg.Add(1)
go s.buildInitialGraph(observationStore)
// Start periodic graph rebuild timer
if s.config.GraphRebuildIntervalMin > 0 {
s.wg.Add(1)
go s.startGraphRebuildTimer(observationStore)
}
}
// Start queue processor if SDK processor is available
if processor != nil {
s.wg.Add(1)
@@ -1136,6 +1165,10 @@ func (s *Service) setupRoutes() {
r.Get("/api/observations/{id}/relations", s.handleGetRelations)
r.Get("/api/observations/{id}/graph", s.handleGetRelationGraph)
r.Get("/api/observations/{id}/related", s.handleGetRelatedObservations)
// LEANN Phase 2: Graph-based search and hybrid vector storage
r.Get("/api/graph/stats", s.handleGetGraphStats)
r.Get("/api/vector/metrics", s.handleGetVectorMetrics)
})
}
@@ -1346,6 +1379,87 @@ func (s *Service) processAllSessions() {
s.broadcastProcessingStatus()
}
// buildInitialGraph builds the observation relationship graph in the background.
func (s *Service) buildInitialGraph(observationStore *gorm.ObservationStore) {
defer s.wg.Done()
log.Info().Msg("Building initial observation graph...")
start := time.Now()
// Fetch all observations
observations, err := observationStore.GetAllObservations(s.ctx)
if err != nil {
log.Error().Err(err).Msg("Failed to fetch observations for graph building")
return
}
if len(observations) == 0 {
log.Info().Msg("No observations to build graph from")
return
}
// Build graph using RebuildGraph method
if err := s.graphSearchClient.RebuildGraph(s.ctx, observations); err != nil {
log.Error().Err(err).Msg("Failed to build observation graph")
return
}
elapsed := time.Since(start)
stats := s.graphSearchClient.GetGraphStats()
log.Info().
Int("observations", len(observations)).
Int("nodes", stats.NodeCount).
Int("edges", stats.EdgeCount).
Float64("avg_degree", stats.AvgDegree).
Int("max_degree", stats.MaxDegree).
Dur("elapsed", elapsed).
Msg("Initial observation graph built successfully")
}
// startGraphRebuildTimer starts a periodic ticker to rebuild the observation graph.
func (s *Service) startGraphRebuildTimer(observationStore *gorm.ObservationStore) {
defer s.wg.Done()
interval := time.Duration(s.config.GraphRebuildIntervalMin) * time.Minute
s.graphRebuildTicker = time.NewTicker(interval)
log.Info().
Dur("interval", interval).
Msg("Started periodic graph rebuild timer")
for {
select {
case <-s.ctx.Done():
s.graphRebuildTicker.Stop()
log.Info().Msg("Stopped graph rebuild timer")
return
case <-s.graphRebuildTicker.C:
log.Info().Msg("Periodic graph rebuild triggered")
start := time.Now()
observations, err := observationStore.GetAllObservations(s.ctx)
if err != nil {
log.Error().Err(err).Msg("Failed to fetch observations for graph rebuild")
continue
}
if err := s.graphSearchClient.RebuildGraph(s.ctx, observations); err != nil {
log.Error().Err(err).Msg("Failed to rebuild observation graph")
continue
}
stats := s.graphSearchClient.GetGraphStats()
log.Info().
Int("nodes", stats.NodeCount).
Int("edges", stats.EdgeCount).
Dur("elapsed", time.Since(start)).
Msg("Periodic graph rebuild complete")
}
}
}
// Shutdown gracefully shuts down the service.
func (s *Service) Shutdown(ctx context.Context) error {
s.cancel()
internal/worker/session/manager.go
+20 -23
View File
@@ -21,11 +21,11 @@ const (
// ObservationData contains data for a tool observation.
type ObservationData struct {
ToolName string
ToolInput interface{}
ToolResponse interface{}
PromptNumber int
ToolName string
CWD string
PromptNumber int
}
// SummarizeData contains data for a summarize request.
@@ -36,30 +36,28 @@ type SummarizeData struct {
// PendingMessage represents a message queued for SDK processing.
type PendingMessage struct {
Type MessageType
Observation *ObservationData
Summarize *SummarizeData
Type MessageType
}
// ActiveSession represents an in-memory active session being processed.
type ActiveSession struct {
SessionDBID int64
ClaudeSessionID string
SDKSessionID string
StartTime time.Time
ctx context.Context
cancel context.CancelFunc
notify chan struct{}
Project string
UserPrompt string
SDKSessionID string
ClaudeSessionID string
pendingMessages []PendingMessage
LastPromptNumber int
StartTime time.Time
CumulativeInputTokens int64
CumulativeOutputTokens int64
// Concurrency control
pendingMessages []PendingMessage
messageMu sync.Mutex
notify chan struct{}
ctx context.Context
cancel context.CancelFunc
generatorActive atomic.Bool
SessionDBID int64
messageMu sync.Mutex
generatorActive atomic.Bool
}
// SessionTimeout is how long an inactive session can exist before cleanup.
@@ -70,15 +68,14 @@ const CleanupInterval = 5 * time.Minute
// Manager manages active session lifecycles.
type Manager struct {
sessionStore *gorm.SessionStore
sessions map[int64]*ActiveSession
mu sync.RWMutex
onCreated func(int64)
onDeleted func(int64)
ctx context.Context
cancel context.CancelFunc
// Global notification channel for immediate processing
ctx context.Context
sessionStore *gorm.SessionStore
sessions map[int64]*ActiveSession
onCreated func(int64)
onDeleted func(int64)
cancel context.CancelFunc
ProcessNotify chan struct{}
mu sync.RWMutex
}
// NewManager creates a new session manager.
internal/worker/session/manager_test.go
+7 -7
View File
@@ -669,16 +669,16 @@ func TestActiveSessionCWD(t *testing.T) {
// TestToolInputResponse tests various tool input/response types.
func TestToolInputResponse(t *testing.T) {
tests := []struct {
name string
input interface{}
response interface{}
name string
}{
{"nil_values", nil, nil},
{"string_values", "input string", "response string"},
{"map_values", map[string]string{"key": "value"}, map[string]interface{}{"result": true}},
{"slice_values", []string{"a", "b"}, []int{1, 2, 3}},
{"int_values", 42, 100},
{"bool_values", true, false},
{name: "nil_values", input: nil, response: nil},
{name: "string_values", input: "input string", response: "response string"},
{name: "map_values", input: map[string]string{"key": "value"}, response: map[string]interface{}{"result": true}},
{name: "slice_values", input: []string{"a", "b"}, response: []int{1, 2, 3}},
{name: "int_values", input: 42, response: 100},
{name: "bool_values", input: true, response: false},
}
for _, tt := range tests {
internal/worker/sse/broadcaster.go
+1 -1
View File
@@ -19,10 +19,10 @@ const (
// Client represents a connected SSE client.
type Client struct {
ID string
Writer http.ResponseWriter
Flusher http.Flusher
Done chan struct{}
ID string
}
// Broadcaster manages SSE client connections and message broadcasting.
internal/worker/sse/broadcaster_test.go
+1 -1
View File
@@ -256,8 +256,8 @@ func TestHandleSSE(t *testing.T) {
// TestBroadcastJSON tests broadcasting various JSON types.
func TestBroadcastJSON(t *testing.T) {
tests := []struct {
name string
data interface{}
name string
wantErr bool
}{
{