scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

fix(code-assist): offload async oauth lookup and validate token

Browse source
This commit is contained in:
balazss 2026-03-17 22:37:55 -07:00
parent 9c3d36743a
commit 6714590c9c
4 changed files with 96 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/llms/gemini/fallback_handler.py
View file

@ -4,6 +4,8 @@ from litellm._logging import verbose_logger
from litellm.llms.gemini.common_utils import should_fallback_to_google_code_assist
from litellm.llms.google_code_assist.chat import GoogleCodeAssistChat
_google_code_assist_chat = GoogleCodeAssistChat()
async def run_gemini_acompletion_with_code_assist_fallback(
primary_call: Awaitable[Any],
@ -27,7 +29,7 @@ async def run_gemini_acompletion_with_code_assist_fallback(
"Gemini request failed with ACCESS_TOKEN_SCOPE_INSUFFICIENT. "
"Falling back to google_code_assist."
)
return await GoogleCodeAssistChat().acompletion(**fallback_kwargs)
return await _google_code_assist_chat.acompletion(**fallback_kwargs)
def run_gemini_completion_with_code_assist_fallback(
@ -52,4 +54,4 @@ def run_gemini_completion_with_code_assist_fallback(
"Gemini request failed with ACCESS_TOKEN_SCOPE_INSUFFICIENT. "
"Falling back to google_code_assist."
)
return GoogleCodeAssistChat().completion(**fallback_kwargs)
return _google_code_assist_chat.completion(**fallback_kwargs)

13
litellm/llms/google_code_assist/chat.py
View file

@ -1,3 +1,4 @@
import asyncio
import httpx
from typing import Any, Optional
@ -43,6 +44,11 @@ class GoogleCodeAssistChat:
token = gemini_auth_data.get("token")
initial_project_id = gemini_auth_data.get("project_id")
if not token:
raise GoogleCodeAssistError(
status_code=401,
message="Missing Gemini OAuth token value. Re-run 'litellm-proxy gemini login'.",
)
client = _get_httpx_client()
@ -102,7 +108,7 @@ class GoogleCodeAssistChat:
try:
from litellm.llms.gemini.common_utils import get_gemini_oauth_token
gemini_auth_data = get_gemini_oauth_token()
gemini_auth_data = await asyncio.to_thread(get_gemini_oauth_token)
if not gemini_auth_data:
raise GoogleCodeAssistError(
status_code=401,
@ -111,6 +117,11 @@ class GoogleCodeAssistChat:
token = gemini_auth_data.get("token")
initial_project_id = gemini_auth_data.get("project_id")
if not token:
raise GoogleCodeAssistError(
status_code=401,
message="Missing Gemini OAuth token value. Re-run 'litellm-proxy gemini login'.",
)
async_handler = AsyncHTTPHandler()
try:

7
litellm/main.py
View file

@ -161,6 +161,8 @@ from litellm.utils import (
validate_openai_optional_params,
)
_google_code_assist_chat = GoogleCodeAssistChat()
from ._logging import verbose_logger
from .caching.caching import disable_cache, enable_cache, update_cache
from .litellm_core_utils.core_helpers import safe_deep_copy
@ -3690,9 +3692,8 @@ def completion( # type: ignore # noqa: PLR0915
response = model_response
elif custom_llm_provider == "google_code_assist":
google_code_assist_chat = GoogleCodeAssistChat()
if acompletion is True:
response = google_code_assist_chat.acompletion(
response = _google_code_assist_chat.acompletion(
model=model,
messages=messages,
model_response=model_response,
@ -3703,7 +3704,7 @@ def completion( # type: ignore # noqa: PLR0915
logger_fn=logger_fn,
)
else:
response = google_code_assist_chat.completion(
response = _google_code_assist_chat.completion(
model=model,
messages=messages,
model_response=model_response,

76
tests/test_litellm/llms/google_code_assist/test_google_code_assist.py
View file

@ -70,6 +70,22 @@ class TestGoogleCodeAssist:
assert response.choices[0].message.content == "Hello world"
assert response.usage.total_tokens == 7
@patch("litellm.llms.gemini.common_utils.get_gemini_oauth_token")
def test_completion_raises_when_token_missing(self, mock_get_token):
mock_get_token.return_value = {"token": None}
handler = GoogleCodeAssistChat()
with pytest.raises(Exception, match="Missing Gemini OAuth token value"):
handler.completion(
model="google_code_assist/gemini-1.5-flash",
messages=[{"role": "user", "content": "hi"}],
model_response=ModelResponse(),
print_verbose=False,
logging_obj=MagicMock(),
optional_params={},
litellm_params={},
)
@pytest.mark.asyncio
@patch("litellm.llms.google_code_assist.chat.AsyncHTTPHandler.post")
@patch(
@ -132,3 +148,63 @@ class TestGoogleCodeAssist:
assert response.choices[0].message.content == "Async success"
mock_async_close.assert_awaited_once()
@pytest.mark.asyncio
@patch(
"litellm.llms.google_code_assist.chat.asyncio.to_thread", new_callable=AsyncMock
)
async def test_acompletion_uses_thread_for_oauth_lookup(self, mock_to_thread):
mock_to_thread.return_value = {"token": "test-token"}
with (
patch(
"litellm.llms.google_code_assist.chat.GoogleCodeAssistChat._ahandle_handshake",
new_callable=AsyncMock,
return_value="final-project",
),
patch(
"litellm.llms.google_code_assist.chat.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_async_post,
patch(
"litellm.llms.google_code_assist.chat.AsyncHTTPHandler.close",
new_callable=AsyncMock,
),
):
completion_data = {
"response": {
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "Async success"}],
},
"finishReason": "STOP",
}
],
"usageMetadata": {
"promptTokenCount": 5,
"candidatesTokenCount": 2,
"totalTokenCount": 7,
},
}
}
mock_completion_resp = httpx.Response(
status_code=200,
content=json.dumps(completion_data).encode(),
request=httpx.Request("POST", "https://completion"),
)
mock_async_post.return_value = mock_completion_resp
handler = GoogleCodeAssistChat()
await handler.acompletion(
model="google_code_assist/gemini-1.5-flash",
messages=[{"role": "user", "content": "hi"}],
model_response=ModelResponse(),
print_verbose=False,
logging_obj=MagicMock(),
optional_params={},
litellm_params={},
)
mock_to_thread.assert_awaited_once()