diff --git a/pkg/aiproxy/handlers/chat_completions.go b/pkg/aiproxy/handlers/chat_completions.go index fac2948870..e52cd4c0a9 100644 --- a/pkg/aiproxy/handlers/chat_completions.go +++ b/pkg/aiproxy/handlers/chat_completions.go @@ -36,6 +36,7 @@ import ( ) const headerAiVirtualKey = "X-Ai-Virtual-Key" +const headerAiRoutingId = "X-Ai-Routing-Id" func extractVirtualKey(r *http.Request) string { if v := strings.TrimSpace(r.Header.Get(headerAiVirtualKey)); v != "" { @@ -49,6 +50,10 @@ func extractVirtualKey(r *http.Request) string { return "" } +func extractRoutingId(r *http.Request) string { + return strings.TrimSpace(r.Header.Get(headerAiRoutingId)) +} + func upstreamErrorStatusCode(uerr *upstream.Error) int { if uerr == nil || uerr.StatusCode <= 0 { return 0 @@ -141,7 +146,7 @@ func chatCompletionsHandler(ctx context.Context, w http.ResponseWriter, r *http. vk := extractVirtualKey(r) userCred := auth.AdminCredential() - up, err := models.ResolveChatUpstream(ctx, userCred, vk, dict) + up, err := models.ResolveChatUpstream(ctx, userCred, vk, dict, extractRoutingId(r)) if err != nil { dbg.Error("resolve upstream: %v", err) failAPILogRecord(rec, http.StatusInternalServerError, "resolve_upstream", err) diff --git a/pkg/aiproxy/handlers/completions.go b/pkg/aiproxy/handlers/completions.go index 83e8403b5c..ed8edff470 100644 --- a/pkg/aiproxy/handlers/completions.go +++ b/pkg/aiproxy/handlers/completions.go @@ -71,7 +71,7 @@ func completionsHandler(ctx context.Context, w http.ResponseWriter, r *http.Requ vk := extractVirtualKey(r) userCred := auth.AdminCredential() - up, err := models.ResolveChatUpstream(ctx, userCred, vk, dict) + up, err := models.ResolveChatUpstream(ctx, userCred, vk, dict, extractRoutingId(r)) if err != nil { dbg.Error("resolve upstream: %v", err) failAPILogRecord(rec, http.StatusInternalServerError, "resolve_upstream", err) diff --git a/pkg/aiproxy/handlers/embeddings.go b/pkg/aiproxy/handlers/embeddings.go index ce420a1a8a..bc27a325b8 100644 --- a/pkg/aiproxy/handlers/embeddings.go +++ b/pkg/aiproxy/handlers/embeddings.go @@ -69,7 +69,7 @@ func embeddingsHandler(ctx context.Context, w http.ResponseWriter, r *http.Reque vk := extractVirtualKey(r) userCred := auth.AdminCredential() - up, err := models.ResolveChatUpstream(ctx, userCred, vk, dict) + up, err := models.ResolveChatUpstream(ctx, userCred, vk, dict, extractRoutingId(r)) if err != nil { dbg.Error("resolve upstream: %v", err) failAPILogRecord(rec, http.StatusInternalServerError, "resolve_upstream", err) diff --git a/pkg/aiproxy/handlers/images_generations.go b/pkg/aiproxy/handlers/images_generations.go index c8cd7351e7..48fe6a886d 100644 --- a/pkg/aiproxy/handlers/images_generations.go +++ b/pkg/aiproxy/handlers/images_generations.go @@ -69,7 +69,7 @@ func imagesGenerationsHandler(ctx context.Context, w http.ResponseWriter, r *htt vk := extractVirtualKey(r) userCred := auth.AdminCredential() - up, err := models.ResolveChatUpstream(ctx, userCred, vk, dict) + up, err := models.ResolveChatUpstream(ctx, userCred, vk, dict, extractRoutingId(r)) if err != nil { dbg.Error("resolve upstream: %v", err) failAPILogRecord(rec, http.StatusInternalServerError, "resolve_upstream", err) diff --git a/pkg/aiproxy/handlers/messages.go b/pkg/aiproxy/handlers/messages.go index a5267c1f2a..d23f6758ed 100644 --- a/pkg/aiproxy/handlers/messages.go +++ b/pkg/aiproxy/handlers/messages.go @@ -77,7 +77,7 @@ func messagesHandler(ctx context.Context, w http.ResponseWriter, r *http.Request vk := extractVirtualKey(r) userCred := auth.AdminCredential() - up, err := models.ResolveChatUpstream(ctx, userCred, vk, dict) + up, err := models.ResolveChatUpstream(ctx, userCred, vk, dict, extractRoutingId(r)) if err != nil { dbg.Error("resolve upstream: %v", err) failAPILogRecord(rec, http.StatusInternalServerError, "resolve_upstream", err) diff --git a/pkg/aiproxy/handlers/responses.go b/pkg/aiproxy/handlers/responses.go index 79a245d51d..705b14ce8e 100644 --- a/pkg/aiproxy/handlers/responses.go +++ b/pkg/aiproxy/handlers/responses.go @@ -97,7 +97,7 @@ func handleResponsesCreate(ctx context.Context, w http.ResponseWriter, r *http.R vk := extractVirtualKey(r) userCred := auth.AdminCredential() - up, err := models.ResolveChatUpstream(ctx, userCred, vk, dict) + up, err := models.ResolveChatUpstream(ctx, userCred, vk, dict, extractRoutingId(r)) if err != nil { dbg.Error("resolve upstream: %v", err) failAPILogRecord(rec, http.StatusInternalServerError, "resolve_upstream", err) @@ -293,7 +293,7 @@ func handleResponsesSubResource(ctx context.Context, w http.ResponseWriter, r *h writeResponsesError(ctx, w, http.StatusBadRequest, "invalid_request_error", "model query parameter is required") return } - up, err := models.ResolveChatUpstream(ctx, userCred, vk, probe) + up, err := models.ResolveChatUpstream(ctx, userCred, vk, probe, extractRoutingId(r)) if err != nil { dbg.Error("resolve upstream: %v", err) httperrors.GeneralServerError(ctx, w, err) diff --git a/pkg/aiproxy/models/chat_upstream.go b/pkg/aiproxy/models/chat_upstream.go index 3e09d7d422..4873149e17 100644 --- a/pkg/aiproxy/models/chat_upstream.go +++ b/pkg/aiproxy/models/chat_upstream.go @@ -191,7 +191,8 @@ func listProjectRoutingsForVirtualKey(ctx context.Context, userCred mcclient.Tok // pickRoutingForRequest chooses the best matching ai_routing on the current aiproxy instance. // Hierarchical refs (routingKey/catalogPart) match only ai_routing.model_key on routeKey. -// Flat refs: Phase 1 exact ai_routing.model_key match, Phase 2 ai_routing.model_pattern match. +// Flat refs: Phase 1 exact ai_routing.model_key match, Phase 2 non-empty ai_routing.model_pattern match. +// Empty model_pattern is not a wildcard; unmatched models return nil (ResolveChatUpstream maps that to ErrNotFound). func pickRoutingForRequest(routings []SAiRouting, reqModel, currentNodeId string) (*SAiRouting, error) { ref := parseClientModelRef(reqModel) if ref.hierarchical { @@ -205,6 +206,9 @@ func pickRoutingForRequest(routings []SAiRouting, reqModel, currentNodeId string return picked, err } return pickRoutingByMatch(routings, reqModel, currentNodeId, func(r *SAiRouting, reqModel string) bool { + if strings.TrimSpace(r.ModelPattern) == "" { + return false + } return modelPatternMatches(r.ModelPattern, reqModel) }) } @@ -238,6 +242,33 @@ func pickRoutingByMatch(routings []SAiRouting, reqModel, currentNodeId string, m return nil, nil } +func pickRoutingById(routings []SAiRouting, routingId, currentNodeId string) (*SAiRouting, error) { + routingId = strings.TrimSpace(routingId) + if routingId == "" { + return nil, nil + } + var boundElsewhere *SAiRouting + for i := range routings { + r := &routings[i] + if r.Id != routingId { + continue + } + if !proxyNodeScopeMatches(r.AiProxyNodeId, currentNodeId) { + if strings.TrimSpace(r.AiProxyNodeId) != "" { + boundElsewhere = r + } + continue + } + return r, nil + } + if boundElsewhere != nil { + return nil, errors.Wrapf(httperrors.ErrForbidden, + "ai_routing %q is bound to ai_proxy_node %q; use that instance endpoint", + boundElsewhere.Name, boundElsewhere.AiProxyNodeId) + } + return nil, nil +} + type resolvedCatalogModel struct { provider *SAiProvider model *SAiModel @@ -291,10 +322,10 @@ func resolveCatalogModelFromRouting( // ResolveChatUpstream resolves upstream URL, API key, and catalog model_key for a chat request: // 1. ai_virtual_key (auth + project scope) -// 2. ai_routing in that project (model_key exact match first, then model_pattern / optional proxy-node scope, priority) +// 2. optional preferredRoutingId (X-Ai-Routing-Id) pins the ai_routing; otherwise model_key then non-empty model_pattern // 3. ai_routing_model -> ai_provider + ai_model // 4. ai_key rows for that provider matching the catalog model_key (weight) -func ResolveChatUpstream(ctx context.Context, userCred mcclient.TokenCredential, virtualKey string, body *jsonutils.JSONDict) (*ChatUpstream, error) { +func ResolveChatUpstream(ctx context.Context, userCred mcclient.TokenCredential, virtualKey string, body *jsonutils.JSONDict, preferredRoutingId string) (*ChatUpstream, error) { vk, err := loadEnabledVirtualKey(virtualKey) if err != nil { return nil, err @@ -309,7 +340,12 @@ func ResolveChatUpstream(ctx context.Context, userCred mcclient.TokenCredential, if err != nil { return nil, err } - routing, err := pickRoutingForRequest(routings, reqModel, CurrentProxyNodeId()) + var routing *SAiRouting + if strings.TrimSpace(preferredRoutingId) != "" { + routing, err = pickRoutingById(routings, preferredRoutingId, CurrentProxyNodeId()) + } else { + routing, err = pickRoutingForRequest(routings, reqModel, CurrentProxyNodeId()) + } if err != nil { return nil, err } diff --git a/pkg/aiproxy/models/chat_upstream_test.go b/pkg/aiproxy/models/chat_upstream_test.go index 1045413dd8..94928b5316 100644 --- a/pkg/aiproxy/models/chat_upstream_test.go +++ b/pkg/aiproxy/models/chat_upstream_test.go @@ -103,11 +103,41 @@ func TestRouterRequestModel(t *testing.T) { } } -func TestPickRoutingForRequestBoundElsewhereHierarchical(t *testing.T) { - routings := []SAiRouting{ - {Priority: 10, ModelKey: "claude", AiProxyNodeId: "node-b"}, +func TestPickRoutingByIdHitsPreferred(t *testing.T) { + qwen := SAiRouting{ModelKey: "qwen38-Qwen3.8-27B-NVFP4"} + qwen.Id = "rid-qwen" + qwen.Name = "qwen" + ds := SAiRouting{ModelKey: "deepseek-v4-flash"} + ds.Id = "rid-ds" + ds.Name = "deepseek" + routings := []SAiRouting{qwen, ds} + + picked, err := pickRoutingById(routings, "rid-qwen", "primary") + if err != nil { + t.Fatalf("unexpected err: %v", err) } - _, err := pickRoutingForRequest(routings, "claude/gpt-4", "node-a") + if picked == nil || picked.Id != "rid-qwen" { + t.Fatalf("expected pinned qwen routing, got %#v", picked) + } +} + +func TestPickRoutingByIdNotInList(t *testing.T) { + ds := SAiRouting{ModelKey: "deepseek-v4-flash"} + ds.Id = "rid-ds" + picked, err := pickRoutingById([]SAiRouting{ds}, "rid-missing", "primary") + if err != nil { + t.Fatalf("unexpected err: %v", err) + } + if picked != nil { + t.Fatalf("expected nil when routing id is not visible, got %#v", picked) + } +} + +func TestPickRoutingByIdBoundElsewhere(t *testing.T) { + r := SAiRouting{ModelKey: "qwen", AiProxyNodeId: "node-b"} + r.Id = "rid-qwen" + r.Name = "qwen" + _, err := pickRoutingById([]SAiRouting{r}, "rid-qwen", "node-a") if err == nil { t.Fatal("expected forbidden error") } diff --git a/pkg/aiproxy/models/list_models_test.go b/pkg/aiproxy/models/list_models_test.go index 28b19f6842..40d7be7302 100644 --- a/pkg/aiproxy/models/list_models_test.go +++ b/pkg/aiproxy/models/list_models_test.go @@ -169,6 +169,20 @@ func TestPickRoutingForRequestFallbackPattern(t *testing.T) { } } +func TestPickRoutingForRequestEmptyPatternDoesNotMatch(t *testing.T) { + routings := []SAiRouting{ + {Priority: 10, ModelPattern: ""}, + {Priority: 20, ModelKey: "deepseek-v4-flash"}, + } + picked, err := pickRoutingForRequest(routings, "qwen38-Qwen3.8-27B-NVFP4", "primary") + if err != nil { + t.Fatalf("unexpected err: %v", err) + } + if picked != nil { + t.Fatalf("empty model_pattern must not catch unmatched models, got %#v", picked) + } +} + func TestModelKeyMatches(t *testing.T) { if !modelKeyMatches("Foo", "foo") { t.Fatal("expected case-insensitive match") diff --git a/pkg/apis/llm/mcp_agent.go b/pkg/apis/llm/mcp_agent.go index f695fe0fe8..0af6c782cd 100644 --- a/pkg/apis/llm/mcp_agent.go +++ b/pkg/apis/llm/mcp_agent.go @@ -76,33 +76,39 @@ type MCPAgentListInput struct { type MCPAgentCreateInput struct { apis.SharableVirtualResourceCreateInput - LLMId string `json:"llm_id" help:"LLM 实例 ID，如果提供则自动获取 llm_url"` - LLMUrl string `json:"llm_url" help:"后端大模型的 base 请求地址"` - LLMDriver string `json:"llm_driver" help:"使用的大模型驱动，可以是 ollama 或 openai"` - Model string `json:"model" help:"使用的模型名称"` - ApiKey string `json:"api_key" help:"在 llm_driver 中需要用到的认证"` - McpServer string `json:"mcp_server" help:"mcp 服务器的后端地址"` - DefaultAgent *bool `json:"default_agent,omitempty" help:"set as default MCP agent (only one can be true globally)"` + LLMId string `json:"llm_id" help:"deprecated; ignored"` + LLMUrl string `json:"llm_url" help:"AI 网关 OpenAI 兼容 base 请求地址"` + LLMDriver string `json:"llm_driver" help:"使用的大模型驱动，固定为 openai"` + Model string `json:"model" help:"使用的模型名称"` + ApiKey string `json:"api_key" help:"deprecated; ignored"` + McpServer string `json:"mcp_server" help:"mcp 服务器的后端地址"` + AiProxyRoutingId string `json:"aiproxy_routing_id" help:"关联的 AI 网关路由规则 ID"` + AiproxyVirtualKeyId string `json:"aiproxy_virtual_key_id" help:"关联的 AI 网关 API Key ID"` + DefaultAgent *bool `json:"default_agent,omitempty" help:"set as default MCP agent (only one can be true globally)"` } type MCPAgentUpdateInput struct { apis.SharableVirtualResourceCreateInput - LLMId *string `json:"llm_id,omitempty" help:"LLM 实例 ID，如果提供则自动获取 llm_url"` - LLMUrl *string `json:"llm_url,omitempty" help:"后端大模型的 base 请求地址"` - LLMDriver *string `json:"llm_driver,omitempty" help:"使用的大模型驱动，可以是 ollama 或 openai"` - Model *string `json:"model,omitempty" help:"使用的模型名称"` - ApiKey *string `json:"api_key,omitempty" help:"在 llm_driver 中需要用到的认证"` - McpServer *string `json:"mcp_server,omitempty" help:"mcp 服务器的后端地址"` - DefaultAgent *bool `json:"default_agent,omitempty" help:"set as default MCP agent (only one can be true globally)"` + LLMId *string `json:"llm_id,omitempty" help:"deprecated; ignored"` + LLMUrl *string `json:"llm_url,omitempty" help:"AI 网关 OpenAI 兼容 base 请求地址"` + LLMDriver *string `json:"llm_driver,omitempty" help:"使用的大模型驱动，固定为 openai"` + Model *string `json:"model,omitempty" help:"使用的模型名称"` + ApiKey *string `json:"api_key,omitempty" help:"deprecated; ignored"` + McpServer *string `json:"mcp_server,omitempty" help:"mcp 服务器的后端地址"` + AiProxyRoutingId *string `json:"aiproxy_routing_id,omitempty" help:"关联的 AI 网关路由规则 ID"` + AiproxyVirtualKeyId *string `json:"aiproxy_virtual_key_id,omitempty" help:"关联的 AI 网关 API Key ID"` + DefaultAgent *bool `json:"default_agent,omitempty" help:"set as default MCP agent (only one can be true globally)"` } type MCPAgentDetails struct { apis.SharableVirtualResourceDetails - LLMId string `json:"llm_id"` - LLMName string `json:"llm_name"` - DefaultAgent bool `json:"default_agent"` + LLMId string `json:"llm_id"` + LLMName string `json:"llm_name"` + AiProxyRoutingId string `json:"aiproxy_routing_id"` + AiproxyVirtualKeyId string `json:"aiproxy_virtual_key_id"` + DefaultAgent bool `json:"default_agent"` } type LLMToolRequestInput struct { diff --git a/pkg/llm/drivers/llm_client/openai.go b/pkg/llm/drivers/llm_client/openai.go index 324c98d4cd..0e06434c47 100644 --- a/pkg/llm/drivers/llm_client/openai.go +++ b/pkg/llm/drivers/llm_client/openai.go @@ -29,6 +29,20 @@ func newOpenAI() models.ILLMClient { return new(openai) } +func setAiproxyRequestHeaders(ctx context.Context, httpReq *http.Request, mcpAgent *models.SMCPAgent) error { + apiKey, err := mcpAgent.GetAiproxyVirtualKey(ctx) + if err != nil { + return err + } + if apiKey != "" { + httpReq.Header.Set("Authorization", "Bearer "+apiKey) + } + if rid := strings.TrimSpace(mcpAgent.AiproxyRoutingId); rid != "" { + httpReq.Header.Set("X-Ai-Routing-Id", rid) + } + return nil +} + func (o *openai) GetType() api.LLMClientType { return api.LLM_CLIENT_OPENAI } @@ -283,13 +297,9 @@ func (o *openai) doChatStreamRequest(ctx context.Context, mcpAgent *models.SMCPA return errors.Wrap(err, "create request") } httpReq.Header.Set("Content-Type", "application/json") - apiKey, err := mcpAgent.GetApiKey() - if err != nil { + if err := setAiproxyRequestHeaders(ctx, httpReq, mcpAgent); err != nil { return err } - if apiKey != "" { - httpReq.Header.Set("Authorization", "Bearer "+apiKey) - } client := &http.Client{ // Stream request no timeout @@ -372,12 +382,8 @@ func (o *openai) doChatRequest(ctx context.Context, mcpAgent *models.SMCPAgent, return nil, errors.Wrap(err, "create request") } httpReq.Header.Set("Content-Type", "application/json") - apiKey, err := mcpAgent.GetApiKey() - if err != nil { - return nil, errors.Wrap(err, "get apiKey") - } - if apiKey != "" { - httpReq.Header.Set("Authorization", "Bearer "+apiKey) + if err := setAiproxyRequestHeaders(ctx, httpReq, mcpAgent); err != nil { + return nil, err } client := &http.Client{ diff --git a/pkg/llm/models/mcp_agent.go b/pkg/llm/models/mcp_agent.go index add48c0132..76301a2510 100644 --- a/pkg/llm/models/mcp_agent.go +++ b/pkg/llm/models/mcp_agent.go @@ -12,7 +12,6 @@ import ( "yunion.io/x/jsonutils" "yunion.io/x/log" "yunion.io/x/pkg/errors" - seclib "yunion.io/x/pkg/utils" "yunion.io/x/sqlchemy" api "yunion.io/x/onecloud/pkg/apis/llm" @@ -23,6 +22,7 @@ import ( "yunion.io/x/onecloud/pkg/llm/options" "yunion.io/x/onecloud/pkg/llm/utils" "yunion.io/x/onecloud/pkg/mcclient" + apmodules "yunion.io/x/onecloud/pkg/mcclient/modules/aiproxy" "yunion.io/x/onecloud/pkg/util/stringutils2" ) @@ -119,17 +119,21 @@ func (man *SMCPAgentManager) GetDefaultMcpServerTools(ctx context.Context, userC type SMCPAgent struct { db.SSharableVirtualResourceBase - // LLMId 关联的 LLM 实例 ID - LLMId string `width:"128" charset:"ascii" nullable:"true" list:"user" create:"optional" update:"user"` + // LLMId 旧字段，新建不再写入 + LLMId string `width:"128" charset:"ascii" nullable:"true" list:"user"` + // AiproxyRoutingId 关联的 AI 网关路由规则 ID + AiproxyRoutingId string `width:"128" charset:"ascii" nullable:"true" list:"user" create:"required" update:"user" json:"aiproxy_routing_id"` + // AiproxyVirtualKeyId 关联的 AI 网关 API Key ID（密钥只存在 aiproxy，聊天时按 ID 读取） + AiproxyVirtualKeyId string `width:"128" charset:"ascii" nullable:"true" list:"user" create:"required" update:"user" json:"aiproxy_virtual_key_id"` - // LLMUrl 对应后端大模型的 base 请求地址 + // LLMUrl 对应 aiproxy OpenAI 兼容 base 请求地址 LLMUrl string `width:"512" charset:"utf8" nullable:"false" list:"user" create:"required" update:"user"` - // LLMDriver 对应使用的大模型驱动（llm_client），现在可以被设置为 ollama 或 openai + // LLMDriver 固定为 openai（走 AI 网关） LLMDriver string `width:"64" charset:"ascii" nullable:"false" list:"user" create:"required" update:"user"` - // Model 使用的模型名称 - Model string `width:"128" charset:"ascii" nullable:"false" list:"user" create:"required" update:"user"` - // ApiKey 即在 llm_driver 中需要用到的认证 - ApiKey string `width:"512" charset:"utf8" nullable:"true" create:"optional" update:"user"` + // Model 使用的模型名称（可为 aiproxy 扁平或层级 client model id） + Model string `width:"256" charset:"ascii" nullable:"false" list:"user" create:"required" update:"user"` + // ApiKey 旧字段，新建不再写入 + ApiKey string `width:"512" charset:"utf8" nullable:"true"` // McpServer 即 mcp 服务器的后端地址 McpServer string `width:"512" charset:"utf8" nullable:"false" list:"user" create:"optional" update:"user"` // DefaultAgent 是否为默认 Agent，全局仅允许一条为 true @@ -140,32 +144,11 @@ func (mcp *SMCPAgent) BeforeInsert() { if len(mcp.Id) == 0 { mcp.Id = db.DefaultUUIDGenerator() } - if len(mcp.ApiKey) > 0 { - sec, err := seclib.EncryptAESBase64(mcp.Id, mcp.ApiKey) - if err != nil { - log.Errorf("EncryptAESBase64 fail %s", err) - } else { - mcp.ApiKey = sec - } - } + mcp.ApiKey = "" + mcp.LLMId = "" mcp.SSharableVirtualResourceBase.BeforeInsert() } -func (mcp *SMCPAgent) BeforeUpdate() { - if len(mcp.ApiKey) > 0 { - // heuristic to check if it is plaintext - _, err := seclib.DescryptAESBase64(mcp.Id, mcp.ApiKey) - if err != nil { - sec, err := seclib.EncryptAESBase64(mcp.Id, mcp.ApiKey) - if err != nil { - log.Errorf("EncryptAESBase64 fail %s", err) - } else { - mcp.ApiKey = sec - } - } - } -} - func (mcp *SMCPAgent) PostCreate(ctx context.Context, userCred mcclient.TokenCredential, ownerId mcclient.IIdentityProvider, query jsonutils.JSONObject, data jsonutils.JSONObject) { mcp.SSharableVirtualResourceBase.PostCreate(ctx, userCred, ownerId, query, data) if mcp.DefaultAgent { @@ -182,18 +165,36 @@ func (mcp *SMCPAgent) PostUpdate(ctx context.Context, userCred mcclient.TokenCre log.Errorf("unsetOtherDefaultAgents after update: %v", err) } } + if strings.TrimSpace(mcp.ApiKey) != "" || strings.TrimSpace(mcp.LLMId) != "" { + if _, err := db.Update(mcp, func() error { + mcp.ApiKey = "" + mcp.LLMId = "" + return nil + }); err != nil { + log.Errorf("clear mcp agent llm_id/api_key: %v", err) + } + } } -func (mcp *SMCPAgent) GetApiKey() (string, error) { - if len(mcp.ApiKey) == 0 { - return "", nil +func (mcp *SMCPAgent) GetAiproxyVirtualKey(ctx context.Context) (string, error) { + id := strings.TrimSpace(mcp.AiproxyVirtualKeyId) + if id == "" { + return "", errors.Wrap(httperrors.ErrInvalidStatus, "mcp agent has no aiproxy_virtual_key_id; update the agent to bind an AI gateway API key") } - // try decrypt - key, err := seclib.DescryptAESBase64(mcp.Id, mcp.ApiKey) - if err == nil { - return key, nil + session := aiproxyAdminSession(ctx) + if session == nil { + return "", errors.Wrap(httperrors.ErrInvalidStatus, "aiproxy admin session is nil") } - return mcp.ApiKey, nil + resp, err := apmodules.AiVirtualKeys.Get(session, id, nil) + if err != nil { + return "", errors.Wrapf(err, "get ai_virtual_key %s", id) + } + key, _ := resp.GetString("virtual_key") + key = strings.TrimSpace(key) + if key == "" { + return "", errors.Wrapf(httperrors.ErrInvalidStatus, "ai_virtual_key %s has empty virtual_key", id) + } + return key, nil } func (man *SMCPAgentManager) CustomizeHandlerInfo(info *appsrv.SHandlerInfo) { @@ -228,53 +229,33 @@ func (man *SMCPAgentManager) ValidateCreateData(ctx context.Context, userCred mc return input, errors.Wrap(err, "validate SharableVirtualResourceCreateInput") } - // 如果提供了 llm_id，则通过 LLM 获取 llm_url 和 model - if len(input.LLMId) > 0 { - llm, err := FetchAccessibleLLM(ctx, userCred, input.LLMId) - if err != nil { - return input, errors.Wrapf(err, "fetch LLM by id %s", input.LLMId) - } - input.LLMId = llm.Id - llmUrl, err := llm.GetLLMAccessUrlInfo(ctx, userCred, query) - if err != nil { - return input, errors.Wrapf(err, "get LLM URL from LLM %s", input.LLMId) - } - input.LLMUrl = llmUrl.LoginUrl - - if len(input.Model) == 0 { - mdlInfos, err := llm.getProbedInstantModelsExt(ctx, userCred) - if err != nil { - return input, errors.Wrap(err, "get probed models from LLM instance") - } - if len(mdlInfos) == 0 { - return input, httperrors.NewBadRequestError("no available models found in LLM instance %s", input.LLMId) - } - var firstModel api.LLMInternalInstantMdlInfo - for _, mdlInfo := range mdlInfos { - firstModel = mdlInfo - break - } - input.Model = fmt.Sprintf("%s:%s", firstModel.Name, firstModel.Tag) - } + input.LLMId = "" + input.ApiKey = "" + input.AiProxyRoutingId = strings.TrimSpace(input.AiProxyRoutingId) + input.AiproxyVirtualKeyId = strings.TrimSpace(input.AiproxyVirtualKeyId) + if input.AiProxyRoutingId == "" { + return input, errors.Wrap(httperrors.ErrInputParameter, "aiproxy_routing_id is required") + } + if input.AiproxyVirtualKeyId == "" { + return input, errors.Wrap(httperrors.ErrInputParameter, "aiproxy_virtual_key_id is required") } - // 验证 llm_url 不为空 if len(input.LLMUrl) == 0 { - return input, errors.Wrap(httperrors.ErrInputParameter, "llm_url is required (or provide llm_id to auto-fetch)") + return input, errors.Wrap(httperrors.ErrInputParameter, "llm_url is required") } - // 验证 llm_driver 必须是 ollama 或 openai input.LLMDriver = strings.ToLower(strings.TrimSpace(input.LLMDriver)) - if !api.IsLLMClientType(input.LLMDriver) { - return input, errors.Wrapf(httperrors.ErrInputParameter, "llm_driver must be one of: %s, got: %s", api.LLM_CLIENT_TYPES.List(), input.LLMDriver) + if input.LLMDriver == "" { + input.LLMDriver = string(api.LLM_CLIENT_OPENAI) + } + if input.LLMDriver != string(api.LLM_CLIENT_OPENAI) { + return input, errors.Wrapf(httperrors.ErrInputParameter, "llm_driver must be %s", api.LLM_CLIENT_OPENAI) } - // 验证 model 不为空 if len(input.Model) == 0 { return input, errors.Wrap(httperrors.ErrInputParameter, "model is required") } - // 验证 mcp_server 不为空 if len(input.McpServer) == 0 { input.McpServer = options.Options.MCPServerURL } @@ -282,11 +263,6 @@ func (man *SMCPAgentManager) ValidateCreateData(ctx context.Context, userCred mc return input, httperrors.NewInputParameterError("%s", err.Error()) } - // 对于 openai 驱动，api_key 是必需的 - if input.LLMDriver == string(api.LLM_CLIENT_OPENAI) && len(input.ApiKey) == 0 { - return input, errors.Wrap(httperrors.ErrInputParameter, "api_key is required when llm_driver is openai") - } - input.Status = api.STATUS_READY return input, nil } @@ -298,43 +274,32 @@ func (man *SMCPAgentManager) ValidateUpdateData(ctx context.Context, userCred mc return input, errors.Wrap(err, "validate SharableVirtualResourceCreateInput") } - // 如果提供了 llm_id，则通过 LLM 获取 llm_url 和 model - if input.LLMId != nil && len(*input.LLMId) > 0 { - llm, err := FetchAccessibleLLM(ctx, userCred, *input.LLMId) - if err != nil { - return input, errors.Wrapf(err, "fetch LLM by id %s", *input.LLMId) - } - llmId := llm.Id - input.LLMId = &llmId - llmUrl, err := llm.GetLLMAccessUrlInfo(ctx, userCred, query) - if err != nil { - return input, errors.Wrapf(err, "get LLM URL from LLM %s", *input.LLMId) - } - input.LLMUrl = &llmUrl.LoginUrl + input.LLMId = nil + input.ApiKey = nil - if input.Model == nil || len(*input.Model) == 0 { - mdlInfos, err := llm.getProbedInstantModelsExt(ctx, userCred) - if err != nil { - return input, errors.Wrap(err, "get probed models from LLM instance") - } - if len(mdlInfos) == 0 { - return input, httperrors.NewBadRequestError("no available models found in LLM instance %s", *input.LLMId) - } - var firstModel api.LLMInternalInstantMdlInfo - for _, mdlInfo := range mdlInfos { - firstModel = mdlInfo - break - } - modelStr := fmt.Sprintf("%s:%s", firstModel.Name, firstModel.Tag) - input.Model = &modelStr + if input.AiProxyRoutingId != nil { + trimmed := strings.TrimSpace(*input.AiProxyRoutingId) + if trimmed == "" { + return input, errors.Wrap(httperrors.ErrInputParameter, "aiproxy_routing_id is required") } + input.AiProxyRoutingId = &trimmed + } + if input.AiproxyVirtualKeyId != nil { + trimmed := strings.TrimSpace(*input.AiproxyVirtualKeyId) + if trimmed == "" { + return input, errors.Wrap(httperrors.ErrInputParameter, "aiproxy_virtual_key_id is required") + } + input.AiproxyVirtualKeyId = &trimmed } - // 如果更新 llm_driver，验证其值 if input.LLMDriver != nil { *input.LLMDriver = strings.ToLower(strings.TrimSpace(*input.LLMDriver)) - if !api.IsLLMClientType(*input.LLMDriver) { - return input, errors.Wrapf(httperrors.ErrInputParameter, "llm_driver must be one of: %s, got: %s", api.LLM_CLIENT_TYPES.List(), *input.LLMDriver) + if *input.LLMDriver == "" { + openai := string(api.LLM_CLIENT_OPENAI) + input.LLMDriver = &openai + } + if *input.LLMDriver != string(api.LLM_CLIENT_OPENAI) { + return input, errors.Wrapf(httperrors.ErrInputParameter, "llm_driver must be %s", api.LLM_CLIENT_OPENAI) } } @@ -385,29 +350,11 @@ func (manager *SMCPAgentManager) FetchCustomizeColumns( agents := []SMCPAgent{} jsonutils.Update(&agents, objs) - llmIds := make([]string, 0) - for i := range agents { - if len(agents[i].LLMId) > 0 { - llmIds = append(llmIds, agents[i].LLMId) - } - } - - var llmIdNameMap map[string]string - if len(llmIds) > 0 { - var err error - llmIdNameMap, err = db.FetchIdNameMap2(GetLLMManager(), llmIds) - if err != nil { - log.Errorf("FetchIdNameMap2 for LLMs failed: %v", err) - } - } - for i := range rows { rows[i].SharableVirtualResourceDetails = vrows[i] if i < len(agents) { - rows[i].LLMId = agents[i].LLMId - if name, ok := llmIdNameMap[agents[i].LLMId]; ok { - rows[i].LLMName = name - } + rows[i].AiProxyRoutingId = agents[i].AiproxyRoutingId + rows[i].AiproxyVirtualKeyId = agents[i].AiproxyVirtualKeyId rows[i].DefaultAgent = agents[i].DefaultAgent } } @@ -577,6 +524,12 @@ func friendlyChatStreamError(err error) string { // process 处理用户请求（多轮工具调用，直到模型不再发 tool_calls 或达到上限） func (mcp *SMCPAgent) process(ctx context.Context, userCred mcclient.TokenCredential, req *api.LLMMCPAgentRequestInput, onStream func(string) error) (*api.MCPAgentResponse, error) { + if strings.TrimSpace(mcp.AiproxyRoutingId) == "" { + return nil, errors.Wrap(httperrors.ErrInvalidStatus, "mcp agent has no aiproxy_routing_id; update the agent to bind an AI gateway routing rule") + } + if strings.TrimSpace(mcp.AiproxyVirtualKeyId) == "" { + return nil, errors.Wrap(httperrors.ErrInvalidStatus, "mcp agent has no aiproxy_virtual_key_id; update the agent to bind an AI gateway API key") + } mcpServerUrl, err := mcp.GetMcpServerUrl(ctx, userCred) if err != nil { return nil, errors.Wrap(err, "GetMcpServerUrl") diff --git a/pkg/mcclient/options/llm/mcp_agent.go b/pkg/mcclient/options/llm/mcp_agent.go index 2e0d3cc412..f29533b1e9 100644 --- a/pkg/mcclient/options/llm/mcp_agent.go +++ b/pkg/mcclient/options/llm/mcp_agent.go @@ -33,13 +33,13 @@ func (o *MCPAgentShowOptions) Params() (jsonutils.JSONObject, error) { type MCPAgentCreateOptions struct { apis.SharableVirtualResourceCreateInput - LlmId string `help:"LLM 实例 ID，如果提供则自动获取 llm_url" json:"llm_id"` - LLM_URL string `help:"后端大模型的 base 请求地址" json:"llm_url"` - LLM_DRIVER string `help:"使用的大模型驱动，可以是 ollama 或 openai" json:"llm_driver" choices:"ollama|openai"` - MODEL string `help:"使用的模型名称" json:"model"` - API_KEY string `help:"访问大模型的密钥" json:"api_key"` - McpServer string `help:"mcp 服务器的后端地址" json:"mcp_server"` - DefaultAgent bool `help:"set as default MCP agent (only one can be true globally)" json:"default_agent"` + LLM_URL string `help:"AI 网关 OpenAI 兼容 base 请求地址" json:"llm_url"` + LLM_DRIVER string `help:"使用的大模型驱动，固定为 openai" json:"llm_driver" choices:"openai"` + MODEL string `help:"使用的模型名称" json:"model"` + McpServer string `help:"mcp 服务器的后端地址" json:"mcp_server"` + AiProxyRoutingId string `help:"关联的 AI 网关路由规则 ID" json:"aiproxy_routing_id"` + AiproxyVirtualKeyId string `help:"关联的 AI 网关 API Key ID" json:"aiproxy_virtual_key_id"` + DefaultAgent bool `help:"set as default MCP agent (only one can be true globally)" json:"default_agent"` } func (o *MCPAgentCreateOptions) Params() (jsonutils.JSONObject, error) { @@ -49,14 +49,14 @@ func (o *MCPAgentCreateOptions) Params() (jsonutils.JSONObject, error) { type MCPAgentUpdateOptions struct { apis.SharableVirtualResourceBaseUpdateInput - ID string - LlmId *string `help:"LLM 实例 ID，如果提供则自动获取 llm_url" json:"llm_id,omitempty"` - LlmUrl *string `help:"后端大模型的 base 请求地址" json:"llm_url,omitempty"` - LlmDriver *string `help:"使用的大模型驱动，可以是 ollama 或 openai" json:"llm_driver,omitempty" choices:"ollama|openai"` - Model *string `help:"使用的模型名称" json:"model,omitempty"` - ApiKey *string `help:"访问大模型的密钥" json:"api_key,omitempty"` - McpServer *string `help:"mcp 服务器的后端地址" json:"mcp_server,omitempty"` - DefaultAgent *bool `help:"set as default MCP agent (only one can be true globally)" json:"default_agent,omitempty"` + ID string + LlmUrl *string `help:"AI 网关 OpenAI 兼容 base 请求地址" json:"llm_url,omitempty"` + LlmDriver *string `help:"使用的大模型驱动，固定为 openai" json:"llm_driver,omitempty" choices:"openai"` + Model *string `help:"使用的模型名称" json:"model,omitempty"` + McpServer *string `help:"mcp 服务器的后端地址" json:"mcp_server,omitempty"` + AiProxyRoutingId *string `help:"关联的 AI 网关路由规则 ID" json:"aiproxy_routing_id,omitempty"` + AiproxyVirtualKeyId *string `help:"关联的 AI 网关 API Key ID" json:"aiproxy_virtual_key_id,omitempty"` + DefaultAgent *bool `help:"set as default MCP agent (only one can be true globally)" json:"default_agent,omitempty"` } func (o *MCPAgentUpdateOptions) GetId() string { @@ -66,9 +66,6 @@ func (o *MCPAgentUpdateOptions) GetId() string { func (o *MCPAgentUpdateOptions) Params() (jsonutils.JSONObject, error) { // 只包含非空字段 params := jsonutils.NewDict() - if o.LlmId != nil && len(*o.LlmId) > 0 { - params.Set("llm_id", jsonutils.NewString(*o.LlmId)) - } if o.LlmUrl != nil && len(*o.LlmUrl) > 0 { params.Set("llm_url", jsonutils.NewString(*o.LlmUrl)) } @@ -78,12 +75,15 @@ func (o *MCPAgentUpdateOptions) Params() (jsonutils.JSONObject, error) { if o.Model != nil && len(*o.Model) > 0 { params.Set("model", jsonutils.NewString(*o.Model)) } - if o.ApiKey != nil && len(*o.ApiKey) > 0 { - params.Set("api_key", jsonutils.NewString(*o.ApiKey)) - } if o.McpServer != nil && len(*o.McpServer) > 0 { params.Set("mcp_server", jsonutils.NewString(*o.McpServer)) } + if o.AiProxyRoutingId != nil && len(*o.AiProxyRoutingId) > 0 { + params.Set("aiproxy_routing_id", jsonutils.NewString(*o.AiProxyRoutingId)) + } + if o.AiproxyVirtualKeyId != nil && len(*o.AiproxyVirtualKeyId) > 0 { + params.Set("aiproxy_virtual_key_id", jsonutils.NewString(*o.AiproxyVirtualKeyId)) + } if o.DefaultAgent != nil { params.Set("default_agent", jsonutils.NewBool(*o.DefaultAgent)) }