Fixes across 25 files addressing 46+ review comments: - pkg/ai/metrics.go: handle error from Close() on writable file handle - pkg/ansible: restore loop vars after loop, restore become settings, fix Upload with become=true and no password (use sudo -n), honour SSH timeout config, use E() helper for contextual errors, quote git refs in checkout commands - pkg/rag: validate chunk config, guard negative-to-uint64 conversion, use E() helper for errors, add context timeout to Ollama HTTP calls - pkg/deploy/python: fix exec.ExitError type assertion (was os.PathError), handle os.UserHomeDir() error - pkg/build/buildcmd: use cmd.Context() instead of context.Background() for proper Ctrl+C cancellation - install.bat: add curl timeouts, CRLF line endings, use --connect-timeout for archive downloads - install.sh: use absolute path for version check in CI mode - tools/rag: fix broken ingest.py function def, escape HTML in query.py, pin qdrant-client version, add markdown code block languages - internal/cmd/rag: add chunk size validation, env override handling Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
110 lines
2.4 KiB
Go
110 lines
2.4 KiB
Go
package rag
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/host-uk/core/pkg/i18n"
|
|
"github.com/host-uk/core/pkg/rag"
|
|
"github.com/spf13/cobra"
|
|
)
|
|
|
|
var (
|
|
queryCollection string
|
|
limit int
|
|
threshold float32
|
|
category string
|
|
format string
|
|
)
|
|
|
|
var queryCmd = &cobra.Command{
|
|
Use: "query [question]",
|
|
Short: i18n.T("cmd.rag.query.short"),
|
|
Long: i18n.T("cmd.rag.query.long"),
|
|
Args: cobra.ExactArgs(1),
|
|
RunE: runQuery,
|
|
}
|
|
|
|
func runQuery(cmd *cobra.Command, args []string) error {
|
|
question := args[0]
|
|
ctx := context.Background()
|
|
|
|
// Connect to Qdrant
|
|
qdrantClient, err := rag.NewQdrantClient(rag.QdrantConfig{
|
|
Host: qdrantHost,
|
|
Port: qdrantPort,
|
|
UseTLS: false,
|
|
})
|
|
if err != nil {
|
|
return fmt.Errorf("failed to connect to Qdrant: %w", err)
|
|
}
|
|
defer qdrantClient.Close()
|
|
|
|
// Connect to Ollama
|
|
ollamaClient, err := rag.NewOllamaClient(rag.OllamaConfig{
|
|
Host: ollamaHost,
|
|
Port: ollamaPort,
|
|
Model: model,
|
|
})
|
|
if err != nil {
|
|
return fmt.Errorf("failed to connect to Ollama: %w", err)
|
|
}
|
|
|
|
// Configure query
|
|
if limit < 0 {
|
|
limit = 0
|
|
}
|
|
cfg := rag.QueryConfig{
|
|
Collection: queryCollection,
|
|
Limit: uint64(limit),
|
|
Threshold: threshold,
|
|
Category: category,
|
|
}
|
|
|
|
// Run query
|
|
results, err := rag.Query(ctx, qdrantClient, ollamaClient, question, cfg)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Format output
|
|
switch format {
|
|
case "json":
|
|
fmt.Println(rag.FormatResultsJSON(results))
|
|
case "context":
|
|
fmt.Println(rag.FormatResultsContext(results))
|
|
default:
|
|
fmt.Println(rag.FormatResultsText(results))
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// QueryDocs is exported for use by other packages (e.g., MCP).
|
|
func QueryDocs(ctx context.Context, question, collectionName string, topK int) ([]rag.QueryResult, error) {
|
|
qdrantClient, err := rag.NewQdrantClient(rag.DefaultQdrantConfig())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer qdrantClient.Close()
|
|
|
|
ollamaClient, err := rag.NewOllamaClient(rag.DefaultOllamaConfig())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
cfg := rag.DefaultQueryConfig()
|
|
cfg.Collection = collectionName
|
|
cfg.Limit = uint64(topK)
|
|
|
|
return rag.Query(ctx, qdrantClient, ollamaClient, question, cfg)
|
|
}
|
|
|
|
// QueryDocsContext is exported and returns context-formatted results.
|
|
func QueryDocsContext(ctx context.Context, question, collectionName string, topK int) (string, error) {
|
|
results, err := QueryDocs(ctx, question, collectionName, topK)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return rag.FormatResultsContext(results), nil
|
|
}
|