Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions internal/aistudio/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,11 @@ func (c *Client) do(ctx context.Context, method string, accountID string, reques
func (c *Client) doProtected(ctx context.Context, request GenerateRequest, body []byte) (*RPCResponse, error) {
rpc := newRPCRequest("GenerateContent", request.AccountID, request.ID, body, true)
c.applyBenefitTier(rpc.Method, request.AccountID, rpc.Header)
if request.ImageRoute {
// 该头随 bootstrap 页面(gemini-flash-latest)采样,图像模型页面取值不同;
// 置空使其在头合并时被删除,官网实测缺省可通过(错值会被 Code 7 拒绝)
rpc.Header["X-Goog-Ext-519733851-Bin"] = nil
}
response, err := c.protected.DoProtected(ctx, request, rpc)
if err != nil {
return nil, fmt.Errorf("发送 AI Studio GenerateContent: %w", err)
Expand Down
43 changes: 40 additions & 3 deletions internal/aistudio/generate.go
Original file line number Diff line number Diff line change
Expand Up @@ -42,13 +42,19 @@ func EncodeGenerateContentRequest(request GenerateRequest, defaults GenerationDe
wire := make([]any, length)
wire[0] = wireModelName(request.Model)
wire[1] = contents
wire[2] = observedSafetySettings()
if !defaults.ImageRoute {
wire[2] = observedSafetySettings()
}
wire[3] = config
if request.System != "" {
wire[5] = encodeSystemInstruction(request.System)
}
if explicitTools {
switch {
case explicitTools:
wire[6] = tools
case defaults.OutputResolution:
// 官网对可设置分辨率的图像模型固定下发该槽位;缺失会被上游以 Code 7 拒绝
wire[6] = []any{[]any{nil, nil, nil, []any{nil, []any{}}}}
}
wire[10] = int64(1)
if runtime.Timezone != "" {
Expand Down Expand Up @@ -134,6 +140,13 @@ func encodeGenerationConfig(config GenerationConfig, defaults GenerationDefaults
return nil, err
}
imageConfig := encodeImageConfig(config.ImageConfig)
if defaults.OutputResolution && imageConfig == nil {
// 官网对可设置分辨率的图像模型固定带 imageConfig;缺失会被上游以 Code 7 拒绝
imageConfig = []any{nil, "1K"}
}
if !defaults.OutputResolution {
imageConfig = nil
}
speechConfig, err := encodeSpeechConfig(config.SpeechConfig)
if err != nil {
return nil, err
Expand Down Expand Up @@ -330,18 +343,41 @@ func encodeSpeechConfig(config *SpeechConfig) ([]any, error) {
}

func applyModelMediaDefaults(config GenerationConfig, model Model) GenerationConfig {
if model.Capabilities["image_route"] && imageModalityNeedsText(config.ResponseModalities) {
// 官网对图像模型总是下发 [IMAGE, TEXT];纯 IMAGE 会被上游以 Code 7 拒绝。
config.ResponseModalities = []ResponseModality{ResponseModalityImage, ResponseModalityText}
return config
}
if config.ResponseModalities != nil {
return config
}
switch {
case model.Capabilities["speech_route"], model.Capabilities["music_route"]:
config.ResponseModalities = []ResponseModality{ResponseModalityAudio}
case model.Capabilities["image_route"]:
config.ResponseModalities = []ResponseModality{ResponseModalityImage}
config.ResponseModalities = []ResponseModality{ResponseModalityImage, ResponseModalityText}
}
return config
}

// imageModalityNeedsText 判断图像模型请求是否缺 TEXT 模态
func imageModalityNeedsText(modalities []ResponseModality) bool {
if modalities == nil {
return true
}
hasImage := false
hasText := false
for _, modality := range modalities {
switch ResponseModality(strings.ToUpper(strings.TrimSpace(string(modality)))) {
case ResponseModalityImage:
hasImage = true
case ResponseModalityText:
hasText = true
}
}
return hasImage && !hasText
}

func applySpeechTranscript(contents []Content, model Model, config GenerationConfig) []Content {
if !model.Capabilities["speech_route"] || config.SpeechConfig == nil {
return contents
Expand Down Expand Up @@ -382,6 +418,7 @@ func (c *Client) Generate(ctx context.Context, request GenerateRequest) (<-chan
return nil, fmt.Errorf("%w: %v", ErrInvalidArgument, err)
}
request.Config = applyModelMediaDefaults(request.Config, entry.model)
request.ImageRoute = entry.defaults.ImageRoute
request.Contents = applySpeechTranscript(request.Contents, entry.model, request.Config)
runtime := RequestContext{}
if c.contextProvider != nil {
Expand Down
5 changes: 5 additions & 0 deletions internal/aistudio/models.go
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,10 @@ type GenerationDefaults struct {
ThinkingLevel bool
DefaultThinkingLevel int64
ThinkingLevels []int64
// ImageRoute 表示图像生成模型(官网 GenerateContent 走扩展约束)
ImageRoute bool
// OutputResolution 表示支持输出分辨率(官网会带 imageConfig 与扩展槽位)
OutputResolution bool
}

type modelEntry struct {
Expand Down Expand Up @@ -253,6 +257,7 @@ func decodeGenerationDefaults(row []json.RawMessage, path string, evidence json.
defaults := GenerationDefaults{
Thinking: capabilities["thinking"], ThinkingBudget: capabilities["thinking_budget"],
ThinkingLevel: capabilities["thinking_level"], DefaultThinkingLevel: 2,
ImageRoute: capabilities["image_route"], OutputResolution: capabilities["output_resolution"],
}
maxOutput, err := optionalIntField(row, 6, path, evidence)
if err != nil {
Expand Down
2 changes: 2 additions & 0 deletions internal/aistudio/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,8 @@ type GenerateRequest struct {
Config GenerationConfig `json:"config,omitempty"`
Tools Tools `json:"tools,omitempty"`
AccountID string `json:"account_id,omitempty"`
// ImageRoute 内部标记:图像生成模型(由模型目录能力推导)
ImageRoute bool `json:"-"`
}

// TokenCountRequest 表示计数请求
Expand Down