Update thinking budget configs

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "uv_build"
 [project]
 name = "draive"
 description = "Framework designed to simplify and accelerate the development of LLM-based applications."
-version = "0.106.3"
+version = "0.106.4"
 readme = "README.md"
 maintainers = [
   { name = "Kacper Kaliński", email = "kacper.kalinski@miquido.com" },

src/draive/gemini/config.py

@@ -40,6 +40,8 @@ class GeminiConfig(Configuration):
     speech_language_code: str | Missing = MISSING
     media_resolution: Literal["low", "medium", "high"] | Missing = MISSING
     thinking_budget: int | Missing = MISSING
+    thinking_level: Literal["minimal", "low", "medium", "high"] | Missing = MISSING
+    context_window_compression: bool | Missing = MISSING
     safety: GeminiSafetyConfig | Missing = MISSING
 
 

src/draive/gemini/live.py

@@ -23,6 +23,7 @@
     Modality,
     Part,
     PartDict,
+    ThinkingLevel,
     Transcription,
 )
 from haiway import MISSING, Meta, as_dict, ctx
@@ -416,6 +417,33 @@ async def close_session(
         )
 
 
+def _resolve_thinking_level(level: str) -> ThinkingLevel:
+    match level:
+        case "minimal":
+            return ThinkingLevel.MINIMAL
+        case "low":
+            return ThinkingLevel.LOW
+        case "medium":
+            return ThinkingLevel.MEDIUM
+        case "high":
+            return ThinkingLevel.HIGH
+        case _:
+            raise ValueError(f"Unsupported thinking level: {level}")
+
+
+def _resolve_media_resolution(config: GeminiConfig) -> MediaResolution | None:
+    if config.media_resolution is MISSING:
+        return None
+    elif config.media_resolution == "low":
+        return MediaResolution.MEDIA_RESOLUTION_LOW
+    elif config.media_resolution == "medium":
+        return MediaResolution.MEDIA_RESOLUTION_MEDIUM
+    elif config.media_resolution == "high":
+        return MediaResolution.MEDIA_RESOLUTION_HIGH
+    else:
+        raise ValueError(f"Unsupported media resolution: {config.media_resolution}")
+
+
 def _live_connect_config(
     *,
     instructions: ModelInstructions,
@@ -431,12 +459,12 @@ def _live_connect_config(
         "history_config": {
             "initial_history_in_client_content": True,
         },
-        "context_window_compression": {
-            "sliding_window": {},
-        },
         "seed": unwrap_missing(config.seed),
     }
 
+    if config.context_window_compression is True:
+        live_config["context_window_compression"] = {"sliding_window": {}}
+
     if instructions:
         live_config["system_instruction"] = instructions
 
@@ -461,25 +489,17 @@ def _live_connect_config(
             }
         ]
 
-    if config.media_resolution is MISSING:
-        pass  # skip missing
-
-    elif config.media_resolution == "low":
-        live_config["media_resolution"] = MediaResolution.MEDIA_RESOLUTION_LOW
-
-    elif config.media_resolution == "medium":
-        live_config["media_resolution"] = MediaResolution.MEDIA_RESOLUTION_MEDIUM
-
-    elif config.media_resolution == "high":
-        live_config["media_resolution"] = MediaResolution.MEDIA_RESOLUTION_HIGH
-
-    else:
-        raise ValueError(f"Unsupported media resolution: {config.media_resolution}")
+    if resolution := _resolve_media_resolution(config):
+        live_config["media_resolution"] = resolution
 
     if speech := speech_config(config):
         live_config["speech_config"] = speech
 
-    if config.thinking_budget is not MISSING:
+    if config.thinking_level is not MISSING:
+        live_config["thinking_config"] = {
+            "thinking_level": _resolve_thinking_level(cast(str, config.thinking_level)),
+        }
+    elif config.thinking_budget is not MISSING:
         live_config["thinking_config"] = {
             "include_thoughts": True,
             "thinking_budget": cast(int, config.thinking_budget),