Skip to content
Open
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
4 changes: 3 additions & 1 deletion Sources/AgentRunKit/LLM/GeminiClient.swift
Original file line number Diff line number Diff line change
Expand Up @@ -133,10 +133,12 @@ extension GeminiClient {
let level = effortToLevel(config.effort)
let budget = config.budgetTokens

// Gemini requires exactly one of thinkingBudget or thinkingLevel.
// If an explicit budget is set, prefer it; otherwise use the level.
return GeminiThinkingConfig(
includeThoughts: true,
thinkingBudget: budget,
thinkingLevel: level
thinkingLevel: budget == nil ? level : nil
)
}

Expand Down
11 changes: 11 additions & 0 deletions Sources/AgentRunKit/LLM/GeminiClientTypes.swift
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,17 @@ struct GeminiThinkingConfig: Encodable {
let includeThoughts: Bool
let thinkingBudget: Int?
let thinkingLevel: String?

private enum CodingKeys: String, CodingKey {
case includeThoughts, thinkingBudget, thinkingLevel
}

func encode(to encoder: any Encoder) throws {
var container = encoder.container(keyedBy: CodingKeys.self)
try container.encode(includeThoughts, forKey: .includeThoughts)
try container.encodeIfPresent(thinkingBudget, forKey: .thinkingBudget)
try container.encodeIfPresent(thinkingLevel, forKey: .thinkingLevel)
}
}

struct GeminiResponse: Decodable {
Expand Down
17 changes: 17 additions & 0 deletions Tests/AgentRunKitTests/GeminiClientTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,7 @@ struct GeminiRequestSerializationTests {
let thinking = genConfig?["thinkingConfig"] as? [String: Any]
#expect(thinking?["includeThoughts"] as? Bool == true)
#expect(thinking?["thinkingLevel"] as? String == "HIGH")
#expect(thinking?["thinkingBudget"] == nil)
}

@Test
Expand Down Expand Up @@ -167,6 +168,22 @@ struct GeminiRequestSerializationTests {
let thinking = genConfig?["thinkingConfig"] as? [String: Any]
#expect(thinking?["includeThoughts"] as? Bool == true)
#expect(thinking?["thinkingBudget"] as? Int == 10000)
#expect(thinking?["thinkingLevel"] == nil)
}

@Test
func effortWithBudgetSendsOnlyBudget() throws {
// Simulates what ProviderService produces: effort + budgetTokens.
let config = ReasoningConfig(effort: .medium, budgetTokens: 32000)
let client = makeClient(reasoningConfig: config)
let request = try client.buildRequest(messages: [.user("Hi")], tools: [])
let json = try encodeRequest(request)

let genConfig = json["generationConfig"] as? [String: Any]
let thinking = genConfig?["thinkingConfig"] as? [String: Any]
#expect(thinking?["includeThoughts"] as? Bool == true)
#expect(thinking?["thinkingBudget"] as? Int == 32000)
#expect(thinking?["thinkingLevel"] == nil)
}

@Test
Expand Down