126 lines
3.2 KiB
Go
126 lines
3.2 KiB
Go
package channel
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
app_errors "gpt-load/internal/errors"
|
|
"gpt-load/internal/models"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func init() {
|
|
Register("gemini", newGeminiChannel)
|
|
}
|
|
|
|
type GeminiChannel struct {
|
|
*BaseChannel
|
|
}
|
|
|
|
func newGeminiChannel(f *Factory, group *models.Group) (ChannelProxy, error) {
|
|
base, err := f.newBaseChannel("gemini", group)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &GeminiChannel{
|
|
BaseChannel: base,
|
|
}, nil
|
|
}
|
|
|
|
// ModifyRequest adds the API key as a query parameter for Gemini requests.
|
|
func (ch *GeminiChannel) ModifyRequest(req *http.Request, apiKey *models.APIKey, group *models.Group) {
|
|
q := req.URL.Query()
|
|
q.Set("key", apiKey.KeyValue)
|
|
req.URL.RawQuery = q.Encode()
|
|
}
|
|
|
|
|
|
// IsStreamRequest checks if the request is for a streaming response.
|
|
func (ch *GeminiChannel) IsStreamRequest(c *gin.Context, bodyBytes []byte) bool {
|
|
path := c.Request.URL.Path
|
|
if strings.HasSuffix(path, ":streamGenerateContent") {
|
|
return true
|
|
}
|
|
|
|
// Also check for standard streaming indicators as a fallback.
|
|
if strings.Contains(c.GetHeader("Accept"), "text/event-stream") {
|
|
return true
|
|
}
|
|
if c.Query("stream") == "true" {
|
|
return true
|
|
}
|
|
|
|
return false
|
|
}
|
|
|
|
// ExtractKey extracts the API key from the X-Goog-Api-Key header or the "key" query parameter.
|
|
func (ch *GeminiChannel) ExtractKey(c *gin.Context) string {
|
|
// 1. Check X-Goog-Api-Key header
|
|
if key := c.GetHeader("X-Goog-Api-Key"); key != "" {
|
|
return key
|
|
}
|
|
|
|
// 2. Check "key" query parameter
|
|
if key := c.Query("key"); key != "" {
|
|
return key
|
|
}
|
|
|
|
return ""
|
|
}
|
|
|
|
// ValidateKey checks if the given API key is valid by making a generateContent request.
|
|
func (ch *GeminiChannel) ValidateKey(ctx context.Context, key string) (bool, error) {
|
|
upstreamURL := ch.getUpstreamURL()
|
|
if upstreamURL == nil {
|
|
return false, fmt.Errorf("no upstream URL configured for channel %s", ch.Name)
|
|
}
|
|
|
|
reqURL := fmt.Sprintf("%s/v1beta/models/%s:generateContent?key=%s", upstreamURL.String(), ch.TestModel, key)
|
|
|
|
payload := gin.H{
|
|
"contents": []gin.H{
|
|
{"parts": []gin.H{
|
|
{"text": "hi"},
|
|
}},
|
|
},
|
|
}
|
|
body, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return false, fmt.Errorf("failed to marshal validation payload: %w", err)
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, "POST", reqURL, bytes.NewBuffer(body))
|
|
if err != nil {
|
|
return false, fmt.Errorf("failed to create validation request: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
resp, err := ch.HTTPClient.Do(req)
|
|
if err != nil {
|
|
return false, fmt.Errorf("failed to send validation request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
// A 200 OK status code indicates the key is valid.
|
|
if resp.StatusCode == http.StatusOK {
|
|
return true, nil
|
|
}
|
|
|
|
// For non-200 responses, parse the body to provide a more specific error reason.
|
|
errorBody, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return false, fmt.Errorf("key is invalid (status %d), but failed to read error body: %w", resp.StatusCode, err)
|
|
}
|
|
|
|
// Use the new parser to extract a clean error message.
|
|
parsedError := app_errors.ParseUpstreamError(errorBody)
|
|
|
|
return false, fmt.Errorf("[status %d] %s", resp.StatusCode, parsedError)
|
|
}
|