Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
13 changes: 13 additions & 0 deletions .kokoro/presubmit/unit_agentplatform_llama_index_py311.cfg
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
# Format: //devtools/kokoro/config/proto/build.proto

# Run unit tests for LlamaIndex on Python 3.11
env_vars: {
key: "NOX_SESSION"
value: "unit_agentplatform_llama_index-3.11"
}

# Run unit tests in parallel, splitting up by file
env_vars: {
key: "PYTEST_ADDOPTS"
value: "-n=auto --dist=loadscope"
}
36 changes: 35 additions & 1 deletion agentplatform/agent_engines/templates/adk.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,15 @@
except (ImportError, AttributeError):
BaseMemoryService = Any

try:
from google.adk.auth.credential_service.base_credential_service import (
BaseCredentialService,
)

BaseCredentialService = BaseCredentialService
except (ImportError, AttributeError):
BaseCredentialService = Any

try:
from opentelemetry.sdk import trace

Expand Down Expand Up @@ -682,6 +691,9 @@ def __init__(
session_service_builder: Optional[Callable[..., "BaseSessionService"]] = None,
artifact_service_builder: Optional[Callable[..., "BaseArtifactService"]] = None,
memory_service_builder: Optional[Callable[..., "BaseMemoryService"]] = None,
credential_service_builder: Optional[
Callable[..., "BaseCredentialService"]
] = None,
instrumentor_builder: Optional[Callable[..., Any]] = None,
):
"""An ADK Application.
Expand Down Expand Up @@ -715,6 +727,9 @@ def __init__(
Defaults to a callable that returns InMemoryMemoryService
when running locally and VertexAiMemoryBankService when running
on Agent Engine.
credential_service_builder (Callable[..., BaseCredentialService]):
Optional. A callable that returns an ADK credential service.
Defaults to a callable that returns InMemoryCredentialService.
instrumentor_builder (Callable[..., Any]):
Optional. Callable that returns a new instrumentor. This can be
used for customizing the instrumentation logic of the Agent.
Expand Down Expand Up @@ -759,6 +774,7 @@ def __init__(
"session_service_builder": session_service_builder,
"artifact_service_builder": artifact_service_builder,
"memory_service_builder": memory_service_builder,
"credential_service_builder": credential_service_builder,
"instrumentor_builder": instrumentor_builder,
"express_mode_api_key": (
initializer.global_config.api_key or os.environ.get("GOOGLE_API_KEY")
Expand Down Expand Up @@ -912,6 +928,9 @@ def clone(self):
session_service_builder=self._tmpl_attrs.get("session_service_builder"),
artifact_service_builder=self._tmpl_attrs.get("artifact_service_builder"),
memory_service_builder=self._tmpl_attrs.get("memory_service_builder"),
credential_service_builder=self._tmpl_attrs.get(
"credential_service_builder"
),
instrumentor_builder=self._tmpl_attrs.get("instrumentor_builder"),
)

Expand All @@ -924,6 +943,9 @@ def set_up(self):
InMemoryArtifactService,
)
from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
from google.adk.auth.credential_service.in_memory_credential_service import (
InMemoryCredentialService,
)

os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1"
project = self._tmpl_attrs.get("project")
Expand Down Expand Up @@ -1078,6 +1100,12 @@ def set_up(self):
else:
self._tmpl_attrs["memory_service"] = InMemoryMemoryService()

credential_service_builder = self._tmpl_attrs.get("credential_service_builder")
if credential_service_builder:
self._tmpl_attrs["credential_service"] = credential_service_builder()
else:
self._tmpl_attrs["credential_service"] = InMemoryCredentialService()

self._tmpl_attrs["runner"] = Runner(
app=self._tmpl_attrs.get("app"),
agent=(
Expand Down Expand Up @@ -1114,6 +1142,7 @@ def set_up(self):
session_service=self._tmpl_attrs.get("in_memory_session_service"),
artifact_service=self._tmpl_attrs.get("in_memory_artifact_service"),
memory_service=self._tmpl_attrs.get("in_memory_memory_service"),
credential_service=self._tmpl_attrs.get("credential_service"),
)

async def async_stream_query(
Expand Down Expand Up @@ -1183,11 +1212,16 @@ async def async_stream_query(
from google.adk.events.event import Event

session_service = self._tmpl_attrs.get("session_service")
session_obj = await session_service.get_session(
app_name=self._app_name(),
user_id=user_id,
session_id=session_id,
)
for event in session_events:
if not isinstance(event, Event):
event = Event.model_validate(event)
await session_service.append_event(
session=session,
session=session_obj,
event=event,
)

Expand Down
Loading
Loading