fix mypy check issues

This commit is contained in:
Henry Wang 2025-09-30 10:37:21 +08:00
parent 4e5db9476c
commit 33218606b8
3 changed files with 18 additions and 11 deletions

View file

@ -9,6 +9,7 @@ from litellm.types.llms.openai import (
ChatCompletionAssistantMessage,
ChatCompletionAssistantToolCall,
ChatCompletionRequest,
ChatCompletionSystemMessage,
ChatCompletionToolCallFunctionChunk,
ChatCompletionToolChoiceValues,
ChatCompletionToolMessage,
@ -33,7 +34,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
sent_first_chunk: bool = False
# State tracking for accumulating partial tool calls
gccumulated_tool_calls: Dict[str, Dict[str, Any]]
accumulated_tool_calls: Dict[str, Dict[str, Any]]
def __init__(self, completion_stream: Any):
self.sent_first_chunk = False
@ -373,7 +374,7 @@ class GoogleGenAIAdapter:
system_parts = system_instruction.get("parts", [])
if system_parts and "text" in system_parts[0]:
messages.append(
ChatCompletionUserMessage(
ChatCompletionSystemMessage(
role="system", content=system_parts[0]["text"]
)
)
@ -466,10 +467,7 @@ class GoogleGenAIAdapter:
Returns:
Dict in Google GenAI generate_content response format
"""
if isinstance(response, AdapterCompletionStreamWrapper):
return self.translate_streaming_completion_to_generate_content(
response, wrapper=response
)
# Extract the main response content
choice = response.choices[0] if response.choices else None

View file

@ -24,6 +24,7 @@ from functools import partial
from typing import (
TYPE_CHECKING,
Any,
AsyncIterator,
Callable,
Coroutine,
Dict,
@ -5141,12 +5142,18 @@ async def aadapter_completion(
async def aadapter_generate_content(
**kwargs,
) -> Union[ModelResponse, CustomStreamWrapper]:
) -> Union[Dict[str, Any], AsyncIterator[bytes]]:
from litellm.google_genai.adapters.handler import (
GenerateContentToCompletionHandler,
)
return await GenerateContentToCompletionHandler.async_generate_content_handler(**kwargs, _is_async=True)
coro = cast(
Coroutine[Any, Any, Union[Dict[str, Any], AsyncIterator[bytes]]],
GenerateContentToCompletionHandler.generate_content_handler(
**kwargs, _is_async=True
),
)
return await coro
def adapter_completion(

View file

@ -1,4 +1,4 @@
from fastapi import APIRouter, Depends, Request, Response
from fastapi import APIRouter, Depends, Request, Response, HTTPException
from fastapi.responses import StreamingResponse
from litellm.proxy._types import *
@ -30,9 +30,9 @@ async def google_generate_content(
data = await _read_request_body(request=request)
if "model" not in data:
data["model"] = model_name
data["stream"] = False
# call router
if llm_router is None:
raise HTTPException(status_code=500, detail="Router not initialized")
response = await llm_router.agenerate_content(**data)
return response
@ -61,6 +61,8 @@ async def google_stream_generate_content(
data["stream"] = True # enforce streaming for this endpoint
# call router
if llm_router is None:
raise HTTPException(status_code=500, detail="Router not initialized")
response = await llm_router.agenerate_content(**data)
# Check if response is an async iterator (streaming response)