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