This article is co-authored by Łukasz Olejniczak and Rafał Knapik @raknapik (rafal.knapik@allegro.com).
All the code needed to reproduce this tutorial can be found in my repo: https://github.com/lolejniczak-shared/adkagents/tree/master/04-adk-samples/adk_sample_01_authenticated_vais
In our last blog post, we learned how to add custom datastores to our Agentspace application. In this post, we’ll go a step further and explore Agentspace’s knowledge activation layer.
We will demonstrate how to build a custom, “high-code” AI agent that utilizes our custom datastore as a tool. Critically, because this datastore is enabled with access control, our AI agent will be able to retrieve only data that the end user is authorized to access.
Our AI agent will be implemented using the Google Agent Development Kit (ADK), an open-source library for building production-ready, multi-agent applications.
The core building blocks of AI agents are a Large Language Model (LLM), along with components for planning, memory, and tools:
Our agent will be implemented with only an LLM and a tool, but I am sure this example will serve as a great starting point for you to create a true specialist AI Agents that work with Agentspace datastores to meet your unique requirements.
Agentspace Datastores can be queried via an API. To simplify this process, we’ll create a dedicated Python class called DatastoreService.
This auxiliary class will handle two primary responsibilities:
- Authentication: It ensures we have a valid OAuth2 access token.
- Browsing: It performs the search query on the datastore.
class DatastoreService:
def __init__(self, access_token: str):
self.access_token = None
if access_token:
self.access_token = access_token
else:
creds, project_id = default()
auth_req = transport.requests.Request() # Use google.auth here
creds.refresh(auth_req)
access_token = creds.token
self.access_token = access_tokendef search_datastore(self, project_id, location, datastore_id, query):
# Define API endpoint and headers
url = f"https://{location}-discoveryengine.googleapis.com/v1alpha/projects/{project_id}/locations/{location}/collections/default_collection/dataStores/{datastore_id}/servingConfigs/default_search:search"
headers = {
"Authorization": f"Bearer {self.access_token}",
"Content-Type": "application/json"
}
# Define request data with placeholders for query
data = {
"query": f"{query}",
"pageSize":10,
"queryExpansionSpec":{"condition":"AUTO"},
"spellCorrectionSpec":{"mode":"AUTO"},
"relevanceScoreSpec":{"returnRelevanceScore":True},
"languageCode":"en-US",
"contentSearchSpec":{"snippetSpec":{"returnSnippet":True}},
"naturalLanguageQueryUnderstandingSpec":{"filterExtractionCondition":"ENABLED"},
"userInfo":{"timeZone":"Europe/Warsaw"}
}
# Make POST request
response = requests.post(url, headers=headers, json=data)
resp = response.json()
print(resp)
return resp
When we query our custom datastore for “Employee Handbook,” we should get at least one result corresponding to a document with the ID ‘TSK-1003’. Why not all items we loaded to custom datastore? Because this datastore is configured to respect access control, and we will only get results that authenticated user is privileged to retrieve.
The following is the code snippet demonstrating how to use the DatastoreService class to perform a search query:
service = DatastoreService()query_for_user1 = "Employee Handbook"
service.search_datastore(PROJECT_ID, LOCATION, DATA_STORE_ID, query_for_user1)
The OAuth2 access token needed to search the datastore is generated by the DatastoreService constructor when it’s instantiated.
service = DatastoreService()
The generated access token will represent user1.
When we run the search_datastore function with the query "Employee Handbook", here is the response we get:
{'results': [{'id': 'TSK-1003', 'document': {'name': 'projects/680248386202/locations/global/collections/default_collection/dataStores/auth-datastore-test00/branches/0/documents/TSK-1003', 'id': 'TSK-1003', 'structData': {'date': '2025-08-26T11:00:00Z', 'tags': ['policy', 'documentation', 'compliance'], 'author': 'Ewa Jankowska', 'url': 'http://localhost', 'body': {'priority': 'LOW', 'assigned_to': 'Krzysztof Wójcik', 'creation_date': '2025-08-26T11:00:00Z', 'status': 'To Do', 'assigned_to_email': 'user1@gmail.com', 'description': "Review and update the company's employee handbook with the new remote work and benefits policies. Circulate draft to management for approval by end of Q4."}}, 'derivedStructData': {'snippets': [{'snippet': 'No snippet is available for this page.', 'snippet_status': 'NO_SNIPPET_AVAILABLE'}], 'can_fetch_raw_content': 'true', 'is_exact_match_query': 0}}, 'modelScores': {'relevance_score': {'values': [0.6]}}, 'rankSignals': {'semanticSimilarityScore': 0.7028787, 'topicalityRank': 1, 'boostingFactor': 0, 'defaultRank': 1}}], 'totalSize': 1, 'attributionToken': 'kAL0DwEKDAi01__FBhDzs7GGAxIkNjkwNzAwYWQtMDAwMC0yZmMzLTkyZjQtZDRmNTQ3ZWVhZGUwIgdHRU5FUklDKkjQ2ok309qJN9SynRXC8J4Vt7eMLZvWty2VksUw4-uQN7uR-jG-kfoxtqqiMrmqojKQ97IwxcvzF46-nRXg65A3mNa3LY6RyTAwAVKEAXByb2plY3RzLzY4MDI0ODM4NjIwMi9sb2NhdGlvbnMvZ2xvYmFsL2NvbGxlY3Rpb25zL2RlZmF1bHRfY29sbGVjdGlvbi9kYXRhU3RvcmVzL2F1dGgtZGF0YXN0b3JlLXRlc3QwMC9zZXJ2aW5nQ29uZmlncy9kZWZhdWx0X3NlYXJjaA', 'guidedSearchResult': {}, 'summary': {}, 'queryExpansionInfo': {}, 'naturalLanguageQueryUnderstandingInfo': {'rewrittenQuery': 'Employee Handbook'}}
It’s a good sign that TSK-1003 is on the list, as it confirms that the access control is working correctly. Now, let’s proceed with the next step: defining a Python function to act as an ADK Tool.
def search_tasks(query: str, tool_context: ToolContext):
"""
Searches the tsk registry using the DatastoreService.Args:
query (str): The search query string.
Returns:
dict: The search results from the DatastoreService in JSON format.
"""
datastore_service = None
auth_name= f"temp:{AUTH_NAME}"
access_token = tool_context.state.get(auth_name)
if access_token:
datastore_service = DatastoreService(access_token)
else:
access_token = ""
datastore_service = DatastoreService(access_token)
# Call the search method of the DatastoreService with the project ID, App Engine ID, and query
results = datastore_service.search_datastore(PROJECT_ID, LOCATION, DATA_STORE_ID, query)
# Return the search results
return results
This function takes two parameters: the user’s query and a ToolContext object. The ToolContext acts as a bridge to the ADK runtime, giving us access to the agent’s session state.
The session state can be thought of as a key-value store for the agent, scoped to the current conversation. We will use this to retrieve the access token that represents the user interacting with our agent.
The agent expects this token to be available in the session state under a static key: temp:{AUTH_NAME}. While the key’s name never changes, its value is dynamic. Every user has a unique access token, and each token has a limited lifetime, such as one hour.
You can officially turn any Python function into an ADK tool with just a single line of code. You do this by wrapping the function in a FunctionTool class.
task_search_tool = FunctionTool(func=search_tasks)
The ADK framework automatically inspects the function’s signature and docstring to understand what the tool does, what parameters it requires, and when it should be used. This makes it easy for the agent’s LLM to intelligently select and utilize the tool when needed.
The final step is to attach our new tool to the agent.
We’ll define our AI Agent by creating an instance of the Agent class. This is where we configure its behavior with a prompt, specify the model it will use, and, most importantly, provide it with the tools it can access.
root_agent = Agent(
model="gemini-2.5-flash",
name="custom_datastore_assistant",
instruction="You are helpful assitant answering user queries using available tools",
tools = [search_tasks]
)
The agent’s LLM can now intelligently choose to execute search_tasks tool whenever a user’s query requires it to answer access knowledge available in the datastore.
With this setup complete, the agent is ready to be run locally:
adk web
We can now ask it to: list my tasks
How was the agent able to identify us as ‘user1’?
There is no magic here. This happened because I executed:
gcloud auth application-default login
to authenticate as user1 and only then started ADK Web UI:
adk web
Because we started our agent with an empty session state:
access_token = tool_context.state.get(auth_name)
it was unable to find an access token for the key temp:{AUTH_NAME}. As a result of the missing access token, the DatastoreService was instantiated to handle authentication. It used the following code to get new credentials and a token:
creds, project_id = default()
auth_req = transport.requests.Request() # Use google.auth here
creds.refresh(auth_req)
access_token = creds.token
self.access_token = access_token
This authentication process uses Application Default Credentials (ADC), a strategy employed by Google authentication libraries that automatically finds and uses credentials based on the application’s environment.
This access token then is used when calling DatastoreService and as a result our Agent generates answers only from DATA that user1 can access.
How would this work if our Agent would be able to find access token in session state? It would simply use the token directly to make API calls, rather than checking local environment for user credentials to turn it into access token.
So, how one could put access token into session state? The key to this is an external component — like Agentspace — that handles the OAuth2 process and injects the token into the agent’s session state before the conversation begins.
To get this setup working, a few steps are required. The first is to deploy our ADK agent to GCP Agent Engine, a serverless runtime for AI agents. The ADK library makes this process very easy — just run the following command:
adk deploy agent_engine \
--project=${GOOGLE_CLOUD_PROJECT} \
--region=${GOOGLE_CLOUD_LOCATION} \
--staging_bucket=gs://${STAGING_BUCKET} \
--display_name=${ADK_APP_NAME} \
agent
We will use Agentspace as a user interface for our Agent but we will also delegate the user authorization process to it. As a result Agentspace will become responsible for the OAuth2 flow, generating the necessary access token and placing it directly into our ADK agent’s session state whenever a new conversation with our agent is started.
For this flow to work, we need to provide Agentspace with details about the OAuth2 process we want it to own. We do this by registering an authorization resource object within the Agentspace application.
This object has a unique identifier, auth_id. This identifier is crucial because it is used to construct the name of the session state key where Agentspace will place the access token, using the format temp:{auth_id}.
The authorization resource object contains all the necessary details for an OAuth2 flow to authorize a user and obtain a corresponding access token. This also includes scopes, which are groups of permissions that our agent will need to act on the user’s behalf within the underlying system and will therefore ask the end user to authorize this.
client = AgentspaceManager(project_id=PROJECT_ID, app_id=AGENTSPACE_APP_ID)auth_uri = client.generate_auth_uri(
base_auth_uri="https://accounts.google.com/o/oauth2/v2/auth",
client_id=GOOGLE_CLIENT_ID,
scopes=["openid",
"https://www.googleapis.com/auth/cloud-platform"
]
)
resp = client.create_authorization(
auth_id=auth_object_id,
client_id = GOOGLE_CLIENT_ID,
client_secret= GOOGLE_CLIENT_SECRET,
auth_uri=auth_uri,
token_uri="https://oauth2.googleapis.com/token"
)
print(resp)
AgentspaceManager is not official Agentspace SDK. It is part of my package describing the process. https://github.com/lolejniczak-shared/adkagents/tree/master/04-adk-samples/adk_sample_01_authenticated_vais
We are now ready to register our ADK agent to Agentspace Agent Garden. A single agent may need multiple authorization resources. When registering the agent, we provide a list of these resources, which describe the various authorization flows that have been outsourced to Agentspace.
client = AgentspaceManager(project_id=PROJECT_ID,
app_id=AGENTSPACE_APP_ID, location="global"
)resp = client.register_agent(
auth_ids=[AGENT_AUTH_OBJECT_ID, ],
display_name = AGENTSPACE_APP_NAME,
description = "Agent that answers questions about user tasks.",
tool_description = "Agent that answers questions about user tasks",
adk_deployment_id = AGENT_ENGINE_ID,
adk_deployment_location = LOCATION,
icon_uri="https://raw.githubusercontent.com/google/material-design-icons/master/src/action/android/materialicons/24px.svg"
)
print(resp)
Now, navigate to the Agentspace application. Upon visiting the Agent Garden, you will see our newly registered agent.
After you select the agent, you will first be asked to authorize it so that it is able to act on your behalf on Datastores.
Once the OAuth2 flow is complete, you are ready to chat with your agent.
This is precisely the behavior we would expect from our agent. Agentspace successfully passed the access token to the session state of our ADK agent, our ADK agent was able to use it, which enabled us to receive answers generated from data to which we have access.
Source Credit: https://medium.com/google-cloud/adk-agents-with-agentspace-authenticated-datastores-64a49466074d?source=rss—-e52cf94d98af—4
