8000 ADK changes · ognis1205/adk-python@ca99327 · GitHub
[go: up one dir, main page]
Skip to content

Commit ca99327

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
ADK changes
PiperOrigin-RevId: 750763037
1 parent a49d339 commit ca99327

File tree

6 files changed

+435
-105
lines changed

6 files changed

+435
-105
lines changed

‎src/google/adk/tools/application_integration_tool/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,9 @@
1313
# limitations under the License.
1414

1515
from .application_integration_toolset import ApplicationIntegrationToolset
16+
from .integration_connector_tool import IntegrationConnectorTool
1617

1718
__all__ = [
1819
'ApplicationIntegrationToolset',
20+
'IntegrationConnectorTool',
1921
]

‎src/google/adk/tools/application_integration_tool/application_integration_toolset.py

Lines changed: 48 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -12,21 +12,21 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
from typing import Dict
16-
from typing import List
17-
from typing import Optional
15+
from typing import Dict, List, Optional
1816

1917
from fastapi.openapi.models import HTTPBearer
20-
from google.adk.tools.application_integration_tool.clients.connections_client import ConnectionsClient
21-
from google.adk.tools.application_integration_tool.clients.integration_client import IntegrationClient
22-
from google.adk.tools.openapi_tool.auth.auth_helpers import service_account_scheme_credential
23-
from google.adk.tools.openapi_tool.openapi_spec_parser.openapi_toolset import OpenAPIToolset
24-
from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
2518

2619
from ...auth.auth_credential import AuthCredential
2720
from ...auth.auth_credential import AuthCredentialTypes
2821
from ...auth.auth_credential import ServiceAccount
2922
from ...auth.auth_credential import ServiceAccountCredential
23+
from ..openapi_tool.auth.auth_helpers import service_account_scheme_credential
24+
from ..openapi_tool.openapi_spec_parser.openapi_spec_parser import OpenApiSpecParser
25+
from ..openapi_tool.openapi_spec_parser.openapi_toolset import OpenAPIToolset
26+
from ..openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
27+
from .clients.connections_client import ConnectionsClient
28+
from .clients.integration_client import IntegrationClient
29+
from .integration_connector_tool import IntegrationConnectorTool
3030

3131

3232
# TODO(cheliu): Apply a common toolset interface
@@ -168,23 +168,14 @@ def __init__(
168168
actions,
169169
service_account_json,
170170
)
171+
connection_details = {}
171172
if integration and trigger:
172173
spec = integration_client.get_openapi_spec_for_integration()
173174
elif connection and (entity_operations or actions):
174175
connections_client = ConnectionsClient(
175176
project, location, connection, service_account_json
176177
)
177178
connection_details = connections_client.get_connection_details()
178-
tool_instructions += (
179-
"ALWAYS use serviceName = "
180-
+ connection 628C _details["serviceName"]
181-
+ ", host = "
182-
+ connection_details["host"]
183-
+ " and the connection name = "
184-
+ f"projects/{project}/locations/{location}/connections/{connection} when"
185-
" using this tool"
186-
+ ". DONOT ask the user for these values as you already have those."
187-
)
188179
spec = integration_client.get_openapi_spec_for_connection(
189180
tool_name,
190181
tool_instructions,
@@ -194,9 +185,9 @@ def __init__(
194185
"Either (integration and trigger) or (connection and"
195186
" (entity_operations or actions)) should be provided."
196187
)
197-
self._parse_spec_to_tools(spec)
188+
self._parse_spec_to_tools(spec, connection_details)
198189

199-
def _parse_spec_to_tools(self, spec_dict):
190+
def _parse_spec_to_tools(self, spec_dict, connection_details):
200191
"""Parses the spec dict to a list of RestApiTool."""
201192
if self.service_account_json:
202193
sa_credential = ServiceAccountCredential.model_validate_json(
@@ -218,12 +209,43 @@ def _parse_spec_to_tools(self, spec_dict):
218209
),
219210
)
220211
auth_scheme = HTTPBearer(bearerFormat="JWT")
221-
tools = OpenAPIToolset(
222-
spec_dict=spec_dict,
223-
auth_credential=auth_credential,
224-
auth_scheme=auth_scheme,
225-
).get_tools()
226-
for tool in tools:
212+
213+
if self.integration and self.trigger:
214+
tools = OpenAPIToolset(
215+
spec_dict=spec_dict,
216+
auth_credential=auth_credential,
217+
auth_scheme=auth_scheme,
218+
).get_tools()
219+
for tool in tools:
220+
self.generated_tools[tool.name] = tool
221+
return
222+
223+
operations = OpenApiSpecParser().parse(spec_dict)
224+
225+
for open_api_operation in operations:
226+
operation = getattr(open_api_operation.operation, "x-operation")
227+
entity = None
228+
action = None
229+
if hasattr(open_api_operation.operation, "x-entity"):
230+
entity = getattr(open_api_operation.operation, "x-entity")
231+
elif hasattr(open_api_operation.operation, "x-action"):
232+
action = getattr(open_api_operation.operation, "x-action")
233+
rest_api_tool = RestApiTool.from_parsed_operation(open_api_operation)
234+
if auth_scheme:
235+
rest_api_tool.configure_auth_scheme(auth_scheme)
236+
if auth_credential:
237+
rest_api_tool.configure_auth_credential(auth_credential)
238+
tool = IntegrationConnectorTool(
239+
name=rest_api_tool.name,
240+
description=rest_api_tool.description,
241+
connection_name=connection_details["name"],
242+
connection_host=connection_details["host"],
243+
connection_service_name=connection_details["serviceName"],
244+
entity=entity,
245+
action=action,
246+
operation=operation,
247+
rest_api_tool=rest_api_tool,
248+
)
227249
self.generated_tools[tool.name] = tool
228250

229251
def get_tools(self) -> List[RestApiTool]:

‎src/google/adk/tools/application_integration_tool/clients/connections_client.py

Lines changed: 26 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -68,12 +68,14 @@ def get_connection_details(self) -> Dict[str, Any]:
6868
response = self._execute_api_call(url)
6969

7070
connection_data = response.json()
71+
connection_name = connection_data.get("name", "")
7172
service_name = connection_data.get("serviceDirectory", "")
7273
host = connection_data.get("host", "")
7374
if host:
7475
service_name = connection_data.get("tlsServiceDirectory", "")
7576
auth_override_enabled = connection_data.get("authOverrideEnabled", False)
7677
return {
78+
"name": connection_name,
7779
"serviceName": service_name,
7880
"host": host,
7981
"authOverrideEnabled": auth_override_enabled,
@@ -291,13 +293,9 @@ def get_action_operation(
291293
tool_name: str = "",
292294
tool_instructions: str = "",
293295
) -> Dict[str, Any]:
294-
description = (
295-
f"Use this tool with" f' action = "{action}" and'
296-
) + f' operation = "{operation}" only. Dont ask these values from user.'
296+
description = f"Use this tool to execute {action}"
297297
if operation == "EXECUTE_QUERY":
298-
description = (
299-
(f"Use this tool with" f' action = "{action}" and')
300-
+ f' operation = "{operation}" only. Dont ask these values from user.'
298+
description += (
301299
" Use pageSize = 50 and timeout = 120 until user specifies a"
302300
" different value otherwise. If user provides a query in natural"
303301
" language, convert it to SQL query and then execute it using the"
@@ -308,6 +306,8 @@ def get_action_operation(
308306
"summary": f"{action_display_name}",
309307
"description": f"{description} {tool_instructions}",
310308
"operationId": f"{tool_name}_{action_display_name}",
309+
"x-action": f"{action}",
310+
"x-operation": f"{operation}",
311311
"requestBody": {
312312
"content": {
313313
"application/json": {
@@ -347,16 +347,12 @@ def list_operation(
347347
"post": {
348348
"summary": f"List {entity}",
349349
"description": (
350-
f"Returns all entities of type {entity}. Use this tool with"
351-
+ f' entity = "{entity}" and'
352-
+ ' operation = "LIST_ENTITIES" only. Dont ask these values'
353-
" from"
354-
+ ' user. Always use ""'
355-
+ ' as filter clause and ""'
356-
+ " as page token and 50 as page size until user specifies a"
357-
" different value otherwise. Use single quotes for strings in"
358-
f" filter clause. {tool_instructions}"
350+
f"""Returns the list of {entity} data. If the page token was available in the response, let users know there are more records available. Ask if the user wants to fetch the next page of results. When passing filter use the
351+
following format: `field_name1='value1' AND field_name2='value2'
352+
`. {tool_instructions}"""
359353
),
354+
"x-operation": "LIST_ENTITIES",
355+
"x-entity": f"{entity}",
360356
"operationId": f"{tool_name}_list_{entity}",
361357
"requestBody": {
362358
"content": {
@@ -401,14 +397,11 @@ def get_operation(
401397
"post": {
402398
"summary": f"Get {entity}",
403399
"description": (
404-
(
405-
f"Returns the details of the {entity}. Use this tool with"
406-
f' entity = "{entity}" and'
407-
)
408-
+ ' operation = "GET_ENTITY" only. Dont ask these values from'
409-
f" user. {tool_instructions}"
400+
f"Returns the details of the {entity}. {tool_instructions}"
410401
),
411402
"operationId": f"{tool_name}_get_{entity}",
403+
"x-operation": "GET_ENTITY",
404+
"x-entity": f"{entity}",
412405
"requestBody": {
413406
"content": {
414407
"application/json": {
@@ -445,17 +438,10 @@ def create_operation(
445438
) -> Dict[str, Any]:
446439
return {
447440
"post": {
448-
"summary": f"Create {entity}",
449-
"description": (
450-
(
451-
f"Creates a new entity of type {entity}. Use this tool with"
452-
f' entity = "{entity}" and'
453-
)
454-
+ ' operation = "CREATE_ENTITY" only. Dont ask these values'
455-
" from"
456-
+ " user. Follow the schema of the entity provided in the"
457-
f" instructions to create {entity}. {tool_instructions}"
458-
),
441+
"summary": f"Creates a new {entity}",
442+
"description": f"Creates a new {entity}. {tool_instructions}",
443+
"x-operation": "CREATE_ENTITY",
444+
"x-entity": f"{entity}",
459445
"operationId": f"{tool_name}_create_{entity}",
460446
"requestBody": {
461447
"content": {
@@ -491,18 +477,10 @@ def update_operation(
491477
) -> Dict[str, Any]:
492478
return {
493479
"post": {
494-
"summary": f"Update {entity}",
495-
"description": (
496-
(
497-
f"Updates an entity of type {entity}. Use this tool with"
498-
f' entity = "{entity}" and'
499-
)
500-
+ ' operation = "UPDATE_ENTITY" only. Dont ask these values'
501-
" from"
502-
+ " user. Use entityId to uniquely identify the entity to"
503-
" update. Follow the schema of the entity provided in the"
504-
f" instructions to update {entity}. {tool_instructions}"
505-
),
480+
"summary": f"Updates the {entity}",
481+
"description": f"Updates the {entity}. {tool_instructions}",
482+
"x-operation": "UPDATE_ENTITY",
483+
"x-entity": f"{entity}",
506484
"operationId": f"{tool_name}_update_{entity}",
507485
"requestBody": {
508486
"content": {
@@ -538,16 +516,10 @@ def delete_operation(
538516
) -> Dict[str, Any]:
539517
return {
540518
"post": {
541-
"summary": f"Delete {entity}",
542-
"description": (
543-
(
544-
f"Deletes an entity of type {entity}. Use this tool with"
545-
f' entity = "{entity}" and'
546-
)
547-
+ ' operation = "DELETE_ENTITY" only. Dont ask these values'
548-
" from"
549-
f" user. {tool_instructions}"
550-
),
519+
"summary": f"Delete the {entity}",
520+
"description": f"Deletes the {entity}. {tool_instructions}",
521+
"x-operation": "DELETE_ENTITY",
522+
"x-entity": f"{entity}",
551523
"operationId": f"{tool_name}_delete_{entity}",
552524
"requestBody": {
553525
"content": {

0 commit comments

Comments
 (0)
0