Skip to content

feat: Add injectable time and UUID providers #1255

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
DABH wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: Add injectable time and UUID providers #1255

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
6 changes: 4 additions & 2 deletions core/src/main/java/com/google/adk/agents/BaseAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -418,7 +418,8 @@ private Maybe<Event> callCallback(
content -> {
invocationContext.setEndInvocation(true);
return Event.builder()
.id(Event.generateEventId())
.id(invocationContext.newUuid())
.timestamp(invocationContext.now().toEpochMilli())
.invocationId(invocationContext.invocationId())
.author(name())
.branch(invocationContext.branch().orElse(null))
Expand All @@ -435,7 +436,8 @@ private Maybe<Event> callCallback(
if (callbackContext.state().hasDelta()) {
Event.Builder eventBuilder =
Event.builder()
.id(Event.generateEventId())
.id(invocationContext.newUuid())
.timestamp(invocationContext.now().toEpochMilli())
.invocationId(invocationContext.invocationId())
.author(name())
.branch(invocationContext.branch().orElse(null))
Expand Down
71 changes: 67 additions & 4 deletions core/src/main/java/com/google/adk/agents/InvocationContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,17 +21,19 @@
import com.google.adk.artifacts.BaseArtifactService;
Comment thread
DABH marked this conversation as resolved.
Comment thread
DABH marked this conversation as resolved.
import com.google.adk.memory.BaseMemoryService;
import com.google.adk.models.LlmCallsLimitExceededException;
import com.google.adk.platform.TimeProvider;
import com.google.adk.platform.UuidProvider;
import com.google.adk.plugins.Plugin;
import com.google.adk.plugins.PluginManager;
import com.google.adk.sessions.BaseSessionService;
import com.google.adk.sessions.Session;
import com.google.adk.summarizer.EventsCompactionConfig;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.genai.types.Content;
import java.time.Instant;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import org.jspecify.annotations.Nullable;

Expand All @@ -52,6 +54,8 @@ public class InvocationContext {
@Nullable private final ContextCacheConfig contextCacheConfig;
private final InvocationCostManager invocationCostManager;
private final Map<String, Object> callbackContextData;
private final TimeProvider timeProvider;
private final UuidProvider uuidProvider;

@Nullable private String branch;
private BaseAgent agent;
Expand All @@ -78,6 +82,8 @@ protected InvocationContext(Builder builder) {
// invocation invocation so that Plugins can access the same data it during the invocation
// across all types of callbacks.
this.callbackContextData = builder.callbackContextData;
this.timeProvider = builder.timeProvider;
this.uuidProvider = builder.uuidProvider;
}

/** Returns a new {@link Builder} for creating {@link InvocationContext} instances. */
Expand Down Expand Up @@ -192,9 +198,34 @@ public String userId() {
return session.userId();
}

/** Returns the {@link TimeProvider} for this invocation. */
public TimeProvider timeProvider() {
return timeProvider;
}

/** Returns the {@link UuidProvider} for this invocation. */
public UuidProvider uuidProvider() {
return uuidProvider;
}

/** Returns the current time from this invocation's {@link TimeProvider}. */
public Instant now() {
return timeProvider.now();
}

/** Returns a new unique identifier from this invocation's {@link UuidProvider}. */
public String newUuid() {
return uuidProvider.newUuid();
}

/** Generates a new unique ID for an invocation context. */
public static String newInvocationContextId() {
return "e-" + UUID.randomUUID();
return newInvocationContextId(UuidProvider.SYSTEM);
}

/** Generates a new unique ID for an invocation context using the given {@link UuidProvider}. */
public static String newInvocationContextId(UuidProvider uuidProvider) {
return "e-" + uuidProvider.newUuid();
}

/**
Expand Down Expand Up @@ -275,6 +306,8 @@ private Builder(InvocationContext context) {
// invocation invocation so that Plugins can access the same data it during the invocation
// across all types of callbacks.
this.callbackContextData = context.callbackContextData;
this.timeProvider = context.timeProvider;
this.uuidProvider = context.uuidProvider;
}

private BaseSessionService sessionService;
Expand All @@ -294,6 +327,8 @@ private Builder(InvocationContext context) {
@Nullable private ContextCacheConfig contextCacheConfig;
private InvocationCostManager invocationCostManager = new InvocationCostManager();
private Map<String, Object> callbackContextData = new ConcurrentHashMap<>();
private TimeProvider timeProvider = TimeProvider.SYSTEM;
private UuidProvider uuidProvider = UuidProvider.SYSTEM;

/**
* Sets the session service for managing session state.
Expand Down Expand Up @@ -475,6 +510,30 @@ public Builder callbackContextData(Map<String, Object> callbackContextData) {
return this;
}

/**
* Sets the time provider for the invocation. Defaults to {@link TimeProvider#SYSTEM}.
*
* @param timeProvider the provider for the current time.
* @return this builder instance for chaining.
*/
@CanIgnoreReturnValue
public Builder timeProvider(TimeProvider timeProvider) {
this.timeProvider = timeProvider;
return this;
}

/**
* Sets the UUID provider for the invocation. Defaults to {@link UuidProvider#SYSTEM}.
*
* @param uuidProvider the provider for new unique identifiers.
* @return this builder instance for chaining.
*/
@CanIgnoreReturnValue
public Builder uuidProvider(UuidProvider uuidProvider) {
this.uuidProvider = uuidProvider;
return this;
}

/**
* Builds the {@link InvocationContext} instance.
*
Expand Down Expand Up @@ -531,7 +590,9 @@ public boolean equals(Object o) {
&& Objects.equals(eventsCompactionConfig, that.eventsCompactionConfig)
&& Objects.equals(contextCacheConfig, that.contextCacheConfig)
&& Objects.equals(invocationCostManager, that.invocationCostManager)
&& Objects.equals(callbackContextData, that.callbackContextData);
&& Objects.equals(callbackContextData, that.callbackContextData)
&& Objects.equals(timeProvider, that.timeProvider)
&& Objects.equals(uuidProvider, that.uuidProvider);
}

@Override
Expand All @@ -553,6 +614,8 @@ public int hashCode() {
eventsCompactionConfig,
contextCacheConfig,
invocationCostManager,
callbackContextData);
callbackContextData,
timeProvider,
uuidProvider);
}
}
8 changes: 4 additions & 4 deletions core/src/main/java/com/google/adk/events/Event.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
import com.google.adk.JsonBaseModel;
import com.google.adk.platform.TimeProvider;
import com.google.adk.platform.UuidProvider;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
Expand All @@ -34,12 +36,10 @@
import com.google.genai.types.GenerateContentResponseUsageMetadata;
import com.google.genai.types.GroundingMetadata;
import com.google.genai.types.Transcription;
import java.time.Instant;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import org.jspecify.annotations.Nullable;

// TODO - b/413761119 update Agent.java when resolved.
Expand Down Expand Up @@ -74,7 +74,7 @@ public class Event extends JsonBaseModel {
private Event() {}

public static String generateEventId() {
return UUID.randomUUID().toString();
return UuidProvider.SYSTEM.newUuid();
}

/** The event id. */
Expand Down Expand Up @@ -587,7 +587,7 @@ public Event build() {
event.setCustomMetadata(customMetadata);
event.setModelVersion(modelVersion);
event.setActions(actions().orElseGet(() -> EventActions.builder().build()));
event.setTimestamp(timestamp().orElseGet(() -> Instant.now().toEpochMilli()));
event.setTimestamp(timestamp().orElseGet(() -> TimeProvider.SYSTEM.now().toEpochMilli()));
event.setInputTranscription(inputTranscription);
event.setOutputTranscription(outputTranscription);
return event;
Expand Down
20 changes: 12 additions & 8 deletions core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -438,7 +438,7 @@ private Flowable<Event> runOneStep(Context spanContext, InvocationContext contex

final Event mutableEventTemplate =
Event.builder()
.id(Event.generateEventId())
Comment thread
DABH marked this conversation as resolved.
.id(context.newUuid())
.invocationId(context.invocationId())
.author(context.agent().name())
.branch(context.branch().orElse(null))
Expand All @@ -453,15 +453,15 @@ private Flowable<Event> runOneStep(Context spanContext, InvocationContext contex
.doFinally(
() -> {
String oldId = mutableEventTemplate.id();
String newId = Event.generateEventId();
String newId = context.newUuid();
logger.debug("Resetting event ID from {} to {}", oldId, newId);
mutableEventTemplate.setId(newId);
})
.concatMap(
event -> {
// Update event ID for the new resulting events
String oldId = event.id();
String newId = Event.generateEventId();
String newId = context.newUuid();
logger.debug("Resetting event ID from {} to {}", oldId, newId);
event = event.toBuilder().id(newId).build();
Flowable<Event> postProcessedEvents = Flowable.just(event);
Expand Down Expand Up @@ -555,7 +555,7 @@ public Flowable<Event> runLive(InvocationContext invocationContext) {
return Flowable.empty();
}

String eventIdForSendData = Event.generateEventId();
String eventIdForSendData = invocationContext.newUuid();
LlmAgent agent = (LlmAgent) invocationContext.agent();
BaseLlm llm =
agent.resolvedModel().model().isPresent()
Expand Down Expand Up @@ -647,7 +647,7 @@ public void onError(Throwable e) {
.flatMap(
llmResponse -> {
Event baseEventForThisLlmResponse =
liveEventBuilderTemplate.id(Event.generateEventId()).build();
liveEventBuilderTemplate.id(invocationContext.newUuid()).build();
return postprocess(
invocationContext,
baseEventForThisLlmResponse,
Expand Down Expand Up @@ -727,7 +727,7 @@ private Flowable<Event> buildPostprocessingEvents(
}

Event modelResponseEvent =
buildModelResponseEvent(baseEventForLlmResponse, llmRequest, updatedResponse);
buildModelResponseEvent(context, baseEventForLlmResponse, llmRequest, updatedResponse);
if (modelResponseEvent.functionCalls().isEmpty()) {
return processorEvents.concatWith(Flowable.just(modelResponseEvent));
}
Expand Down Expand Up @@ -772,9 +772,13 @@ private void traceCallLlm(
}

private Event buildModelResponseEvent(
Event baseEventForLlmResponse, LlmRequest llmRequest, LlmResponse llmResponse) {
InvocationContext context,
Event baseEventForLlmResponse,
LlmRequest llmRequest,
LlmResponse llmResponse) {
Event.Builder eventBuilder =
baseEventForLlmResponse.toBuilder()
.timestamp(context.now().toEpochMilli())
.content(llmResponse.content().orElse(null))
.partial(llmResponse.partial().orElse(null))
.errorCode(llmResponse.errorCode().orElse(null))
Expand All @@ -794,7 +798,7 @@ private Event buildModelResponseEvent(
logger.debug("event: {} functionCalls: {}", event, event.functionCalls());

if (!event.functionCalls().isEmpty()) {
Functions.populateClientFunctionCallId(event);
Functions.populateClientFunctionCallId(event, context.uuidProvider());
Set<String> longRunningToolIds =
Functions.getLongRunningFunctionCalls(event.functionCalls(), llmRequest.tools());
logger.debug("longRunningToolIds: {}", longRunningToolIds);
Expand Down
6 changes: 6 additions & 0 deletions core/src/main/java/com/google/adk/flows/llmflows/CodeExecution.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,8 @@ private static Flowable<Event> runPreProcessor(
llmRequest.contents().add(codeContent);
Event codeEvent =
Event.builder()
.id(invocationContext.newUuid())
.timestamp(invocationContext.now().toEpochMilli())
.invocationId(invocationContext.invocationId())
.author(llmAgent.name())
.content(codeContent)
Expand Down Expand Up @@ -307,6 +309,8 @@ private static Flowable<Event> runPostProcessor(

Event codeEvent =
Event.builder()
.id(invocationContext.newUuid())
.timestamp(invocationContext.now().toEpochMilli())
.invocationId(invocationContext.invocationId())
.author(llmAgent.name())
.content(responseContent)
Expand Down Expand Up @@ -456,6 +460,8 @@ private static Single<Event> postProcessCodeExecutionResult(
}
eventActionsBuilder.artifactDelta(artifactDelta);
return Event.builder()
.id(invocationContext.newUuid())
.timestamp(invocationContext.now().toEpochMilli())
.invocationId(invocationContext.invocationId())
.author(invocationContext.agent().name())
.content(resultContent)
Expand Down
Loading