mirror of
https://github.com/Warky-Devs/vecna.git
synced 2026-05-05 01:26:58 +00:00
feat(query): add support for extra maps in query command
This commit is contained in:
@@ -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,9 +117,11 @@ func runQuery(_ *cobra.Command, args []string) error {
|
||||
vec := resp.Embeddings[0]
|
||||
|
||||
if !queryRaw {
|
||||
adp, err := buildAdapter(cfg)
|
||||
if err != nil {
|
||||
return fmt.Errorf("build adapter: %w", err)
|
||||
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 {
|
||||
|
||||
Reference in New Issue
Block a user