兼容pydantic v1 和 v2

Author: su-kaka

src/antigravity_api.py

@@ -21,7 +21,7 @@
 
 from .credential_manager import CredentialManager
 from .httpx_client import create_streaming_client_with_kwargs, http_client
-from .models import Model
+from .models import Model, model_to_dict
 from .utils import ANTIGRAVITY_USER_AGENT, parse_quota_reset_timestamp
 
 async def _check_should_auto_ban(status_code: int) -> bool:
@@ -453,7 +453,7 @@ async def fetch_available_models(
                             created=current_timestamp,
                             owned_by='google'
                         )
-                        model_list.append(model.model_dump())
+                        model_list.append(model_to_dict(model))
 
                 # 添加额外的 claude-opus-4-5 模型
                 claude_opus_model = Model(
@@ -462,7 +462,7 @@ async def fetch_available_models(
                     created=current_timestamp,
                     owned_by='google'
                 )
-                model_list.append(claude_opus_model.model_dump())
+                model_list.append(model_to_dict(claude_opus_model))
 
                 log.info(f"[ANTIGRAVITY] Fetched {len(model_list)} available models")
                 return model_list

src/antigravity_router.py

@@ -27,6 +27,7 @@
     GeminiGenerationConfig,
     Model,
     ModelList,
+    model_to_dict,
     OpenAIChatCompletionChoice,
     OpenAIChatCompletionResponse,
     OpenAIChatMessage,
@@ -265,10 +266,8 @@ def clean_parameters(obj):
                 func_params = function.get("parameters", {})
 
                 # 转换为字典（如果是 Pydantic 模型）
-                if hasattr(func_params, "dict"):
-                    func_params = func_params.dict()
-                elif hasattr(func_params, "model_dump"):
-                    func_params = func_params.model_dump()
+                if hasattr(func_params, "dict") or hasattr(func_params, "model_dump"):
+                    func_params = model_to_dict(func_params)
 
                 # 清理参数
                 cleaned_params = clean_parameters(func_params)
@@ -355,7 +354,7 @@ def convert_to_openai_tool_call(function_call: Dict[str, Any]) -> Dict[str, Any]
             arguments=json.dumps(function_call.get("args", {}))
         )
     )
-    return tool_call.model_dump()
+    return model_to_dict(tool_call)
 
 
 async def convert_antigravity_stream_to_openai(
@@ -652,7 +651,7 @@ def convert_antigravity_response_to_openai(
         usage=usage
     )
 
-    return response.model_dump()
+    return model_to_dict(response)
 
 
 def convert_antigravity_response_to_gemini(

src/models.py

@@ -3,6 +3,21 @@
 from pydantic import BaseModel, Field
 
 
+# Pydantic v1/v2 兼容性辅助函数
+def model_to_dict(model: BaseModel) -> Dict[str, Any]:
+    """
+    兼容 Pydantic v1 和 v2 的模型转字典方法
+    - v1: model.dict()
+    - v2: model.model_dump()
+    """
+    if hasattr(model, 'model_dump'):
+        # Pydantic v2
+        return model.model_dump()
+    else:
+        # Pydantic v1
+        return model.dict()
+
+
 # Common Models
 class Model(BaseModel):
     id: str

src/openai_transfer.py

@@ -21,8 +21,9 @@
     should_include_thoughts,
 )
 from log import log
+from pydantic import BaseModel
 
-from .models import ChatCompletionRequest
+from .models import ChatCompletionRequest, model_to_dict
 
 
 async def openai_request_to_gemini_payload(
@@ -778,10 +779,8 @@ def convert_openai_tools_to_gemini(openai_tools: List) -> List[Dict[str, Any]]:
 
     for tool in openai_tools:
         # 处理 Pydantic 模型
-        if hasattr(tool, "model_dump"):
-            tool_dict = tool.model_dump()
-        elif hasattr(tool, "dict"):
-            tool_dict = tool.dict()
+        if hasattr(tool, "model_dump") or hasattr(tool, "dict"):
+            tool_dict = model_to_dict(tool)
         else:
             tool_dict = tool