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
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

package com.google.adk.artifacts;

import com.google.adk.sessions.SessionKey;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Completable;
Expand All @@ -39,6 +40,11 @@ public interface BaseArtifactService {
Single<Integer> saveArtifact(
String appName, String userId, String sessionId, String filename, Part artifact);

default Single<Integer> saveArtifact(SessionKey sessionKey, String filename, Part artifact) {
return saveArtifact(
sessionKey.appName(), sessionKey.userId(), sessionKey.id(), filename, artifact);
}

/**
* Saves an artifact and returns it with fileData if available.
*
Expand All @@ -58,18 +64,33 @@ default Single<Part> saveAndReloadArtifact(
.flatMap(version -> loadArtifact(appName, userId, sessionId, filename, version).toSingle());
}

default Single<Part> saveAndReloadArtifact(
SessionKey sessionKey, String filename, Part artifact) {
return saveAndReloadArtifact(
sessionKey.appName(), sessionKey.userId(), sessionKey.id(), filename, artifact);
}

/** Loads the latest version of an artifact from the service. */
default Maybe<Part> loadArtifact(
String appName, String userId, String sessionId, String filename) {
return loadArtifact(appName, userId, sessionId, filename, Optional.empty());
}

default Maybe<Part> loadArtifact(SessionKey sessionKey, String filename) {
return loadArtifact(sessionKey.appName(), sessionKey.userId(), sessionKey.id(), filename);
}

/** Loads a specific version of an artifact from the service. */
default Maybe<Part> loadArtifact(
String appName, String userId, String sessionId, String filename, int version) {
return loadArtifact(appName, userId, sessionId, filename, Optional.of(version));
}

default Maybe<Part> loadArtifact(SessionKey sessionKey, String filename, int version) {
return loadArtifact(
sessionKey.appName(), sessionKey.userId(), sessionKey.id(), filename, version);
}

/**
* @deprecated Use {@link #loadArtifact(String, String, String, String)} or {@link
* #loadArtifact(String, String, String, String, int)} instead.
Expand All @@ -78,6 +99,12 @@ default Maybe<Part> loadArtifact(
Maybe<Part> loadArtifact(
String appName, String userId, String sessionId, String filename, Optional<Integer> version);

default Maybe<Part> loadArtifact(
SessionKey sessionKey, String filename, Optional<Integer> version) {
return loadArtifact(
sessionKey.appName(), sessionKey.userId(), sessionKey.id(), filename, version);
}

/**
* Lists all the artifact filenames within a session.
*
Expand All @@ -88,6 +115,10 @@ Maybe<Part> loadArtifact(
*/
Single<ListArtifactsResponse> listArtifactKeys(String appName, String userId, String sessionId);

default Single<ListArtifactsResponse> listArtifactKeys(SessionKey sessionKey) {
return listArtifactKeys(sessionKey.appName(), sessionKey.userId(), sessionKey.id());
}

/**
* Deletes an artifact.
*
Expand All @@ -98,6 +129,10 @@ Maybe<Part> loadArtifact(
*/
Completable deleteArtifact(String appName, String userId, String sessionId, String filename);

default Completable deleteArtifact(SessionKey sessionKey, String filename) {
return deleteArtifact(sessionKey.appName(), sessionKey.userId(), sessionKey.id(), filename);
}

/**
* Lists all the versions (as revision IDs) of an artifact.
*
Expand All @@ -109,4 +144,8 @@ Maybe<Part> loadArtifact(
*/
Single<ImmutableList<Integer>> listVersions(
String appName, String userId, String sessionId, String filename);

default Single<ImmutableList<Integer>> listVersions(SessionKey sessionKey, String filename) {
return listVersions(sessionKey.appName(), sessionKey.userId(), sessionKey.id(), filename);
}
}
31 changes: 31 additions & 0 deletions core/src/main/java/com/google/adk/runner/Runner.java
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import com.google.adk.sessions.BaseSessionService;
import com.google.adk.sessions.InMemorySessionService;
import com.google.adk.sessions.Session;
import com.google.adk.sessions.SessionKey;
import com.google.adk.summarizer.EventsCompactionConfig;
import com.google.adk.summarizer.LlmEventSummarizer;
import com.google.adk.summarizer.SlidingWindowEventCompactor;
Expand Down Expand Up @@ -383,6 +384,25 @@ public Flowable<Event> runAsync(
.flatMapPublisher(session -> this.runAsyncImpl(session, newMessage, runConfig, stateDelta));
}

/** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */
public Flowable<Event> runAsync(
SessionKey sessionKey,
Content newMessage,
RunConfig runConfig,
@Nullable Map<String, Object> stateDelta) {
return runAsync(sessionKey.userId(), sessionKey.id(), newMessage, runConfig, stateDelta);
}

/** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */
public Flowable<Event> runAsync(SessionKey sessionKey, Content newMessage, RunConfig runConfig) {
return runAsync(sessionKey, newMessage, runConfig, /* stateDelta= */ null);
}

/** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */
public Flowable<Event> runAsync(SessionKey sessionKey, Content newMessage) {
return runAsync(sessionKey, newMessage, RunConfig.builder().build());
}

/** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */
public Flowable<Event> runAsync(String userId, String sessionId, Content newMessage) {
return runAsync(userId, sessionId, newMessage, RunConfig.builder().build());
Expand Down Expand Up @@ -671,6 +691,17 @@ public Flowable<Event> runLive(
.flatMapPublisher(session -> this.runLive(session, liveRequestQueue, runConfig));
}

/**
* Retrieves the session and runs the agent in live mode.
*
* @return stream of events from the agent.
* @throws IllegalArgumentException if the session is not found.
*/
public Flowable<Event> runLive(
SessionKey sessionKey, LiveRequestQueue liveRequestQueue, RunConfig runConfig) {
return runLive(sessionKey.userId(), sessionKey.id(), liveRequestQueue, runConfig);
}

/**
* Runs the agent asynchronously with a default user ID.
*
Expand Down
57 changes: 44 additions & 13 deletions core/src/main/java/com/google/adk/sessions/BaseSessionService.java
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,18 @@ default Single<Session> createSession(
return createSession(appName, userId, ensureConcurrentMap(state), sessionId);
}

/**
* Creates a new session with the specified parameters.
*
* @param sessionKey The session key containing appName, userId and sessionId.
* @param state An optional map representing the initial state of the session. Can be null or
* empty.
*/
default Single<Session> createSession(
SessionKey sessionKey, @Nullable Map<String, Object> state) {
return createSession(sessionKey.appName(), sessionKey.userId(), state, sessionKey.id());
}

/**
* Creates a new session with the specified application name and user ID, using a default state
* (null) and allowing the service to generate a unique session ID.
Expand All @@ -94,6 +106,10 @@ default Single<Session> createSession(String appName, String userId) {
return createSession(appName, userId, null, null);
}

default Single<Session> createSession(SessionKey sessionKey) {
return createSession(sessionKey.appName(), sessionKey.userId(), null, sessionKey.id());
}

/**
* Retrieves a specific session, optionally filtering the events included.
*
Expand All @@ -110,6 +126,11 @@ default Single<Session> createSession(String appName, String userId) {
Maybe<Session> getSession(
String appName, String userId, String sessionId, Optional<GetSessionConfig> config);

default Maybe<Session> getSession(SessionKey sessionKey, @Nullable GetSessionConfig config) {
return getSession(
sessionKey.appName(), sessionKey.userId(), sessionKey.id(), Optional.ofNullable(config));
}

/**
* Lists sessions associated with a specific application and user.
*
Expand All @@ -123,6 +144,10 @@ Maybe<Session> getSession(
*/
Single<ListSessionsResponse> listSessions(String appName, String userId);

default Single<ListSessionsResponse> listSessions(SessionKey sessionKey) {
return listSessions(sessionKey.appName(), sessionKey.userId());
}

/**
* Deletes a specific session.
*
Expand All @@ -134,6 +159,10 @@ Maybe<Session> getSession(
*/
Completable deleteSession(String appName, String userId, String sessionId);

default Completable deleteSession(SessionKey sessionKey) {
return deleteSession(sessionKey.appName(), sessionKey.userId(), sessionKey.id());
}

/**
* Lists the events within a specific session. Supports pagination via the response object.
*
Expand All @@ -147,6 +176,10 @@ Maybe<Session> getSession(
*/
Single<ListEventsResponse> listEvents(String appName, String userId, String sessionId);

default Single<ListEventsResponse> listEvents(SessionKey sessionKey) {
return listEvents(sessionKey.appName(), sessionKey.userId(), sessionKey.id());
}

/**
* Closes a session. This is currently a placeholder and may involve finalizing session state or
* performing cleanup actions in future implementations. The default implementation does nothing.
Expand Down Expand Up @@ -190,20 +223,18 @@ default Single<Event> appendEvent(Session session, Event event) {
EventActions actions = event.actions();
if (actions != null) {
Map<String, Object> stateDelta = actions.stateDelta();
if (stateDelta != null && !stateDelta.isEmpty()) {
Map<String, Object> sessionState = session.state();
if (sessionState != null) {
stateDelta.forEach(
(key, value) -> {
if (!key.startsWith(State.TEMP_PREFIX)) {
if (value == State.REMOVED) {
sessionState.remove(key);
} else {
sessionState.put(key, value);
}
Map<String, Object> sessionState = session.state();
if (stateDelta != null && !stateDelta.isEmpty() && sessionState != null) {
stateDelta.forEach(
(key, value) -> {
if (!key.startsWith(State.TEMP_PREFIX)) {
if (value == State.REMOVED) {
sessionState.remove(key);
} else {
sessionState.put(key, value);
}
});
}
}
});
}
}

Expand Down
22 changes: 22 additions & 0 deletions core/src/main/java/com/google/adk/sessions/Session.java
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,10 @@ public static Builder builder(String id) {
return new Builder(id);
}

public static Builder builder(SessionKey sessionKey) {
return new Builder(sessionKey);
}

/** Builder for {@link Session}. */
public static final class Builder {
private String id;
Expand All @@ -62,6 +66,12 @@ public Builder(String id) {
this.id = id;
}

public Builder(SessionKey sessionKey) {
this.id = sessionKey.id();
this.appName = sessionKey.appName();
this.userId = sessionKey.userId();
}

@JsonCreator
private Builder() {}

Expand All @@ -72,6 +82,14 @@ public Builder id(String id) {
return this;
}

@CanIgnoreReturnValue
public Builder sessionKey(SessionKey sessionKey) {
this.id = sessionKey.id();
this.appName = sessionKey.appName();
this.userId = sessionKey.userId();
return this;
}

@CanIgnoreReturnValue
public Builder state(State state) {
this.state = state;
Expand Down Expand Up @@ -130,6 +148,10 @@ public Session build() {
}
}

public SessionKey sessionKey() {
return new SessionKey(appName, userId, id);
}

@JsonProperty("id")
public String id() {
return id;
Expand Down
82 changes: 82 additions & 0 deletions core/src/main/java/com/google/adk/sessions/SessionKey.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
/*
* Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package com.google.adk.sessions;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.adk.JsonBaseModel;
import java.util.Objects;

/** Key for a session, composed of appName, userId and session id. */
public final class SessionKey extends JsonBaseModel {
private final String appName;
private final String userId;
private final String id;

@JsonCreator
public SessionKey(
@JsonProperty("appName") String appName,
@JsonProperty("userId") String userId,
@JsonProperty("id") String id) {
this.appName = appName;
this.userId = userId;
this.id = id;
}

@JsonProperty("appName")
public String appName() {
return appName;
}

@JsonProperty("userId")
public String userId() {
return userId;
}

@JsonProperty("id")
public String id() {
return id;
}

@Override
public boolean equals(Object o) {
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
SessionKey that = (SessionKey) o;
return Objects.equals(appName, that.appName)
&& Objects.equals(userId, that.userId)
&& Objects.equals(id, that.id);
}

@Override
public int hashCode() {
return Objects.hash(appName, userId, id);
}

@Override
public String toString() {
return toJson();
}

public static SessionKey fromJson(String json) {
return fromJsonString(json, SessionKey.class);
}
}
Loading