mirror of
https://github.com/lukaszraczylo/claude-mnemonic.git
synced 2026-10-19 06:18:13 +00:00
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:
@@ -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")
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user