fix

Author: m1rl0k

main.go

@@ -99,6 +99,7 @@ var (
 	// CLI flags
 	enableLLM           = flag.Bool("llm", false, "Enable LLM-powered verification")
 	modelPath           = flag.String("model-path", ".gosecretscanner/models/granite-4.0-micro-Q4_K_M.gguf", "Path to LLM model")
+	embeddingsPath      = flag.String("embeddings-path", "", "Path to embeddings models directory (defaults to .gosecretscanner/models)")
 	dbPath              = flag.String("db-path", ".gosecretscanner/findings.db", "Path to vector store database")
 	similarityThreshold = flag.Float64("similarity", 0.8, "Similarity threshold for vector search")
 	keepVectorStore     = flag.Bool("keep-vector-store", false, "Keep the vector store database after the run")
@@ -127,6 +128,7 @@ func main() {
 			Enabled:             true,
 			DBPath:              *dbPath,
 			ModelPath:           *modelPath,
+			EmbeddingsPath:      *embeddingsPath,
 			SimilarityThreshold: float32(*similarityThreshold),
 			EphemeralStore:      !*keepVectorStore,
 			LLMEndpoint:         *llmEndpoint,

pkg/verification/pipeline.go

@@ -27,6 +27,7 @@ type Config struct {
 	Enabled             bool
 	DBPath              string
 	ModelPath           string
+	EmbeddingsPath      string  // Path to embeddings models directory
 	SimilarityThreshold float32
 	EphemeralStore      bool
 	LLMEndpoint         string
@@ -44,8 +45,17 @@ func NewPipeline(config *Config) (*Pipeline, error) {
 	embeddingGen := embeddings.NewEmbeddingGenerator(config.Enabled)
 
 	// Initialize BGE embeddings with ONNX model
-	if config.Enabled && config.ModelPath != "" {
-		modelsDir := filepath.Dir(config.ModelPath)
+	// Use EmbeddingsPath if specified, otherwise try to derive from ModelPath, or default to .gosecretscanner/models
+	if config.Enabled {
+		var modelsDir string
+		if config.EmbeddingsPath != "" {
+			modelsDir = config.EmbeddingsPath
+		} else if config.ModelPath != "" {
+			modelsDir = filepath.Dir(config.ModelPath)
+		} else {
+			modelsDir = ".gosecretscanner/models"
+		}
+
 		if err := embeddingGen.Initialize(modelsDir); err != nil {
 			// Non-fatal: fall back to hash-based embeddings
 			fmt.Printf("Warning: failed to initialize BGE embeddings, using hash-based fallback: %v\n", err)
@@ -105,12 +115,15 @@ func (p *Pipeline) VerifyFinding(
 		return nil, fmt.Errorf("failed to parse file: %w", err)
 	}
 
+	// Get surrounding code context (±5 lines)
+	surroundingCode := p.parser.GetContext(content, lineNumber, 5)
+
 	// Create code context
 	codeContext := &llm.CodeContext{
 		FilePath:        filePath,
 		Language:        parsed.Language,
 		Function:        p.findFunction(parsed, lineNumber),
-		SurroundingCode: line,
+		SurroundingCode: surroundingCode,
 		Imports:         parsed.Imports,
 		IsTest:          parsed.IsTest,
 	}