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

- [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:
2026-01-07 20:43:10 +00:00
parent 7ab4b07cf2
commit 74ae8ed4c1
83 changed files with 5190 additions and 603 deletions
+285
View File
@@ -0,0 +1,285 @@
// Package golang provides AST-aware chunking for Go source files.
package golang
import (
"context"
"fmt"
"go/ast"
"go/parser"
"go/token"
"os"
"strings"
"github.com/lukaszraczylo/claude-mnemonic/internal/chunking"
)
// Chunker implements AST-aware chunking for Go files.
type Chunker struct {
options chunking.ChunkOptions
}
// NewChunker creates a new Go chunker.
func NewChunker(options chunking.ChunkOptions) *Chunker {
return &Chunker{options: options}
}
// Language returns the language this chunker supports.
func (c *Chunker) Language() chunking.Language {
return chunking.LanguageGo
}
// SupportedExtensions returns the file extensions this chunker handles.
func (c *Chunker) SupportedExtensions() []string {
return []string{".go"}
}
// Chunk parses a Go source file and returns semantic code chunks.
func (c *Chunker) Chunk(ctx context.Context, filePath string) ([]chunking.Chunk, error) {
// Read file content
content, err := os.ReadFile(filePath)
if err != nil {
return nil, fmt.Errorf("read file: %w", err)
}
// Parse the Go file
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, filePath, content, parser.ParseComments)
if err != nil {
return nil, fmt.Errorf("parse Go file: %w", err)
}
chunks := make([]chunking.Chunk, 0)
sourceLines := strings.Split(string(content), "\n")
// Extract chunks from declarations
for _, decl := range file.Decls {
switch d := decl.(type) {
case *ast.FuncDecl:
chunk := c.extractFunction(fset, d, sourceLines, filePath)
if chunk != nil {
chunks = append(chunks, *chunk)
}
case *ast.GenDecl:
extracted := c.extractGenDecl(fset, d, sourceLines, filePath)
chunks = append(chunks, extracted...)
}
}
return chunks, nil
}
// extractFunction extracts a function or method declaration as a chunk.
func (c *Chunker) extractFunction(fset *token.FileSet, fn *ast.FuncDecl, sourceLines []string, filePath string) *chunking.Chunk {
// Skip unexported if configured
if !c.options.IncludePrivate && !fn.Name.IsExported() {
return nil
}
startPos := fset.Position(fn.Pos())
endPos := fset.Position(fn.End())
chunk := &chunking.Chunk{
FilePath: filePath,
Language: chunking.LanguageGo,
Name: fn.Name.Name,
StartLine: startPos.Line,
EndLine: endPos.Line,
}
// Determine if this is a method or a function
if fn.Recv != nil && len(fn.Recv.List) > 0 {
chunk.Type = chunking.ChunkTypeMethod
chunk.ParentName = c.extractReceiverType(fn.Recv)
} else {
chunk.Type = chunking.ChunkTypeFunction
}
// Extract content
chunk.Content = c.extractLines(sourceLines, startPos.Line, endPos.Line)
// Extract signature (function declaration without body)
chunk.Signature = c.extractFunctionSignature(fn, fset, sourceLines)
// Extract doc comment
if c.options.IncludeDocComments && fn.Doc != nil {
chunk.DocComment = strings.TrimSpace(fn.Doc.Text())
}
return chunk
}
// extractGenDecl extracts general declarations (type, const, var).
func (c *Chunker) extractGenDecl(fset *token.FileSet, gd *ast.GenDecl, sourceLines []string, filePath string) []chunking.Chunk {
var chunks []chunking.Chunk
for _, spec := range gd.Specs {
switch s := spec.(type) {
case *ast.TypeSpec:
chunk := c.extractTypeSpec(fset, gd, s, sourceLines, filePath)
if chunk != nil {
chunks = append(chunks, *chunk)
}
case *ast.ValueSpec:
// Handle const and var declarations
chunk := c.extractValueSpec(fset, gd, s, sourceLines, filePath)
if chunk != nil {
chunks = append(chunks, *chunk)
}
}
}
return chunks
}
// extractTypeSpec extracts a type declaration (struct, interface, type alias).
func (c *Chunker) extractTypeSpec(fset *token.FileSet, gd *ast.GenDecl, ts *ast.TypeSpec, sourceLines []string, filePath string) *chunking.Chunk {
// Skip unexported if configured
if !c.options.IncludePrivate && !ts.Name.IsExported() {
return nil
}
startPos := fset.Position(gd.Pos())
endPos := fset.Position(gd.End())
chunk := &chunking.Chunk{
FilePath: filePath,
Language: chunking.LanguageGo,
Name: ts.Name.Name,
StartLine: startPos.Line,
EndLine: endPos.Line,
Content: c.extractLines(sourceLines, startPos.Line, endPos.Line),
}
// Determine chunk type based on type expression
switch ts.Type.(type) {
case *ast.StructType:
chunk.Type = chunking.ChunkTypeClass // Treat struct as class
case *ast.InterfaceType:
chunk.Type = chunking.ChunkTypeInterface
default:
chunk.Type = chunking.ChunkTypeType
}
// Extract doc comment
if c.options.IncludeDocComments && gd.Doc != nil {
chunk.DocComment = strings.TrimSpace(gd.Doc.Text())
}
return chunk
}
// extractValueSpec extracts const or var declarations.
func (c *Chunker) extractValueSpec(fset *token.FileSet, gd *ast.GenDecl, vs *ast.ValueSpec, sourceLines []string, filePath string) *chunking.Chunk {
// Skip if all names are unexported and we're excluding private
if !c.options.IncludePrivate {
allUnexported := true
for _, name := range vs.Names {
if name.IsExported() {
allUnexported = false
break
}
}
if allUnexported {
return nil
}
}
startPos := fset.Position(gd.Pos())
endPos := fset.Position(gd.End())
// Use first name as the chunk name, join multiple if present
names := make([]string, len(vs.Names))
for i, name := range vs.Names {
names[i] = name.Name
}
chunk := &chunking.Chunk{
FilePath: filePath,
Language: chunking.LanguageGo,
Name: strings.Join(names, ", "),
StartLine: startPos.Line,
EndLine: endPos.Line,
Content: c.extractLines(sourceLines, startPos.Line, endPos.Line),
}
// Set type based on token
if gd.Tok == token.CONST {
chunk.Type = chunking.ChunkTypeConst
} else {
chunk.Type = chunking.ChunkTypeVar
}
// Extract doc comment
if c.options.IncludeDocComments && gd.Doc != nil {
chunk.DocComment = strings.TrimSpace(gd.Doc.Text())
}
return chunk
}
// extractReceiverType extracts the receiver type name from a method.
func (c *Chunker) extractReceiverType(recv *ast.FieldList) string {
if len(recv.List) == 0 {
return ""
}
field := recv.List[0]
switch t := field.Type.(type) {
case *ast.Ident:
return t.Name
case *ast.StarExpr:
if ident, ok := t.X.(*ast.Ident); ok {
return ident.Name
}
}
return ""
}
// extractFunctionSignature extracts the function signature without the body.
func (c *Chunker) extractFunctionSignature(fn *ast.FuncDecl, fset *token.FileSet, sourceLines []string) string {
if fn.Body == nil {
// No body, return entire declaration
startPos := fset.Position(fn.Pos())
endPos := fset.Position(fn.End())
return c.extractLines(sourceLines, startPos.Line, endPos.Line)
}
// Extract from start of function to just before body
startPos := fset.Position(fn.Pos())
bodyPos := fset.Position(fn.Body.Pos())
// If body is on the same line, extract just that line up to the opening brace
if startPos.Line == bodyPos.Line {
line := sourceLines[startPos.Line-1]
// Find the opening brace position
if idx := strings.Index(line[startPos.Column-1:], "{"); idx >= 0 {
return strings.TrimSpace(line[startPos.Column-1 : startPos.Column-1+idx])
}
return strings.TrimSpace(line[startPos.Column-1:])
}
// Get lines from start to the line containing the opening brace
sig := c.extractLines(sourceLines, startPos.Line, bodyPos.Line)
// Remove the opening brace and anything after it
if idx := strings.Index(sig, "{"); idx >= 0 {
sig = sig[:idx]
}
return strings.TrimSpace(sig)
}
// extractLines extracts a range of lines from source (1-indexed, inclusive).
func (c *Chunker) extractLines(lines []string, start, end int) string {
if start < 1 || end < start || start > len(lines) {
return ""
}
// Adjust for 0-indexed array (start and end are 1-indexed)
startIdx := start - 1
endIdx := end
if endIdx > len(lines) {
endIdx = len(lines)
}
return strings.Join(lines[startIdx:endIdx], "\n")
}
+214
View File
@@ -0,0 +1,214 @@
package golang
import (
"context"
"os"
"path/filepath"
"testing"
"github.com/lukaszraczylo/claude-mnemonic/internal/chunking"
)
func TestGoChunker_BasicFunctions(t *testing.T) {
// Create temp test file
tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "test.go")
testCode := `package main
import "fmt"
// Greet prints a greeting message
func Greet(name string) {
fmt.Printf("Hello, %s!\n", name)
}
// Add adds two numbers
func Add(a, b int) int {
return a + b
}
// unexported function should be included by default
func helper() string {
return "helper"
}
`
if err := os.WriteFile(testFile, []byte(testCode), 0600); err != nil {
t.Fatalf("Failed to create test file: %v", err)
}
// Create chunker with default options
chunker := NewChunker(chunking.DefaultChunkOptions())
// Chunk the file
chunks, err := chunker.Chunk(context.Background(), testFile)
if err != nil {
t.Fatalf("Chunk() failed: %v", err)
}
// Verify we got all functions
if len(chunks) != 3 {
t.Errorf("Expected 3 chunks (Greet, Add, helper), got %d", len(chunks))
}
// Verify chunk details
expectedNames := map[string]bool{
"Greet": false,
"Add": false,
"helper": false,
}
for _, chunk := range chunks {
if chunk.Type != chunking.ChunkTypeFunction {
t.Errorf("Expected chunk type 'function', got '%s'", chunk.Type)
}
if chunk.Language != chunking.LanguageGo {
t.Errorf("Expected language 'go', got '%s'", chunk.Language)
}
if _, ok := expectedNames[chunk.Name]; !ok {
t.Errorf("Unexpected function name: %s", chunk.Name)
} else {
expectedNames[chunk.Name] = true
}
// Verify content is non-empty
if chunk.Content == "" {
t.Errorf("Chunk %s has empty content", chunk.Name)
}
// Verify signature is present for functions
if chunk.Signature == "" {
t.Errorf("Chunk %s has empty signature", chunk.Name)
}
}
// Verify all expected functions were found
for name, found := range expectedNames {
if !found {
t.Errorf("Expected function %s not found", name)
}
}
}
func TestGoChunker_StructsAndMethods(t *testing.T) {
tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "test.go")
testCode := `package main
// User represents a user
type User struct {
ID int
Name string
}
// GetName returns the user's name
func (u *User) GetName() string {
return u.Name
}
// SetName sets the user's name
func (u *User) SetName(name string) {
u.Name = name
}
`
if err := os.WriteFile(testFile, []byte(testCode), 0600); err != nil {
t.Fatalf("Failed to create test file: %v", err)
}
chunker := NewChunker(chunking.DefaultChunkOptions())
chunks, err := chunker.Chunk(context.Background(), testFile)
if err != nil {
t.Fatalf("Chunk() failed: %v", err)
}
// Should have 1 struct + 2 methods = 3 chunks
if len(chunks) != 3 {
t.Errorf("Expected 3 chunks (User struct, GetName, SetName), got %d", len(chunks))
}
// Find the struct and methods
var structChunk, getNameChunk, setNameChunk *chunking.Chunk
for i := range chunks {
switch chunks[i].Name {
case "User":
structChunk = &chunks[i]
case "GetName":
getNameChunk = &chunks[i]
case "SetName":
setNameChunk = &chunks[i]
}
}
// Verify struct
if structChunk == nil {
t.Fatal("User struct not found")
}
if structChunk.Type != chunking.ChunkTypeClass {
t.Errorf("Expected User to be ChunkTypeClass, got %s", structChunk.Type)
}
// Verify methods
if getNameChunk == nil {
t.Fatal("GetName method not found")
}
if getNameChunk.Type != chunking.ChunkTypeMethod {
t.Errorf("Expected GetName to be ChunkTypeMethod, got %s", getNameChunk.Type)
}
if getNameChunk.ParentName != "User" {
t.Errorf("Expected GetName parent to be 'User', got '%s'", getNameChunk.ParentName)
}
if setNameChunk == nil {
t.Fatal("SetName method not found")
}
if setNameChunk.Type != chunking.ChunkTypeMethod {
t.Errorf("Expected SetName to be ChunkTypeMethod, got %s", setNameChunk.Type)
}
if setNameChunk.ParentName != "User" {
t.Errorf("Expected SetName parent to be 'User', got '%s'", setNameChunk.ParentName)
}
}
func TestGoChunker_DocComments(t *testing.T) {
tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "test.go")
testCode := `package main
// Calculate performs a calculation.
// It takes two integers and returns their sum.
func Calculate(a, b int) int {
return a + b
}
`
if err := os.WriteFile(testFile, []byte(testCode), 0600); err != nil {
t.Fatalf("Failed to create test file: %v", err)
}
chunker := NewChunker(chunking.DefaultChunkOptions())
chunks, err := chunker.Chunk(context.Background(), testFile)
if err != nil {
t.Fatalf("Chunk() failed: %v", err)
}
if len(chunks) != 1 {
t.Fatalf("Expected 1 chunk, got %d", len(chunks))
}
chunk := chunks[0]
if chunk.DocComment == "" {
t.Error("Expected doc comment to be present")
}
// Doc comment should contain the comment text
expectedComment := "Calculate performs a calculation.\nIt takes two integers and returns their sum."
if chunk.DocComment != expectedComment {
t.Errorf("Expected doc comment '%s', got '%s'", expectedComment, chunk.DocComment)
}
}