mirror of
https://github.com/EvolutionAPI/adk-python.git
synced 2025-07-16 04:02:55 -06:00
chore: Add the name of the agent as a billing label. This will help consumers slice their billing reports by Agent.
Fixes https://github.com/google/adk-python/issues/452 PiperOrigin-RevId: 765402643
This commit is contained in:
parent
63822c3fa8
commit
e06e6753ad
@ -23,6 +23,7 @@ from typing import cast
|
|||||||
from typing import Optional
|
from typing import Optional
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from google.genai import types
|
||||||
from websockets.exceptions import ConnectionClosedOK
|
from websockets.exceptions import ConnectionClosedOK
|
||||||
|
|
||||||
from . import functions
|
from . import functions
|
||||||
@ -50,6 +51,8 @@ if TYPE_CHECKING:
|
|||||||
|
|
||||||
logger = logging.getLogger('google_adk.' + __name__)
|
logger = logging.getLogger('google_adk.' + __name__)
|
||||||
|
|
||||||
|
_ADK_AGENT_NAME_LABEL_KEY = 'adk_agent_name'
|
||||||
|
|
||||||
|
|
||||||
class BaseLlmFlow(ABC):
|
class BaseLlmFlow(ABC):
|
||||||
"""A basic flow that calls the LLM in a loop until a final response is generated.
|
"""A basic flow that calls the LLM in a loop until a final response is generated.
|
||||||
@ -499,6 +502,16 @@ class BaseLlmFlow(ABC):
|
|||||||
yield response
|
yield response
|
||||||
return
|
return
|
||||||
|
|
||||||
|
llm_request.config = llm_request.config or types.GenerateContentConfig()
|
||||||
|
llm_request.config.labels = llm_request.config.labels or {}
|
||||||
|
|
||||||
|
# Add agent name as a label to the llm_request. This will help with slicing
|
||||||
|
# the billing reports on a per-agent basis.
|
||||||
|
if _ADK_AGENT_NAME_LABEL_KEY not in llm_request.config.labels:
|
||||||
|
llm_request.config.labels[_ADK_AGENT_NAME_LABEL_KEY] = (
|
||||||
|
invocation_context.agent.name
|
||||||
|
)
|
||||||
|
|
||||||
# Calls the LLM.
|
# Calls the LLM.
|
||||||
llm = self.__get_llm(invocation_context)
|
llm = self.__get_llm(invocation_context)
|
||||||
with tracer.start_as_current_span('call_llm'):
|
with tracer.start_as_current_span('call_llm'):
|
||||||
|
@ -44,3 +44,4 @@ def test_output_schema():
|
|||||||
assert len(mockModel.requests) == 1
|
assert len(mockModel.requests) == 1
|
||||||
assert mockModel.requests[0].config.response_schema == CustomOutput
|
assert mockModel.requests[0].config.response_schema == CustomOutput
|
||||||
assert mockModel.requests[0].config.response_mime_type == 'application/json'
|
assert mockModel.requests[0].config.response_mime_type == 'application/json'
|
||||||
|
assert mockModel.requests[0].config.labels == {'adk_agent_name': 'root_agent'}
|
||||||
|
Loading…
Reference in New Issue
Block a user