Skip to content

Add monkey patching for ADK's litellm import #1660

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
saglave wants to merge 1 commit into
base: develop
Choose a base branch
from
Draft

Add monkey patching for ADK's litellm import #1660

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion packages/nvidia_nat_adk/src/nat/plugins/adk/callback_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ def __init__(self):
# Original references to Google ADK Tool and LLM methods (for uninstrumenting if needed)
self._original_tool_call = None
self._original_llm_call = None
self._original_adk_llm_call = None
self._instrumented = False

def instrument(self) -> None:
Expand All @@ -77,6 +78,7 @@ def instrument(self) -> None:
logger.exception("litellm import failed; skipping instrumentation")
return
try:
import google.adk.models.lite_llm as adk_lite_llm
from google.adk.tools.function_tool import FunctionTool
except Exception as _e:
logger.exception("ADK import failed; skipping instrumentation")
Expand All @@ -85,9 +87,12 @@ def instrument(self) -> None:
# Save the originals
self._original_tool_call = FunctionTool.run_async
self._original_llm_call = litellm.acompletion
self._original_adk_llm_call = adk_lite_llm.acompletion

wrapped_llm = self._llm_call_monkey_patch()
FunctionTool.run_async = self._tool_use_monkey_patch()
litellm.acompletion = self._llm_call_monkey_patch()
litellm.acompletion = wrapped_llm
adk_lite_llm.acompletion = wrapped_llm

logger.debug("ADKProfilerHandler instrumentation applied successfully.")
self._instrumented = True
Expand All @@ -97,6 +102,7 @@ def uninstrument(self) -> None:
Add an explicit unpatch to avoid side-effects across tests/process lifetime.
"""
try:
import google.adk.models.lite_llm as adk_lite_llm
import litellm
from google.adk.tools.function_tool import FunctionTool
if self._original_tool_call is not None:
Expand All @@ -107,6 +113,10 @@ def uninstrument(self) -> None:
litellm.acompletion = self._original_llm_call
self._original_llm_call = None

if self._original_adk_llm_call is not None:
adk_lite_llm.acompletion = self._original_adk_llm_call
self._original_adk_llm_call = None

self._instrumented = False
self.last_call_ts = 0.0
logger.debug("ADKProfilerHandler uninstrumented successfully.")
Expand Down