feat(query): add support for extra maps in query command

This commit is contained in:
2026-04-11 22:36:08 +02:00
parent 058aaa73e4
commit 200b0dae95

View File

@@ -4,19 +4,20 @@ import (
"context"
"encoding/json"
"fmt"
"net/http"
"os"
"strings"
"time"
"github.com/spf13/cobra"
"github.com/Warky-Devs/vecna.git/pkg/adapter"
"github.com/Warky-Devs/vecna.git/pkg/config"
"github.com/Warky-Devs/vecna.git/pkg/embedclient"
)
var (
queryTarget string
queryMap string
queryRaw bool
queryCompact bool
)
@@ -35,6 +36,8 @@ Text can be supplied as a positional argument or via stdin (use - as the argumen
func init() {
queryCmd.Flags().StringVar(&queryTarget, "target", "",
"forward target to use (default: forward.default from config)")
queryCmd.Flags().StringVar(&queryMap, "map", "",
"extra_map key to use; overrides both adapter and forward target")
queryCmd.Flags().BoolVar(&queryRaw, "raw", false,
"skip the adapter — output the raw vector from the backing model")
queryCmd.Flags().BoolVar(&queryCompact, "compact", false,
@@ -54,8 +57,24 @@ func runQuery(_ *cobra.Command, args []string) error {
return err
}
// Resolve target.
// Resolve adapter and forward target — --map overrides both.
var adp adapter.Adapter
targetName := queryTarget
if queryMap != "" {
mc, ok := cfg.ExtraMaps[queryMap]
if !ok {
return fmt.Errorf("extra_map %q not found in config", queryMap)
}
extraMaps, buildErr := buildExtraMapAdapters(cfg)
if buildErr != nil {
return fmt.Errorf("build extra_maps: %w", buildErr)
}
adp = extraMaps[queryMap].Adapter
if mc.ForwardTarget != "" && targetName == "" {
targetName = mc.ForwardTarget
}
}
if targetName == "" {
targetName = cfg.Forward.Default
}
@@ -67,30 +86,24 @@ func runQuery(_ *cobra.Command, args []string) error {
return fmt.Errorf("target %q has no endpoints", targetName)
}
// Build client (use first endpoint directly — no router needed for a one-shot query).
ep := target.Endpoints[0]
apiKey := ep.APIKey
if apiKey == "" {
apiKey = target.APIKey
// Build a full router so retry/demerit logic applies identically to the server path.
client, err := buildClients(cfg)
if err != nil {
return fmt.Errorf("build clients: %w", err)
}
embedClient, ok := client[targetName]
if !ok {
return fmt.Errorf("target %q not found after building clients", targetName)
}
timeout := time.Duration(target.TimeoutSecs) * time.Second
if timeout == 0 {
timeout = 30 * time.Second
}
httpClient := &http.Client{Timeout: timeout}
var client embedclient.Client
switch target.APIType {
case "google":
client = embedclient.NewGoogle(ep.URL, apiKey, target.Model, httpClient)
default:
client = embedclient.NewOpenAI(ep.URL, apiKey, httpClient)
}
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
resp, err := client.Embed(ctx, embedclient.Request{
resp, err := embedClient.Embed(ctx, embedclient.Request{
Texts: []string{text},
Model: target.Model,
})
@@ -104,10 +117,12 @@ func runQuery(_ *cobra.Command, args []string) error {
vec := resp.Embeddings[0]
if !queryRaw {
adp, err := buildAdapter(cfg)
if adp == nil {
adp, err = buildAdapter(cfg)
if err != nil {
return fmt.Errorf("build adapter: %w", err)
}
}
vec, err = adp.Adapt(vec)
if err != nil {
return fmt.Errorf("adapt: %w", err)