Skip to content

Commit 7407e37

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: Implement Trace management, add HITL support
This change also contains such changes: - Extended functionality of callbacks in BigQueryAgentAnalyticsPlugin to capture model response metrics (latency, time to first token) and token usage. - Additional plugging configuration features PiperOrigin-RevId: 894177714
1 parent 9d6cc80 commit 7407e37

14 files changed

Lines changed: 2025 additions & 148 deletions

File tree

a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,14 +28,14 @@
2828
import com.google.adk.agents.Callbacks;
2929
import com.google.adk.agents.InvocationContext;
3030
import com.google.adk.events.Event;
31+
import com.google.adk.utils.AgentEnums.AgentOrigin;
3132
import com.google.common.collect.ImmutableList;
3233
import com.google.errorprone.annotations.CanIgnoreReturnValue;
3334
import com.google.genai.types.Content;
3435
import com.google.genai.types.CustomMetadata;
3536
import com.google.genai.types.Part;
3637
import io.a2a.client.Client;
3738
import io.a2a.client.ClientEvent;
38-
import io.a2a.client.MessageEvent;
3939
import io.a2a.client.TaskEvent;
4040
import io.a2a.client.TaskUpdateEvent;
4141
import io.a2a.spec.A2AClientException;
@@ -541,6 +541,11 @@ protected Flowable<Event> runLiveImpl(InvocationContext invocationContext) {
541541
"runLiveImpl for " + getClass() + " via A2A is not implemented.");
542542
}
543543

544+
@Override
545+
public AgentOrigin toolOrigin() {
546+
return AgentOrigin.A2A;
547+
}
548+
544549
/** Exception thrown when the agent card cannot be resolved. */
545550
public static class AgentCardResolutionError extends RuntimeException {
546551
public AgentCardResolutionError(String message) {

core/src/main/java/com/google/adk/agents/BaseAgent.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import com.google.adk.events.Event;
2626
import com.google.adk.plugins.Plugin;
2727
import com.google.adk.telemetry.Tracing;
28+
import com.google.adk.utils.AgentEnums.AgentOrigin;
2829
import com.google.common.collect.ImmutableList;
2930
import com.google.errorprone.annotations.CanIgnoreReturnValue;
3031
import com.google.errorprone.annotations.DoNotCall;
@@ -256,6 +257,15 @@ public ImmutableList<? extends AfterAgentCallback> afterAgentCallback() {
256257
return afterAgentCallback;
257258
}
258259

260+
/**
261+
* Returns the origin of the tool when this agent is used as a tool.
262+
*
263+
* @return the tool origin, defaults to "BASE_AGENT".
264+
*/
265+
public AgentOrigin toolOrigin() {
266+
return AgentOrigin.BASE_AGENT;
267+
}
268+
259269
/**
260270
* The resolved beforeAgentCallback field as a list.
261271
*

core/src/main/java/com/google/adk/plugins/agentanalytics/BatchProcessor.java

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,12 @@ public void flush() {
147147
} else {
148148
logger.fine("Successfully wrote " + batch.size() + " rows to BigQuery.");
149149
}
150-
} catch (AppendSerializationError ase) {
150+
}
151+
} catch (Exception e) {
152+
if (e instanceof InterruptedException) {
153+
Thread.currentThread().interrupt();
154+
}
155+
if (e.getCause() instanceof AppendSerializationError ase) {
151156
logger.log(
152157
Level.SEVERE, "Failed to write batch to BigQuery due to serialization error", ase);
153158
Map<Integer, String> rowIndexToErrorMessage = ase.getRowIndexToErrorMessage();
@@ -161,12 +166,9 @@ public void flush() {
161166
logger.severe(
162167
"AppendSerializationError occurred, but no row-specific errors were provided.");
163168
}
169+
} else {
170+
logger.log(Level.SEVERE, "Failed to write batch to BigQuery", e);
164171
}
165-
} catch (Exception e) {
166-
if (e instanceof InterruptedException) {
167-
Thread.currentThread().interrupt();
168-
}
169-
logger.log(Level.SEVERE, "Failed to write batch to BigQuery", e);
170172
} finally {
171173
// Clear the vectors to release the memory.
172174
root.clear();
@@ -185,7 +187,12 @@ private void populateVector(FieldVector vector, int index, Object value) {
185187
return;
186188
}
187189
if (vector instanceof VarCharVector varCharVector) {
188-
String strValue = (value instanceof JsonNode jsonNode) ? jsonNode.asText() : value.toString();
190+
String strValue;
191+
if (value instanceof JsonNode jsonNode) {
192+
strValue = jsonNode.isTextual() ? jsonNode.asText() : jsonNode.toString();
193+
} else {
194+
strValue = value.toString();
195+
}
189196
varCharVector.setSafe(index, strValue.getBytes(UTF_8));
190197
} else if (vector instanceof BigIntVector bigIntVector) {
191198
long longValue;

0 commit comments

Comments
 (0)