diff --git a/src/khoj/processor/conversation/google/gemini_chat.py b/src/khoj/processor/conversation/google/gemini_chat.py index aed144c0..e6923fad 100644 --- a/src/khoj/processor/conversation/google/gemini_chat.py +++ b/src/khoj/processor/conversation/google/gemini_chat.py @@ -41,7 +41,7 @@ def gemini_send_message_to_model( if tools: model_kwargs["tools"] = tools # Monitor for flakiness in 1.5+ models. This would cause unwanted behavior and terminate response early in 1.5 models. - elif response_type == "json_object" and not model.startswith("gemini-1.5"): + elif response_type == "json_object": model_kwargs["response_mime_type"] = "application/json" if response_schema: model_kwargs["response_schema"] = response_schema diff --git a/src/khoj/processor/conversation/google/utils.py b/src/khoj/processor/conversation/google/utils.py index f2b0dc84..fbfebc13 100644 --- a/src/khoj/processor/conversation/google/utils.py +++ b/src/khoj/processor/conversation/google/utils.py @@ -4,7 +4,7 @@ import os import random from copy import deepcopy from time import perf_counter -from typing import AsyncGenerator, AsyncIterator, Dict, List +from typing import Any, AsyncGenerator, AsyncIterator, Dict, List import httpx from google import genai @@ -456,13 +456,31 @@ def is_reasoning_model(model_name: str) -> bool: def to_gemini_tools(tools: List[ToolDefinition]) -> List[gtypes.ToolDict] | None: "Transform tool definitions from standard format to Gemini format." + + def clean_schema(schema: Dict[str, Any]) -> Dict[str, Any]: + """Remove additionalProperties from schema as Gemini doesn't accept it.""" + if not isinstance(schema, dict): + return schema + + cleaned: Dict[str, Any] = {} + for key, value in schema.items(): + if key == "additionalProperties": + continue + if isinstance(value, dict): + cleaned[key] = clean_schema(value) + elif isinstance(value, list): + cleaned[key] = [clean_schema(item) if isinstance(item, dict) else item for item in value] + else: + cleaned[key] = value + return cleaned + gemini_tools = [ gtypes.ToolDict( function_declarations=[ gtypes.FunctionDeclarationDict( name=tool.name, description=tool.description, - parameters=tool.schema, + parameters=clean_schema(tool.schema), ) for tool in tools ]