diff --git a/plugins/llm/openai_compat.go b/plugins/llm/openai_compat.go index 965c48ecc5..59669e0ace 100644 --- a/plugins/llm/openai_compat.go +++ b/plugins/llm/openai_compat.go @@ -6,6 +6,7 @@ import ( "encoding/json" "fmt" "io" + "maps" "net/http" "time" @@ -15,11 +16,13 @@ import ( ) type compatConfig struct { - APIKey string `json:"api_key"` - Model string `json:"model"` - BaseURL string `json:"base_url"` - AgentName string `json:"agent_name"` - HistoryLen int `json:"history_length"` + APIKey string `json:"api_key"` + Model string `json:"model"` + BaseURL string `json:"base_url"` + AgentName string `json:"agent_name"` + HistoryLen int `json:"history_length"` + ToolChoice string `json:"tool_choice"` + ExtraBody map[string]any `json:"extra_params"` } type openAICompatLLM struct { @@ -44,7 +47,10 @@ func newOpenAICompat(provider string, configMap map[string]any, defaultModel, de if cfg.BaseURL == "" { cfg.BaseURL = defaultBaseURL } - return &openAICompatLLM{provider: provider, config: cfg, toolChoice: toolChoice}, nil + if cfg.ToolChoice != "" { + toolChoice = cfg.ToolChoice + } + return &openAICompatLLM{provider: provider, config: cfg, toolChoice: toolChoice, extraBody: cfg.ExtraBody}, nil } func (c *openAICompatLLM) FunctionSchemas() []map[string]any { return c.schemas } @@ -52,20 +58,17 @@ func (c *openAICompatLLM) FunctionSchemas() []map[string]any { return c.schemas func (c *openAICompatLLM) SetSchemas(schemas []map[string]any) { c.schemas = schemas } func (c *openAICompatLLM) Call(ctx context.Context, prompt string, history []llm.Message) (*llm.Response, error) { - requestBody := map[string]any{ - "model": c.config.Model, - "messages": buildMessages(prompt, history), - } + requestBody := make(map[string]any, len(c.extraBody)+4) + maps.Copy(requestBody, c.extraBody) + + requestBody["model"] = c.config.Model + requestBody["messages"] = buildMessages(prompt, history) if len(c.schemas) > 0 { requestBody["tools"] = c.schemas requestBody["tool_choice"] = c.toolChoice } - for k, v := range c.extraBody { - requestBody[k] = v - } - body, err := c.doRequest(ctx, requestBody) if err != nil { return nil, err diff --git a/plugins/llm/openai_compat_test.go b/plugins/llm/openai_compat_test.go index 643470ae07..d84dc62993 100644 --- a/plugins/llm/openai_compat_test.go +++ b/plugins/llm/openai_compat_test.go @@ -72,6 +72,18 @@ func TestConfigOverridesDefaults(t *testing.T) { require.Equal(t, "https://example.com/v1", c.config.BaseURL) } +func TestConfigOverridesToolChoiceAndExtraParams(t *testing.T) { + instance, err := llm.Load("XAILLM", map[string]any{ + "api_key": "k", + "tool_choice": "required", + "extra_params": map[string]any{"temperature": 0.8}, + }) + require.NoError(t, err) + c := instance.(*openAICompatLLM) + require.Equal(t, "required", c.toolChoice) + require.Equal(t, 0.8, c.extraBody["temperature"]) +} + func newTestCompat(t *testing.T, baseURL, toolChoice string) *openAICompatLLM { t.Helper() c, err := newOpenAICompat("TestLLM", map[string]any{ @@ -142,6 +154,26 @@ func TestOpenAICompatCallMergesExtraBody(t *testing.T) { require.Equal(t, float64(10), cap.body["max_tokens"]) } +func TestOpenAICompatCallExtraBodyCannotOverrideCoreFields(t *testing.T) { + srv, cap := captureServer(t, http.StatusOK, `{"choices":[{"message":{"content":"ok"}}]}`) + + c := newTestCompat(t, srv.URL, "required") + c.SetSchemas([]map[string]any{{"type": "function", "function": map[string]any{"name": "speak"}}}) + c.extraBody = map[string]any{ + "model": "hacked-model", + "tool_choice": "none", + "messages": "nonsense", + "temperature": 0.7, + } + + _, err := c.Call(context.Background(), "hello", nil) + require.NoError(t, err) + require.Equal(t, "test-model", cap.body["model"], "core model must win") + require.Equal(t, "required", cap.body["tool_choice"], "dedicated tool_choice must win") + require.IsType(t, []any{}, cap.body["messages"], "messages must stay the built array") + require.Equal(t, 0.7, cap.body["temperature"], "non-reserved params still pass through") +} + func TestOpenAICompatCallErrorStatus(t *testing.T) { srv, _ := captureServer(t, http.StatusInternalServerError, `{"error":"boom"}`) diff --git a/plugins/llm/router.go b/plugins/llm/router.go new file mode 100644 index 0000000000..54a8078ea3 --- /dev/null +++ b/plugins/llm/router.go @@ -0,0 +1,166 @@ +package llm + +import ( + "context" + "fmt" + "regexp" + "strings" + + "github.com/openmind/om1/internal/llm" + "github.com/openmind/om1/internal/logger" + "go.uber.org/zap" +) + +func init() { + llm.Register("RouterLLM", NewRouter) +} + +type routeConfig struct { + Name string `json:"name"` + LLMType string `json:"llm_type"` + LLMConfig map[string]any `json:"llm_config"` + Keywords []string `json:"keywords"` + Patterns []string `json:"patterns"` +} + +type routerConfig struct { + Routes []routeConfig `json:"routes"` + DefaultRoute string `json:"default_route"` + APIKey string `json:"api_key"` +} + +type route struct { + name string + llm llm.LLM + keywords []string + regexes []*regexp.Regexp +} + +func (r *route) score(text string) int { + score := 0 + for _, kw := range r.keywords { + if kw != "" && strings.Contains(text, kw) { + score++ + } + } + for _, re := range r.regexes { + if re.MatchString(text) { + score++ + } + } + return score +} + +type routerLLM struct { + routes []*route + def *route + log *zap.Logger +} + +func NewRouter(configMap map[string]any) (llm.LLM, error) { + var cfg routerConfig + if err := remarshal(configMap, &cfg); err != nil { + return nil, fmt.Errorf("RouterLLM config: %w", err) + } + if len(cfg.Routes) == 0 { + return nil, fmt.Errorf("RouterLLM: at least one route is required") + } + + router := &routerLLM{log: logger.Get().Named("RouterLLM")} + + for _, rc := range cfg.Routes { + if rc.Name == "" { + return nil, fmt.Errorf("RouterLLM: every route needs a name") + } + if rc.LLMType == "" { + return nil, fmt.Errorf("RouterLLM: route %q has no llm_type", rc.Name) + } + + subCfg := cloneStringAnyMap(rc.LLMConfig) + if cfg.APIKey != "" { + if _, ok := subCfg["api_key"]; !ok { + subCfg["api_key"] = cfg.APIKey + } + } + + sub, err := llm.Load(rc.LLMType, subCfg) + if err != nil { + return nil, fmt.Errorf("RouterLLM: load route %q (%s): %w", rc.Name, rc.LLMType, err) + } + + regexes := make([]*regexp.Regexp, 0, len(rc.Patterns)) + for _, p := range rc.Patterns { + re, err := regexp.Compile("(?i)" + p) + if err != nil { + return nil, fmt.Errorf("RouterLLM: route %q pattern %q: %w", rc.Name, p, err) + } + regexes = append(regexes, re) + } + + keywords := make([]string, 0, len(rc.Keywords)) + for _, kw := range rc.Keywords { + keywords = append(keywords, strings.ToLower(kw)) + } + + router.routes = append(router.routes, &route{ + name: rc.Name, + llm: sub, + keywords: keywords, + regexes: regexes, + }) + } + + router.def = router.routes[0] + if cfg.DefaultRoute != "" { + found := false + for _, r := range router.routes { + if r.name == cfg.DefaultRoute { + router.def = r + found = true + break + } + } + if !found { + return nil, fmt.Errorf("RouterLLM: default_route %q is not a defined route", cfg.DefaultRoute) + } + } + + return router, nil +} + +func (r *routerLLM) SetSchemas(schemas []map[string]any) { + for _, rt := range r.routes { + rt.llm.SetSchemas(schemas) + } +} + +func (r *routerLLM) FunctionSchemas() []map[string]any { return r.def.llm.FunctionSchemas() } + +func (r *routerLLM) logger() *zap.Logger { + if r.log != nil { + return r.log + } + return logger.Get().Named("RouterLLM") +} + +func (r *routerLLM) pick(prompt string) *route { + text := strings.ToLower(extractVoiceInput(prompt)) + if text == "" { + return r.def + } + + best := r.def + bestScore := 0 + for _, rt := range r.routes { + if s := rt.score(text); s > bestScore { + best, bestScore = rt, s + } + } + return best +} + +func (r *routerLLM) Call(ctx context.Context, prompt string, history []llm.Message) (*llm.Response, error) { + chosen := r.pick(prompt) + r.logger().Info("routed", zap.String("route", chosen.name)) + return chosen.llm.Call(ctx, prompt, history) +} diff --git a/plugins/llm/router_test.go b/plugins/llm/router_test.go new file mode 100644 index 0000000000..aae6125182 --- /dev/null +++ b/plugins/llm/router_test.go @@ -0,0 +1,106 @@ +package llm + +import ( + "context" + "regexp" + "testing" + + "github.com/openmind/om1/internal/llm" + "github.com/stretchr/testify/require" +) + +func newTestRouter() (*routerLLM, *stubLLM, *stubLLM) { + chat := &stubLLM{resp: respWithCalls("speak")} + command := &stubLLM{resp: respWithCalls("move")} + chatRoute := &route{name: "chat", llm: chat, regexes: []*regexp.Regexp{regexp.MustCompile(`(?i)\?`)}} + cmdRoute := &route{ + name: "command", + llm: command, + keywords: []string{"sit", "come", "fetch", "spin"}, + } + r := &routerLLM{routes: []*route{cmdRoute, chatRoute}, def: chatRoute} + return r, chat, command +} + +func TestRouterPicksCommandOnKeyword(t *testing.T) { + r, _, _ := newTestRouter() + require.Equal(t, "command", r.pick("INPUT Voice: come here and sit").name) +} + +func TestRouterPicksChatOnQuestion(t *testing.T) { + r, _, _ := newTestRouter() + require.Equal(t, "chat", r.pick("Voice: how are you today?").name) +} + +func TestRouterFallsBackToDefault(t *testing.T) { + r, _, _ := newTestRouter() + require.Equal(t, "chat", r.pick("Voice: the weather is nice").name, "no match → default route") +} + +func TestRouterEmptyVoiceUsesDefault(t *testing.T) { + r, _, _ := newTestRouter() + require.Equal(t, "chat", r.pick("no voice line here").name) +} + +func TestRouterHighestScoreWins(t *testing.T) { + r, _, _ := newTestRouter() + // Two command keywords beat a single chat question mark. + require.Equal(t, "command", r.pick("Voice: sit and fetch, ok?").name) +} + +func TestRouterCallDispatchesToChosen(t *testing.T) { + r, _, command := newTestRouter() + resp, err := r.Call(context.Background(), "Voice: fetch the ball", nil) + require.NoError(t, err) + require.Same(t, command.resp, resp) +} + +func TestRouterSetSchemasPropagates(t *testing.T) { + r, chat, command := newTestRouter() + schemas := []map[string]any{{"name": "speak"}} + r.SetSchemas(schemas) + require.Equal(t, schemas, chat.schemas) + require.Equal(t, schemas, command.schemas) +} + +func TestNewRouterRequiresRoutes(t *testing.T) { + _, err := NewRouter(map[string]any{}) + require.Error(t, err) +} + +func TestNewRouterUnknownDefaultRoute(t *testing.T) { + _, err := NewRouter(map[string]any{ + "default_route": "nope", + "routes": []map[string]any{ + {"name": "chat", "llm_type": "GeminiLLM", "llm_config": map[string]any{"api_key": "k"}}, + }, + }) + require.Error(t, err) +} + +func TestNewRouterBuildsRoutes(t *testing.T) { + got, err := NewRouter(map[string]any{ + "api_key": "shared-key", + "default_route": "chat", + "routes": []map[string]any{ + { + "name": "command", + "llm_type": "GeminiLLM", + "keywords": []string{"sit", "come"}, + "patterns": []string{`\bstop\b`}, + }, + { + "name": "chat", + "llm_type": "OpenRouter", + "llm_config": map[string]any{"temperature": 0.9}, + }, + }, + }) + require.NoError(t, err) + rt := got.(*routerLLM) + require.Len(t, rt.routes, 2) + require.Equal(t, "chat", rt.def.name) + require.Equal(t, "command", rt.pick("Voice: STOP now").name) +} + +var _ llm.LLM = (*routerLLM)(nil)