mirror of
https://gitee.com/wanwujie/sub2api
synced 2026-04-27 18:04:47 +08:00
sub2api: add bedrock support
This commit is contained in:
@@ -207,14 +207,14 @@ func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account
|
||||
testModelID = claude.DefaultTestModel
|
||||
}
|
||||
|
||||
// For API Key accounts with model mapping, map the model
|
||||
// API Key 账号测试连接时也需要应用通配符模型映射。
|
||||
if account.Type == "apikey" {
|
||||
mapping := account.GetModelMapping()
|
||||
if len(mapping) > 0 {
|
||||
if mappedModel, exists := mapping[testModelID]; exists {
|
||||
testModelID = mappedModel
|
||||
}
|
||||
}
|
||||
testModelID = account.GetMappedModel(testModelID)
|
||||
}
|
||||
|
||||
// Bedrock accounts use a separate test path
|
||||
if account.IsBedrock() {
|
||||
return s.testBedrockAccountConnection(c, ctx, account, testModelID)
|
||||
}
|
||||
|
||||
// Determine authentication method and API URL
|
||||
@@ -312,6 +312,109 @@ func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account
|
||||
return s.processClaudeStream(c, resp.Body)
|
||||
}
|
||||
|
||||
// testBedrockAccountConnection tests a Bedrock (SigV4 or API Key) account using non-streaming invoke
|
||||
func (s *AccountTestService) testBedrockAccountConnection(c *gin.Context, ctx context.Context, account *Account, testModelID string) error {
|
||||
region := bedrockRuntimeRegion(account)
|
||||
resolvedModelID, ok := ResolveBedrockModelID(account, testModelID)
|
||||
if !ok {
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported Bedrock model: %s", testModelID))
|
||||
}
|
||||
testModelID = resolvedModelID
|
||||
|
||||
// Set SSE headers (test UI expects SSE)
|
||||
c.Writer.Header().Set("Content-Type", "text/event-stream")
|
||||
c.Writer.Header().Set("Cache-Control", "no-cache")
|
||||
c.Writer.Header().Set("Connection", "keep-alive")
|
||||
c.Writer.Header().Set("X-Accel-Buffering", "no")
|
||||
c.Writer.Flush()
|
||||
|
||||
// Create a minimal Bedrock-compatible payload (no stream, no cache_control)
|
||||
bedrockPayload := map[string]any{
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"messages": []map[string]any{
|
||||
{
|
||||
"role": "user",
|
||||
"content": []map[string]any{
|
||||
{
|
||||
"type": "text",
|
||||
"text": "hi",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"max_tokens": 256,
|
||||
"temperature": 1,
|
||||
}
|
||||
bedrockBody, _ := json.Marshal(bedrockPayload)
|
||||
|
||||
// Use non-streaming endpoint (response is standard Claude JSON)
|
||||
apiURL := BuildBedrockURL(region, testModelID, false)
|
||||
|
||||
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(bedrockBody))
|
||||
if err != nil {
|
||||
return s.sendErrorAndEnd(c, "Failed to create request")
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
// Sign or set auth based on account type
|
||||
if account.IsBedrockAPIKey() {
|
||||
apiKey := account.GetCredential("api_key")
|
||||
if apiKey == "" {
|
||||
return s.sendErrorAndEnd(c, "No API key available")
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
} else {
|
||||
signer, err := NewBedrockSignerFromAccount(account)
|
||||
if err != nil {
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to create Bedrock signer: %s", err.Error()))
|
||||
}
|
||||
if err := signer.SignRequest(ctx, req, bedrockBody); err != nil {
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to sign request: %s", err.Error()))
|
||||
}
|
||||
}
|
||||
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
proxyURL = account.Proxy.URL()
|
||||
}
|
||||
|
||||
resp, err := s.httpUpstream.DoWithTLS(req, proxyURL, account.ID, account.Concurrency, false)
|
||||
if err != nil {
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Request failed: %s", err.Error()))
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)))
|
||||
}
|
||||
|
||||
// Bedrock non-streaming response is standard Claude JSON, extract the text
|
||||
var result struct {
|
||||
Content []struct {
|
||||
Text string `json:"text"`
|
||||
} `json:"content"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &result); err != nil {
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to parse response: %s", err.Error()))
|
||||
}
|
||||
|
||||
text := ""
|
||||
if len(result.Content) > 0 {
|
||||
text = result.Content[0].Text
|
||||
}
|
||||
if text == "" {
|
||||
text = "(empty response)"
|
||||
}
|
||||
|
||||
s.sendEvent(c, TestEvent{Type: "content", Text: text})
|
||||
s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
|
||||
return nil
|
||||
}
|
||||
|
||||
// testOpenAIAccountConnection tests an OpenAI account's connection
|
||||
func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account *Account, modelID string) error {
|
||||
ctx := c.Request.Context()
|
||||
|
||||
Reference in New Issue
Block a user