8000 chore: Add credential service to runner and invocation context · stufflebeam/adk-python@5f89a46 · GitHub
[go: up one dir, main page]
Skip to content

Commit 5f89a46

Browse files
seanzhougooglecopybara-github
seanzhougoogle
authored and
copybara-github
committed
chore: Add credential service to runner and invocation context
PiperOrigin-RevId: 772697298
1 parent 6d174eb commit 5f89a46

File tree

6 files changed

+33
-6
lines changed

6 files changed

+33
-6
lines changed

‎src/google/adk/agents/invocation_context.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from pydantic import ConfigDict
2323

2424
from ..artifacts.base_artifact_service import BaseArtifactService
25+
from ..auth.credential_service.base_credential_service import BaseCredentialService
2526
from ..memory.base_memory_service import BaseMemoryService
2627
from ..sessions.base_session_service import BaseSessionService
2728
from ..sessions.session import Session
@@ -115,6 +116,7 @@ class InvocationContext(BaseModel):
115116
artifact_service: Optional[BaseArtifactService] = None
116117
session_service: BaseSessionService
117118
memory_service: Optional[BaseMemoryService] = None
119+
credential_service: Optional[BaseCredentialService] = None
118120

119121
invocation_id: str
120122
"""The id of this invocation context. Readonly."""

‎src/google/adk/auth/credential_service/base_credential_service.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,12 @@
1919
from typing import Optional
2020

2121
from ...tools.tool_context import ToolContext
22-
from ...utils.feature_decorator import working_in_progress
22+
from ...utils.feature_decorator import experimental
2323
from ..auth_credential import AuthCredential
2424
from ..auth_tool import AuthConfig
2525

2626

27-
@working_in_progress("Implementation are in progress. Don't use it for now.")
27+
@experimental
2828
class BaseCredentialService(ABC):
2929
"""Abstract class for Service that loads / saves tool credentials from / to
3030
the backend credential store."""

‎src/google/adk/cli/cli.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424
from ..agents.llm_agent import LlmAgent
2525
from ..artifacts import BaseArtifactService
2626
from ..artifacts import InMemoryArtifactService
27+
from ..auth.credential_service.base_credential_service import BaseCredentialService
28+
from ..auth.credential_service.in_memory_credential_service import InMemoryCredentialService
2729
from ..runners import Runner
2830
from ..sessions.base_session_service import BaseSessionService
2931
from ..sessions.in_memory_session_service import InMemorySessionService
@@ -43,13 +45,15 @@ async def run_input_file(
4345
root_agent: LlmAgent,
4446
artifact_service: BaseArtifactService,
4547
session_service: BaseSessionService,
48+
credential_service: BaseCredentialService,
4649
input_path: str,
4750
) -> Session:
4851
runner = Runner(
4952
app_name=app_name,
5053
agent=root_agent,
5154
artifact_service=artifact_service,
5255
session_service=session_service,
56+
credential_service=credential_service,
5357
)
5458
with open(input_path, 'r', encoding='utf-8') as f:
5559
input_file = InputFile.model_validate_json(f.read())
@@ -75,12 +79,14 @@ async def run_interactively(
7579
artifact_service: BaseArtifactService,
7680
session: Session,
7781
session_service: BaseSessionService,
82+
credential_service: BaseCredentialService,
7883
) -> None:
7984
runner = Runner(
8085
app_name=session.app_name,
8186
agent=root_agent,
8287
artifact_service=artifact_service,
8388
session_service=session_service,
89+
credential_service=credential_service,
8490
)
8591
while True:
8692
query = input('[user]: ')
@@ -125,6 +131,7 @@ async def run_cli(
125131

126132
artifact_service = InMemoryArtifactService()
127133
session_service = InMemorySessionService()
134+
credential_service = InMemoryCredentialService()
128135

129136
user_id = 'test_user'
130137
session = await session_service.create_session(
@@ -141,6 +148,7 @@ async def run_cli(
141148
root_agent=root_agent,
142149
artifact_service=artifact_service,
143150
session_service=session_service,
151+
credential_service=credential_service,
144152
input_path=input_file,
145153
)
146154
elif saved_session_file:
@@ -163,6 +171,7 @@ async def run_cli(
163171
artifact_service,
164172
session,
165173
session_service,
174+
credential_service,
166175
)
167176
else:
168177
click.echo(f'Running agent {root_agent.name}, type exit to exit.')
@@ -171,6 +180,7 @@ async def run_cli(
171180
artifact_service,
172181
session,
173182
session_service,
183+
credential_service,
174184
)
175185

176186
if save_session:

‎src/google/adk/cli/fast_api.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@
5757
from ..agents.run_config import StreamingMode
5858
from ..artifacts.gcs_artifact_service import GcsArtifactService
5959
from ..artifacts.in_memory_artifact_service import InMemoryArtifactService
60+
from ..auth.credential_service.in_memory_credential_service import InMemoryCredentialService
6061
from ..errors.not_found_error import NotFoundError
6162
from ..evaluation.eval_case import EvalCase
6263
from ..evaluation.eval_case import SessionInput
@@ -305,6 +306,9 @@ async def internal_lifespan(app: FastAPI):
305306
else:
306307
artifact_service = InMemoryArtifactService()
307308

309+
# Build the Credential service
310+
credential_service = InMemoryCredentialService()
311+
308312
# initialize Agent Loader
309313
agent_loader = AgentLoader(agents_dir)
310314

@@ -929,6 +933,7 @@ async def _get_runner_async(app_name: str) -> Runner:
929933
artifact_service=artifact_service,
930934
session_service=session_service,
931935
memory_service=memory_service,
936+
credential_service=credential_service,
932937
)
933938
runner_dict[app_name] = runner
934939
return runner

‎src/google/adk/runners.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
import asyncio
1818
import logging
1919
import queue
20-
import threading
2120
from typing import AsyncGenerator
2221
from typing import Generator
2322
from typing import Optional
@@ -34,6 +33,7 @@
3433
from .agents.run_config import RunConfig
3534
from .artifacts.base_artifact_service import BaseArtifactService
3635
from .artifacts.in_memory_artifact_service import InMemoryArtifactService
36+
from .auth.credential_service.base_credential_service import BaseCredentialService
3737
from .code_executors.built_in_code_executor import BuiltInCodeExecutor
3838
from .events.event import Event
3939
from .memory.base_memory_service import BaseMemoryService
@@ -73,6 +73,8 @@ class Runner:
7373
"""The session service for the runner."""
7474
memory_service: Optional[BaseMemoryService] = None
7575
"""The memory service for the runner."""
76+
credential_service: Optional[BaseCredentialService] = None
77+
"""The credential service for the runner."""
7678

7779
def __init__(
7880
self,
@@ -82,6 +84,7 @@ def __init__(
8284
artifact_service: Optional[BaseArtifactService] = None,
8385
session_service: BaseSessionService,
8486
memory_service: Optional[BaseMemoryService] = None,
87+
credential_service: Optional[BaseCredentialService] = None,
8588
):
8689
"""Initializes the Runner.
8790
@@ -97,6 +100,7 @@ def __init__(
97100
self.artifact_service = artifact_service
98101
self.session_service = session_service
99102
self.memory_service = memory_service
103+
self.credential_service = credential_service
100104

101105
def run(
102106
self,
@@ -418,6 +422,7 @@ def _new_invocation_context(
418422
artifact_service=self.artifact_service,
419423
session_service=self.session_service,
420424
memory_service=self.memory_service,
425+
credential_service=self.credential_service,
421426
invocation_id=invocation_id,
422427
agent=self.agent,
423428
session=session,

‎tests/unittests/cli/utils/test_cli.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,7 @@ def _echo(msg: str) -> None:
129129

130130
artifact_service = cli.InMemoryArtifactService()
131131
session_service = cli.InMemorySessionService()
132+
credential_service = cli.InMemoryCredentialService()
132133
dummy_root = types.SimpleNamespace(name="root")
133134

134135
session = await cli.run_input_file(
@@ -137,6 +138,7 @@ def _echo(msg: str) -> None:
137138
root_agent=dummy_root,
138139
artifact_service=artifact_service,
139140
session_service=session_service,
141+
credential_service=credential_service,
140142
input_path=str(input_path),
141143
)
142144

@@ -199,9 +201,10 @@ async def test_run_interactively_whitespace_and_exit(
199201
) -> None:
200202
"""run_interactively should skip blank input, echo once, then exit."""
201203
# make a session that belongs to dummy agent
202-
svc = cli.InMemorySessionService()
203-
sess = await svc.create_session(app_name="dummy", user_id="u")
204+
session_service = cli.InMemorySessionService()
205+
sess = await session_service.create_session(app_name="dummy", user_id="u")
204206
artifact_service = cli.InMemoryArtifactService()
207+
credential_service = cli.InMemoryCredentialService()
205208
root_agent = types.SimpleNamespace(name="root")
206209

207210
# fake user input: blank -> 'hello' -> 'exit'
@@ -212,7 +215,9 @@ async def test_run_interactively_whitespace_and_exit(
212215
echoed: list[str] = []
213216
monkeypatch.setattr(click, "echo", lambda msg: echoed.append(msg))
214217

215-
await cli.run_interactively(root_agent, artifact_service, sess, svc)
218+
await cli.run_interactively(
219+
root_agent, artifact_service, sess, session_service, credential_service
220+
)
216221

217222
# verify: assistant echoed once with 'echo:hello'
218223
assert any("echo:hello" in m for m in echoed)

0 commit comments

Comments
 (0)
0