/*
 * 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.flows.llmflows;

import com.google.adk.agents.BaseAgent;
import com.google.adk.agents.InvocationContext;
import com.google.adk.agents.LlmAgent;
import com.google.adk.events.EventActions;
import com.google.adk.models.LlmRequest;
import com.google.adk.tools.FunctionTool;
import com.google.adk.tools.ToolContext;
import com.google.common.collect.ImmutableList;
import io.reactivex.rxjava3.core.Single;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.List;

/** {@link RequestProcessor} that handles agent transfer for LLM flow. */

public final class AgentTransfer implements RequestProcessor {
  public AgentTransfer() {}
  @Override
  public Single<RequestProcessor.RequestProcessingResult> processRequest(
      InvocationContext context, LlmRequest request) {
    BaseAgent baseAgent = context.agent();
    if (!(baseAgent instanceof LlmAgent agent)) {
      throw new IllegalArgumentException(
          "Base agent in InvocationContext is not an instance of Agent.");
    }
    List<BaseAgent> transferTargets = getTransferTargets(agent);
    if (transferTargets.isEmpty()) {
      return Single.just(
          RequestProcessor.RequestProcessingResult.create(request, ImmutableList.of()));
    }
    LlmRequest.Builder builder =
        request.toBuilder()
            .appendInstructions(
                ImmutableList.of(buildTargetAgentsInstructions(agent, transferTargets)));
    Method transferToAgentMethod;
    try {
      transferToAgentMethod =
          AgentTransfer.class.getMethod("transferToAgent", String.class, ToolContext.class);
    } catch (NoSuchMethodException e) {
      throw new IllegalStateException(e);
    }
    FunctionTool agentTransferTool = FunctionTool.create(transferToAgentMethod);
    agentTransferTool.processLlmRequest(builder, ToolContext.builder(context).build());
    return Single.just(
        RequestProcessor.RequestProcessingResult.create(builder.build(), ImmutableList.of()));
  }
  private String buildTargetAgentsInfo(BaseAgent targetAgent) {
    return String.format(
        "Agent name: %s\nAgent description: %s", targetAgent.name(), targetAgent.description());
  }
  private String buildTargetAgentsInstructions(LlmAgent agent, List<BaseAgent> transferTargets) {
    StringBuilder sb = new StringBuilder();
    sb.append("You have a list of other agents to transfer to:\n");
    for (BaseAgent targetAgent : transferTargets) {
      sb.append(buildTargetAgentsInfo(targetAgent));
      sb.append("\n");
    }
    sb.append(
        "If you are the best to answer the question according to your description, you can answer"
            + " it.\n");
    sb.append(
        "If another agent is better for answering the question according to its description, call"
            + " `transferToAgent` function to transfer the question to that agent. When"
            + " transferring, do not generate any text other than the function call.\n");
    if (agent.parentAgent() != null) {
      sb.append("Your parent agent is ");
      sb.append(agent.parentAgent().name());
      sb.append(
          ".If neither the other agents nor you are best for answering the question according to"
              + " the descriptions, transfer to your parent agent. If you don't have parent agent,"
              + " try answer by yourself.\n");
    }
    return sb.toString();
  }
  private List<BaseAgent> getTransferTargets(LlmAgent agent) {
    List<BaseAgent> transferTargets = new ArrayList<>();
    transferTargets.addAll(agent.subAgents());
    BaseAgent parent = agent.parentAgent();
    if (!(parent instanceof LlmAgent)) {
      return transferTargets;
    }
    if (!agent.disallowTransferToParent()) {
      transferTargets.add(parent);
    }
    if (!agent.disallowTransferToPeers()) {
      for (BaseAgent peerAgent : parent.subAgents()) {
        if (!peerAgent.name().equals(agent.name())) {
          transferTargets.add(peerAgent);
        }
      }
    }
    return transferTargets;
  }
  public static void transferToAgent(String agentName, ToolContext toolContext) {
    EventActions eventActions = toolContext.eventActions();
    toolContext.setActions(eventActions.toBuilder().transferToAgent(agentName).build());
  }
}

/*
 * 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.flows.llmflows;

import com.google.adk.Telemetry;
import com.google.adk.agents.BaseAgent;
import com.google.adk.agents.CallbackContext;
import com.google.adk.agents.Callbacks.AfterModelCallback;
import com.google.adk.agents.Callbacks.BeforeModelCallback;
import com.google.adk.agents.InvocationContext;
import com.google.adk.agents.LiveRequest;
import com.google.adk.agents.LlmAgent;
import com.google.adk.agents.ReadonlyContext;
import com.google.adk.agents.RunConfig.StreamingMode;
import com.google.adk.events.Event;
import com.google.adk.exceptions.LlmCallsLimitExceededException;
import com.google.adk.flows.BaseFlow;
import com.google.adk.flows.llmflows.RequestProcessor.RequestProcessingResult;
import com.google.adk.flows.llmflows.ResponseProcessor.ResponseProcessingResult;
import com.google.adk.models.BaseLlm;
import com.google.adk.models.BaseLlmConnection;
import com.google.adk.models.LlmRegistry;
import com.google.adk.models.LlmRequest;
import com.google.adk.models.LlmResponse;
import com.google.adk.tools.ToolContext;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.genai.types.FunctionResponse;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.StatusCode;
import io.opentelemetry.context.Scope;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import io.reactivex.rxjava3.disposables.Disposable;
import io.reactivex.rxjava3.observers.DisposableCompletableObserver;
import io.reactivex.rxjava3.schedulers.Schedulers;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** A basic flow that calls the LLM in a loop until a final response is generated. */
public abstract class BaseLlmFlow implements BaseFlow {
  private static final Logger logger = LoggerFactory.getLogger(BaseLlmFlow.class);

  protected final List<RequestProcessor> requestProcessors;
  protected final List<ResponseProcessor> responseProcessors;

  // Warning: This is local, in-process state that won't be preserved if the runtime is restarted.
  // "Max steps" is experimental and may evolve in the future (e.g., to support persistence).
  protected int stepsCompleted = 0;
  protected final int maxSteps;

  public BaseLlmFlow(
      List<RequestProcessor> requestProcessors, List<ResponseProcessor> responseProcessors) {
    this(requestProcessors, responseProcessors, /* maxSteps= */ Optional.empty());
  }

  public BaseLlmFlow(
      List<RequestProcessor> requestProcessors,
      List<ResponseProcessor> responseProcessors,
      Optional<Integer> maxSteps) {
    this.requestProcessors = requestProcessors;
    this.responseProcessors = responseProcessors;
    this.maxSteps = maxSteps.orElse(Integer.MAX_VALUE);
  }

  /**
   * Pre-processes the LLM request before sending it to the LLM. Executes all registered {@link
   * RequestProcessor}.
   */
  protected Single<RequestProcessingResult> preprocess(
      InvocationContext context, LlmRequest llmRequest) {

    List<Iterable<Event>> eventIterables = new ArrayList<>();
    LlmAgent agent = (LlmAgent) context.agent();

    Single<LlmRequest> currentLlmRequest = Single.just(llmRequest);
    for (RequestProcessor processor : requestProcessors) {
      currentLlmRequest =
          currentLlmRequest
              .flatMap(request -> processor.processRequest(context, request))
              .doOnSuccess(
                  result -> {
                    if (result.events() != null) {
                      eventIterables.add(result.events());
                    }
                  })
              .map(RequestProcessingResult::updatedRequest);
    }

    return currentLlmRequest.flatMap(
        processedRequest -> {
          LlmRequest.Builder updatedRequestBuilder = processedRequest.toBuilder();

          return agent
              .canonicalTools(new ReadonlyContext(context))
              .concatMapCompletable(
                  tool ->
                      tool.processLlmRequest(
                          updatedRequestBuilder, ToolContext.builder(context).build()))
              .andThen(
                  Single.fromCallable(
                      () -> {
                        Iterable<Event> combinedEvents = Iterables.concat(eventIterables);
                        return RequestProcessingResult.create(
                            updatedRequestBuilder.build(), combinedEvents);
                      }));
        });
  }

  /**
   * Post-processes the LLM response after receiving it from the LLM. Executes all registered {@link
   * ResponseProcessor} instances. Handles function calls if present in the response.
   */
  protected Single<ResponseProcessingResult> postprocess(
      InvocationContext context,
      Event baseEventForLlmResponse,
      LlmRequest llmRequest,
      LlmResponse llmResponse) {

    List<Iterable<Event>> eventIterables = new ArrayList<>();
    Single<LlmResponse> currentLlmResponse = Single.just(llmResponse);
    for (ResponseProcessor processor : responseProcessors) {
      currentLlmResponse =
          currentLlmResponse
              .flatMap(response -> processor.processResponse(context, response))
              .doOnSuccess(
                  result -> {
                    if (result.events() != null) {
                      eventIterables.add(result.events());
                    }
                  })
              .map(ResponseProcessingResult::updatedResponse);
    }

    return currentLlmResponse.flatMap(
        updatedResponse -> {
          if (updatedResponse.content().isEmpty()
              && updatedResponse.errorCode().isEmpty()
              && !updatedResponse.interrupted().orElse(false)
              && !updatedResponse.turnComplete().orElse(false)) {
            return Single.just(
                ResponseProcessingResult.create(
                    updatedResponse, Iterables.concat(eventIterables), Optional.empty()));
          }

          Event modelResponseEvent =
              buildModelResponseEvent(baseEventForLlmResponse, llmRequest, updatedResponse);
          eventIterables.add(Collections.singleton(modelResponseEvent));

          Maybe<Event> maybeFunctionCallEvent =
              modelResponseEvent.functionCalls().isEmpty()
                  ? Maybe.empty()
                  : Functions.handleFunctionCalls(context, modelResponseEvent, llmRequest.tools());

          return maybeFunctionCallEvent
              .map(Optional::of)
              .defaultIfEmpty(Optional.empty())
              .map(
                  functionCallEventOpt -> {
                    Optional<String> transferToAgent = Optional.empty();
                    if (functionCallEventOpt.isPresent()) {
                      Event functionCallEvent = functionCallEventOpt.get();
                      eventIterables.add(Collections.singleton(functionCallEvent));
                      transferToAgent = functionCallEvent.actions().transferToAgent();
                    }
                    Iterable<Event> combinedEvents = Iterables.concat(eventIterables);
                    return ResponseProcessingResult.create(
                        updatedResponse, combinedEvents, transferToAgent);
                  });
        });
  }

  /**
   * Sends a request to the LLM and returns its response.
   *
   * @param context The invocation context.
   * @param llmRequest The LLM request.
   * @param eventForCallbackUsage An Event object primarily for providing context (like actions) to
   *     callbacks. Callbacks should not rely on its ID if they create their own separate events.
   */
  private Flowable<LlmResponse> callLlm(
      InvocationContext context, LlmRequest llmRequest, Event eventForCallbackUsage) {
    LlmAgent agent = (LlmAgent) context.agent();

    return handleBeforeModelCallback(context, llmRequest, eventForCallbackUsage)
        .flatMapPublisher(
            beforeResponse -> {
              if (beforeResponse.isPresent()) {
                return Flowable.just(beforeResponse.get());
              }
              BaseLlm llm =
                  agent.resolvedModel().model().isPresent()
                      ? agent.resolvedModel().model().get()
                      : LlmRegistry.getLlm(agent.resolvedModel().modelName().get());
              return Flowable.defer(
                      () -> {
                        Span llmCallSpan =
                            Telemetry.getTracer().spanBuilder("call_llm").startSpan();

                        try (Scope scope = llmCallSpan.makeCurrent()) {
                          return llm.generateContent(
                                  llmRequest,
                                  context.runConfig().streamingMode() == StreamingMode.SSE)
                              .doOnNext(
                                  llmResp -> {
                                    try (Scope innerScope = llmCallSpan.makeCurrent()) {
                                      Telemetry.traceCallLlm(
                                          context, eventForCallbackUsage.id(), llmRequest, llmResp);
                                    }
                                  })
                              .doOnError(
                                  error -> {
                                    llmCallSpan.setStatus(StatusCode.ERROR, error.getMessage());
                                    llmCallSpan.recordException(error);
                                  })
                              .doFinally(llmCallSpan::end);
                        }
                      })
                  .concatMap(
                      llmResp ->
                          handleAfterModelCallback(context, llmResp, eventForCallbackUsage)
                              .toFlowable());
            });
  }

  /**
   * Invokes {@link BeforeModelCallback}s. If any returns a response, it's used instead of calling
   * the LLM.
   *
   * @return A {@link Single} with the callback result or {@link Optional#empty()}.
   */
  private Single<Optional<LlmResponse>> handleBeforeModelCallback(
      InvocationContext context, LlmRequest llmRequest, Event modelResponseEvent) {
    LlmAgent agent = (LlmAgent) context.agent();

    Optional<List<BeforeModelCallback>> callbacksOpt = agent.beforeModelCallback();
    if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {
      return Single.just(Optional.empty());
    }

    Event callbackEvent = modelResponseEvent.toBuilder().build();
    List<BeforeModelCallback> callbacks = callbacksOpt.get();

    return Flowable.fromIterable(callbacks)
        .concatMapSingle(
            callback -> {
              CallbackContext callbackContext =
                  new CallbackContext(context, callbackEvent.actions());
              return callback
                  .call(callbackContext, llmRequest)
                  .map(Optional::of)
                  .defaultIfEmpty(Optional.empty());
            })
        .filter(Optional::isPresent)
        .firstElement()
        .switchIfEmpty(Single.just(Optional.empty()));
  }

  /**
   * Invokes {@link AfterModelCallback}s after an LLM response. If any returns a response, it
   * replaces the original.
   *
   * @return A {@link Single} with the final {@link LlmResponse}.
   */
  private Single<LlmResponse> handleAfterModelCallback(
      InvocationContext context, LlmResponse llmResponse, Event modelResponseEvent) {
    LlmAgent agent = (LlmAgent) context.agent();
    Optional<List<AfterModelCallback>> callbacksOpt = agent.afterModelCallback();

    if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {
      return Single.just(llmResponse);
    }

    Event callbackEvent = modelResponseEvent.toBuilder().content(llmResponse.content()).build();
    List<AfterModelCallback> callbacks = callbacksOpt.get();

    return Flowable.fromIterable(callbacks)
        .concatMapSingle(
            callback -> {
              CallbackContext callbackContext =
                  new CallbackContext(context, callbackEvent.actions());
              return callback
                  .call(callbackContext, llmResponse)
                  .map(Optional::of)
                  .defaultIfEmpty(Optional.empty());
            })
        .filter(Optional::isPresent)
        .firstElement()
        .map(Optional::get)
        .switchIfEmpty(Single.just(llmResponse));
  }

  /**
   * Executes a single iteration of the LLM flow: preprocessing → LLM call → postprocessing.
   *
   *
Handles early termination, LLM call limits, and agent transfer if needed.
   *
   * @return A {@link Flowable} of {@link Event} objects from this step.
   * @throws LlmCallsLimitExceededException if the agent exceeds allowed LLM invocations.
   * @throws IllegalStateException if a transfer agent is specified but not found.
   */
  private Flowable<Event> runOneStep(InvocationContext context) {
    LlmRequest initialLlmRequest = LlmRequest.builder().build();

    return preprocess(context, initialLlmRequest)
        .flatMapPublisher(
            preResult -> {
              LlmRequest llmRequestAfterPreprocess = preResult.updatedRequest();
              Iterable<Event> preEvents = preResult.events();

              if (context.endInvocation()) {
                logger.debug("End invocation requested during preprocessing.");
                return Flowable.fromIterable(preEvents);
              }

              try {
                context.incrementLlmCallsCount();
              } catch (LlmCallsLimitExceededException e) {
                logger.error("LLM calls limit exceeded.", e);
                return Flowable.fromIterable(preEvents).concatWith(Flowable.error(e));
              }

              final Event mutableEventTemplate =
                  Event.builder()
                      .id(Event.generateEventId())
                      .invocationId(context.invocationId())
                      .author(context.agent().name())
                      .branch(context.branch())
                      .build();

              Flowable<Event> restOfFlow =
                  callLlm(context, llmRequestAfterPreprocess, mutableEventTemplate)
                      .concatMap(
                          llmResponse -> {
                            Single<ResponseProcessingResult> postResultSingle =
                                postprocess(
                                    context,
                                    mutableEventTemplate,
                                    llmRequestAfterPreprocess,
                                    llmResponse);

                            return postResultSingle
                                .doOnSuccess(
                                    ignored -> {
                                      String oldId = mutableEventTemplate.id();
                                      mutableEventTemplate.setId(Event.generateEventId());
                                      logger.debug(
                                          "Updated mutableEventTemplate ID from {} to {} for next"
                                              + " LlmResponse",
                                          oldId,
                                          mutableEventTemplate.id());
                                    })
                                .toFlowable();
                          })
                      .concatMap(
                          postResult -> {
                            Flowable<Event> postProcessedEvents =
                                Flowable.fromIterable(postResult.events());
                            if (postResult.transferToAgent().isPresent()) {
                              String agentToTransfer = postResult.transferToAgent().get();
                              logger.debug("Transferring to agent: {}", agentToTransfer);
                              BaseAgent rootAgent = context.agent().rootAgent();
                              BaseAgent nextAgent = rootAgent.findAgent(agentToTransfer);
                              if (nextAgent == null) {
                                String errorMsg =
                                    "Agent not found for transfer: " + agentToTransfer;
                                logger.error(errorMsg);
                                return postProcessedEvents.concatWith(
                                    Flowable.error(new IllegalStateException(errorMsg)));
                              }
                              return postProcessedEvents.concatWith(
                                  Flowable.defer(() -> nextAgent.runAsync(context)));
                            }
                            return postProcessedEvents;
                          });

              return restOfFlow.startWithIterable(preEvents);
            });
  }

  /**
   * Executes the full LLM flow by repeatedly calling {@link #runOneStep} until a final response is
   * produced.
   *
   * @return A {@link Flowable} of all {@link Event}s generated during the flow.
   */
  @Override
  public Flowable<Event> run(InvocationContext invocationContext) {
    Flowable<Event> currentStepEvents = runOneStep(invocationContext).cache();
    if (++stepsCompleted >= maxSteps) {
      logger.debug("Ending flow execution because max steps reached.");
      return currentStepEvents;
    }

    return currentStepEvents.concatWith(
        currentStepEvents
            .toList()
            .flatMapPublisher(
                eventList -> {
                  if (eventList.isEmpty()
                      || Iterables.getLast(eventList).finalResponse()
                      || Iterables.getLast(eventList).actions().endInvocation().orElse(false)) {
                    logger.debug(
                        "Ending flow execution based on final response, endInvocation action or"
                            + " empty event list.");
                    return Flowable.empty();
                  } else {
                    logger.debug("Continuing to next step of the flow.");
                    return Flowable.defer(() -> run(invocationContext));
                  }
                }));
  }

  /**
   * Executes the LLM flow in streaming mode.
   *
   *
Handles sending history and live requests to the LLM, receiving responses, processing them,
   * and managing agent transfers.
   *
   * @return A {@link Flowable} of {@link Event}s streamed in real-time.
   */
  @Override
  public Flowable<Event> runLive(InvocationContext invocationContext) {
    LlmRequest llmRequest = LlmRequest.builder().build();

    return preprocess(invocationContext, llmRequest)
        .flatMapPublisher(
            preResult -> {
              LlmRequest llmRequestAfterPreprocess = preResult.updatedRequest();
              if (invocationContext.endInvocation()) {
                return Flowable.fromIterable(preResult.events());
              }

              String eventIdForSendData = Event.generateEventId();
              LlmAgent agent = (LlmAgent) invocationContext.agent();
              BaseLlm llm =
                  agent.resolvedModel().model().isPresent()
                      ? agent.resolvedModel().model().get()
                      : LlmRegistry.getLlm(agent.resolvedModel().modelName().get());
              BaseLlmConnection connection = llm.connect(llmRequestAfterPreprocess);
              Completable historySent =
                  llmRequestAfterPreprocess.contents().isEmpty()
                      ? Completable.complete()
                      : Completable.defer(
                          () -> {
                            Span sendDataSpan =
                                Telemetry.getTracer().spanBuilder("send_data").startSpan();
                            try (Scope scope = sendDataSpan.makeCurrent()) {
                              return connection
                                  .sendHistory(llmRequestAfterPreprocess.contents())
                                  .doOnComplete(
                                      () -> {
                                        try (Scope innerScope = sendDataSpan.makeCurrent()) {
                                          Telemetry.traceSendData(
                                              invocationContext,
                                              eventIdForSendData,
                                              llmRequestAfterPreprocess.contents());
                                        }
                                      })
                                  .doOnError(
                                      error -> {
                                        sendDataSpan.setStatus(
                                            StatusCode.ERROR, error.getMessage());
                                        sendDataSpan.recordException(error);
                                        try (Scope innerScope = sendDataSpan.makeCurrent()) {
                                          Telemetry.traceSendData(
                                              invocationContext,
                                              eventIdForSendData,
                                              llmRequestAfterPreprocess.contents());
                                        }
                                      })
                                  .doFinally(sendDataSpan::end);
                            }
                          });

              Flowable<LiveRequest> liveRequests = invocationContext.liveRequestQueue().get().get();
              Disposable sendTask =
                  historySent
                      .observeOn(agent.executor().map(Schedulers::from).orElse(Schedulers.io()))
                      .andThen(
                          liveRequests
                              .onBackpressureBuffer()
                              .concatMapCompletable(
                                  request -> {
                                    if (request.content().isPresent()) {
                                      return connection.sendContent(request.content().get());
                                    } else if (request.blob().isPresent()) {
                                      return connection.sendRealtime(request.blob().get());
                                    }
                                    return Completable.fromAction(connection::close);
                                  }))
                      .subscribeWith(
                          new DisposableCompletableObserver() {
                            @Override
                            public void onComplete() {
                              connection.close();
                            }

                            @Override
                            public void onError(Throwable e) {
                              connection.close(e);
                            }
                          });

              Event.Builder liveEventBuilderTemplate =
                  Event.builder()
                      .invocationId(invocationContext.invocationId())
                      .author(invocationContext.agent().name())
                      .branch(invocationContext.branch());

              Flowable<Event> receiveFlow =
                  connection
                      .receive()
                      .flatMapSingle(
                          llmResponse -> {
                            Event baseEventForThisLlmResponse =
                                liveEventBuilderTemplate.id(Event.generateEventId()).build();
                            return postprocess(
                                invocationContext,
                                baseEventForThisLlmResponse,
                                llmRequestAfterPreprocess,
                                llmResponse);
                          })
                      .flatMap(
                          postResult -> {
                            Flowable<Event> events = Flowable.fromIterable(postResult.events());
                            if (postResult.transferToAgent().isPresent()) {
                              BaseAgent rootAgent = invocationContext.agent().rootAgent();
                              BaseAgent nextAgent =
                                  rootAgent.findAgent(postResult.transferToAgent().get());
                              if (nextAgent == null) {
                                throw new IllegalStateException(
                                    "Agent not found: " + postResult.transferToAgent().get());
                              }
                              Flowable<Event> nextAgentEvents =
                                  nextAgent.runLive(invocationContext);
                              events = Flowable.concat(events, nextAgentEvents);
                            }
                            return events;
                          })
                      .doOnNext(
                          event -> {
                            ImmutableList<FunctionResponse> functionResponses =
                                event.functionResponses();
                            if (!functionResponses.isEmpty()) {
                              invocationContext
                                  .liveRequestQueue()
                                  .get()
                                  .content(event.content().get());
                            }
                            if (functionResponses.stream()
                                    .anyMatch(
                                        functionResponse ->
                                            functionResponse
                                                .name()
                                                .orElse("")
                                                .equals("transferToAgent"))
                                || event.actions().endInvocation().orElse(false)) {
                              sendTask.dispose();
                              connection.close();
                            }
                          });

              return receiveFlow
                  .takeWhile(event -> !event.actions().endInvocation().orElse(false))
                  .startWithIterable(preResult.events());
            });
  }

  /**
   * Builds an {@link Event} from LLM response, request, and base event data.
   *
   *
Populates the event with LLM output and tool function call metadata.
   *
   * @return A fully constructed {@link Event} representing the LLM response.
   */
  private Event buildModelResponseEvent(
      Event baseEventForLlmResponse, LlmRequest llmRequest, LlmResponse llmResponse) {
    Event.Builder eventBuilder =
        baseEventForLlmResponse.toBuilder()
            .content(llmResponse.content())
            .partial(llmResponse.partial())
            .errorCode(llmResponse.errorCode())
            .errorMessage(llmResponse.errorMessage())
            .interrupted(llmResponse.interrupted())
            .turnComplete(llmResponse.turnComplete())
            .groundingMetadata(llmResponse.groundingMetadata());

    Event event = eventBuilder.build();

    if (!event.functionCalls().isEmpty()) {
      Functions.populateClientFunctionCallId(event);
      Set<String> longRunningToolIds =
          Functions.getLongRunningFunctionCalls(event.functionCalls(), llmRequest.tools());
      if (!longRunningToolIds.isEmpty()) {
        event.setLongRunningToolIds(Optional.of(longRunningToolIds));
      }
    }
    return event;
  }
}

/*
 * 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.agents;

import static java.util.stream.Collectors.joining;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.google.adk.SchemaUtils;
import com.google.adk.agents.Callbacks.AfterAgentCallback;
import com.google.adk.agents.Callbacks.AfterAgentCallbackBase;
import com.google.adk.agents.Callbacks.AfterAgentCallbackSync;
import com.google.adk.agents.Callbacks.AfterModelCallback;
import com.google.adk.agents.Callbacks.AfterModelCallbackBase;
import com.google.adk.agents.Callbacks.AfterModelCallbackSync;
import com.google.adk.agents.Callbacks.AfterToolCallback;
import com.google.adk.agents.Callbacks.AfterToolCallbackBase;
import com.google.adk.agents.Callbacks.AfterToolCallbackSync;
import com.google.adk.agents.Callbacks.BeforeAgentCallback;
import com.google.adk.agents.Callbacks.BeforeAgentCallbackBase;
import com.google.adk.agents.Callbacks.BeforeAgentCallbackSync;
import com.google.adk.agents.Callbacks.BeforeModelCallback;
import com.google.adk.agents.Callbacks.BeforeModelCallbackBase;
import com.google.adk.agents.Callbacks.BeforeModelCallbackSync;
import com.google.adk.agents.Callbacks.BeforeToolCallback;
import com.google.adk.agents.Callbacks.BeforeToolCallbackBase;
import com.google.adk.agents.Callbacks.BeforeToolCallbackSync;
import com.google.adk.events.Event;
import com.google.adk.examples.BaseExampleProvider;
import com.google.adk.examples.Example;
import com.google.adk.flows.llmflows.AutoFlow;
import com.google.adk.flows.llmflows.BaseLlmFlow;
import com.google.adk.flows.llmflows.SingleFlow;
import com.google.adk.models.BaseLlm;
import com.google.adk.models.Model;
import com.google.adk.tools.BaseTool;
import com.google.adk.tools.BaseToolset;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.genai.types.Content;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.Schema;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.Executor;
import javax.annotation.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** The LLM-based agent. */
public class LlmAgent extends BaseAgent {

  private static final Logger logger = LoggerFactory.getLogger(LlmAgent.class);

  /**
   * Enum to define if contents of previous events should be included in requests to the underlying
   * LLM.
   */
  public enum IncludeContents {
    DEFAULT,
    NONE
  }

  private final Optional<Model> model;
  private final Instruction instruction;
  private final Instruction globalInstruction;
  private final List<Object> toolsUnion;
  private final Optional<GenerateContentConfig> generateContentConfig;
  private final Optional<BaseExampleProvider> exampleProvider;
  private final IncludeContents includeContents;

  private final boolean planning;
  private final Optional<Integer> maxSteps;
  private final boolean disallowTransferToParent;
  private final boolean disallowTransferToPeers;
  private final Optional<List<BeforeModelCallback>> beforeModelCallback;
  private final Optional<List<AfterModelCallback>> afterModelCallback;
  private final Optional<List<BeforeToolCallback>> beforeToolCallback;
  private final Optional<List<AfterToolCallback>> afterToolCallback;
  private final Optional<Schema> inputSchema;
  private final Optional<Schema> outputSchema;
  private final Optional<Executor> executor;
  private final Optional<String> outputKey;

  private volatile Model resolvedModel;
  private final BaseLlmFlow llmFlow;

  protected LlmAgent(Builder builder) {
    super(
        builder.name,
        builder.description,
        builder.subAgents,
        builder.beforeAgentCallback,
        builder.afterAgentCallback);
    this.model = Optional.ofNullable(builder.model);
    this.instruction =
        builder.instruction == null ? new Instruction.Static("") : builder.instruction;
    this.globalInstruction =
        builder.globalInstruction == null ? new Instruction.Static("") : builder.globalInstruction;
    this.generateContentConfig = Optional.ofNullable(builder.generateContentConfig);
    this.exampleProvider = Optional.ofNullable(builder.exampleProvider);
    this.includeContents =
        builder.includeContents != null ? builder.includeContents : IncludeContents.DEFAULT;
    this.planning = builder.planning != null && builder.planning;
    this.maxSteps = Optional.ofNullable(builder.maxSteps);
    this.disallowTransferToParent = builder.disallowTransferToParent;
    this.disallowTransferToPeers = builder.disallowTransferToPeers;
    this.beforeModelCallback = Optional.ofNullable(builder.beforeModelCallback);
    this.afterModelCallback = Optional.ofNullable(builder.afterModelCallback);
    this.beforeToolCallback = Optional.ofNullable(builder.beforeToolCallback);
    this.afterToolCallback = Optional.ofNullable(builder.afterToolCallback);
    this.inputSchema = Optional.ofNullable(builder.inputSchema);
    this.outputSchema = Optional.ofNullable(builder.outputSchema);
    this.executor = Optional.ofNullable(builder.executor);
    this.outputKey = Optional.ofNullable(builder.outputKey);
    this.toolsUnion = builder.toolsUnion != null ? builder.toolsUnion : ImmutableList.of();

    this.llmFlow = determineLlmFlow();

    // Validate name not empty.
    Preconditions.checkArgument(!this.name().isEmpty(), "Agent name cannot be empty.");
  }

  /** Returns a {@link Builder} for {@link LlmAgent}. */
  public static Builder builder() {
    return new Builder();
  }

  /** Builder for {@link LlmAgent}. */
  public static class Builder {
    private String name;
    private String description;

    private Model model;

    private Instruction instruction;
    private Instruction globalInstruction;
    private ImmutableList<BaseAgent> subAgents;
    private ImmutableList<Object> toolsUnion;
    private GenerateContentConfig generateContentConfig;
    private BaseExampleProvider exampleProvider;
    private IncludeContents includeContents;
    private Boolean planning;
    private Integer maxSteps;
    private Boolean disallowTransferToParent;
    private Boolean disallowTransferToPeers;
    private ImmutableList<BeforeModelCallback> beforeModelCallback;
    private ImmutableList<AfterModelCallback> afterModelCallback;
    private ImmutableList<BeforeAgentCallback> beforeAgentCallback;
    private ImmutableList<AfterAgentCallback> afterAgentCallback;
    private ImmutableList<BeforeToolCallback> beforeToolCallback;
    private ImmutableList<AfterToolCallback> afterToolCallback;
    private Schema inputSchema;
    private Schema outputSchema;
    private Executor executor;
    private String outputKey;

    @CanIgnoreReturnValue
    public Builder name(String name) {
      this.name = name;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder description(String description) {
      this.description = description;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder model(String model) {
      this.model = Model.builder().modelName(model).build();
      return this;
    }

    @CanIgnoreReturnValue
    public Builder model(BaseLlm model) {
      this.model = Model.builder().model(model).build();
      return this;
    }

    @CanIgnoreReturnValue
    public Builder instruction(Instruction instruction) {
      this.instruction = instruction;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder instruction(String instruction) {
      this.instruction = (instruction == null) ? null : new Instruction.Static(instruction);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder globalInstruction(Instruction globalInstruction) {
      this.globalInstruction = globalInstruction;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder globalInstruction(String globalInstruction) {
      this.globalInstruction =
          (globalInstruction == null) ? null : new Instruction.Static(globalInstruction);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder subAgents(List<? extends BaseAgent> subAgents) {
      this.subAgents = ImmutableList.copyOf(subAgents);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder subAgents(BaseAgent... subAgents) {
      this.subAgents = ImmutableList.copyOf(subAgents);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder tools(List<?> tools) {
      this.toolsUnion = ImmutableList.copyOf(tools);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder tools(Object... tools) {
      this.toolsUnion = ImmutableList.copyOf(tools);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder generateContentConfig(GenerateContentConfig generateContentConfig) {
      this.generateContentConfig = generateContentConfig;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder exampleProvider(BaseExampleProvider exampleProvider) {
      this.exampleProvider = exampleProvider;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder exampleProvider(List<Example> examples) {
      this.exampleProvider =
          new BaseExampleProvider() {
            @Override
            public List<Example> getExamples(String query) {
              return examples;
            }
          };
      return this;
    }

    @CanIgnoreReturnValue
    public Builder exampleProvider(Example... examples) {
      this.exampleProvider =
          new BaseExampleProvider() {
            @Override
            public ImmutableList<Example> getExamples(String query) {
              return ImmutableList.copyOf(examples);
            }
          };
      return this;
    }

    @CanIgnoreReturnValue
    public Builder includeContents(IncludeContents includeContents) {
      this.includeContents = includeContents;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder planning(boolean planning) {
      this.planning = planning;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder maxSteps(int maxSteps) {
      this.maxSteps = maxSteps;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder disallowTransferToParent(boolean disallowTransferToParent) {
      this.disallowTransferToParent = disallowTransferToParent;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder disallowTransferToPeers(boolean disallowTransferToPeers) {
      this.disallowTransferToPeers = disallowTransferToPeers;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeModelCallback(BeforeModelCallback beforeModelCallback) {
      this.beforeModelCallback = ImmutableList.of(beforeModelCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeModelCallback(List<BeforeModelCallbackBase> beforeModelCallback) {
      if (beforeModelCallback == null) {
        this.beforeModelCallback = null;
      } else if (beforeModelCallback.isEmpty()) {
        this.beforeModelCallback = ImmutableList.of();
      } else {
        ImmutableList.Builder<BeforeModelCallback> builder = ImmutableList.builder();
        for (BeforeModelCallbackBase callback : beforeModelCallback) {
          if (callback instanceof BeforeModelCallback beforeModelCallbackInstance) {
            builder.add(beforeModelCallbackInstance);
          } else if (callback instanceof BeforeModelCallbackSync beforeModelCallbackSyncInstance) {
            builder.add(
                (BeforeModelCallback)
                    (callbackContext, llmRequest) ->
                        Maybe.fromOptional(
                            beforeModelCallbackSyncInstance.call(callbackContext, llmRequest)));
          } else {
            logger.warn(
                "Invalid beforeModelCallback callback type: %s. Ignoring this callback.",
                callback.getClass().getName());
          }
        }
        this.beforeModelCallback = builder.build();
      }

      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeModelCallbackSync(BeforeModelCallbackSync beforeModelCallbackSync) {
      this.beforeModelCallback =
          ImmutableList.of(
              (callbackContext, llmRequest) ->
                  Maybe.fromOptional(beforeModelCallbackSync.call(callbackContext, llmRequest)));
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterModelCallback(AfterModelCallback afterModelCallback) {
      this.afterModelCallback = ImmutableList.of(afterModelCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterModelCallback(List<AfterModelCallbackBase> afterModelCallback) {
      if (afterModelCallback == null) {
        this.afterModelCallback = null;
      } else if (afterModelCallback.isEmpty()) {
        this.afterModelCallback = ImmutableList.of();
      } else {
        ImmutableList.Builder<AfterModelCallback> builder = ImmutableList.builder();
        for (AfterModelCallbackBase callback : afterModelCallback) {
          if (callback instanceof AfterModelCallback afterModelCallbackInstance) {
            builder.add(afterModelCallbackInstance);
          } else if (callback instanceof AfterModelCallbackSync afterModelCallbackSyncInstance) {
            builder.add(
                (AfterModelCallback)
                    (callbackContext, llmResponse) ->
                        Maybe.fromOptional(
                            afterModelCallbackSyncInstance.call(callbackContext, llmResponse)));
          } else {
            logger.warn(
                "Invalid afterModelCallback callback type: %s. Ignoring this callback.",
                callback.getClass().getName());
          }
        }
        this.afterModelCallback = builder.build();
      }

      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterModelCallbackSync(AfterModelCallbackSync afterModelCallbackSync) {
      this.afterModelCallback =
          ImmutableList.of(
              (callbackContext, llmResponse) ->
                  Maybe.fromOptional(afterModelCallbackSync.call(callbackContext, llmResponse)));
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeAgentCallback(BeforeAgentCallback beforeAgentCallback) {
      this.beforeAgentCallback = ImmutableList.of(beforeAgentCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeAgentCallback(List<BeforeAgentCallbackBase> beforeAgentCallback) {
      this.beforeAgentCallback = CallbackUtil.getBeforeAgentCallbacks(beforeAgentCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeAgentCallbackSync(BeforeAgentCallbackSync beforeAgentCallbackSync) {
      this.beforeAgentCallback =
          ImmutableList.of(
              (callbackContext) ->
                  Maybe.fromOptional(beforeAgentCallbackSync.call(callbackContext)));
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterAgentCallback(AfterAgentCallback afterAgentCallback) {
      this.afterAgentCallback = ImmutableList.of(afterAgentCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterAgentCallback(List<AfterAgentCallbackBase> afterAgentCallback) {
      this.afterAgentCallback = CallbackUtil.getAfterAgentCallbacks(afterAgentCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterAgentCallbackSync(AfterAgentCallbackSync afterAgentCallbackSync) {
      this.afterAgentCallback =
          ImmutableList.of(
              (callbackContext) ->
                  Maybe.fromOptional(afterAgentCallbackSync.call(callbackContext)));
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeToolCallback(BeforeToolCallback beforeToolCallback) {
      this.beforeToolCallback = ImmutableList.of(beforeToolCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeToolCallback(@Nullable List<BeforeToolCallbackBase> beforeToolCallbacks) {
      if (beforeToolCallbacks == null) {
        this.beforeToolCallback = null;
      } else if (beforeToolCallbacks.isEmpty()) {
        this.beforeToolCallback = ImmutableList.of();
      } else {
        ImmutableList.Builder<BeforeToolCallback> builder = ImmutableList.builder();
        for (BeforeToolCallbackBase callback : beforeToolCallbacks) {
          if (callback instanceof BeforeToolCallback beforeToolCallbackInstance) {
            builder.add(beforeToolCallbackInstance);
          } else if (callback instanceof BeforeToolCallbackSync beforeToolCallbackSyncInstance) {
            builder.add(
                (invocationContext, baseTool, input, toolContext) ->
                    Maybe.fromOptional(
                        beforeToolCallbackSyncInstance.call(
                            invocationContext, baseTool, input, toolContext)));
          } else {
            logger.warn(
                "Invalid beforeToolCallback callback type: {}. Ignoring this callback.",
                callback.getClass().getName());
          }
        }
        this.beforeToolCallback = builder.build();
      }
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeToolCallbackSync(BeforeToolCallbackSync beforeToolCallbackSync) {
      this.beforeToolCallback =
          ImmutableList.of(
              (invocationContext, baseTool, input, toolContext) ->
                  Maybe.fromOptional(
                      beforeToolCallbackSync.call(
                          invocationContext, baseTool, input, toolContext)));
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterToolCallback(AfterToolCallback afterToolCallback) {
      this.afterToolCallback = ImmutableList.of(afterToolCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterToolCallback(@Nullable List<AfterToolCallbackBase> afterToolCallbacks) {
      if (afterToolCallbacks == null) {
        this.afterToolCallback = null;
      } else if (afterToolCallbacks.isEmpty()) {
        this.afterToolCallback = ImmutableList.of();
      } else {
        ImmutableList.Builder<AfterToolCallback> builder = ImmutableList.builder();
        for (AfterToolCallbackBase callback : afterToolCallbacks) {
          if (callback instanceof AfterToolCallback afterToolCallbackInstance) {
            builder.add(afterToolCallbackInstance);
          } else if (callback instanceof AfterToolCallbackSync afterToolCallbackSyncInstance) {
            builder.add(
                (invocationContext, baseTool, input, toolContext, response) ->
                    Maybe.fromOptional(
                        afterToolCallbackSyncInstance.call(
                            invocationContext, baseTool, input, toolContext, response)));
          } else {
            logger.warn(
                "Invalid afterToolCallback callback type: {}. Ignoring this callback.",
                callback.getClass().getName());
          }
        }
        this.afterToolCallback = builder.build();
      }
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterToolCallbackSync(AfterToolCallbackSync afterToolCallbackSync) {
      this.afterToolCallback =
          ImmutableList.of(
              (invocationContext, baseTool, input, toolContext, response) ->
                  Maybe.fromOptional(
                      afterToolCallbackSync.call(
                          invocationContext, baseTool, input, toolContext, response)));
      return this;
    }

    @CanIgnoreReturnValue
    public Builder inputSchema(Schema inputSchema) {
      this.inputSchema = inputSchema;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder outputSchema(Schema outputSchema) {
      this.outputSchema = outputSchema;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder executor(Executor executor) {
      this.executor = executor;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder outputKey(String outputKey) {
      this.outputKey = outputKey;
      return this;
    }

    protected void validate() {
      this.disallowTransferToParent =
          this.disallowTransferToParent != null && this.disallowTransferToParent;
      this.disallowTransferToPeers =
          this.disallowTransferToPeers != null && this.disallowTransferToPeers;

      if (this.outputSchema != null) {
        if (!this.disallowTransferToParent || !this.disallowTransferToPeers) {
          System.err.println(
              "Warning: Invalid config for agent "
                  + this.name
                  + ": outputSchema cannot co-exist with agent transfer"
                  + " configurations. Setting disallowTransferToParent=true and"
                  + " disallowTransferToPeers=true.");
          this.disallowTransferToParent = true;
          this.disallowTransferToPeers = true;
        }

        if (this.subAgents != null && !this.subAgents.isEmpty()) {
          throw new IllegalArgumentException(
              "Invalid config for agent "
                  + this.name
                  + ": if outputSchema is set, subAgents must be empty to disable agent"
                  + " transfer.");
        }
        if (this.toolsUnion != null && !this.toolsUnion.isEmpty()) {
          throw new IllegalArgumentException(
              "Invalid config for agent "
                  + this.name
                  + ": if outputSchema is set, tools must be empty.");
        }
      }
    }

    public LlmAgent build() {
      validate();
      return new LlmAgent(this);
    }
  }

  protected BaseLlmFlow determineLlmFlow() {
    if (disallowTransferToParent() && disallowTransferToPeers() && subAgents().isEmpty()) {
      return new SingleFlow(maxSteps);
    } else {
      return new AutoFlow(maxSteps);
    }
  }

  private void maybeSaveOutputToState(Event event) {
    if (outputKey().isPresent() && event.finalResponse() && event.content().isPresent()) {
      // Concatenate text from all parts.
      Object output;
      String rawResult =
          event.content().flatMap(Content::parts).orElse(ImmutableList.of()).stream()
              .map(part -> part.text().orElse(""))
              .collect(joining());

      Optional<Schema> outputSchema = outputSchema();
      if (outputSchema.isPresent()) {
        try {
          Map<String, Object> validatedMap =
              SchemaUtils.validateOutputSchema(rawResult, outputSchema.get());
          output = validatedMap;
        } catch (JsonProcessingException e) {
          System.err.println(
              "Error: LlmAgent output for outputKey '"
                  + outputKey().get()
                  + "' was not valid JSON, despite an outputSchema being present."
                  + " Saving raw output to state. Error: "
                  + e.getMessage());
          output = rawResult;
        } catch (IllegalArgumentException e) {
          System.err.println(
              "Error: LlmAgent output for outputKey '"
                  + outputKey().get()
                  + "' did not match the outputSchema."
                  + " Saving raw output to state. Error: "
                  + e.getMessage());
          output = rawResult;
        }
      } else {
        output = rawResult;
      }
      event.actions().stateDelta().put(outputKey().get(), output);
    }
  }

  @Override
  protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
    return llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState);
  }

  @Override
  protected Flowable<Event> runLiveImpl(InvocationContext invocationContext) {
    return llmFlow.runLive(invocationContext).doOnNext(this::maybeSaveOutputToState);
  }

  /**
   * Constructs the text instruction for this agent based on the {@link #instruction} field.
   *
   * <p>This method is only for use by Agent Development Kit.
   *
   * @param context The context to retrieve the session state.
   * @return The resolved instruction as a {@link Single} wrapped string.
   */
  public Single<String> canonicalInstruction(ReadonlyContext context) {
    if (instruction instanceof Instruction.Static staticInstr) {
      return Single.just(staticInstr.instruction());
    } else if (instruction instanceof Instruction.Provider provider) {
      return provider.getInstruction().apply(context);
    }
    throw new IllegalStateException("Unknown Instruction subtype: " + instruction.getClass());
  }

  /**
   * Constructs the text global instruction for this agent based on the {@link #globalInstruction}
   * field.
   *
   * <p>This method is only for use by Agent Development Kit.
   *
   * @param context The context to retrieve the session state.
   * @return The resolved global instruction as a {@link Single} wrapped string.
   */
  public Single<String> canonicalGlobalInstruction(ReadonlyContext context) {
    if (globalInstruction instanceof Instruction.Static staticInstr) {
      return Single.just(staticInstr.instruction());
    } else if (globalInstruction instanceof Instruction.Provider provider) {
      return provider.getInstruction().apply(context);
    }
    throw new IllegalStateException("Unknown Instruction subtype: " + instruction.getClass());
  }

  /**
   * Constructs the list of tools for this agent based on the {@link #tools} field.
   *
   * <p>This method is only for use by Agent Development Kit.
   *
   * @param context The context to retrieve the session state.
   * @return The resolved list of tools as a {@link Single} wrapped list of {@link BaseTool}.
   */
  public Flowable<BaseTool> canonicalTools(Optional<ReadonlyContext> context) {
    List<Flowable<BaseTool>> toolFlowables = new ArrayList<>();
    for (Object toolOrToolset : toolsUnion) {
      if (toolOrToolset instanceof BaseTool baseTool) {
        toolFlowables.add(Flowable.just(baseTool));
      } else if (toolOrToolset instanceof BaseToolset baseToolset) {
        toolFlowables.add(baseToolset.getTools(context.orElse(null)));
      } else {
        throw new IllegalArgumentException(
            "Object in tools list is not of a supported type: "
                + toolOrToolset.getClass().getName());
      }
    }
    return Flowable.concat(toolFlowables);
  }

  /** Overload of canonicalTools that defaults to an empty context. */
  public Flowable<BaseTool> canonicalTools() {
    return canonicalTools(Optional.empty());
  }

  /** Convenience overload of canonicalTools that accepts a non-optional ReadonlyContext. */
  public Flowable<BaseTool> canonicalTools(ReadonlyContext context) {
    return canonicalTools(Optional.ofNullable(context));
  }

  public Instruction instruction() {
    return instruction;
  }

  public Instruction globalInstruction() {
    return globalInstruction;
  }

  public Optional<Model> model() {
    return model;
  }

  public boolean planning() {
    return planning;
  }

  public Optional<Integer> maxSteps() {
    return maxSteps;
  }

  public Optional<GenerateContentConfig> generateContentConfig() {
    return generateContentConfig;
  }

  public Optional<BaseExampleProvider> exampleProvider() {
    return exampleProvider;
  }

  public IncludeContents includeContents() {
    return includeContents;
  }

  public List<BaseTool> tools() {
    return canonicalTools().toList().blockingGet();
  }

  public List<Object> toolsUnion() {
    return toolsUnion;
  }

  public boolean disallowTransferToParent() {
    return disallowTransferToParent;
  }

  public boolean disallowTransferToPeers() {
    return disallowTransferToPeers;
  }

  public Optional<List<BeforeModelCallback>> beforeModelCallback() {
    return beforeModelCallback;
  }

  public Optional<List<AfterModelCallback>> afterModelCallback() {
    return afterModelCallback;
  }

  public Optional<List<BeforeToolCallback>> beforeToolCallback() {
    return beforeToolCallback;
  }

  public Optional<List<AfterToolCallback>> afterToolCallback() {
    return afterToolCallback;
  }

  public Optional<Schema> inputSchema() {
    return inputSchema;
  }

  public Optional<Schema> outputSchema() {
    return outputSchema;
  }

  public Optional<Executor> executor() {
    return executor;
  }

  public Optional<String> outputKey() {
    return outputKey;
  }

  public Model resolvedModel() {
    if (resolvedModel == null) {
      synchronized (this) {
        if (resolvedModel == null) {
          resolvedModel = resolveModelInternal();
        }
      }
    }
    return resolvedModel;
  }

  /**
   * Resolves the model for this agent, checking first if it is defined locally, then searching
   * through ancestors.
   *
   * <p>This method is only for use by Agent Development Kit.
   *
   * @return The resolved {@link Model} for this agent.
   * @throws IllegalStateException if no model is found for this agent or its ancestors.
   */
  private Model resolveModelInternal() {
    if (this.model.isPresent()) {
      if (this.model().isPresent()) {
        return this.model.get();
      }
    }
    BaseAgent current = this.parentAgent();
    while (current != null) {
      if (current instanceof LlmAgent) {
        return ((LlmAgent) current).resolvedModel();
      }
      current = current.parentAgent();
    }
    throw new IllegalStateException("No model found for agent " + name() + " or its ancestors.");
  }
}

/*
 * 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.runner;

import com.google.adk.Telemetry;
import com.google.adk.agents.BaseAgent;
import com.google.adk.agents.InvocationContext;
import com.google.adk.agents.LiveRequestQueue;
import com.google.adk.agents.LlmAgent;
import com.google.adk.agents.RunConfig;
import com.google.adk.artifacts.BaseArtifactService;
import com.google.adk.events.Event;
import com.google.adk.sessions.BaseSessionService;
import com.google.adk.sessions.Session;
import com.google.adk.utils.CollectionUtils;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.AudioTranscriptionConfig;
import com.google.genai.types.Content;
import com.google.genai.types.Modality;
import com.google.genai.types.Part;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.StatusCode;
import io.opentelemetry.context.Scope;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;

/** The main class for the GenAI Agents runner. */
public class Runner {
  private final BaseAgent agent;
  private final String appName;
  private final BaseArtifactService artifactService;
  private final BaseSessionService sessionService;

  /** Creates a new {@code Runner}. */
  public Runner(
      BaseAgent agent,
      String appName,
      BaseArtifactService artifactService,
      BaseSessionService sessionService) {
    this.agent = agent;
    this.appName = appName;
    this.artifactService = artifactService;
    this.sessionService = sessionService;
  }

  public BaseAgent agent() {
    return this.agent;
  }

  public String appName() {
    return this.appName;
  }

  public BaseArtifactService artifactService() {
    return this.artifactService;
  }

  public BaseSessionService sessionService() {
    return this.sessionService;
  }

  /**
   * Appends a new user message to the session history.
   *
   * @throws IllegalArgumentException if message has no parts.
   */
  private void appendNewMessageToSession(
      Session session,
      Content newMessage,
      InvocationContext invocationContext,
      boolean saveInputBlobsAsArtifacts) {
    if (newMessage.parts().isEmpty()) {
      throw new IllegalArgumentException("No parts in the new_message.");
    }

    if (this.artifactService != null && saveInputBlobsAsArtifacts) {
      // The runner directly saves the artifacts (if applicable) in the
      // user message and replaces the artifact data with a file name
      // placeholder.
      for (int i = 0; i < newMessage.parts().get().size(); i++) {
        Part part = newMessage.parts().get().get(i);
        if (part.inlineData().isEmpty()) {
          continue;
        }
        String fileName = "artifact_" + invocationContext.invocationId() + "_" + i;
        var unused =
            this.artifactService.saveArtifact(
                this.appName, session.userId(), session.id(), fileName, part);

        newMessage
            .parts()
            .get()
            .set(
                i,
                Part.fromText(
                    "Uploaded file: " + fileName + ". It has been saved to the artifacts"));
      }
    }
    // Appends only. We do not yield the event because it's not from the model.
    Event event =
        Event.builder()
            .id(Event.generateEventId())
            .invocationId(invocationContext.invocationId())
            .author("user")
            .content(Optional.of(newMessage))
            .build();
    this.sessionService.appendEvent(session, event);
  }

  /**
   * Runs the agent in the standard mode.
   *
   * @param userId The ID of the user for the session.
   * @param sessionId The ID of the session to run the agent in.
   * @param newMessage The new message from the user to process.
   * @param runConfig Configuration for the agent run.
   * @return A Flowable stream of {@link Event} objects generated by the agent during execution.
   */
  public Flowable<Event> runAsync(
      String userId, String sessionId, Content newMessage, RunConfig runConfig) {
    Maybe<Session> maybeSession =
        this.sessionService.getSession(appName, userId, sessionId, Optional.empty());
    return maybeSession
        .switchIfEmpty(
            Single.error(
                new IllegalArgumentException(
                    String.format("Session not found: %s for user %s", sessionId, userId))))
        .flatMapPublisher(session -> this.runAsync(session, newMessage, runConfig));
  }

  /**
   * Asynchronously runs the agent for a given user and session, processing a new message and using
   * a default {@link RunConfig}.
   *
   * <p>This method initiates an agent execution within the specified session, appending the
   * provided new message to the session's history. It utilizes a default {@code RunConfig} to
   * control execution parameters. The method returns a stream of {@link Event} objects representing
   * the agent's activity during the run.
   *
   * @param userId The ID of the user initiating the session.
   * @param sessionId The ID of the session in which the agent will run.
   * @param newMessage The new {@link Content} message to be processed by the agent.
   * @return A {@link Flowable} emitting {@link Event} objects generated by the agent.
   */
  public Flowable<Event> runAsync(String userId, String sessionId, Content newMessage) {
    return runAsync(userId, sessionId, newMessage, RunConfig.builder().build());
  }

  /**
   * Runs the agent in the standard mode using a provided Session object.
   *
   * @param session The session to run the agent in.
   * @param newMessage The new message from the user to process.
   * @param runConfig Configuration for the agent run.
   * @return A Flowable stream of {@link Event} objects generated by the agent during execution.
   */
  public Flowable<Event> runAsync(Session session, Content newMessage, RunConfig runConfig) {
    Span span = Telemetry.getTracer().spanBuilder("invocation").startSpan();
    try (Scope scope = span.makeCurrent()) {
      return Flowable.just(session)
          .flatMap(
              sess -> {
                BaseAgent rootAgent = this.agent;
                InvocationContext invocationContext =
                    InvocationContext.create(
                        this.sessionService,
                        this.artifactService,
                        InvocationContext.newInvocationContextId(),
                        rootAgent,
                        sess,
                        newMessage,
                        runConfig);

                if (newMessage != null) {
                  appendNewMessageToSession(
                      sess, newMessage, invocationContext, runConfig.saveInputBlobsAsArtifacts());
                }

                invocationContext.agent(this.findAgentToRun(sess, rootAgent));
                Flowable<Event> events = invocationContext.agent().runAsync(invocationContext);
                return events.doOnNext(event -> this.sessionService.appendEvent(sess, event));
              })
          .doOnError(
              throwable -> {
                span.setStatus(StatusCode.ERROR, "Error in runAsync Flowable execution");
                span.recordException(throwable);
              })
          .doFinally(span::end);
    } catch (Throwable t) {
      span.setStatus(StatusCode.ERROR, "Error during runAsync synchronous setup");
      span.recordException(t);
      span.end();
      return Flowable.error(t);
    }
  }

  /**
   * Creates an {@link InvocationContext} for a live (streaming) run.
   *
   * @return invocation context configured for a live run.
   */
  private InvocationContext newInvocationContextForLive(
      Session session, Optional<LiveRequestQueue> liveRequestQueue, RunConfig runConfig) {
    RunConfig.Builder runConfigBuilder = RunConfig.builder(runConfig);
    if (!CollectionUtils.isNullOrEmpty(runConfig.responseModalities())
        && liveRequestQueue.isPresent()) {
      // Default to AUDIO modality if not specified.
      if (CollectionUtils.isNullOrEmpty(runConfig.responseModalities())) {
        runConfigBuilder.setResponseModalities(
            ImmutableList.of(new Modality(Modality.Known.AUDIO)));
        if (runConfig.outputAudioTranscription() == null) {
          runConfigBuilder.setOutputAudioTranscription(AudioTranscriptionConfig.builder().build());
        }
      } else if (!runConfig.responseModalities().contains(new Modality(Modality.Known.TEXT))) {
        if (runConfig.outputAudioTranscription() == null) {
          runConfigBuilder.setOutputAudioTranscription(AudioTranscriptionConfig.builder().build());
        }
      }
    }
    return newInvocationContext(session, liveRequestQueue, runConfigBuilder.build());
  }

  /**
   * Creates an {@link InvocationContext} for the given session, request queue, and config.
   *
   * @return a new {@link InvocationContext}.
   */
  private InvocationContext newInvocationContext(
      Session session, Optional<LiveRequestQueue> liveRequestQueue, RunConfig runConfig) {
    BaseAgent rootAgent = this.agent;
    InvocationContext invocationContext =
        InvocationContext.create(
            this.sessionService,
            this.artifactService,
            rootAgent,
            session,
            liveRequestQueue.orElse(null),
            runConfig);
    invocationContext.agent(this.findAgentToRun(session, rootAgent));
    return invocationContext;
  }

  /**
   * Runs the agent in live mode, appending generated events to the session.
   *
   * @return stream of events from the agent.
   */
  public Flowable<Event> runLive(
      Session session, LiveRequestQueue liveRequestQueue, RunConfig runConfig) {
    Span span = Telemetry.getTracer().spanBuilder("invocation").startSpan();
    try (Scope scope = span.makeCurrent()) {
      InvocationContext invocationContext =
          newInvocationContextForLive(session, Optional.of(liveRequestQueue), runConfig);
      return invocationContext
          .agent()
          .runLive(invocationContext)
          .doOnNext(event -> this.sessionService.appendEvent(session, event))
          .doOnError(
              throwable -> {
                span.setStatus(StatusCode.ERROR, "Error in runLive Flowable execution");
                span.recordException(throwable);
              })
          .doFinally(span::end);
    } catch (Throwable t) {
      span.setStatus(StatusCode.ERROR, "Error during runLive synchronous setup");
      span.recordException(t);
      span.end();
      return Flowable.error(t);
    }
  }

  /**
   * 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(
      String userId, String sessionId, LiveRequestQueue liveRequestQueue, RunConfig runConfig) {
    return this.sessionService
        .getSession(appName, userId, sessionId, Optional.empty())
        .flatMapPublisher(
            session -> {
              if (session == null) {
                return Flowable.error(
                    new IllegalArgumentException(
                        String.format("Session not found: %s for user %s", sessionId, userId)));
              }
              return this.runLive(session, liveRequestQueue, runConfig);
            });
  }

  /**
   * Runs the agent asynchronously with a default user ID.
   *
   * @return stream of generated events.
   */
  public Flowable<Event> runWithSessionId(
      String sessionId, Content newMessage, RunConfig runConfig) {
    // TODO(b/410859954): Add user_id to getter or method signature. Assuming "tmp-user" for now.
    return this.runAsync("tmp-user", sessionId, newMessage, runConfig);
  }

  /**
   * Checks if the agent and its parent chain allow transfer up the tree.
   *
   * @return true if transferable, false otherwise.
   */
  private boolean isTransferableAcrossAgentTree(BaseAgent agentToRun) {
    BaseAgent current = agentToRun;
    while (current != null) {
      // Agents eligible to transfer must have an LLM-based agent parent.
      if (!(current instanceof LlmAgent)) {
        return false;
      }
      // If any agent can't transfer to its parent, the chain is broken.
      LlmAgent agent = (LlmAgent) current;
      if (agent.disallowTransferToParent()) {
        return false;
      }
      current = current.parentAgent();
    }
    return true;
  }

  /**
   * Returns the agent that should handle the next request based on session history.
   *
   * @return agent to run.
   */
  private BaseAgent findAgentToRun(Session session, BaseAgent rootAgent) {
    List<Event> events = new ArrayList<>(session.events());
    Collections.reverse(events);

    for (Event event : events) {
      String author = event.author();
      if (author.equals("user")) {
        continue;
      }

      if (author.equals(rootAgent.name())) {
        return rootAgent;
      }

      BaseAgent agent = rootAgent.findSubAgent(author);

      if (agent == null) {
        continue;
      }

      if (this.isTransferableAcrossAgentTree(agent)) {
        return agent;
      }
    }

    return rootAgent;
  }

  // TODO: run statelessly
}

/*
 * 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.agents;

import static com.google.common.collect.ImmutableList.toImmutableList;

import com.google.adk.Telemetry;
import com.google.adk.agents.Callbacks.AfterAgentCallback;
import com.google.adk.agents.Callbacks.BeforeAgentCallback;
import com.google.adk.events.Event;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.Content;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Scope;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.util.List;
import java.util.Optional;
import java.util.function.Function;
import org.jspecify.annotations.Nullable;

/** Base class for all agents. */
public abstract class BaseAgent {

  /** The agent's name. Must be a unique identifier within the agent tree. */
  private final String name;

  /**
   * One line description about the agent's capability. The system can use this for decision-making
   * when delegating control to different agents.
   */
  private final String description;

  /**
   * The parent agent in the agent tree. Note that one agent cannot be added to two different
   * parents' sub-agents lists.
   */
  private BaseAgent parentAgent;

  private final List<? extends BaseAgent> subAgents;

  private final Optional<List<BeforeAgentCallback>> beforeAgentCallback;
  private final Optional<List<AfterAgentCallback>> afterAgentCallback;

  /**
   * Creates a new BaseAgent.
   *
   * @param name Unique agent name. Cannot be "user" (reserved).
   * @param description Agent purpose.
   * @param subAgents Agents managed by this agent.
   * @param beforeAgentCallback Callbacks before agent execution. Invoked in order until one doesn't
   *     return null.
   * @param afterAgentCallback Callbacks after agent execution. Invoked in order until one doesn't
   *     return null.
   */
  public BaseAgent(
      String name,
      String description,
      List<? extends BaseAgent> subAgents,
      List<BeforeAgentCallback> beforeAgentCallback,
      List<AfterAgentCallback> afterAgentCallback) {
    this.name = name;
    this.description = description;
    this.parentAgent = null;
    this.subAgents = subAgents != null ? subAgents : ImmutableList.of();
    this.beforeAgentCallback = Optional.ofNullable(beforeAgentCallback);
    this.afterAgentCallback = Optional.ofNullable(afterAgentCallback);

    // Establish parent relationships for all sub-agents if needed.
    for (BaseAgent subAgent : this.subAgents) {
      subAgent.parentAgent(this);
    }
  }

  /**
   * Gets the agent's unique name.
   *
   * @return the unique name of the agent.
   */
  public final String name() {
    return name;
  }

  /**
   * Gets the one-line description of the agent's capability.
   *
   * @return the description of the agent.
   */
  public final String description() {
    return description;
  }

  /**
   * Retrieves the parent agent in the agent tree.
   *
   * @return the parent agent, or {@code null} if this agent does not have a parent.
   */
  public BaseAgent parentAgent() {
    return parentAgent;
  }

  /**
   * Sets the parent agent.
   *
   * @param parentAgent The parent agent to set.
   */
  protected void parentAgent(BaseAgent parentAgent) {
    this.parentAgent = parentAgent;
  }

  /**
   * Returns the root agent for this agent by traversing up the parent chain.
   *
   * @return the root agent.
   */
  public BaseAgent rootAgent() {
    BaseAgent agent = this;
    while (agent.parentAgent() != null) {
      agent = agent.parentAgent();
    }
    return agent;
  }

  /**
   * Finds an agent (this or descendant) by name.
   *
   * @return the agent or descendant with the given name, or {@code null} if not found.
   */
  public BaseAgent findAgent(String name) {
    if (this.name().equals(name)) {
      return this;
    }
    return findSubAgent(name);
  }

  /** Recursively search sub agent by name. */
  public @Nullable BaseAgent findSubAgent(String name) {
    for (BaseAgent subAgent : subAgents) {
      if (subAgent.name().equals(name)) {
        return subAgent;
      }
      BaseAgent result = subAgent.findSubAgent(name);
      if (result != null) {
        return result;
      }
    }
    return null;
  }

  public List<? extends BaseAgent> subAgents() {
    return subAgents;
  }

  public Optional<List<BeforeAgentCallback>> beforeAgentCallback() {
    return beforeAgentCallback;
  }

  public Optional<List<AfterAgentCallback>> afterAgentCallback() {
    return afterAgentCallback;
  }

  /**
   * Creates a shallow copy of the parent context with the agent properly being set to this
   * instance.
   *
   * @param parentContext Parent context to copy.
   * @return new context with updated branch name.
   */
  private InvocationContext createInvocationContext(InvocationContext parentContext) {
    InvocationContext invocationContext = InvocationContext.copyOf(parentContext);
    invocationContext.agent(this);
    // Check for branch to be truthy (not None, not empty string),
    if (parentContext.branch().filter(s -> !s.isEmpty()).isPresent()) {
      invocationContext.branch(parentContext.branch().get() + "." + name());
    }
    return invocationContext;
  }

  /**
   * Runs the agent asynchronously.
   *
   * @param parentContext Parent context to inherit.
   * @return stream of agent-generated events.
   */
  public Flowable<Event> runAsync(InvocationContext parentContext) {
    Tracer tracer = Telemetry.getTracer();
    return Flowable.defer(
        () -> {
          Span span = tracer.spanBuilder("agent_run [" + name() + "]").startSpan();
          try (Scope scope = span.makeCurrent()) {
            InvocationContext invocationContext = createInvocationContext(parentContext);

            Flowable<Event> executionFlowable =
                beforeAgentCallback
                    .map(
                        callback ->
                            callCallback(beforeCallbacksToFunctions(callback), invocationContext))
                    .orElse(Single.just(Optional.empty()))
                    .flatMapPublisher(
                        beforeEventOpt -> {
                          if (invocationContext.endInvocation()) {
                            return Flowable.fromOptional(beforeEventOpt);
                          }

                          Flowable<Event> beforeEvents = Flowable.fromOptional(beforeEventOpt);
                          Flowable<Event> mainEvents =
                              Flowable.defer(() -> runAsyncImpl(invocationContext));
                          Flowable<Event> afterEvents =
                              afterAgentCallback
                                  .map(
                                      callback ->
                                          Flowable.defer(
                                              () ->
                                                  callCallback(
                                                          afterCallbacksToFunctions(callback),
                                                          invocationContext)
                                                      .flatMapPublisher(Flowable::fromOptional)))
                                  .orElse(Flowable.empty());

                          return Flowable.concat(beforeEvents, mainEvents, afterEvents);
                        });
            return executionFlowable.doFinally(span::end);
          }
        });
  }

  /**
   * Converts before-agent callbacks to functions.
   *
   * @param callbacks Before-agent callbacks.
   * @return callback functions.
   */
  private ImmutableList<Function<CallbackContext, Maybe<Content>>> beforeCallbacksToFunctions(
      List<BeforeAgentCallback> callbacks) {
    return callbacks.stream()
        .map(callback -> (Function<CallbackContext, Maybe<Content>>) callback::call)
        .collect(toImmutableList());
  }

  /**
   * Converts after-agent callbacks to functions.
   *
   * @param callbacks After-agent callbacks.
   * @return callback functions.
   */
  private ImmutableList<Function<CallbackContext, Maybe<Content>>> afterCallbacksToFunctions(
      List<AfterAgentCallback> callbacks) {
    return callbacks.stream()
        .map(callback -> (Function<CallbackContext, Maybe<Content>>) callback::call)
        .collect(toImmutableList());
  }

  /**
   * Calls agent callbacks and returns the first produced event, if any.
   *
   * @param agentCallbacks Callback functions.
   * @param invocationContext Current invocation context.
   * @return single emitting first event, or empty if none.
   */
  private Single<Optional<Event>> callCallback(
      List<Function<CallbackContext, Maybe<Content>>> agentCallbacks,
      InvocationContext invocationContext) {
    if (agentCallbacks == null || agentCallbacks.isEmpty()) {
      return Single.just(Optional.empty());
    }

    CallbackContext callbackContext =
        new CallbackContext(invocationContext, /* eventActions= */ null);

    return Flowable.fromIterable(agentCallbacks)
        .concatMap(
            callback -> {
              Maybe<Content> maybeContent = callback.apply(callbackContext);

              return maybeContent
                  .map(
                      content -> {
                        Event.Builder eventBuilder =
                            Event.builder()
                                .id(Event.generateEventId())
                                .invocationId(invocationContext.invocationId())
                                .author(name())
                                .branch(invocationContext.branch())
                                .actions(callbackContext.eventActions());

                        eventBuilder.content(Optional.of(content));
                        invocationContext.setEndInvocation(true);
                        return Optional.of(eventBuilder.build());
                      })
                  .toFlowable();
            })
        .firstElement()
        .switchIfEmpty(
            Single.defer(
                () -> {
                  if (callbackContext.state().hasDelta()) {
                    Event.Builder eventBuilder =
                        Event.builder()
                            .id(Event.generateEventId())
                            .invocationId(invocationContext.invocationId())
                            .author(name())
                            .branch(invocationContext.branch())
                            .actions(callbackContext.eventActions());

                    return Single.just(Optional.of(eventBuilder.build()));
                  } else {
                    return Single.just(Optional.empty());
                  }
                }));
  }

  /**
   * Runs the agent synchronously.
   *
   * @param parentContext Parent context to inherit.
   * @return stream of agent-generated events.
   */
  public Flowable<Event> runLive(InvocationContext parentContext) {
    Tracer tracer = Telemetry.getTracer();
    return Flowable.defer(
        () -> {
          Span span = tracer.spanBuilder("agent_run [" + name() + "]").startSpan();
          try (Scope scope = span.makeCurrent()) {
            InvocationContext invocationContext = createInvocationContext(parentContext);
            Flowable<Event> executionFlowable = runLiveImpl(invocationContext);
            return executionFlowable.doFinally(span::end);
          }
        });
  }

  /**
   * Agent-specific asynchronous logic.
   *
   * @param invocationContext Current invocation context.
   * @return stream of agent-generated events.
   */
  protected abstract Flowable<Event> runAsyncImpl(InvocationContext invocationContext);

  /**
   * Agent-specific synchronous logic.
   *
   * @param invocationContext Current invocation context.
   * @return stream of agent-generated events.
   */
  protected abstract Flowable<Event> runLiveImpl(InvocationContext invocationContext);
}

/*
 * 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.flows.llmflows;

import com.google.adk.agents.InvocationContext;
import com.google.adk.agents.LlmAgent;
import com.google.adk.agents.ReadonlyContext;
import com.google.adk.models.LlmRequest;
import com.google.adk.utils.InstructionUtils;
import com.google.common.collect.ImmutableList;
import io.reactivex.rxjava3.core.Single;

/** {@link RequestProcessor} that handles instructions and global instructions for LLM flows. */
public final class Instructions implements RequestProcessor {
  public Instructions() {}

  @Override
  public Single<RequestProcessor.RequestProcessingResult> processRequest(
      InvocationContext context, LlmRequest request) {
    if (!(context.agent() instanceof LlmAgent)) {
      return Single.error(
          new IllegalArgumentException(
              "Agent in InvocationContext is not an instance of LlmAgent."));
    }
    LlmAgent agent = (LlmAgent) context.agent();
    ReadonlyContext readonlyContext = new ReadonlyContext(context);
    Single<LlmRequest.Builder> builderSingle = Single.just(request.toBuilder());
    if (agent.rootAgent() instanceof LlmAgent) {
      LlmAgent rootAgent = (LlmAgent) agent.rootAgent();
      builderSingle =
          builderSingle.flatMap(
              builder ->
                  rootAgent
                      .canonicalGlobalInstruction(readonlyContext)
                      .flatMap(
                          globalInstr -> {
                            if (!globalInstr.isEmpty()) {
                              return InstructionUtils.injectSessionState(context, globalInstr)
                                  .map(
                                      resolvedGlobalInstr ->
                                          builder.appendInstructions(
                                              ImmutableList.of(resolvedGlobalInstr)));
                            }
                            return Single.just(builder);
                          }));
    }

    builderSingle =
        builderSingle.flatMap(
            builder ->
                agent
                    .canonicalInstruction(readonlyContext)
                    .flatMap(
                        agentInstr -> {
                          if (!agentInstr.isEmpty()) {
                            return InstructionUtils.injectSessionState(context, agentInstr)
                                .map(
                                    resolvedAgentInstr ->
                                        builder.appendInstructions(
                                            ImmutableList.of(resolvedAgentInstr)));
                          }
                          return Single.just(builder);
                        }));

    return builderSingle.map(
        finalBuilder ->
            RequestProcessor.RequestProcessingResult.create(
                finalBuilder.build(), ImmutableList.of()));
  }
}

/*
 * 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.flows.llmflows;

import static com.google.common.collect.ImmutableList.toImmutableList;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.google.adk.JsonBaseModel;
import com.google.adk.agents.InvocationContext;
import com.google.adk.agents.LlmAgent;
import com.google.adk.events.Event;
import com.google.adk.models.LlmRequest;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.genai.types.Content;
import com.google.genai.types.FunctionCall;
import com.google.genai.types.FunctionResponse;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Single;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

/** {@link RequestProcessor} that populates content in request for LLM flows. */
public final class Contents implements RequestProcessor {
  public Contents() {}

  @Override
  public Single<RequestProcessor.RequestProcessingResult> processRequest(
      InvocationContext context, LlmRequest request) {
    if (!(context.agent() instanceof LlmAgent)) {
      return Single.just(
          RequestProcessor.RequestProcessingResult.create(request, context.session().events()));
    }
    LlmAgent llmAgent = (LlmAgent) context.agent();

    if (llmAgent.includeContents() == LlmAgent.IncludeContents.NONE) {
      return Single.just(
          RequestProcessor.RequestProcessingResult.create(
              request.toBuilder().contents(ImmutableList.of()).build(), ImmutableList.of()));
    }

    ImmutableList<Content> contents =
        getContents(context.branch(), context.session().events(), context.agent().name());

    return Single.just(
        RequestProcessor.RequestProcessingResult.create(
            request.toBuilder().contents(contents).build(), ImmutableList.of()));
  }

  private ImmutableList<Content> getContents(
      Optional<String> currentBranch, List<Event> events, String agentName) {
    List<Event> filteredEvents = new ArrayList<>();

    // Filter the events, leaving the contents and the function calls and responses from the current
    // agent.
    for (Event event : events) {
      // Skip events without content, or generated neither by user nor by model or has empty text.
      // E.g. events purely for mutating session states.
      if (event.content().isEmpty()) {
        continue;
      }
      var content = event.content().get();
      if (content.role().isEmpty()
          || content.role().get().isEmpty()
          || content.parts().isEmpty()
          || content.parts().get().isEmpty()
          || content.parts().get().get(0).text().map(String::isEmpty).orElse(false)) {
        continue;
      }

      if (!isEventBelongsToBranch(currentBranch, event)) {
        continue;
      }

      // TODO: Skip auth events.

      if (isOtherAgentReply(agentName, event)) {
        filteredEvents.add(convertForeignEvent(event));
      } else {
        filteredEvents.add(event);
      }
    }

    List<Event> resultEvents = rearrangeEventsForLatestFunctionResponse(filteredEvents);
    resultEvents = rearrangeEventsForAsyncFunctionResponsesInHistory(resultEvents);

    return resultEvents.stream()
        .map(Event::content)
        .flatMap(Optional::stream)
        .collect(toImmutableList());
  }

  /** Whether the event is a reply from another agent. */
  private static boolean isOtherAgentReply(String agentName, Event event) {
    return !agentName.isEmpty()
        && !event.author().equals(agentName)
        && !event.author().equals("user");
  }

  /** Converts an {@code event} authored by another agent to a 'contextual-only' event. */
  private static Event convertForeignEvent(Event event) {
    if (event.content().isEmpty()
        || event.content().get().parts().isEmpty()
        || event.content().get().parts().get().isEmpty()) {
      return event;
    }

    List<Part> parts = new ArrayList<>();
    parts.add(Part.fromText("For context:"));

    String originalAuthor = event.author();

    for (Part part : event.content().get().parts().get()) {
      if (part.text().isPresent() && !part.text().get().isEmpty()) {
        parts.add(Part.fromText(String.format("[%s] said: %s", originalAuthor, part.text().get())));
      } else if (part.functionCall().isPresent()) {
        FunctionCall functionCall = part.functionCall().get();
        parts.add(
            Part.fromText(
                String.format(
                    "[%s] called tool `%s` with parameters: %s",
                    originalAuthor,
                    functionCall.name().orElse("unknown_tool"),
                    functionCall.args().map(Contents::convertMapToJson).orElse("{}"))));
      } else if (part.functionResponse().isPresent()) {
        FunctionResponse functionResponse = part.functionResponse().get();
        parts.add(
            Part.fromText(
                String.format(
                    "[%s] `%s` tool returned result: %s",
                    originalAuthor,
                    functionResponse.name().orElse("unknown_tool"),
                    functionResponse.response().map(Contents::convertMapToJson).orElse("{}"))));
      } else {
        parts.add(part);
      }
    }

    Content content = Content.builder().role("user").parts(parts).build();
    return event.toBuilder().author("user").content(content).build();
  }

  private static String convertMapToJson(Map<String, Object> struct) {
    try {
      return JsonBaseModel.getMapper().writeValueAsString(struct);
    } catch (JsonProcessingException e) {
      throw new IllegalStateException("Failed to serialize the object to JSON.", e);
    }
  }

  private static boolean isEventBelongsToBranch(Optional<String> invocationBranchOpt, Event event) {
    Optional<String> eventBranchOpt = event.branch();

    if (invocationBranchOpt.isEmpty() || invocationBranchOpt.get().isEmpty()) {
      return true;
    }
    if (eventBranchOpt.isEmpty() || eventBranchOpt.get().isEmpty()) {
      return true;
    }
    return invocationBranchOpt.get().startsWith(eventBranchOpt.get());
  }

  /**
   * Rearranges the events for the latest function response. If the latest function response is for
   * an async function call, all events between the initial function call and the latest function
   * response will be removed.
   *
   * @param events The list of events.
   * @return A new list of events with the appropriate rearrangement.
   */
  private static List<Event> rearrangeEventsForLatestFunctionResponse(List<Event> events) {
    // TODO: b/412663475 - Handle parallel function calls within the same event. Currently, this
    // throws an error.
    if (events.isEmpty() || Iterables.getLast(events).functionResponses().isEmpty()) {
      // No need to process if the list is empty or the last event is not a function response
      return events;
    }

    Event latestEvent = Iterables.getLast(events);
    // Extract function response IDs from the latest event
    Set<String> functionResponseIds = new HashSet<>();
    latestEvent
        .content()
        .flatMap(Content::parts)
        .ifPresent(
            parts -> {
              for (Part part : parts) {
                part.functionResponse()
                    .flatMap(FunctionResponse::id)
                    .ifPresent(functionResponseIds::add);
              }
            });

    if (functionResponseIds.isEmpty()) {
      return events;
    }

    // Check if the second to last event contains the corresponding function call
    if (events.size() >= 2) {
      Event penultimateEvent = events.get(events.size() - 2);
      boolean matchFound =
          penultimateEvent
              .content()
              .flatMap(Content::parts)
              .map(
                  parts -> {
                    for (Part part : parts) {
                      if (part.functionCall()
                          .flatMap(FunctionCall::id)
                          .map(functionResponseIds::contains)
                          .orElse(false)) {
                        return true; // Found a matching function call ID
                      }
                    }
                    return false;
                  })
              .orElse(false);
      if (matchFound) {
        // The latest function response is already matched with the immediately preceding event
        return events;
      }
    }

    // Look for the corresponding function call event by iterating backwards
    int functionCallEventIndex = -1;
    for (int i = events.size() - 3; i >= 0; i--) { // Start from third-to-last
      Event event = events.get(i);
      Optional<List<Part>> partsOptional = event.content().flatMap(Content::parts);
      if (partsOptional.isPresent()) {
        List<Part> parts = partsOptional.get();
        for (Part part : parts) {
          Optional<String> callIdOpt = part.functionCall().flatMap(FunctionCall::id);
          if (callIdOpt.isPresent() && functionResponseIds.contains(callIdOpt.get())) {
            functionCallEventIndex = i;
            // Add all function call IDs from this event to the set
            parts.forEach(
                p ->
                    p.functionCall().flatMap(FunctionCall::id).ifPresent(functionResponseIds::add));
            break; // Found the matching event
          }
        }
      }
      if (functionCallEventIndex != -1) {
        break; // Exit outer loop once found
      }
    }

    if (functionCallEventIndex == -1) {
      if (!functionResponseIds.isEmpty()) {
        throw new IllegalStateException(
            "No function call event found for function response IDs: " + functionResponseIds);
      } else {
        return events; // No IDs to match, no rearrangement based on this logic.
      }
    }

    List<Event> resultEvents = new ArrayList<>(events.subList(0, functionCallEventIndex + 1));

    // Collect all function response events between the call and the latest response
    List<Event> functionResponseEventsToMerge = new ArrayList<>();
    for (int i = functionCallEventIndex + 1; i < events.size() - 1; i++) {
      Event intermediateEvent = events.get(i);
      boolean hasMatchingResponse =
          intermediateEvent
              .content()
              .flatMap(Content::parts)
              .map(
                  parts -> {
                    for (Part part : parts) {
                      if (part.functionResponse()
                          .flatMap(FunctionResponse::id)
                          .map(functionResponseIds::contains)
                          .orElse(false)) {
                        return true;
                      }
                    }
                    return false;
                  })
              .orElse(false);
      if (hasMatchingResponse) {
        functionResponseEventsToMerge.add(intermediateEvent);
      }
    }
    functionResponseEventsToMerge.add(latestEvent);

    if (!functionResponseEventsToMerge.isEmpty()) {
      resultEvents.add(mergeFunctionResponseEvents(functionResponseEventsToMerge));
    }

    return resultEvents;
  }

  private static List<Event> rearrangeEventsForAsyncFunctionResponsesInHistory(List<Event> events) {
    Map<String, Integer> functionCallIdToResponseEventIndex = new HashMap<>();
    for (int i = 0; i < events.size(); i++) {
      final int index = i;
      Event event = events.get(index);
      event
          .content()
          .flatMap(Content::parts)
          .ifPresent(
              parts -> {
                for (Part part : parts) {
                  part.functionResponse()
                      .ifPresent(
                          response ->
                              response
                                  .id()
                                  .ifPresent(
                                      functionCallId ->
                                          functionCallIdToResponseEventIndex.put(
                                              functionCallId, index)));
                }
              });
    }

    List<Event> resultEvents = new ArrayList<>();
    // Keep track of response events already added to avoid duplicates when merging
    Set<Integer> processedResponseIndices = new HashSet<>();

    for (int i = 0; i < events.size(); i++) {
      Event event = events.get(i);

      // Skip response events that have already been processed and added alongside their call event
      if (processedResponseIndices.contains(i)) {
        continue;
      }

      Optional<List<Part>> partsOptional = event.content().flatMap(Content::parts);
      boolean hasFunctionCalls =
          partsOptional
              .map(parts -> parts.stream().anyMatch(p -> p.functionCall().isPresent()))
              .orElse(false);

      if (hasFunctionCalls) {
        Set<Integer> responseEventIndices = new HashSet<>();
        // Iterate through parts again to get function call IDs
        partsOptional
            .get()
            .forEach(
                part ->
                    part.functionCall()
                        .ifPresent(
                            call ->
                                call.id()
                                    .ifPresent(
                                        functionCallId -> {
                                          if (functionCallIdToResponseEventIndex.containsKey(
                                              functionCallId)) {
                                            responseEventIndices.add(
                                                functionCallIdToResponseEventIndex.get(
                                                    functionCallId));
                                          }
                                        })));

        resultEvents.add(event); // Add the function call event

        if (!responseEventIndices.isEmpty()) {
          List<Event> responseEventsToAdd = new ArrayList<>();
          List<Integer> sortedIndices = new ArrayList<>(responseEventIndices);
          Collections.sort(sortedIndices); // Process in chronological order

          for (int index : sortedIndices) {
            if (processedResponseIndices.add(index)) { // Add index and check if it was newly added
              responseEventsToAdd.add(events.get(index));
            }
          }

          if (responseEventsToAdd.size() == 1) {
            resultEvents.add(responseEventsToAdd.get(0));
          } else if (responseEventsToAdd.size() > 1) {
            resultEvents.add(mergeFunctionResponseEvents(responseEventsToAdd));
          }
        }
      } else {
        resultEvents.add(event);
      }
    }

    return resultEvents;
  }

  /**
   * Merges a list of function response events into one event.
   *
   * <p>The key goal is to ensure: 1. functionCall and functionResponse are always of the same
   * number. 2. The functionCall and functionResponse are consecutively in the content.
   *
   * @param functionResponseEvents A list of function response events. NOTE: functionResponseEvents
   *     must fulfill these requirements: 1. The list is in increasing order of timestamp; 2. the
   *     first event is the initial function response event; 3. all later events should contain at
   *     least one function response part that related to the function call event. Caveat: This
   *     implementation doesn't support when a parallel function call event contains async function
   *     call of the same name.
   * @return A merged event, that is 1. All later function_response will replace function response
   *     part in the initial function response event. 2. All non-function response parts will be
   *     appended to the part list of the initial function response event.
   */
  private static Event mergeFunctionResponseEvents(List<Event> functionResponseEvents) {
    if (functionResponseEvents.isEmpty()) {
      throw new IllegalArgumentException("At least one functionResponse event is required.");
    }
    if (functionResponseEvents.size() == 1) {
      return functionResponseEvents.get(0);
    }

    Event baseEvent = functionResponseEvents.get(0);
    Content baseContent =
        baseEvent
            .content()
            .orElseThrow(() -> new IllegalArgumentException("Base event must have content."));
    List<Part> baseParts =
        baseContent
            .parts()
            .orElseThrow(() -> new IllegalArgumentException("Base event content must have parts."));

    if (baseParts.isEmpty()) {
      throw new IllegalArgumentException(
          "There should be at least one functionResponse part in the base event.");
    }
    List<Part> partsInMergedEvent = new ArrayList<>(baseParts);

    Map<String, Integer> partIndicesInMergedEvent = new HashMap<>();
    for (int i = 0; i < partsInMergedEvent.size(); i++) {
      final int index = i;
      Part part = partsInMergedEvent.get(i);
      if (part.functionResponse().isPresent()) {
        part.functionResponse()
            .get()
            .id()
            .ifPresent(functionCallId -> partIndicesInMergedEvent.put(functionCallId, index));
      }
    }

    for (Event event : functionResponseEvents.subList(1, functionResponseEvents.size())) {
      if (!hasContentWithNonEmptyParts(event)) {
        continue;
      }

      for (Part part : event.content().get().parts().get()) {
        if (part.functionResponse().isPresent()) {
          Optional<String> functionCallIdOpt = part.functionResponse().get().id();
          if (functionCallIdOpt.isPresent()) {
            String functionCallId = functionCallIdOpt.get();
            if (partIndicesInMergedEvent.containsKey(functionCallId)) {
              partsInMergedEvent.set(partIndicesInMergedEvent.get(functionCallId), part);
            } else {
              partsInMergedEvent.add(part);
              partIndicesInMergedEvent.put(functionCallId, partsInMergedEvent.size() - 1);
            }
          } else {
            partsInMergedEvent.add(part);
          }
        } else {
          partsInMergedEvent.add(part);
        }
      }
    }

    return baseEvent.toBuilder()
        .content(
            Optional.of(
                Content.builder().role(baseContent.role().get()).parts(partsInMergedEvent).build()))
        .build();
  }

  private static boolean hasContentWithNonEmptyParts(Event event) {
    return event
        .content() // Optional<Content>
        .flatMap(Content::parts) // Optional<List<Part>>
        .map(list -> !list.isEmpty()) // Optional<Boolean>
        .orElse(false);
  }
}

/*
 * 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.flows.llmflows;

import com.google.adk.agents.BaseAgent;
import com.google.adk.agents.InvocationContext;
import com.google.adk.models.LlmRequest;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
import io.reactivex.rxjava3.core.Single;

/** {@link RequestProcessor} that gives the agent identity from the framework */
public final class Identity implements RequestProcessor {

  public Identity() {}

  @Override
  public Single<RequestProcessor.RequestProcessingResult> processRequest(
      InvocationContext context, LlmRequest request) {
    BaseAgent agent = context.agent();
    StringBuilder builder =
        new StringBuilder()
            .append("You are an agent. Your internal name is ")
            .append(agent.name())
            .append(".");
    if (!Strings.isNullOrEmpty(agent.description())) {
      builder.append(" The description about you is ").append(agent.description());
    }
    return Single.just(
        RequestProcessor.RequestProcessingResult.create(
            request.toBuilder().appendInstructions(ImmutableList.of(builder.toString())).build(),
            ImmutableList.of()));
  }
}

/*
 * 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.flows.llmflows;

import com.google.adk.agents.InvocationContext;
import com.google.adk.agents.LlmAgent;
import com.google.adk.examples.ExampleUtils;
import com.google.adk.models.LlmRequest;
import com.google.common.collect.ImmutableList;
import io.reactivex.rxjava3.core.Single;

/** {@link RequestProcessor} that populates examples in LLM request. */
public final class Examples implements RequestProcessor {

  public Examples() {}

  @Override
  public Single<RequestProcessor.RequestProcessingResult> processRequest(
      InvocationContext context, LlmRequest request) {
    if (!(context.agent() instanceof LlmAgent)) {
      throw new IllegalArgumentException("Agent in InvocationContext is not an instance of Agent.");
    }
    LlmAgent agent = (LlmAgent) context.agent();
    LlmRequest.Builder builder = request.toBuilder();

    String query =
        context.userContent().isPresent()
            ? context.userContent().get().parts().get().get(0).text().orElse("")
            : "";
    agent
        .exampleProvider()
        .ifPresent(
            exampleProvider ->
                builder.appendInstructions(
                    ImmutableList.of(ExampleUtils.buildExampleSi(exampleProvider, query))));
    return Single.just(
        RequestProcessor.RequestProcessingResult.create(builder.build(), ImmutableList.of()));
  }
}

/*
 * 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.tools;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.google.adk.JsonBaseModel;
import com.google.adk.SchemaUtils;
import com.google.adk.agents.BaseAgent;
import com.google.adk.agents.LlmAgent;
import com.google.adk.events.Event;
import com.google.adk.runner.InMemoryRunner;
import com.google.adk.runner.Runner;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.genai.types.Content;
import com.google.genai.types.FunctionDeclaration;
import com.google.genai.types.Part;
import com.google.genai.types.Schema;
import io.reactivex.rxjava3.core.Single;
import java.util.Map;
import java.util.Optional;

/** AgentTool implements a tool that allows an agent to call another agent. */
public class AgentTool extends BaseTool {

  private final BaseAgent agent;
  private final boolean skipSummarization;

  public static AgentTool create(BaseAgent agent, boolean skipSummarization) {
    return new AgentTool(agent, skipSummarization);
  }

  public static AgentTool create(BaseAgent agent) {
    return new AgentTool(agent, false);
  }

  protected AgentTool(BaseAgent agent, boolean skipSummarization) {
    super(agent.name(), agent.description());
    this.agent = agent;
    this.skipSummarization = skipSummarization;
  }

  @Override
  public Optional<FunctionDeclaration> declaration() {
    FunctionDeclaration.Builder builder =
        FunctionDeclaration.builder().description(this.description()).name(this.name());

    Optional<Schema> agentInputSchema = Optional.empty();
    if (agent instanceof LlmAgent llmAgent) {
      agentInputSchema = llmAgent.inputSchema();
    }

    if (agentInputSchema.isPresent()) {
      builder.parameters(agentInputSchema.get());
    } else {
      builder.parameters(
          Schema.builder()
              .type("OBJECT")
              .properties(ImmutableMap.of("request", Schema.builder().type("STRING").build()))
              .required(ImmutableList.of("request"))
              .build());
    }
    return Optional.of(builder.build());
  }

  @Override
  public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContext toolContext) {

    if (this.skipSummarization) {
      toolContext.actions().setSkipSummarization(true);
    }

    Optional<Schema> agentInputSchema = Optional.empty();
    if (agent instanceof LlmAgent llmAgent) {
      agentInputSchema = llmAgent.inputSchema();
    }

    final Content content;
    if (agentInputSchema.isPresent()) {
      SchemaUtils.validateMapOnSchema(args, agentInputSchema.get(), true);
      try {
        content =
            Content.fromParts(Part.fromText(JsonBaseModel.getMapper().writeValueAsString(args)));
      } catch (JsonProcessingException e) {
        return Single.error(
            new RuntimeException("Error serializing tool arguments to JSON: " + args, e));
      }
    } else {
      Object input = args.get("request");
      content = Content.fromParts(Part.fromText(input.toString()));
    }

    Runner runner = new InMemoryRunner(this.agent, toolContext.agentName());
    // Session state is final, can't update to toolContext state
    // session.toBuilder().setState(toolContext.getState());
    return runner
        .sessionService()
        .createSession(toolContext.agentName(), "tmp-user", toolContext.state(), null)
        .flatMapPublisher(session -> runner.runAsync(session.userId(), session.id(), content))
        .lastElement()
        .map(Optional::of)
        .defaultIfEmpty(Optional.empty())
        .map(
            optionalLastEvent -> {
              if (optionalLastEvent.isEmpty()) {
                return ImmutableMap.of();
              }
              Event lastEvent = optionalLastEvent.get();
              Optional<String> outputText =
                  lastEvent
                      .content()
                      .flatMap(Content::parts)
                      .filter(parts -> !parts.isEmpty())
                      .flatMap(parts -> parts.get(0).text());

              if (outputText.isEmpty()) {
                return ImmutableMap.of();
              }
              String output = outputText.get();

              Optional<Schema> agentOutputSchema = Optional.empty();
              if (agent instanceof LlmAgent llmAgent) {
                agentOutputSchema = llmAgent.outputSchema();
              }

              if (agentOutputSchema.isPresent()) {
                return SchemaUtils.validateOutputSchema(output, agentOutputSchema.get());
              } else {
                return ImmutableMap.of("result", output);
              }
            });
  }
}

/*
 * 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.flows.llmflows;

import com.google.adk.agents.InvocationContext;
import com.google.adk.agents.LlmAgent;
import com.google.adk.models.LlmRequest;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.LiveConnectConfig;
import io.reactivex.rxjava3.core.Single;
import java.util.Optional;

/** {@link RequestProcessor} that handles basic information to build the LLM request. */
public final class Basic implements RequestProcessor {

  public Basic() {}

  @Override
  public Single<RequestProcessor.RequestProcessingResult> processRequest(
      InvocationContext context, LlmRequest request) {
    if (!(context.agent() instanceof LlmAgent)) {
      throw new IllegalArgumentException("Agent in InvocationContext is not an instance of Agent.");
    }
    LlmAgent agent = (LlmAgent) context.agent();
    String modelName =
        agent.resolvedModel().model().isPresent()
            ? agent.resolvedModel().model().get().model()
            : agent.resolvedModel().modelName().get();

    LiveConnectConfig.Builder liveConnectConfigBuilder =
        LiveConnectConfig.builder().responseModalities(context.runConfig().responseModalities());
    Optional.ofNullable(context.runConfig().speechConfig())
        .ifPresent(liveConnectConfigBuilder::speechConfig);
    Optional.ofNullable(context.runConfig().outputAudioTranscription())
        .ifPresent(liveConnectConfigBuilder::outputAudioTranscription);

    LlmRequest.Builder builder =
        request.toBuilder()
            .model(modelName)
            .config(agent.generateContentConfig().orElse(GenerateContentConfig.builder().build()))
            .liveConnectConfig(liveConnectConfigBuilder.build());

    agent.outputSchema().ifPresent(builder::outputSchema);
    return Single.just(
        RequestProcessor.RequestProcessingResult.create(builder.build(), ImmutableList.of()));
  }
}

/*
 * 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.flows.llmflows;

import com.google.adk.Telemetry;
import com.google.adk.agents.Callbacks.AfterToolCallback;
import com.google.adk.agents.Callbacks.BeforeToolCallback;
import com.google.adk.agents.InvocationContext;
import com.google.adk.agents.LlmAgent;
import com.google.adk.events.Event;
import com.google.adk.events.EventActions;
import com.google.adk.tools.BaseTool;
import com.google.adk.tools.ToolContext;
import com.google.common.base.VerifyException;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.Content;
import com.google.genai.types.FunctionCall;
import com.google.genai.types.FunctionResponse;
import com.google.genai.types.Part;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Scope;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.core.Maybe;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import org.jspecify.annotations.Nullable;

/** Utility class for handling function calls. */
public final class Functions {

  private static final String AF_FUNCTION_CALL_ID_PREFIX = "adk-";

  /** Generates a unique ID for a function call. */
  public static String generateClientFunctionCallId() {
    return AF_FUNCTION_CALL_ID_PREFIX + UUID.randomUUID();
  }

  /**
   * Populates missing function call IDs in the provided event's content.
   *
   * <p>If the event contains function calls without an ID, this method generates a unique
   * client-side ID for each and updates the event content.
   *
   * @param modelResponseEvent The event potentially containing function calls.
   */
  public static void populateClientFunctionCallId(Event modelResponseEvent) {
    Optional<Content> originalContentOptional = modelResponseEvent.content();
    if (originalContentOptional.isEmpty()) {
      return;
    }
    Content originalContent = originalContentOptional.get();
    List<Part> originalParts = originalContent.parts().orElse(ImmutableList.of());
    if (originalParts.stream().noneMatch(part -> part.functionCall().isPresent())) {
      return; // No function calls to process
    }

    List<Part> newParts = new ArrayList<>();
    boolean modified = false;
    for (Part part : originalParts) {
      if (part.functionCall().isPresent()) {
        FunctionCall functionCall = part.functionCall().get();
        if (functionCall.id().isEmpty() || functionCall.id().get().isEmpty()) {
          FunctionCall updatedFunctionCall =
              functionCall.toBuilder().id(generateClientFunctionCallId()).build();
          newParts.add(Part.builder().functionCall(updatedFunctionCall).build());
          modified = true;
        } else {
          newParts.add(part); // Keep original part if ID exists
        }
      } else {
        newParts.add(part); // Keep non-function call parts
      }
    }

    if (modified) {
      String role =
          originalContent
              .role()
              .orElseThrow(
                  () ->
                      new IllegalStateException(
                          "Content role is missing in event: " + modelResponseEvent.id()));
      Content newContent = Content.builder().role(role).parts(newParts).build();
      modelResponseEvent.setContent(Optional.of(newContent));
    }
  }

  // TODO - b/413761119 add the remaining methods for function call id.

  public static Maybe<Event> handleFunctionCalls(
      InvocationContext invocationContext, Event functionCallEvent, Map<String, BaseTool> tools) {
    ImmutableList<FunctionCall> functionCalls = functionCallEvent.functionCalls();

    List<Maybe<Event>> functionResponseEvents = new ArrayList<>();

    for (FunctionCall functionCall : functionCalls) {
      if (!tools.containsKey(functionCall.name().get())) {
        throw new VerifyException("Tool not found: " + functionCall.name().get());
      }
      BaseTool tool = tools.get(functionCall.name().get());
      ToolContext toolContext =
          ToolContext.builder(invocationContext)
              .functionCallId(functionCall.id().orElse(""))
              .build();

      Map<String, Object> functionArgs = functionCall.args().orElse(new HashMap<>());

      Maybe<Map<String, Object>> maybeFunctionResult =
          maybeInvokeBeforeToolCall(invocationContext, tool, functionArgs, toolContext)
              .switchIfEmpty(Maybe.defer(() -> callTool(tool, functionArgs, toolContext)));

      Maybe<Event> maybeFunctionResponseEvent =
          maybeFunctionResult
              .map(Optional::of)
              .defaultIfEmpty(Optional.empty())
              .flatMapMaybe(
                  optionalInitialResult -> {
                    Map<String, Object> initialFunctionResult = optionalInitialResult.orElse(null);

                    Maybe<Map<String, Object>> afterToolResultMaybe =
                        maybeInvokeAfterToolCall(
                            invocationContext,
                            tool,
                            functionArgs,
                            toolContext,
                            initialFunctionResult);

                    return afterToolResultMaybe
                        .map(Optional::of)
                        .defaultIfEmpty(Optional.ofNullable(initialFunctionResult))
                        .flatMapMaybe(
                            finalOptionalResult -> {
                              Map<String, Object> finalFunctionResult =
                                  finalOptionalResult.orElse(null);
                              if (tool.longRunning() && finalFunctionResult == null) {
                                return Maybe.empty();
                              }
                              Event functionResponseEvent =
                                  buildResponseEvent(
                                      tool, finalFunctionResult, toolContext, invocationContext);
                              return Maybe.just(functionResponseEvent);
                            });
                  });

      functionResponseEvents.add(maybeFunctionResponseEvent);
    }

    return Maybe.merge(functionResponseEvents)
        .toList()
        .flatMapMaybe(
            events -> {
              if (events.isEmpty()) {
                return Maybe.empty();
              }
              Event mergedEvent = Functions.mergeParallelFunctionResponseEvents(events);
              if (mergedEvent == null) {
                return Maybe.empty();
              }

              if (events.size() > 1) {
                Tracer tracer = Telemetry.getTracer();
                Span mergedSpan = tracer.spanBuilder("tool_response").startSpan();
                try (Scope scope = mergedSpan.makeCurrent()) {
                  Telemetry.traceToolResponse(invocationContext, mergedEvent.id(), mergedEvent);
                } finally {
                  mergedSpan.end();
                }
              }
              return Maybe.just(mergedEvent);
            });
  }

  public static Set<String> getLongRunningFunctionCalls(
      List<FunctionCall> functionCalls, Map<String, BaseTool> tools) {
    Set<String> longRunningFunctionCalls = new HashSet<>();
    for (FunctionCall functionCall : functionCalls) {
      if (!tools.containsKey(functionCall.name().get())) {
        continue;
      }
      BaseTool tool = tools.get(functionCall.name().get());
      if (tool.longRunning()) {
        longRunningFunctionCalls.add(functionCall.id().orElse(""));
      }
    }
    return longRunningFunctionCalls;
  }

  private static @Nullable Event mergeParallelFunctionResponseEvents(
      List<Event> functionResponseEvents) {
    if (functionResponseEvents.isEmpty()) {
      return null;
    }
    if (functionResponseEvents.size() == 1) {
      return functionResponseEvents.get(0);
    }
    // Use the first event as the base for common attributes
    Event baseEvent = functionResponseEvents.get(0);

    List<Part> mergedParts = new ArrayList<>();
    for (Event event : functionResponseEvents) {
      event.content().flatMap(Content::parts).ifPresent(mergedParts::addAll);
    }

    // Merge actions from all events
    // TODO: validate that pending actions are not cleared away
    EventActions.Builder mergedActionsBuilder = EventActions.builder();
    for (Event event : functionResponseEvents) {
      mergedActionsBuilder.merge(event.actions());
    }

    return Event.builder()
        .id(Event.generateEventId())
        .invocationId(baseEvent.invocationId())
        .author(baseEvent.author())
        .branch(baseEvent.branch())
        .content(Optional.of(Content.builder().role("user").parts(mergedParts).build()))
        .actions(mergedActionsBuilder.build())
        .timestamp(baseEvent.timestamp())
        .build();
  }

  private static Maybe<Map<String, Object>> maybeInvokeBeforeToolCall(
      InvocationContext invocationContext,
      BaseTool tool,
      Map<String, Object> functionArgs,
      ToolContext toolContext) {
    if (invocationContext.agent() instanceof LlmAgent) {
      LlmAgent agent = (LlmAgent) invocationContext.agent();

      Optional<List<BeforeToolCallback>> callbacksOpt = agent.beforeToolCallback();
      if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {
        return Maybe.empty();
      }
      List<BeforeToolCallback> callbacks = callbacksOpt.get();

      return Flowable.fromIterable(callbacks)
          .concatMapMaybe(
              callback -> callback.call(invocationContext, tool, functionArgs, toolContext))
          .firstElement();
    }
    return Maybe.empty();
  }

  private static Maybe<Map<String, Object>> maybeInvokeAfterToolCall(
      InvocationContext invocationContext,
      BaseTool tool,
      Map<String, Object> functionArgs,
      ToolContext toolContext,
      Map<String, Object> functionResult) {
    if (invocationContext.agent() instanceof LlmAgent) {
      LlmAgent agent = (LlmAgent) invocationContext.agent();
      Optional<List<AfterToolCallback>> callbacksOpt = agent.afterToolCallback();
      if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {
        return Maybe.empty();
      }
      List<AfterToolCallback> callbacks = callbacksOpt.get();

      return Flowable.fromIterable(callbacks)
          .concatMapMaybe(
              callback ->
                  callback.call(invocationContext, tool, functionArgs, toolContext, functionResult))
          .firstElement();
    }
    return Maybe.empty();
  }

  private static Maybe<Map<String, Object>> callTool(
      BaseTool tool, Map<String, Object> args, ToolContext toolContext) {
    Tracer tracer = Telemetry.getTracer();
    return Maybe.defer(
        () -> {
          Span span = tracer.spanBuilder("tool_call [" + tool.name() + "]").startSpan();
          try (Scope scope = span.makeCurrent()) {
            Telemetry.traceToolCall(args);
            return tool.runAsync(args, toolContext)
                .toMaybe()
                .doOnError(span::recordException)
                .doFinally(span::end);
          } catch (RuntimeException e) {
            span.recordException(e);
            span.end();
            return Maybe.error(new RuntimeException("Failed to call tool: " + tool.name(), e));
          }
        });
  }

  private static Event buildResponseEvent(
      BaseTool tool,
      Map<String, Object> response,
      ToolContext toolContext,
      InvocationContext invocationContext) {
    Tracer tracer = Telemetry.getTracer();
    Span span = tracer.spanBuilder("tool_response [" + tool.name() + "]").startSpan();
    try (Scope scope = span.makeCurrent()) {
      // use a empty placeholder response if tool response is null.
      if (response == null) {
        response = new HashMap<>();
      }

      Part partFunctionResponse =
          Part.builder()
              .functionResponse(
                  FunctionResponse.builder()
                      .id(toolContext.functionCallId().orElse(""))
                      .name(tool.name())
                      .response(response)
                      .build())
              .build();

      Event event =
          Event.builder()
              .id(Event.generateEventId())
              .invocationId(invocationContext.invocationId())
              .author(invocationContext.agent().name())
              .branch(invocationContext.branch())
              .content(
                  Optional.of(
                      Content.builder()
                          .role("user")
                          .parts(Collections.singletonList(partFunctionResponse))
                          .build()))
              .actions(toolContext.eventActions())
              .build();
      Telemetry.traceToolResponse(invocationContext, event.id(), event);
      return event;
    } finally {
      span.end();
    }
  }

  private Functions() {}
}

/*
 * 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.web;

import com.google.adk.agents.BaseAgent;
import com.google.adk.web.config.AgentLoadingProperties;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.PrintWriter;
import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.net.URISyntaxException;
import java.net.URL;
import java.net.URLClassLoader;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.StandardCopyOption;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.eclipse.jdt.core.compiler.batch.BatchCompiler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;

/**
 * Dynamically compiles and loads ADK {@link BaseAgent} implementations from source files. It
 * orchestrates the discovery of the ADK core JAR, compilation of agent sources using the Eclipse
 * JDT (ECJ) compiler, and loading of compiled agents into isolated classloaders. Agents are
 * identified by a public static field named {@code ROOT_AGENT}. Supports agent organization in
 * subdirectories or as individual {@code .java} files.
 */
@Service
public class AgentCompilerLoader {
  private static final Logger logger = LoggerFactory.getLogger(AgentCompilerLoader.class);
  private final AgentLoadingProperties properties;
  private Path compiledAgentsOutputDir;
  private final String adkCoreJarPathForCompilation;

  /**
   * Initializes the loader with agent configuration and proactively attempts to locate the ADK core
   * JAR. This JAR, containing {@link BaseAgent} and other core ADK types, is crucial for agent
   * compilation. The location strategy (see {@link #locateAndPrepareAdkCoreJar()}) includes
   * handling directly available JARs and extracting nested JARs (e.g., in Spring Boot fat JARs) to
   * ensure it's available for the compilation classpath.
   *
   * @param properties Configuration detailing agent source locations and compilation settings.
   */
  public AgentCompilerLoader(AgentLoadingProperties properties) {
    this.properties = properties;
    this.adkCoreJarPathForCompilation = locateAndPrepareAdkCoreJar();
  }

  /**
   * Attempts to find the ADK core JAR, which provides {@link BaseAgent} and essential ADK classes
   * required for compiling dynamically loaded agents.
   *
   *
Strategies include:
   *
   *

   *
Checking if {@code BaseAgent.class} is loaded from a plain {@code .jar} file on the
   *       classpath.
   *
Detecting and extracting the ADK core JAR if it's nested within a "fat JAR" (e.g., {@code
   *       BOOT-INF/lib/} in Spring Boot applications). The extracted JAR is placed in a temporary
   *       file for use during compilation.
   *

   *
   * If located, its absolute path is returned for explicit inclusion in the compiler's classpath.
   * Returns an empty string if the JAR cannot be reliably pinpointed through these specific means,
   * in which case compilation will rely on broader classpath introspection.
   *
   * @return Absolute path to the ADK core JAR if found and prepared, otherwise an empty string.
   */
  private String locateAndPrepareAdkCoreJar() {
    try {
      URL agentClassUrl = BaseAgent.class.getProtectionDomain().getCodeSource().getLocation();
      if (agentClassUrl == null) {
        logger.warn("Could not get location for BaseAgent.class. ADK Core JAR might not be found.");
        return "";
      }
      logger.debug("BaseAgent.class loaded from: {}", agentClassUrl);

      if ("file".equals(agentClassUrl.getProtocol())) {
        Path path = Paths.get(agentClassUrl.toURI());
        if (path.toString().endsWith(".jar") && Files.exists(path)) {
          logger.debug(
              "ADK Core JAR (or where BaseAgent resides) found directly on classpath: {}",
              path.toAbsolutePath());
          return path.toAbsolutePath().toString();
        } else if (Files.isDirectory(path)) {
          logger.debug(
              "BaseAgent.class found in directory (e.g., target/classes): {}. This path will be"
                  + " part of classloader introspection.",
              path.toAbsolutePath());
          return "";
        }
      } else if ("jar".equals(agentClassUrl.getProtocol())) { // Typically for nested JARs
        String urlPath = agentClassUrl.getPath();
        if (urlPath.startsWith("file:")) {
          urlPath = urlPath.substring("file:".length());
        }
        int firstSeparator = urlPath.indexOf("!/");
        if (firstSeparator == -1) {
          logger.warn("Malformed JAR URL for BaseAgent.class: {}", agentClassUrl);
          return "";
        }

        String mainJarPath = urlPath.substring(0, firstSeparator);
        String nestedPath = urlPath.substring(firstSeparator);

        if (nestedPath.startsWith("!/BOOT-INF/lib/") && nestedPath.contains("google-adk-")) {
          int nestedJarStartInPath = "!/BOOT-INF/lib/".length();
          int nestedJarEndInPath = nestedPath.indexOf("!/", nestedJarStartInPath);
          if (nestedJarEndInPath > 0) {
            String nestedJarName = nestedPath.substring(nestedJarStartInPath, nestedJarEndInPath);
            String nestedJarUrlString =
                "jar:file:" + mainJarPath + "!/BOOT-INF/lib/" + nestedJarName;

            Path tempFile = Files.createTempFile("adk-core-extracted-", ".jar");
            try (InputStream is = new URL(nestedJarUrlString).openStream()) {
              Files.copy(is, tempFile, StandardCopyOption.REPLACE_EXISTING);
            }
            tempFile.toFile().deleteOnExit();
            logger.debug(
                "Extracted ADK Core JAR '{}' from nested location to: {}",
                nestedJarName,
                tempFile.toAbsolutePath());
            return tempFile.toAbsolutePath().toString();
          }
        } else if (mainJarPath.contains("google-adk-") && mainJarPath.endsWith(".jar")) {
          File adkJar = new File(mainJarPath);
          if (adkJar.exists()) {
            logger.debug("ADK Core JAR identified as the outer JAR: {}", adkJar.getAbsolutePath());
            return adkJar.getAbsolutePath();
          }
        }
      }
    } catch (Exception e) {
      logger.error("Error trying to locate or extract ADK Core JAR", e);
    }
    logger.warn(
        "ADK Core JAR could not be reliably located for compilation via locateAndPrepareAdkCoreJar."
            + " Relying on classloader introspection.");
    return "";
  }

  /**
   * Discovers, compiles, and loads agents from the configured source directory.
   *
   *
The process for each potential "agent unit" (a subdirectory or a root {@code .java} file):
   *
   *

   *
Collects {@code .java} source files.
   *
Compiles these sources using ECJ (see {@link #compileSourcesWithECJ(List, Path)}) into a
   *       temporary, unit-specific output directory. This directory is cleaned up on JVM exit.
   *
Creates a dedicated {@link URLClassLoader} for the compiled unit, isolating its classes.
   *
Scans compiled classes for a public static field {@code ROOT_AGENT} assignable to {@link
   *       BaseAgent}. This field serves as the designated entry point for an agent.
   *
Instantiates and stores the {@link BaseAgent} if found, keyed by its name.
   *

   *
   * This approach allows for dynamic addition of agents without pre-compilation and supports
   * independent classpaths per agent unit if needed (though current implementation uses a shared
   * parent classloader).
   *
   * @return A map of successfully loaded agent names to their {@link BaseAgent} instances. Returns
   *     an empty map if the source directory isn't configured or no agents are found.
   * @throws IOException If an I/O error occurs (e.g., creating temp directories, reading sources).
   */
  public Map<String, BaseAgent> loadAgents() throws IOException {
    if (properties.getSourceDir() == null || properties.getSourceDir().isEmpty()) {
      logger.info(
          "Agent source directory (adk.agents.source-dir) not configured. No dynamic agents will be"
              + " loaded.");
      return Collections.emptyMap();
    }

    Path agentsSourceRoot = Paths.get(properties.getSourceDir());
    if (!Files.isDirectory(agentsSourceRoot)) {
      logger.warn("Agent source directory does not exist: {}", agentsSourceRoot);
      return Collections.emptyMap();
    }

    this.compiledAgentsOutputDir = Files.createTempDirectory("adk-compiled-agents-");
    this.compiledAgentsOutputDir.toFile().deleteOnExit();
    logger.debug("Compiling agents from {} to {}", agentsSourceRoot, compiledAgentsOutputDir);

    Map<String, BaseAgent> loadedAgents = new HashMap<>();

    try (Stream<Path> stream = Files.list(agentsSourceRoot)) {
      List<Path> entries = stream.collect(Collectors.toList());

      for (Path entry : entries) {
        List<String> javaFilesToCompile = new ArrayList<>();
        String agentUnitName;

        if (Files.isDirectory(entry)) {
          agentUnitName = entry.getFileName().toString();
          logger.debug("Processing agent sources from directory: {}", agentUnitName);
          try (Stream<Path> javaFilesStream =
              Files.walk(entry)
                  .filter(p -> p.toString().endsWith(".java") && Files.isRegularFile(p))) {
            javaFilesToCompile =
                javaFilesStream
                    .map(p -> p.toAbsolutePath().toString())
                    .collect(Collectors.toList());
          }
        } else if (Files.isRegularFile(entry) && entry.getFileName().toString().endsWith(".java")) {
          String fileName = entry.getFileName().toString();
          agentUnitName = fileName.substring(0, fileName.length() - ".java".length());
          logger.debug("Processing agent source file: {}", entry.getFileName());
          javaFilesToCompile.add(entry.toAbsolutePath().toString());
        } else {
          logger.trace("Skipping non-agent entry in agent source root: {}", entry.getFileName());
          continue;
        }

        if (javaFilesToCompile.isEmpty()) {
          logger.debug("No .java files found for agent unit: {}", agentUnitName);
          continue;
        }

        Path unitSpecificOutputDir = compiledAgentsOutputDir.resolve(agentUnitName);
        Files.createDirectories(unitSpecificOutputDir);

        boolean compilationSuccess =
            compileSourcesWithECJ(javaFilesToCompile, unitSpecificOutputDir);

        if (compilationSuccess) {
          try {
            List<URL> classLoaderUrls = new ArrayList<>();
            classLoaderUrls.add(unitSpecificOutputDir.toUri().toURL());

            URLClassLoader agentClassLoader =
                new URLClassLoader(
                    classLoaderUrls.toArray(new URL[0]),
                    AgentCompilerLoader.class.getClassLoader());

            Files.walk(unitSpecificOutputDir)
                .filter(p -> p.toString().endsWith(".class"))
                .forEach(
                    classFile -> {
                      try {
                        String relativePath =
                            unitSpecificOutputDir.relativize(classFile).toString();
                        String className =
                            relativePath
                                .substring(0, relativePath.length() - ".class".length())
                                .replace(File.separatorChar, '.');

                        Class<?> loadedClass = agentClassLoader.loadClass(className);
                        Field rootAgentField = null;
                        try {
                          rootAgentField = loadedClass.getField("ROOT_AGENT");
                        } catch (NoSuchFieldException e) {
                          return;
                        }

                        if (Modifier.isStatic(rootAgentField.getModifiers())
                            && BaseAgent.class.isAssignableFrom(rootAgentField.getType())) {
                          BaseAgent agentInstance = (BaseAgent) rootAgentField.get(null);
                          if (agentInstance != null) {
                            if (loadedAgents.containsKey(agentInstance.name())) {
                              logger.warn(
                                  "Found another agent with name {}. This will overwrite the"
                                      + " original agent loaded with this name from unit {} using"
                                      + " class {}",
                                  agentInstance.name(),
                                  agentUnitName,
                                  className);
                            }
                            loadedAgents.put(agentInstance.name(), agentInstance);
                            logger.debug(
                                "Successfully loaded agent '{}' from unit: {} using class {}",
                                agentInstance.name(),
                                agentUnitName,
                                className);
                          } else {
                            logger.warn(
                                "ROOT_AGENT field in class {} from unit {} was null",
                                className,
                                agentUnitName);
                          }
                        }
                      } catch (ClassNotFoundException | IllegalAccessException e) {
                        logger.error(
                            "Error loading or accessing agent from class file {} for unit {}",
                            classFile,
                            agentUnitName,
                            e);
                      } catch (Exception e) {
                        logger.error(
                            "Unexpected error processing class file {} for unit {}",
                            classFile,
                            agentUnitName,
                            e);
                      }
                    });
          } catch (Exception e) {
            logger.error(
                "Error during class loading setup for unit {}: {}",
                agentUnitName,
                e.getMessage(),
                e);
          }
        } else {
          logger.error("Compilation failed for agent unit: {}", agentUnitName);
        }
      }
    }
    return loadedAgents;
  }

  /**
   * Compiles the given Java source files using the Eclipse JDT (ECJ) batch compiler.
   *
   *
Key aspects of the compilation process:
   *
   *

   *
Sets Java version to 17 (to align with core ADK library) and suppresses warnings by
   *       default.
   *
Constructs the compilation classpath by:
   *

   *
Introspecting the current classloader hierarchy ({@link URLClassLoader} instances)
   *             to gather available JARs and class directories.
   *
Falling back to {@code System.getProperty("java.class.path")} if introspection
   *             yields no results.
   *
Explicitly adding the ADK Core JAR path (determined by {@link
   *             #locateAndPrepareAdkCoreJar()}) to ensure {@link BaseAgent} and related types are
   *             resolvable.
   *
Appending any user-defined classpath entries from {@link
   *             AgentLoadingProperties#getCompileClasspath()}.
   *

   *
Outputs compiled {@code .class} files to the specified {@code outputDir}.
   *

   *
   * This method aims to provide a robust classpath for compiling agents in various runtime
   * environments, including IDEs, standard Java executions, and fat JAR deployments.
   *
   * @param javaFilePaths A list of absolute paths to {@code .java} files to be compiled.
   * @param outputDir The directory where compiled {@code .class} files will be placed.
   * @return {@code true} if compilation succeeds, {@code false} otherwise.
   */
  private boolean compileSourcesWithECJ(List<String> javaFilePaths, Path outputDir) {
    List<String> ecjArgs = new ArrayList<>();
    ecjArgs.add("-17"); // Java version
    ecjArgs.add("-nowarn");
    ecjArgs.add("-d");
    ecjArgs.add(outputDir.toAbsolutePath().toString());

    Set<String> classpathEntries = new LinkedHashSet<>();

    logger.debug("Attempting to derive ECJ classpath from classloader hierarchy...");
    ClassLoader currentClassLoader = AgentCompilerLoader.class.getClassLoader();
    int classLoaderCount = 0;
    while (currentClassLoader != null) {
      classLoaderCount++;
      logger.debug(
          "Inspecting classloader ({}) : {}",
          classLoaderCount,
          currentClassLoader.getClass().getName());
      if (currentClassLoader instanceof java.net.URLClassLoader) {
        URL[] urls = ((java.net.URLClassLoader) currentClassLoader).getURLs();
        logger.debug(
            "  Found {} URLs in URLClassLoader {}",
            urls.length,
            currentClassLoader.getClass().getName());
        for (URL url : urls) {
          try {
            if ("file".equals(url.getProtocol())) {
              String path = Paths.get(url.toURI()).toString();
              classpathEntries.add(path);
              logger.trace("    Added to ECJ classpath: {}", path);
            } else {
              logger.debug(
                  "    Skipping non-file URL from classloader {}: {}",
                  currentClassLoader.getClass().getName(),
                  url);
            }
          } catch (URISyntaxException | IllegalArgumentException e) {
            logger.warn(
                "    Could not convert URL to path or add to classpath from {}: {} (Error: {})",
                currentClassLoader.getClass().getName(),
                url,
                e.getMessage());
          } catch (Exception e) {
            logger.warn(
                "    Unexpected error converting URL to path from {}: {} (Error: {})",
                currentClassLoader.getClass().getName(),
                url,
                e.getMessage(),
                e);
          }
        }
      }
      currentClassLoader = currentClassLoader.getParent();
    }

    if (classpathEntries.isEmpty()) {
      logger.warn(
          "No classpath entries derived from classloader hierarchy. "
              + "Falling back to System.getProperty(\"java.class.path\").");
      String systemClasspath = System.getProperty("java.class.path");
      if (systemClasspath != null && !systemClasspath.isEmpty()) {
        logger.debug("Using system classpath for ECJ (fallback): {}", systemClasspath);
        classpathEntries.addAll(Arrays.asList(systemClasspath.split(File.pathSeparator)));
      } else {
        logger.error("System classpath (java.class.path) is also null or empty.");
      }
    }

    if (this.adkCoreJarPathForCompilation != null && !this.adkCoreJarPathForCompilation.isEmpty()) {
      if (!classpathEntries.contains(this.adkCoreJarPathForCompilation)) {
        logger.debug(
            "Adding ADK Core JAR path explicitly to ECJ classpath: {}",
            this.adkCoreJarPathForCompilation);
        classpathEntries.add(this.adkCoreJarPathForCompilation);
      } else {
        logger.debug(
            "ADK Core JAR path ({}) already found in derived ECJ classpath.",
            this.adkCoreJarPathForCompilation);
      }
    } else if (classpathEntries.stream().noneMatch(p -> p.contains("google-adk"))) {
      logger.error(
          "ADK Core JAR path is missing and no 'google-adk' JAR found in derived classpath. "
              + "Compilation will likely fail to find BaseAgent.");
    }

    if (properties.getCompileClasspath() != null && !properties.getCompileClasspath().isEmpty()) {
      String userClasspath = properties.getCompileClasspath();
      logger.info(
          "Appending user-defined classpath (adk.agents.compile-classpath) to ECJ: {}",
          userClasspath);
      classpathEntries.addAll(Arrays.asList(userClasspath.split(File.pathSeparator)));
    }

    if (!classpathEntries.isEmpty()) {
      String effectiveClasspath =
          classpathEntries.stream().collect(Collectors.joining(File.pathSeparator));
      ecjArgs.add("-cp");
      ecjArgs.add(effectiveClasspath);
      logger.debug("Constructed ECJ classpath with {} entries", classpathEntries.size());
      logger.debug("Final effective ECJ classpath: {}", effectiveClasspath);
    } else {
      logger.error("ECJ Classpath is empty after all attempts. Compilation will fail.");
      return false;
    }

    ecjArgs.addAll(javaFilePaths);

    logger.debug("ECJ Args: {}", String.join(" ", ecjArgs));

    PrintWriter outWriter = new PrintWriter(System.out, true);
    PrintWriter errWriter = new PrintWriter(System.err, true);

    boolean success =
        BatchCompiler.compile(ecjArgs.toArray(new String[0]), outWriter, errWriter, null);
    if (!success) {
      logger.error("ECJ Compilation failed. See console output for details.");
    }
    return success;
  }
}

/*
 * 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.web;

import com.google.adk.agents.BaseAgent;
import com.google.adk.agents.LlmAgent;
import com.google.adk.tools.AgentTool;
import com.google.adk.tools.BaseTool;
import com.google.adk.tools.FunctionTool;
import com.google.adk.tools.retrieval.BaseRetrievalTool;
import guru.nidi.graphviz.attribute.Arrow;
import guru.nidi.graphviz.attribute.Color;
import guru.nidi.graphviz.attribute.Label;
import guru.nidi.graphviz.attribute.Rank;
import guru.nidi.graphviz.attribute.Shape;
import guru.nidi.graphviz.attribute.Style;
import guru.nidi.graphviz.engine.Format;
import guru.nidi.graphviz.engine.Graphviz;
import guru.nidi.graphviz.model.Factory;
import guru.nidi.graphviz.model.Link;
import guru.nidi.graphviz.model.MutableGraph;
import guru.nidi.graphviz.model.MutableNode;
import java.util.HashSet;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Utility class to generate Graphviz DOT representations of Agent structures. */
public class AgentGraphGenerator {

  private static final Logger log = LoggerFactory.getLogger(AgentGraphGenerator.class);

  private static final Color DARK_GREEN = Color.rgb("#0F5223");
  private static final Color LIGHT_GREEN = Color.rgb("#69CB87");
  private static final Color LIGHT_GRAY = Color.rgb("#CCCCCC");
  private static final Color BG_COLOR = Color.rgb("#333537");

  /**
   * Generates the DOT source string for the agent graph.
   *
   * @param rootAgent The root agent of the structure.
   * @param highlightPairs A list where each inner list contains two strings (fromNode, toNode)
   *     representing an edge to highlight. Order matters for direction.
   * @return The DOT source string, or null if graph generation fails.
   */
  public static Optional<String> getAgentGraphDotSource(
      BaseAgent rootAgent, List<List<String>> highlightPairs) {
    log.debug(
        "Building agent graph with root: {}, highlights: {}", rootAgent.name(), highlightPairs);
    try {
      MutableGraph graph =
          Factory.mutGraph("AgentGraph")
              .setDirected(true)
              .graphAttrs()
              .add(Rank.dir(Rank.RankDir.LEFT_TO_RIGHT), BG_COLOR.background());

      Set<String> visitedNodes = new HashSet<>();
      buildGraphRecursive(graph, rootAgent, highlightPairs, visitedNodes);

      String dotSource = Graphviz.fromGraph(graph).render(Format.DOT).toString();
      log.debug("Generated DOT source successfully.");
      return Optional.of(dotSource);
    } catch (Exception e) {
      log.error("Error generating agent graph DOT source", e);
      return Optional.empty();
    }
  }

  /** Recursively builds the graph structure. */
  private static void buildGraphRecursive(
      MutableGraph graph,
      BaseAgent agent,
      List<List<String>> highlightPairs,
      Set<String> visitedNodes) {
    if (agent == null || visitedNodes.contains(getNodeName(agent))) {
      return;
    }

    drawNode(graph, agent, highlightPairs, visitedNodes);

    if (agent.subAgents() != null) {
      for (BaseAgent subAgent : agent.subAgents()) {
        if (subAgent != null) {
          drawEdge(graph, getNodeName(agent), getNodeName(subAgent), highlightPairs);
          buildGraphRecursive(graph, subAgent, highlightPairs, visitedNodes);
        }
      }
    }

    if (agent instanceof LlmAgent) {
      LlmAgent llmAgent = (LlmAgent) agent;
      List<BaseTool> tools = llmAgent.canonicalTools().toList().blockingGet();
      if (tools != null) {
        for (BaseTool tool : tools) {
          if (tool != null) {
            drawNode(graph, tool, highlightPairs, visitedNodes);
            drawEdge(graph, getNodeName(agent), getNodeName(tool), highlightPairs);
          }
        }
      }
    }
  }

  /** Draws a node for an agent or tool, applying highlighting if applicable. */
  private static void drawNode(
      MutableGraph graph,
      Object toolOrAgent,
      List<List<String>> highlightPairs,
      Set<String> visitedNodes) {
    String name = getNodeName(toolOrAgent);
    if (name == null || name.isEmpty() || visitedNodes.contains(name)) {
      return;
    }

    Shape shape = getNodeShape(toolOrAgent);
    String caption = getNodeCaption(toolOrAgent);
    boolean isHighlighted = isNodeHighlighted(name, highlightPairs);

    MutableNode node =
        Factory.mutNode(name).add(Label.of(caption)).add(shape).add(LIGHT_GRAY.font());
    if (isHighlighted) {
      node.add(Style.FILLED);
      node.add(DARK_GREEN);
    } else {
      node.add(Style.ROUNDED);
      node.add(LIGHT_GRAY);
    }
    graph.add(node);
    visitedNodes.add(name);
    log.trace(
        "Added node: name={}, caption={}, shape={}, highlighted={}",
        name,
        caption,
        shape,
        isHighlighted);
  }

  /** Draws an edge between two nodes, applying highlighting if applicable. */
  private static void drawEdge(
      MutableGraph graph, String fromName, String toName, List<List<String>> highlightPairs) {
    if (fromName == null || fromName.isEmpty() || toName == null || toName.isEmpty()) {
      log.warn(
          "Skipping edge draw due to null or empty name: from='{}', to='{}'", fromName, toName);
      return;
    }

    Optional<Boolean> highlightForward = isEdgeHighlighted(fromName, toName, highlightPairs);
    Link link = Factory.to(Factory.mutNode(toName));

    if (highlightForward.isPresent()) {
      link = link.with(LIGHT_GREEN);
      if (!highlightForward.get()) { // If true, means b->a was highlighted, draw reverse arrow
        link = link.with(Arrow.NORMAL.dir(Arrow.DirType.BACK));
      } else {
        link = link.with(Arrow.NORMAL);
      }
    } else {
      link = link.with(LIGHT_GRAY, Arrow.NONE);
    }

    graph.add(Factory.mutNode(fromName).addLink(link));
    log.trace(
        "Added edge: from={}, to={}, highlighted={}",
        fromName,
        toName,
        highlightForward.isPresent());
  }

  private static String getNodeName(Object toolOrAgent) {
    if (toolOrAgent instanceof BaseAgent) {
      return ((BaseAgent) toolOrAgent).name();
    } else if (toolOrAgent instanceof BaseTool) {
      return ((BaseTool) toolOrAgent).name();
    } else {
      log.warn("Unsupported type for getNodeName: {}", toolOrAgent.getClass().getName());
      return "unknown_" + toolOrAgent.hashCode();
    }
  }

  private static String getNodeCaption(Object toolOrAgent) {
    String name = getNodeName(toolOrAgent); // Get name first
    if (toolOrAgent instanceof BaseAgent) {
      return "🤖 " + name;
    } else if (toolOrAgent instanceof BaseRetrievalTool) {
      return "🔎 " + name;
    } else if (toolOrAgent instanceof FunctionTool) {
      return "🔧 " + name;
    } else if (toolOrAgent instanceof AgentTool) {
      return "🤖 " + name;
    } else if (toolOrAgent instanceof BaseTool) {
      return "🔧 " + name;
    } else {
      log.warn("Unsupported type for getNodeCaption: {}", toolOrAgent.getClass().getName());
      return "❓ " + name;
    }
  }

  private static Shape getNodeShape(Object toolOrAgent) {
    if (toolOrAgent instanceof BaseAgent) {
      return Shape.ELLIPSE;
    } else if (toolOrAgent instanceof BaseRetrievalTool) {
      return Shape.CYLINDER;
    } else if (toolOrAgent instanceof FunctionTool) {
      return Shape.BOX;
    } else if (toolOrAgent instanceof BaseTool) {
      return Shape.BOX;
    } else {
      log.warn("Unsupported type for getNodeShape: {}", toolOrAgent.getClass().getName());
      return Shape.EGG;
    }
  }

  private static boolean isNodeHighlighted(String nodeName, List<List<String>> highlightPairs) {
    if (highlightPairs == null || nodeName == null) {
      return false;
    }
    for (List<String> pair : highlightPairs) {
      if (pair != null && pair.contains(nodeName)) {
        return true;
      }
    }
    return false;
  }

  /**
   * Checks if an edge should be highlighted. Returns Optional<Boolean>: empty=no, true=forward,
   * false=backward
   */
  private static Optional<Boolean> isEdgeHighlighted(
      String fromName, String toName, List<List<String>> highlightPairs) {
    if (highlightPairs == null || fromName == null || toName == null) {
      return Optional.empty();
    }
    for (List<String> pair : highlightPairs) {
      if (pair != null && pair.size() == 2) {
        String pairFrom = pair.get(0);
        String pairTo = pair.get(1);
        if (fromName.equals(pairFrom) && toName.equals(pairTo)) {
          return Optional.of(true);
        }
        if (fromName.equals(pairTo) && toName.equals(pairFrom)) {
          return Optional.of(false);
        }
      }
    }
    return Optional.empty();
  }
}

/*
 * 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.events;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.genai.types.Part;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import javax.annotation.Nullable;

/** Represents the actions attached to an event. */
// TODO - b/414081262 make json wire camelCase
@JsonDeserialize(builder = EventActions.Builder.class)
public class EventActions {

  private Optional<Boolean> skipSummarization = Optional.empty();
  private ConcurrentMap<String, Object> stateDelta = new ConcurrentHashMap<>();
  private ConcurrentMap<String, Part> artifactDelta = new ConcurrentHashMap<>();
  private Optional<String> transferToAgent = Optional.empty();
  private Optional<Boolean> escalate = Optional.empty();
  private ConcurrentMap<String, ConcurrentMap<String, Object>> requestedAuthConfigs =
      new ConcurrentHashMap<>();
  private Optional<Boolean> endInvocation = Optional.empty();

  /** Default constructor for Jackson. */
  public EventActions() {}

  @JsonProperty("skipSummarization")
  public Optional<Boolean> skipSummarization() {
    return skipSummarization;
  }

  public void setSkipSummarization(@Nullable Boolean skipSummarization) {
    this.skipSummarization = Optional.ofNullable(skipSummarization);
  }

  public void setSkipSummarization(Optional<Boolean> skipSummarization) {
    this.skipSummarization = skipSummarization;
  }

  public void setSkipSummarization(boolean skipSummarization) {
    this.skipSummarization = Optional.of(skipSummarization);
  }

  @JsonProperty("stateDelta")
  public ConcurrentMap<String, Object> stateDelta() {
    return stateDelta;
  }

  public void setStateDelta(ConcurrentMap<String, Object> stateDelta) {
    this.stateDelta = stateDelta;
  }

  @JsonProperty("artifactDelta")
  public ConcurrentMap<String, Part> artifactDelta() {
    return artifactDelta;
  }

  public void setArtifactDelta(ConcurrentMap<String, Part> artifactDelta) {
    this.artifactDelta = artifactDelta;
  }

  @JsonProperty("transferToAgent")
  public Optional<String> transferToAgent() {
    return transferToAgent;
  }

  public void setTransferToAgent(Optional<String> transferToAgent) {
    this.transferToAgent = transferToAgent;
  }

  public void setTransferToAgent(String transferToAgent) {
    this.transferToAgent = Optional.ofNullable(transferToAgent);
  }

  @JsonProperty("escalate")
  public Optional<Boolean> escalate() {
    return escalate;
  }

  public void setEscalate(Optional<Boolean> escalate) {
    this.escalate = escalate;
  }

  public void setEscalate(boolean escalate) {
    this.escalate = Optional.of(escalate);
  }

  @JsonProperty("requestedAuthConfigs")
  public ConcurrentMap<String, ConcurrentMap<String, Object>> requestedAuthConfigs() {
    return requestedAuthConfigs;
  }

  public void setRequestedAuthConfigs(
      ConcurrentMap<String, ConcurrentMap<String, Object>> requestedAuthConfigs) {
    this.requestedAuthConfigs = requestedAuthConfigs;
  }

  @JsonProperty("endInvocation")
  public Optional<Boolean> endInvocation() {
    return endInvocation;
  }

  public void setEndInvocation(Optional<Boolean> endInvocation) {
    this.endInvocation = endInvocation;
  }

  public void setEndInvocation(boolean endInvocation) {
    this.endInvocation = Optional.of(endInvocation);
  }

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

  public Builder toBuilder() {
    return new Builder(this);
  }

  @Override
  public boolean equals(Object o) {
    if (this == o) {
      return true;
    }
    if (!(o instanceof EventActions that)) {
      return false;
    }
    return Objects.equals(skipSummarization, that.skipSummarization)
        && Objects.equals(stateDelta, that.stateDelta)
        && Objects.equals(artifactDelta, that.artifactDelta)
        && Objects.equals(transferToAgent, that.transferToAgent)
        && Objects.equals(escalate, that.escalate)
        && Objects.equals(requestedAuthConfigs, that.requestedAuthConfigs)
        && Objects.equals(endInvocation, that.endInvocation);
  }

  @Override
  public int hashCode() {
    return Objects.hash(
        skipSummarization,
        stateDelta,
        artifactDelta,
        transferToAgent,
        escalate,
        requestedAuthConfigs,
        endInvocation);
  }

  /** Builder for {@link EventActions}. */
  public static class Builder {
    private Optional<Boolean> skipSummarization = Optional.empty();
    private ConcurrentMap<String, Object> stateDelta = new ConcurrentHashMap<>();
    private ConcurrentMap<String, Part> artifactDelta = new ConcurrentHashMap<>();
    private Optional<String> transferToAgent = Optional.empty();
    private Optional<Boolean> escalate = Optional.empty();
    private ConcurrentMap<String, ConcurrentMap<String, Object>> requestedAuthConfigs =
        new ConcurrentHashMap<>();
    private Optional<Boolean> endInvocation = Optional.empty();

    public Builder() {}

    private Builder(EventActions eventActions) {
      this.skipSummarization = eventActions.skipSummarization();
      this.stateDelta = new ConcurrentHashMap<>(eventActions.stateDelta());
      this.artifactDelta = new ConcurrentHashMap<>(eventActions.artifactDelta());
      this.transferToAgent = eventActions.transferToAgent();
      this.escalate = eventActions.escalate();
      this.requestedAuthConfigs = new ConcurrentHashMap<>(eventActions.requestedAuthConfigs());
      this.endInvocation = eventActions.endInvocation();
    }

    @CanIgnoreReturnValue
    @JsonProperty("skipSummarization")
    public Builder skipSummarization(boolean skipSummarization) {
      this.skipSummarization = Optional.of(skipSummarization);
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("stateDelta")
    public Builder stateDelta(ConcurrentMap<String, Object> value) {
      this.stateDelta = value;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("artifactDelta")
    public Builder artifactDelta(ConcurrentMap<String, Part> value) {
      this.artifactDelta = value;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("transferToAgent")
    public Builder transferToAgent(String agentId) {
      this.transferToAgent = Optional.ofNullable(agentId);
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("escalate")
    public Builder escalate(boolean escalate) {
      this.escalate = Optional.of(escalate);
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("requestedAuthConfigs")
    public Builder requestedAuthConfigs(
        ConcurrentMap<String, ConcurrentMap<String, Object>> value) {
      this.requestedAuthConfigs = value;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("endInvocation")
    public Builder endInvocation(boolean endInvocation) {
      this.endInvocation = Optional.of(endInvocation);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder merge(EventActions other) {
      if (other.skipSummarization().isPresent()) {
        this.skipSummarization = other.skipSummarization();
      }
      if (other.stateDelta() != null) {
        this.stateDelta.putAll(other.stateDelta());
      }
      if (other.artifactDelta() != null) {
        this.artifactDelta.putAll(other.artifactDelta());
      }
      if (other.transferToAgent().isPresent()) {
        this.transferToAgent = other.transferToAgent();
      }
      if (other.escalate().isPresent()) {
        this.escalate = other.escalate();
      }
      if (other.requestedAuthConfigs() != null) {
        this.requestedAuthConfigs.putAll(other.requestedAuthConfigs());
      }
      if (other.endInvocation().isPresent()) {
        this.endInvocation = other.endInvocation();
      }
      return this;
    }

    public EventActions build() {
      EventActions eventActions = new EventActions();
      eventActions.setSkipSummarization(this.skipSummarization);
      eventActions.setStateDelta(this.stateDelta);
      eventActions.setArtifactDelta(this.artifactDelta);
      eventActions.setTransferToAgent(this.transferToAgent);
      eventActions.setEscalate(this.escalate);
      eventActions.setRequestedAuthConfigs(this.requestedAuthConfigs);
      eventActions.setEndInvocation(this.endInvocation);
      return eventActions;
    }
  }
}

/*
 * 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.web;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.adk.JsonBaseModel;
import com.google.adk.agents.BaseAgent;
import com.google.adk.agents.LiveRequest;
import com.google.adk.agents.LiveRequestQueue;
import com.google.adk.agents.RunConfig;
import com.google.adk.agents.RunConfig.StreamingMode;
import com.google.adk.artifacts.BaseArtifactService;
import com.google.adk.artifacts.InMemoryArtifactService;
import com.google.adk.artifacts.ListArtifactsResponse;
import com.google.adk.events.Event;
import com.google.adk.runner.Runner;
import com.google.adk.sessions.BaseSessionService;
import com.google.adk.sessions.InMemorySessionService;
import com.google.adk.sessions.ListSessionsResponse;
import com.google.adk.sessions.Session;
import com.google.adk.web.config.AgentLoadingProperties;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.genai.types.Blob;
import com.google.genai.types.Content;
import com.google.genai.types.FunctionCall;
import com.google.genai.types.FunctionResponse;
import com.google.genai.types.Modality;
import com.google.genai.types.Part;
import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.api.common.AttributeKey;
import io.opentelemetry.api.common.Attributes;
import io.opentelemetry.api.trace.SpanId;
import io.opentelemetry.sdk.OpenTelemetrySdk;
import io.opentelemetry.sdk.common.CompletableResultCode;
import io.opentelemetry.sdk.resources.Resource;
import io.opentelemetry.sdk.trace.SdkTracerProvider;
import io.opentelemetry.sdk.trace.data.SpanData;
import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
import io.opentelemetry.sdk.trace.export.SpanExporter;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import io.reactivex.rxjava3.disposables.Disposable;
import io.reactivex.rxjava3.schedulers.Schedulers;
import java.io.IOException;
import java.net.URI;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.boot.context.properties.ConfigurationPropertiesScan;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.ComponentScan;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Component;
import org.springframework.util.MultiValueMap;
import org.springframework.web.bind.annotation.DeleteMapping;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.server.ResponseStatusException;
import org.springframework.web.servlet.config.annotation.ResourceHandlerRegistry;
import org.springframework.web.servlet.config.annotation.ViewControllerRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.util.UriComponentsBuilder;

/**
 * Single-file Spring Boot application for the Agent Server. Combines configuration, DTOs, and
 * controller logic.
 */
@SpringBootApplication
@ConfigurationPropertiesScan
@ComponentScan(basePackages = {"com.google.adk.web", "com.google.adk.web.config"})
public class AdkWebServer implements WebMvcConfigurer {

  private static final Logger log = LoggerFactory.getLogger(AdkWebServer.class);

  @Value("${adk.web.ui.dir:#{null}}")
  private String webUiDir;

  @Bean
  public BaseSessionService sessionService() {
    // TODO: Add logic to select service based on config (e.g., DB URL)
    log.info("Using InMemorySessionService");
    return new InMemorySessionService();
  }

  /**
   * Provides the singleton instance of the ArtifactService (InMemory). TODO: configure this based
   * on config (e.g., DB URL)
   *
   * @return An instance of BaseArtifactService (currently InMemoryArtifactService).
   */
  @Bean
  public BaseArtifactService artifactService() {
    log.info("Using InMemoryArtifactService");
    return new InMemoryArtifactService();
  }

  @Bean("loadedAgentRegistry")
  public Map<String, BaseAgent> loadedAgentRegistry(
      AgentCompilerLoader loader, AgentLoadingProperties props) {
    if (props.getSourceDir() == null || props.getSourceDir().isEmpty()) {
      log.info("adk.agents.source-dir not set. Initializing with an empty agent registry.");
      return Collections.emptyMap();
    }
    try {
      Map<String, BaseAgent> agents = loader.loadAgents();
      log.info("Loaded {} dynamic agent(s): {}", agents.size(), agents.keySet());
      return agents;
    } catch (IOException e) {
      log.error("Failed to load dynamic agents", e);
      return Collections.emptyMap();
    }
  }

  @Bean
  public ObjectMapper objectMapper() {
    return JsonBaseModel.getMapper();
  }

  /** Service for creating and caching Runner instances. */
  @Component
  public static class RunnerService {
    private static final Logger log = LoggerFactory.getLogger(RunnerService.class);

    private final Map<String, BaseAgent> agentRegistry;
    private final BaseArtifactService artifactService;
    private final BaseSessionService sessionService;
    private final Map<String, Runner> runnerCache = new ConcurrentHashMap<>();

    @Autowired
    public RunnerService(
        @Qualifier("loadedAgentRegistry") Map<String, BaseAgent> agentRegistry,
        BaseArtifactService artifactService,
        BaseSessionService sessionService) {
      this.agentRegistry = agentRegistry;
      this.artifactService = artifactService;
      this.sessionService = sessionService;
    }

    /**
     * Gets the Runner instance for a given application name. Handles potential agent engine ID
     * overrides.
     *
     * @param appName The application name requested by the user.
     * @return A configured Runner instance.
     */
    public Runner getRunner(String appName) {
      return runnerCache.computeIfAbsent(
          appName,
          key -> {
            BaseAgent agent = agentRegistry.get(key);
            if (agent == null) {
              log.error(
                  "Agent/App named '{}' not found in registry. Available apps: {}",
                  key,
                  agentRegistry.keySet());
              throw new ResponseStatusException(
                  HttpStatus.NOT_FOUND, "Agent/App not found: " + key);
            }
            log.info(
                "RunnerService: Creating Runner for appName: {}, using agent" + " definition: {}",
                appName,
                agent.name());
            return new Runner(agent, appName, this.artifactService, this.sessionService);
          });
    }
  }

  /** Configuration class for OpenTelemetry, setting up the tracer provider and span exporter. */
  @Configuration
  public static class OpenTelemetryConfig {
    private static final Logger otelLog = LoggerFactory.getLogger(OpenTelemetryConfig.class);

    @Bean
    public ApiServerSpanExporter apiServerSpanExporter() {
      return new ApiServerSpanExporter();
    }

    @Bean(destroyMethod = "shutdown")
    public SdkTracerProvider sdkTracerProvider(ApiServerSpanExporter apiServerSpanExporter) {
      otelLog.debug("Configuring SdkTracerProvider with ApiServerSpanExporter.");
      Resource resource =
          Resource.getDefault()
              .merge(
                  Resource.create(
                      Attributes.of(AttributeKey.stringKey("service.name"), "adk-web-server")));

      return SdkTracerProvider.builder()
          .addSpanProcessor(SimpleSpanProcessor.create(apiServerSpanExporter))
          .setResource(resource)
          .build();
    }

    @Bean
    public OpenTelemetry openTelemetrySdk(SdkTracerProvider sdkTracerProvider) {
      otelLog.debug("Configuring OpenTelemetrySdk and registering globally.");
      OpenTelemetrySdk otelSdk =
          OpenTelemetrySdk.builder().setTracerProvider(sdkTracerProvider).buildAndRegisterGlobal();

      Runtime.getRuntime().addShutdownHook(new Thread(otelSdk::close));
      return otelSdk;
    }
  }

  /**
   * A custom SpanExporter that stores relevant span data. It handles two types of trace data
   * storage: 1. Event-ID based: Stores attributes of specific spans (call_llm, send_data,
   * tool_response) keyed by `gcp.vertex.agent.event_id`. This is used for debugging individual
   * events. 2. Session-ID based: Stores all exported spans and maintains a mapping from
   * `session_id` (extracted from `call_llm` spans) to a list of `trace_id`s. This is used for
   * retrieving all spans related to a session.
   */
  public static class ApiServerSpanExporter implements SpanExporter {
    private static final Logger exporterLog = LoggerFactory.getLogger(ApiServerSpanExporter.class);

    private final Map<String, Map<String, Object>> eventIdTraceStorage = new ConcurrentHashMap<>();

    // Session ID -> Trace IDs -> Trace Object
    private final Map<String, List<String>> sessionToTraceIdsMap = new ConcurrentHashMap<>();

    private final List<SpanData> allExportedSpans = Collections.synchronizedList(new ArrayList<>());

    public ApiServerSpanExporter() {}

    public Map<String, Object> getEventTraceAttributes(String eventId) {
      return this.eventIdTraceStorage.get(eventId);
    }

    public Map<String, List<String>> getSessionToTraceIdsMap() {
      return this.sessionToTraceIdsMap;
    }

    public List<SpanData> getAllExportedSpans() {
      return this.allExportedSpans;
    }

    @Override
    public CompletableResultCode export(Collection<SpanData> spans) {
      exporterLog.debug("ApiServerSpanExporter received {} spans to export.", spans.size());
      List<SpanData> currentBatch = new ArrayList<>(spans);
      allExportedSpans.addAll(currentBatch);

      for (SpanData span : currentBatch) {
        String spanName = span.getName();
        if ("call_llm".equals(spanName)
            || "send_data".equals(spanName)
            || (spanName != null && spanName.startsWith("tool_response"))) {
          String eventId =
              span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.event_id"));
          if (eventId != null && !eventId.isEmpty()) {
            Map<String, Object> attributesMap = new HashMap<>();
            span.getAttributes().forEach((key, value) -> attributesMap.put(key.getKey(), value));
            attributesMap.put("trace_id", span.getSpanContext().getTraceId());
            attributesMap.put("span_id", span.getSpanContext().getSpanId());
            attributesMap.putIfAbsent("gcp.vertex.agent.event_id", eventId);
            exporterLog.debug("Storing event-based trace attributes for event_id: {}", eventId);
            this.eventIdTraceStorage.put(eventId, attributesMap); // Use internal storage
          } else {
            exporterLog.trace(
                "Span {} for event-based trace did not have 'gcp.vertex.agent.event_id'"
                    + " attribute or it was empty.",
                spanName);
          }
        }

        if ("call_llm".equals(spanName)) {
          String sessionId =
              span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.session_id"));
          if (sessionId != null && !sessionId.isEmpty()) {
            String traceId = span.getSpanContext().getTraceId();
            sessionToTraceIdsMap
                .computeIfAbsent(sessionId, k -> Collections.synchronizedList(new ArrayList<>()))
                .add(traceId);
            exporterLog.trace(
                "Associated trace_id {} with session_id {} for session tracing",
                traceId,
                sessionId);
          } else {
            exporterLog.trace(
                "Span {} for session trace did not have 'gcp.vertex.agent.session_id' attribute.",
                spanName);
          }
        }
      }
      return CompletableResultCode.ofSuccess();
    }

    @Override
    public CompletableResultCode flush() {
      return CompletableResultCode.ofSuccess();
    }

    @Override
    public CompletableResultCode shutdown() {
      exporterLog.debug("Shutting down ApiServerSpanExporter.");
      // no need to clear storage on shutdown, as everything is currently stored in memory.
      return CompletableResultCode.ofSuccess();
    }
  }

  /**
   * Data Transfer Object (DTO) for POST /run and POST /run-sse requests. Contains information
   * needed to execute an agent run.
   */
  public static class AgentRunRequest {
    @JsonProperty("appName")
    public String appName;

    @JsonProperty("userId")
    public String userId;

    @JsonProperty("sessionId")
    public String sessionId;

    @JsonProperty("newMessage")
    public Content newMessage;

    @JsonProperty("streaming")
    public boolean streaming = false;

    public AgentRunRequest() {}

    public String getAppName() {
      return appName;
    }

    public String getUserId() {
      return userId;
    }

    public String getSessionId() {
      return sessionId;
    }

    public Content getNewMessage() {
      return newMessage;
    }

    public boolean getStreaming() {
      return streaming;
    }
  }

  /**
   * DTO for POST /apps/{appName}/eval_sets/{evalSetId}/add-session requests. Contains information
   * to associate a session with an evaluation set.
   */
  public static class AddSessionToEvalSetRequest {
    @JsonProperty("evalId")
    public String evalId;

    @JsonProperty("sessionId")
    public String sessionId;

    @JsonProperty("userId")
    public String userId;

    public AddSessionToEvalSetRequest() {}

    public String getEvalId() {
      return evalId;
    }

    public String getSessionId() {
      return sessionId;
    }

    public String getUserId() {
      return userId;
    }
  }

  /**
   * DTO for POST /apps/{appName}/eval_sets/{evalSetId}/run-eval requests. Contains information for
   * running evaluations.
   */
  public static class RunEvalRequest {
    @JsonProperty("evalIds")
    public List<String> evalIds;

    @JsonProperty("evalMetrics")
    public List<String> evalMetrics;

    public RunEvalRequest() {}

    public List<String> getEvalIds() {
      return evalIds;
    }

    public List<String> getEvalMetrics() {
      return evalMetrics;
    }
  }

  /**
   * DTO for the response of POST /apps/{appName}/eval_sets/{evalSetId}/run-eval. Contains the
   * results of an evaluation run.
   */
  public static class RunEvalResult extends JsonBaseModel {
    @JsonProperty("appName")
    public String appName;

    @JsonProperty("evalSetId")
    public String evalSetId;

    @JsonProperty("evalId")
    public String evalId;

    @JsonProperty("finalEvalStatus")
    public String finalEvalStatus;

    @JsonProperty("evalMetricResults")
    public List<List<Object>> evalMetricResults;

    @JsonProperty("sessionId")
    public String sessionId;

    /**
     * Constructs a RunEvalResult.
     *
     * @param appName The application name.
     * @param evalSetId The evaluation set ID.
     * @param evalId The evaluation ID.
     * @param finalEvalStatus The final status of the evaluation.
     * @param evalMetricResults The results for each metric.
     * @param sessionId The session ID associated with the evaluation.
     */
    public RunEvalResult(
        String appName,
        String evalSetId,
        String evalId,
        String finalEvalStatus,
        List<List<Object>> evalMetricResults,
        String sessionId) {
      this.appName = appName;
      this.evalSetId = evalSetId;
      this.evalId = evalId;
      this.finalEvalStatus = finalEvalStatus;
      this.evalMetricResults = evalMetricResults;
      this.sessionId = sessionId;
    }

    public RunEvalResult() {}
  }

  /**
   * DTO for the response of GET
   * /apps/{appName}/users/{userId}/sessions/{sessionId}/events/{eventId}/graph. Contains the graph
   * representation (e.g., DOT source).
   */
  public static class GraphResponse {
    @JsonProperty("dotSrc")
    public String dotSrc;

    /**
     * Constructs a GraphResponse.
     *
     * @param dotSrc The graph source string (e.g., in DOT format).
     */
    public GraphResponse(String dotSrc) {
      this.dotSrc = dotSrc;
    }

    public GraphResponse() {}

    public String getDotSrc() {
      return dotSrc;
    }
  }

  /**
   * Configures resource handlers for serving static content (like the Dev UI). Maps requests
   * starting with "/dev-ui/" to the directory specified by the 'adk.web.ui.dir' system property.
   */
  @Override
  public void addResourceHandlers(ResourceHandlerRegistry registry) {
    if (webUiDir != null && !webUiDir.isEmpty()) {
      // Ensure the path uses forward slashes and ends with a slash
      String location = webUiDir.replace("\\", "/");
      if (!location.startsWith("file:")) {
        location = "file:" + location; // Ensure file: prefix
      }
      if (!location.endsWith("/")) {
        location += "/";
      }
      log.debug("Mapping URL path /** to static resources at location: {}", location);
      registry
          .addResourceHandler("/**")
          .addResourceLocations(location)
          .setCachePeriod(0)
          .resourceChain(true);

    } else {
      log.debug(
          "System property 'adk.web.ui.dir' or config 'adk.web.ui.dir' is not set. Mapping URL path"
              + " /** to classpath:/browser/");
      registry
          .addResourceHandler("/**")
          .addResourceLocations("classpath:/browser/")
          .setCachePeriod(0)
          .resourceChain(true);
    }
  }

  /**
   * Configures simple automated controllers: - Redirects the root path "/" to "/dev-ui". - Forwards
   * requests to "/dev-ui" to "/dev-ui/index.html" so the ResourceHandler serves it.
   */
  @Override
  public void addViewControllers(ViewControllerRegistry registry) {
    registry.addRedirectViewController("/", "/dev-ui");
    registry.addViewController("/dev-ui").setViewName("forward:/index.html");
    registry.addViewController("/dev-ui/").setViewName("forward:/index.html");
  }

  /** Spring Boot REST Controller handling agent-related API endpoints. */
  @RestController
  public static class AgentController {

    private static final Logger log = LoggerFactory.getLogger(AgentController.class);

    private static final String EVAL_SESSION_ID_PREFIX = "ADK_EVAL_";

    private final BaseSessionService sessionService;
    private final BaseArtifactService artifactService;
    private final Map<String, BaseAgent> agentRegistry;
    private final ApiServerSpanExporter apiServerSpanExporter;
    private final RunnerService runnerService;
    private final ExecutorService sseExecutor = Executors.newCachedThreadPool();

    /**
     * Constructs the AgentController.
     *
     * @param sessionService The service for managing sessions.
     * @param artifactService The service for managing artifacts.
     * @param agentRegistry The registry of loaded agents.
     * @param apiServerSpanExporter The exporter holding all trace data.
     * @param runnerService The service for obtaining Runner instances.
     */
    @Autowired
    public AgentController(
        BaseSessionService sessionService,
        BaseArtifactService artifactService,
        @Qualifier("loadedAgentRegistry") Map<String, BaseAgent> agentRegistry,
        ApiServerSpanExporter apiServerSpanExporter,
        RunnerService runnerService) {
      this.sessionService = sessionService;
      this.artifactService = artifactService;
      this.agentRegistry = agentRegistry;
      this.apiServerSpanExporter = apiServerSpanExporter;
      this.runnerService = runnerService;
      log.info(
          "AgentController initialized with {} dynamic agents: {}",
          agentRegistry.size(),
          agentRegistry.keySet());
      if (agentRegistry.isEmpty()) {
        log.warn(
            "Agent registry is empty. Check 'adk.agents.source-dir' property and compilation"
                + " logs.");
      }
    }

    /**
     * Finds a session by its identifiers or throws a ResponseStatusException if not found or if
     * there's an app/user mismatch.
     *
     * @param appName The application name.
     * @param userId The user ID.
     * @param sessionId The session ID.
     * @return The found Session object.
     * @throws ResponseStatusException with HttpStatus.NOT_FOUND if the session doesn't exist or
     *     belongs to a different app/user.
     */
    private Session findSessionOrThrow(String appName, String userId, String sessionId) {
      Maybe<Session> maybeSession =
          sessionService.getSession(appName, userId, sessionId, Optional.empty());

      Session session = maybeSession.blockingGet();

      if (session == null) {
        log.warn(
            "Session not found for appName={}, userId={}, sessionId={}",
            appName,
            userId,
            sessionId);
        throw new ResponseStatusException(
            HttpStatus.NOT_FOUND,
            String.format(
                "Session not found: appName=%s, userId=%s, sessionId=%s",
                appName, userId, sessionId));
      }

      if (!Objects.equals(session.appName(), appName)
          || !Objects.equals(session.userId(), userId)) {
        log.warn(
            "Session ID {} found but appName/userId mismatch (Expected: {}/{}, Found: {}/{}) -"
                + " Treating as not found.",
            sessionId,
            appName,
            userId,
            session.appName(),
            session.userId());

        throw new ResponseStatusException(
            HttpStatus.NOT_FOUND, "Session found but belongs to a different app/user.");
      }
      log.debug("Found session: {}", sessionId);
      return session;
    }

    /**
     * Lists available applications. Currently returns only the configured root agent's name.
     *
     * @return A list containing the root agent's name.
     */
    @GetMapping("/list-apps")
    public List<String> listApps() {
      log.info("Listing apps from dynamic registry. Found: {}", agentRegistry.keySet());
      List<String> appNames = new ArrayList<>(agentRegistry.keySet());
      Collections.sort(appNames);
      return appNames;
    }

    /**
     * Endpoint for retrieving trace information stored by the ApiServerSpanExporter, based on event
     * ID.
     *
     * @param eventId The ID of the event to trace (expected to be gcp.vertex.agent.event_id).
     * @return A ResponseEntity containing the trace data or NOT_FOUND.
     */
    @GetMapping("/debug/trace/{eventId}")
    public ResponseEntity<?> getTraceDict(@PathVariable String eventId) {
      log.info("Request received for GET /debug/trace/{}", eventId);
      Map<String, Object> traceData = this.apiServerSpanExporter.getEventTraceAttributes(eventId);
      if (traceData == null) {
        log.warn("Trace not found for eventId: {}", eventId);
        return ResponseEntity.status(HttpStatus.NOT_FOUND)
            .body(Collections.singletonMap("message", "Trace not found for eventId: " + eventId));
      }
      log.info("Returning trace data for eventId: {}", eventId);
      return ResponseEntity.ok(traceData);
    }

    /**
     * Retrieves trace spans for a given session ID.
     *
     * @param sessionId The session ID.
     * @return A ResponseEntity containing a list of span data maps for the session, or an empty
     *     list.
     */
    @GetMapping("/debug/trace/session/{sessionId}")
    public ResponseEntity<Object> getSessionTrace(@PathVariable String sessionId) {
      log.info("Request received for GET /debug/trace/session/{}", sessionId);

      List<String> traceIdsForSession =
          this.apiServerSpanExporter.getSessionToTraceIdsMap().get(sessionId);

      if (traceIdsForSession == null || traceIdsForSession.isEmpty()) {
        log.warn("No trace IDs found for session ID: {}", sessionId);
        return ResponseEntity.ok(Collections.emptyList());
      }

      // Iterate over a snapshot of all spans to avoid concurrent modification issues
      // if the exporter is actively adding spans.
      List<SpanData> allSpansSnapshot =
          new ArrayList<>(this.apiServerSpanExporter.getAllExportedSpans());

      if (allSpansSnapshot.isEmpty()) {
        log.warn("No spans have been exported yet overall.");
        return ResponseEntity.ok(Collections.emptyList());
      }

      Set<String> relevantTraceIds = new HashSet<>(traceIdsForSession);
      List<Map<String, Object>> resultSpans = new ArrayList<>();

      for (SpanData span : allSpansSnapshot) {
        if (relevantTraceIds.contains(span.getSpanContext().getTraceId())) {
          Map<String, Object> spanMap = new HashMap<>();
          spanMap.put("name", span.getName());
          spanMap.put("span_id", span.getSpanContext().getSpanId());
          spanMap.put("trace_id", span.getSpanContext().getTraceId());
          spanMap.put("start_time", span.getStartEpochNanos());
          spanMap.put("end_time", span.getEndEpochNanos());

          Map<String, Object> attributesMap = new HashMap<>();
          span.getAttributes().forEach((key, value) -> attributesMap.put(key.getKey(), value));
          spanMap.put("attributes", attributesMap);

          String parentSpanId = span.getParentSpanId();
          if (SpanId.isValid(parentSpanId)) {
            spanMap.put("parent_span_id", parentSpanId);
          } else {
            spanMap.put("parent_span_id", null);
          }
          resultSpans.add(spanMap);
        }
      }

      log.info("Returning {} spans for session ID: {}", resultSpans.size(), sessionId);
      return ResponseEntity.ok(resultSpans);
    }

    /**
     * Retrieves a specific session by its ID.
     *
     * @param appName The application name.
     * @param userId The user ID.
     * @param sessionId The session ID.
     * @return The requested Session object.
     * @throws ResponseStatusException if the session is not found.
     */
    @GetMapping("/apps/{appName}/users/{userId}/sessions/{sessionId}")
    public Session getSession(
        @PathVariable String appName, @PathVariable String userId, @PathVariable String sessionId) {
      log.info(
          "Request received for GET /apps/{}/users/{}/sessions/{}", appName, userId, sessionId);
      return findSessionOrThrow(appName, userId, sessionId);
    }

    /**
     * Lists all non-evaluation sessions for a given app and user.
     *
     * @param appName The name of the application.
     * @param userId The ID of the user.
     * @return A list of sessions, excluding those used for evaluation.
     */
    @GetMapping("/apps/{appName}/users/{userId}/sessions")
    public List<Session> listSessions(@PathVariable String appName, @PathVariable String userId) {
      log.info("Request received for GET /apps/{}/users/{}/sessions", appName, userId);

      Single<ListSessionsResponse> sessionsResponseSingle =
          sessionService.listSessions(appName, userId);

      ListSessionsResponse response = sessionsResponseSingle.blockingGet();
      if (response == null || response.sessions() == null) {
        log.warn(
            "Received null response or null sessions list for listSessions({}, {})",
            appName,
            userId);
        return Collections.emptyList();
      }

      List<Session> filteredSessions =
          response.sessions().stream()
              .filter(s -> !s.id().startsWith(EVAL_SESSION_ID_PREFIX))
              .collect(Collectors.toList());
      log.info(
          "Found {} non-evaluation sessions for app={}, user={}",
          filteredSessions.size(),
          appName,
          userId);
      return filteredSessions;
    }

    /**
     * Creates a new session with a specific ID provided by the client.
     *
     * @param appName The application name.
     * @param userId The user ID.
     * @param sessionId The desired session ID.
     * @param state Optional initial state for the session.
     * @return The newly created Session object.
     * @throws ResponseStatusException if a session with the given ID already exists (BAD_REQUEST)
     *     or if creation fails (INTERNAL_SERVER_ERROR).
     */
    @PostMapping("/apps/{appName}/users/{userId}/sessions/{sessionId}")
    public Session createSessionWithId(
        @PathVariable String appName,
        @PathVariable String userId,
        @PathVariable String sessionId,
        @RequestBody(required = false) Map<String, Object> state) {
      log.info(
          "Request received for POST /apps/{}/users/{}/sessions/{} with state: {}",
          appName,
          userId,
          sessionId,
          state);

      try {
        findSessionOrThrow(appName, userId, sessionId);

        log.warn("Attempted to create session with existing ID: {}", sessionId);
        throw new ResponseStatusException(
            HttpStatus.BAD_REQUEST, "Session already exists: " + sessionId);
      } catch (ResponseStatusException e) {

        if (e.getStatusCode() != HttpStatus.NOT_FOUND) {
          throw e;
        }

        log.info("Session {} not found, proceeding with creation.", sessionId);
      }

      Map<String, Object> initialState = (state != null) ? state : Collections.emptyMap();
      try {
        Session createdSession =
            sessionService
                .createSession(appName, userId, new ConcurrentHashMap<>(initialState), sessionId)
                .blockingGet();

        if (createdSession == null) {

          log.error(
              "Session creation call completed without error but returned null session for {}",
              sessionId);
          throw new ResponseStatusException(
              HttpStatus.INTERNAL_SERVER_ERROR, "Failed to create session (null result)");
        }
        log.info("Session created successfully with id: {}", createdSession.id());
        return createdSession;
      } catch (Exception e) {
        log.error("Error creating session with id {}", sessionId, e);

        throw new ResponseStatusException(
            HttpStatus.INTERNAL_SERVER_ERROR, "Error creating session", e);
      }
    }

    /**
     * Creates a new session where the ID is generated by the service.
     *
     * @param appName The application name.
     * @param userId The user ID.
     * @param state Optional initial state for the session.
     * @return The newly created Session object.
     * @throws ResponseStatusException if creation fails (INTERNAL_SERVER_ERROR).
     */
    @PostMapping("/apps/{appName}/users/{userId}/sessions")
    public Session createSession(
        @PathVariable String appName,
        @PathVariable String userId,
        @RequestBody(required = false) Map<String, Object> state) {

      log.info(
          "Request received for POST /apps/{}/users/{}/sessions (service generates ID) with state:"
              + " {}",
          appName,
          userId,
          state);

      Map<String, Object> initialState = (state != null) ? state : Collections.emptyMap();
      try {

        Session createdSession =
            sessionService
                .createSession(appName, userId, new ConcurrentHashMap<>(initialState), null)
                .blockingGet();

        if (createdSession == null) {
          log.error(
              "Session creation call completed without error but returned null session for user {}",
              userId);
          throw new ResponseStatusException(
              HttpStatus.INTERNAL_SERVER_ERROR, "Failed to create session (null result)");
        }
        log.info("Session created successfully with generated id: {}", createdSession.id());
        return createdSession;
      } catch (Exception e) {
        log.error("Error creating session for user {}", userId, e);
        throw new ResponseStatusException(
            HttpStatus.INTERNAL_SERVER_ERROR, "Error creating session", e);
      }
    }

    /**
     * Deletes a specific session.
     *
     * @param appName The application name.
     * @param userId The user ID.
     * @param sessionId The session ID to delete.
     * @return A ResponseEntity with status NO_CONTENT on success.
     * @throws ResponseStatusException if deletion fails (INTERNAL_SERVER_ERROR).
     */
    @DeleteMapping("/apps/{appName}/users/{userId}/sessions/{sessionId}")
    public ResponseEntity<Void> deleteSession(
        @PathVariable String appName, @PathVariable String userId, @PathVariable String sessionId) {
      log.info(
          "Request received for DELETE /apps/{}/users/{}/sessions/{}", appName, userId, sessionId);
      try {

        sessionService.deleteSession(appName, userId, sessionId).blockingAwait();
        log.info("Session deleted successfully: {}", sessionId);
        return ResponseEntity.noContent().build();
      } catch (Exception e) {

        log.error("Error deleting session {}", sessionId, e);

        throw new ResponseStatusException(
            HttpStatus.INTERNAL_SERVER_ERROR, "Error deleting session", e);
      }
    }

    /**
     * Loads the latest or a specific version of an artifact associated with a session.
     *
     * @param appName The application name.
     * @param userId The user ID.
     * @param sessionId The session ID.
     * @param artifactName The name of the artifact.
     * @param version Optional specific version number. If null, loads the latest.
     * @return The artifact content as a Part object.
     * @throws ResponseStatusException if the artifact is not found (NOT_FOUND).
     */
    @GetMapping("/apps/{appName}/users/{userId}/sessions/{sessionId}/artifacts/{artifactName}")
    public Part loadArtifact(
        @PathVariable String appName,
        @PathVariable String userId,
        @PathVariable String sessionId,
        @PathVariable String artifactName,
        @RequestParam(required = false) Integer version) {
      String versionStr = (version == null) ? "latest" : String.valueOf(version);
      log.info(
          "Request received to load artifact: app={}, user={}, session={}, artifact={}, version={}",
          appName,
          userId,
          sessionId,
          artifactName,
          versionStr);

      Maybe<Part> artifactMaybe =
          artifactService.loadArtifact(
              appName, userId, sessionId, artifactName, Optional.ofNullable(version));

      Part artifact = artifactMaybe.blockingGet();

      if (artifact == null) {
        log.warn(
            "Artifact not found: app={}, user={}, session={}, artifact={}, version={}",
            appName,
            userId,
            sessionId,
            artifactName,
            versionStr);
        throw new ResponseStatusException(HttpStatus.NOT_FOUND, "Artifact not found");
      }
      log.debug("Artifact {} version {} loaded successfully.", artifactName, versionStr);
      return artifact;
    }

    /**
     * Loads a specific version of an artifact.
     *
     * @param appName The application name.
     * @param userId The user ID.
     * @param sessionId The session ID.
     * @param artifactName The name of the artifact.
     * @param versionId The specific version number.
     * @return The artifact content as a Part object.
     * @throws ResponseStatusException if the artifact version is not found (NOT_FOUND).
     */
    @GetMapping(
        "/apps/{appName}/users/{userId}/sessions/{sessionId}/artifacts/{artifactName}/versions/{versionId}")
    public Part loadArtifactVersion(
        @PathVariable String appName,
        @PathVariable String userId,
        @PathVariable String sessionId,
        @PathVariable String artifactName,
        @PathVariable int versionId) {
      log.info(
          "Request received to load artifact version: app={}, user={}, session={}, artifact={},"
              + " version={}",
          appName,
          userId,
          sessionId,
          artifactName,
          versionId);

      Maybe<Part> artifactMaybe =
          artifactService.loadArtifact(
              appName, userId, sessionId, artifactName, Optional.of(versionId));

      Part artifact = artifactMaybe.blockingGet();

      if (artifact == null) {
        log.warn(
            "Artifact version not found: app={}, user={}, session={}, artifact={}, version={}",
            appName,
            userId,
            sessionId,
            artifactName,
            versionId);
        throw new ResponseStatusException(HttpStatus.NOT_FOUND, "Artifact version not found");
      }
      log.debug("Artifact {} version {} loaded successfully.", artifactName, versionId);
      return artifact;
    }

    /**
     * Lists the names of all artifacts associated with a session.
     *
     * @param appName The application name.
     * @param userId The user ID.
     * @param sessionId The session ID.
     * @return A list of artifact names.
     */
    @GetMapping("/apps/{appName}/users/{userId}/sessions/{sessionId}/artifacts")
    public List<String> listArtifactNames(
        @PathVariable String appName, @PathVariable String userId, @PathVariable String sessionId) {
      log.info(
          "Request received to list artifact names for app={}, user={}, session={}",
          appName,
          userId,
          sessionId);

      Single<ListArtifactsResponse> responseSingle =
          artifactService.listArtifactKeys(appName, userId, sessionId);

      ListArtifactsResponse response = responseSingle.blockingGet();
      List<String> filenames =
          (response != null && response.filenames() != null)
              ? response.filenames()
              : Collections.emptyList();
      log.info("Found {} artifact names for session {}", filenames.size(), sessionId);
      return filenames;
    }

    /**
     * Lists the available versions for a specific artifact.
     *
     * @param appName The application name.
     * @param userId The user ID.
     * @param sessionId The session ID.
     * @param artifactName The name of the artifact.
     * @return A list of version numbers (integers).
     */
    @GetMapping(
        "/apps/{appName}/users/{userId}/sessions/{sessionId}/artifacts/{artifactName}/versions")
    public List<Integer> listArtifactVersions(
        @PathVariable String appName,
        @PathVariable String userId,
        @PathVariable String sessionId,
        @PathVariable String artifactName) {
      log.info(
          "Request received to list versions for artifact: app={}, user={}, session={},"
              + " artifact={}",
          appName,
          userId,
          sessionId,
          artifactName);

      Single<ImmutableList<Integer>> versionsSingle =
          artifactService.listVersions(appName, userId, sessionId, artifactName);
      ImmutableList<Integer> versions = versionsSingle.blockingGet();
      log.info(
          "Found {} versions for artifact {}",
          versions != null ? versions.size() : 0,
          artifactName);
      return versions != null ? versions : Collections.emptyList();
    }

    /**
     * Deletes an artifact and all its versions.
     *
     * @param appName The application name.
     * @param userId The user ID.
     * @param sessionId The session ID.
     * @param artifactName The name of the artifact to delete.
     * @return A ResponseEntity with status NO_CONTENT on success.
     * @throws ResponseStatusException if deletion fails (INTERNAL_SERVER_ERROR).
     */
    @DeleteMapping("/apps/{appName}/users/{userId}/sessions/{sessionId}/artifacts/{artifactName}")
    public ResponseEntity<Void> deleteArtifact(
        @PathVariable String appName,
        @PathVariable String userId,
        @PathVariable String sessionId,
        @PathVariable String artifactName) {
      log.info(
          "Request received to delete artifact: app={}, user={}, session={}, artifact={}",
          appName,
          userId,
          sessionId,
          artifactName);

      try {

        artifactService.deleteArtifact(appName, userId, sessionId, artifactName);
        log.info("Artifact deleted successfully: {}", artifactName);
        return ResponseEntity.noContent().build();
      } catch (Exception e) {
        log.error("Error deleting artifact {}", artifactName, e);

        throw new ResponseStatusException(
            HttpStatus.INTERNAL_SERVER_ERROR, "Error deleting artifact", e);
      }
    }

    /**
     * Executes a non-streaming agent run for a given session and message.
     *
     * @param request The AgentRunRequest containing run details.
     * @return A list of events generated during the run.
     * @throws ResponseStatusException if the session is not found or the run fails.
     */
    @PostMapping("/run")
    public List<Event> agentRun(@RequestBody AgentRunRequest request) {
      if (request.appName == null || request.appName.trim().isEmpty()) {
        log.warn("appName cannot be null or empty in POST /run request.");
        throw new ResponseStatusException(
            HttpStatus.BAD_REQUEST, "appName cannot be null or empty");
      }
      if (request.sessionId == null || request.sessionId.trim().isEmpty()) {
        log.warn("sessionId cannot be null or empty in POST /run request.");
        throw new ResponseStatusException(
            HttpStatus.BAD_REQUEST, "sessionId cannot be null or empty");
      }
      log.info("Request received for POST /run for session: {}", request.sessionId);

      Runner runner = this.runnerService.getRunner(request.appName);
      try {

        RunConfig runConfig = RunConfig.builder().setStreamingMode(StreamingMode.NONE).build();
        Flowable<Event> eventStream =
            runner.runAsync(request.userId, request.sessionId, request.newMessage, runConfig);

        List<Event> events = Lists.newArrayList(eventStream.blockingIterable());
        log.info("Agent run for session {} generated {} events.", request.sessionId, events.size());
        return events;
      } catch (Exception e) {
        log.error("Error during agent run for session {}", request.sessionId, e);
        throw new ResponseStatusException(HttpStatus.INTERNAL_SERVER_ERROR, "Agent run failed", e);
      }
    }

    /**
     * Executes an agent run and streams the resulting events using Server-Sent Events (SSE).
     *
     * @param request The AgentRunRequest containing run details.
     * @return A Flux that will stream events to the client.
     */
    @PostMapping(value = "/run_sse", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter agentRunSse(@RequestBody AgentRunRequest request) {
      SseEmitter emitter = new SseEmitter();

      if (request.appName == null || request.appName.trim().isEmpty()) {
        log.warn(
            "appName cannot be null or empty in SseEmitter request for appName: {}, session: {}",
            request.appName,
            request.sessionId);
        emitter.completeWithError(
            new ResponseStatusException(HttpStatus.BAD_REQUEST, "appName cannot be null or empty"));
        return emitter;
      }
      if (request.sessionId == null || request.sessionId.trim().isEmpty()) {
        log.warn(
            "sessionId cannot be null or empty in SseEmitter request for appName: {}, session: {}",
            request.appName,
            request.sessionId);
        emitter.completeWithError(
            new ResponseStatusException(
                HttpStatus.BAD_REQUEST, "sessionId cannot be null or empty"));
        return emitter;
      }

      log.info(
          "SseEmitter Request received for POST /run_sse_emitter for session: {}",
          request.sessionId);

      final String sessionId = request.sessionId;
      sseExecutor.execute(
          () -> {
            Runner runner;
            try {
              runner = this.runnerService.getRunner(request.appName);
            } catch (ResponseStatusException e) {
              log.warn(
                  "Setup failed for SseEmitter request for session {}: {}",
                  sessionId,
                  e.getMessage());
              try {
                emitter.completeWithError(e);
              } catch (Exception ex) {
                log.warn(
                    "Error completing emitter after setup failure for session {}: {}",
                    sessionId,
                    ex.getMessage());
              }
              return;
            }

            final RunConfig runConfig =
                RunConfig.builder()
                    .setStreamingMode(
                        request.getStreaming() ? StreamingMode.SSE : StreamingMode.NONE)
                    .build();

            Flowable<Event> eventFlowable =
                runner.runAsync(request.userId, request.sessionId, request.newMessage, runConfig);

            Disposable disposable =
                eventFlowable
                    .observeOn(Schedulers.io())
                    .subscribe(
                        event -> {
                          try {
                            log.debug(
                                "SseEmitter: Sending event {} for session {}",
                                event.id(),
                                sessionId);
                            emitter.send(SseEmitter.event().data(event.toJson()));
                          } catch (IOException e) {
                            log.error(
                                "SseEmitter: IOException sending event for session {}: {}",
                                sessionId,
                                e.getMessage());
                            throw new RuntimeException("Failed to send event", e);
                          } catch (Exception e) {
                            log.error(
                                "SseEmitter: Unexpected error sending event for session {}: {}",
                                sessionId,
                                e.getMessage(),
                                e);
                            throw new RuntimeException("Unexpected error sending event", e);
                          }
                        },
                        error -> {
                          log.error(
                              "SseEmitter: Stream error for session {}: {}",
                              sessionId,
                              error.getMessage(),
                              error);
                          try {
                            emitter.completeWithError(error);
                          } catch (Exception ex) {
                            log.warn(
                                "Error completing emitter after stream error for session {}: {}",
                                sessionId,
                                ex.getMessage());
                          }
                        },
                        () -> {
                          log.debug(
                              "SseEmitter: Stream completed normally for session: {}", sessionId);
                          try {
                            emitter.complete();
                          } catch (Exception ex) {
                            log.warn(
                                "Error completing emitter after normal completion for session {}:"
                                    + " {}",
                                sessionId,
                                ex.getMessage());
                          }
                        });
            emitter.onCompletion(
                () -> {
                  log.debug(
                      "SseEmitter: onCompletion callback for session: {}. Disposing subscription.",
                      sessionId);
                  if (!disposable.isDisposed()) {
                    disposable.dispose();
                  }
                });
            emitter.onTimeout(
                () -> {
                  log.debug(
                      "SseEmitter: onTimeout callback for session: {}. Disposing subscription and"
                          + " completing.",
                      sessionId);
                  if (!disposable.isDisposed()) {
                    disposable.dispose();
                  }
                  emitter.complete();
                });
          });

      log.debug("SseEmitter: Returning emitter for session: {}", sessionId);
      return emitter;
    }

    /**
     * Endpoint to get a graph representation of an event (currently returns a placeholder).
     * Requires Graphviz or similar tooling for full implementation.
     *
     * @param appName Application name.
     * @param userId User ID.
     * @param sessionId Session ID.
     * @param eventId Event ID.
     * @return ResponseEntity containing a GraphResponse with placeholder DOT source.
     * @throws ResponseStatusException if the session or event is not found.
     */
    @GetMapping("/apps/{appName}/users/{userId}/sessions/{sessionId}/events/{eventId}/graph")
    public ResponseEntity<GraphResponse> getEventGraph(
        @PathVariable String appName,
        @PathVariable String userId,
        @PathVariable String sessionId,
        @PathVariable String eventId) {
      log.info(
          "Request received for GET /apps/{}/users/{}/sessions/{}/events/{}/graph",
          appName,
          userId,
          sessionId,
          eventId);

      BaseAgent currentAppAgent = agentRegistry.get(appName);
      if (currentAppAgent == null) {
        log.warn("Agent app '{}' not found for graph generation.", appName);
        return ResponseEntity.status(HttpStatus.NOT_FOUND)
            .body(new GraphResponse("Agent app not found: " + appName));
      }

      Session session = findSessionOrThrow(appName, userId, sessionId);
      Event event =
          session.events().stream()
              .filter(e -> Objects.equals(e.id(), eventId))
              .findFirst()
              .orElse(null);

      if (event == null) {
        log.warn("Event {} not found in session {}", eventId, sessionId);
        return ResponseEntity.ok(new GraphResponse(null));
      }

      log.debug("Found event {} for graph generation.", eventId);

      List<List<String>> highlightPairs = new ArrayList<>();
      String eventAuthor = event.author();
      List<FunctionCall> functionCalls = event.functionCalls();
      List<FunctionResponse> functionResponses = event.functionResponses();

      if (!functionCalls.isEmpty()) {
        log.debug("Processing {} function calls for highlighting.", functionCalls.size());
        for (FunctionCall fc : functionCalls) {
          Optional<String> toolName = fc.name();
          if (toolName.isPresent() && !toolName.get().isEmpty()) {
            highlightPairs.add(ImmutableList.of(eventAuthor, toolName.get()));
            log.trace("Adding function call highlight: {} -> {}", eventAuthor, toolName.get());
          }
        }
      } else if (!functionResponses.isEmpty()) {
        log.debug("Processing {} function responses for highlighting.", functionResponses.size());
        for (FunctionResponse fr : functionResponses) {
          Optional<String> toolName = fr.name();
          if (toolName.isPresent() && !toolName.get().isEmpty()) {
            highlightPairs.add(ImmutableList.of(toolName.get(), eventAuthor));
            log.trace("Adding function response highlight: {} -> {}", toolName.get(), eventAuthor);
          }
        }
      } else {
        log.debug("Processing simple event, highlighting author: {}", eventAuthor);
        highlightPairs.add(ImmutableList.of(eventAuthor, eventAuthor));
      }

      Optional<String> dotSourceOpt =
          AgentGraphGenerator.getAgentGraphDotSource(currentAppAgent, highlightPairs);

      if (dotSourceOpt.isPresent()) {
        log.debug("Successfully generated graph DOT source for event {}", eventId);
        return ResponseEntity.ok(new GraphResponse(dotSourceOpt.get()));
      } else {
        log.warn(
            "Failed to generate graph DOT source for event {} with agent {}",
            eventId,
            currentAppAgent.name());
        return ResponseEntity.ok(new GraphResponse("Could not generate graph for this event."));
      }
    }

    /** Placeholder for creating an evaluation set. */
    @PostMapping("/apps/{appName}/eval_sets/{evalSetId}")
    public ResponseEntity<Object> createEvalSet(
        @PathVariable String appName, @PathVariable String evalSetId) {
      log.warn("Endpoint /apps/{}/eval_sets/{} (POST) is not implemented", appName, evalSetId);
      return ResponseEntity.status(HttpStatus.NOT_IMPLEMENTED)
          .body(Collections.singletonMap("message", "Eval set creation not implemented"));
    }

    /** Placeholder for listing evaluation sets. */
    @GetMapping("/apps/{appName}/eval_sets")
    public List<String> listEvalSets(@PathVariable String appName) {
      log.warn("Endpoint /apps/{}/eval_sets (GET) is not implemented", appName);
      return Collections.emptyList();
    }

    /** Placeholder for adding a session to an evaluation set. */
    @PostMapping("/apps/{appName}/eval_sets/{evalSetId}/add-session")
    public ResponseEntity<Object> addSessionToEvalSet(
        @PathVariable String appName,
        @PathVariable String evalSetId,
        @RequestBody AddSessionToEvalSetRequest req) {
      log.warn(
          "Endpoint /apps/{}/eval_sets/{}/add-session is not implemented. Request details:"
              + " evalId={}, sessionId={}, userId={}",
          appName,
          evalSetId,
          req.getEvalId(),
          req.getSessionId(),
          req.getUserId());
      return ResponseEntity.status(HttpStatus.NOT_IMPLEMENTED)
          .body(Collections.singletonMap("message", "Adding session to eval set not implemented"));
    }

    /** Placeholder for listing evaluations within an evaluation set. */
    @GetMapping("/apps/{appName}/eval_sets/{evalSetId}/evals")
    public List<String> listEvalsInEvalSet(
        @PathVariable String appName, @PathVariable String evalSetId) {
      log.warn("Endpoint /apps/{}/eval_sets/{}/evals is not implemented", appName, evalSetId);
      return Collections.emptyList();
    }

    /** Placeholder for running evaluations. */
    @PostMapping("/apps/{appName}/eval_sets/{evalSetId}/run-eval")
    public List<RunEvalResult> runEval(
        @PathVariable String appName,
        @PathVariable String evalSetId,
        @RequestBody RunEvalRequest req) {
      log.warn(
          "Endpoint /apps/{}/eval_sets/{}/run-eval is not implemented. Request details: evalIds={},"
              + " evalMetrics={}",
          appName,
          evalSetId,
          req.getEvalIds(),
          req.getEvalMetrics());
      return Collections.emptyList();
    }

    /**
     * Gets a specific evaluation result. (STUB - Not Implemented)
     *
     * @param appName The application name.
     * @param evalResultId The evaluation result ID.
     * @return A ResponseEntity indicating the endpoint is not implemented.
     */
    @GetMapping("/apps/{appName}/eval_results/{evalResultId}")
    public ResponseEntity<Object> getEvalResult(
        @PathVariable String appName, @PathVariable String evalResultId) {
      log.warn("Endpoint /apps/{}/eval_results/{} (GET) is not implemented", appName, evalResultId);
      return ResponseEntity.status(HttpStatus.NOT_IMPLEMENTED)
          .body(Collections.singletonMap("message", "Get evaluation result not implemented"));
    }

    /**
     * Lists all evaluation results for an app. (STUB - Not Implemented)
     *
     * @param appName The application name.
     * @return An empty list, as this endpoint is not implemented.
     */
    @GetMapping("/apps/{appName}/eval_results")
    public List<String> listEvalResults(@PathVariable String appName) {
      log.warn("Endpoint /apps/{}/eval_results (GET) is not implemented", appName);
      return Collections.emptyList();
    }
  }

  /** Configuration class for WebSocket handling. */
  @Configuration
  @EnableWebSocket
  public static class WebSocketConfig implements WebSocketConfigurer {

    private final LiveWebSocketHandler liveWebSocketHandler;

    @Autowired
    public WebSocketConfig(LiveWebSocketHandler liveWebSocketHandler) {
      this.liveWebSocketHandler = liveWebSocketHandler;
    }

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
      registry.addHandler(liveWebSocketHandler, "/run_live").setAllowedOrigins("*");
    }
  }

  /**
   * WebSocket Handler for the /run_live endpoint.
   *
   * <p>Manages bidirectional communication for live agent interactions. Assumes the
   * com.google.adk.runner.Runner class has a method: {@code public Flowable<Event> runLive(Session
   * session, Flowable<LiveRequest> liveRequests, List<String> modalities)}
   */
  @Component
  public static class LiveWebSocketHandler extends TextWebSocketHandler {
    private static final Logger log = LoggerFactory.getLogger(LiveWebSocketHandler.class);
    private static final String LIVE_REQUEST_QUEUE_ATTR = "liveRequestQueue";
    private static final String LIVE_SUBSCRIPTION_ATTR = "liveSubscription";
    private static final int WEBSOCKET_MAX_BYTES_FOR_REASON = 123;

    private final ObjectMapper objectMapper;
    private final BaseSessionService sessionService;
    private final RunnerService runnerService;

    @Autowired
    public LiveWebSocketHandler(
        ObjectMapper objectMapper, BaseSessionService sessionService, RunnerService runnerService) {
      this.objectMapper = objectMapper;
      this.sessionService = sessionService;
      this.runnerService = runnerService;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession wsSession) throws Exception {
      URI uri = wsSession.getUri();
      if (uri == null) {
        log.warn("WebSocket session URI is null, cannot establish connection.");
        wsSession.close(CloseStatus.SERVER_ERROR.withReason("Invalid URI"));
        return;
      }
      String path = uri.getPath();
      log.info("WebSocket connection established: {} from {}", wsSession.getId(), uri);

      MultiValueMap<String, String> queryParams =
          UriComponentsBuilder.fromUri(uri).build().getQueryParams();
      String appName = queryParams.getFirst("app_name");
      String userId = queryParams.getFirst("user_id");
      String sessionId = queryParams.getFirst("session_id");

      if (appName == null || appName.trim().isEmpty()) {
        log.warn(
            "WebSocket connection for session {} rejected: app_name query parameter is required and"
                + " cannot be empty. URI: {}",
            wsSession.getId(),
            uri);
        wsSession.close(
            CloseStatus.POLICY_VIOLATION.withReason(
                "app_name query parameter is required and cannot be empty"));
        return;
      }
      if (sessionId == null || sessionId.trim().isEmpty()) {
        log.warn(
            "WebSocket connection for session {} rejected: session_id query parameter is required"
                + " and cannot be empty. URI: {}",
            wsSession.getId(),
            uri);
        wsSession.close(
            CloseStatus.POLICY_VIOLATION.withReason(
                "session_id query parameter is required and cannot be empty"));
        return;
      }

      log.debug(
          "Extracted params for WebSocket session {}: appName={}, userId={}, sessionId={},",
          wsSession.getId(),
          appName,
          userId,
          sessionId);

      RunConfig runConfig =
          RunConfig.builder()
              .setResponseModalities(ImmutableList.of(new Modality(Modality.Known.AUDIO)))
              .setStreamingMode(StreamingMode.BIDI)
              .build();

      Session session;
      try {
        session =
            sessionService.getSession(appName, userId, sessionId, Optional.empty()).blockingGet();
        if (session == null) {
          log.warn(
              "Session not found for WebSocket: app={}, user={}, id={}. Closing connection.",
              appName,
              userId,
              sessionId);
          wsSession.close(new CloseStatus(1002, "Session not found")); // 1002: Protocol Error
          return;
        }
      } catch (Exception e) {
        log.error(
            "Error retrieving session for WebSocket: app={}, user={}, id={}",
            appName,
            userId,
            sessionId,
            e);
        wsSession.close(CloseStatus.SERVER_ERROR.withReason("Failed to retrieve session"));
        return;
      }

      LiveRequestQueue liveRequestQueue = new LiveRequestQueue();
      wsSession.getAttributes().put(LIVE_REQUEST_QUEUE_ATTR, liveRequestQueue);

      Runner runner;
      try {
        runner = this.runnerService.getRunner(appName);
      } catch (ResponseStatusException e) {
        log.error(
            "Failed to get runner for app {} during WebSocket connection: {}",
            appName,
            e.getMessage());
        wsSession.close(
            CloseStatus.SERVER_ERROR.withReason("Runner unavailable: " + e.getReason()));
        return;
      }

      Flowable<Event> eventStream = runner.runLive(session, liveRequestQueue, runConfig);

      Disposable disposable =
          eventStream
              .subscribeOn(Schedulers.io()) // Offload runner work
              .observeOn(Schedulers.io()) // Send messages on I/O threads
              .subscribe(
                  event -> {
                    try {
                      String jsonEvent = objectMapper.writeValueAsString(event);
                      log.debug(
                          "Sending event via WebSocket session {}: {}",
                          wsSession.getId(),
                          jsonEvent);
                      wsSession.sendMessage(new TextMessage(jsonEvent));
                    } catch (JsonProcessingException e) {
                      log.error(
                          "Error serializing event to JSON for WebSocket session {}",
                          wsSession.getId(),
                          e);
                      // Decide if to close session or just log
                    } catch (IOException e) {
                      log.error(
                          "IOException sending message via WebSocket session {}",
                          wsSession.getId(),
                          e);
                      // This might mean the session is already closed or problematic
                      // Consider closing/disposing here
                      try {
                        wsSession.close(
                            CloseStatus.SERVER_ERROR.withReason("Error sending message"));
                      } catch (IOException ignored) {
                      }
                    }
                  },
                  error -> {
                    log.error(
                        "Error in run_live stream for WebSocket session {}: {}",
                        wsSession.getId(),
                        error.getMessage(),
                        error);
                    String reason =
                        error.getMessage() != null ? error.getMessage() : "Unknown error";
                    try {
                      wsSession.close(
                          new CloseStatus(
                              1011, // Internal Server Error for WebSocket
                              reason.substring(
                                  0, Math.min(reason.length(), WEBSOCKET_MAX_BYTES_FOR_REASON))));
                    } catch (IOException ignored) {
                    }
                  },
                  () -> {
                    log.debug(
                        "run_live stream completed for WebSocket session {}", wsSession.getId());
                    try {
                      wsSession.close(CloseStatus.NORMAL);
                    } catch (IOException ignored) {
                    }
                  });
      wsSession.getAttributes().put(LIVE_SUBSCRIPTION_ATTR, disposable);
      log.debug("Live run started for WebSocket session {}", wsSession.getId());
    }

    @Override
    protected void handleTextMessage(WebSocketSession wsSession, TextMessage message)
        throws Exception {
      LiveRequestQueue liveRequestQueue =
          (LiveRequestQueue) wsSession.getAttributes().get(LIVE_REQUEST_QUEUE_ATTR);

      if (liveRequestQueue == null) {
        log.warn(
            "Received message on WebSocket session {} but LiveRequestQueue is not available (null)."
                + " Message: {}",
            wsSession.getId(),
            message.getPayload());
        return;
      }

      try {
        String payload = message.getPayload();
        log.debug("Received text message on WebSocket session {}: {}", wsSession.getId(), payload);

        JsonNode rootNode = objectMapper.readTree(payload);
        LiveRequest.Builder liveRequestBuilder = LiveRequest.builder();

        if (rootNode.has("content")) {
          Content content = objectMapper.treeToValue(rootNode.get("content"), Content.class);
          liveRequestBuilder.content(content);
        }

        if (rootNode.has("blob")) {
          JsonNode blobNode = rootNode.get("blob");
          Blob.Builder blobBuilder = Blob.builder();
          if (blobNode.has("displayName")) {
            blobBuilder.displayName(blobNode.get("displayName").asText());
          }
          if (blobNode.has("data")) {
            blobBuilder.data(blobNode.get("data").binaryValue());
          }
          // Handle both mime_type and mimeType. Blob states mimeType but we get mime_type from the
          // frontend.
          String mimeType =
              blobNode.has("mimeType")
                  ? blobNode.get("mimeType").asText()
                  : (blobNode.has("mime_type") ? blobNode.get("mime_type").asText() : null);
          if (mimeType != null) {
            blobBuilder.mimeType(mimeType);
          }
          liveRequestBuilder.blob(blobBuilder.build());
        }
        LiveRequest liveRequest = liveRequestBuilder.build();
        liveRequestQueue.send(liveRequest);
      } catch (JsonProcessingException e) {
        log.error(
            "Error deserializing LiveRequest from WebSocket message for session {}: {}",
            wsSession.getId(),
            message.getPayload(),
            e);
        wsSession.sendMessage(
            new TextMessage(
                "{\"error\":\"Invalid JSON format for LiveRequest\", \"details\":\""
                    + e.getMessage()
                    + "\"}"));
      } catch (Exception e) {
        log.error(
            "Unexpected error processing text message for WebSocket session {}: {}",
            wsSession.getId(),
            message.getPayload(),
            e);
        String reason = e.getMessage() != null ? e.getMessage() : "Error processing message";
        wsSession.close(
            new CloseStatus(
                1011,
                reason.substring(0, Math.min(reason.length(), WEBSOCKET_MAX_BYTES_FOR_REASON))));
      }
    }

    @Override
    public void handleTransportError(WebSocketSession wsSession, Throwable exception)
        throws Exception {
      log.error(
          "WebSocket transport error for session {}: {}",
          wsSession.getId(),
          exception.getMessage(),
          exception);
      // Cleanup resources similar to afterConnectionClosed
      cleanupSession(wsSession);
      if (wsSession.isOpen()) {
        String reason = exception.getMessage() != null ? exception.getMessage() : "Transport error";
        wsSession.close(
            CloseStatus.PROTOCOL_ERROR.withReason(
                reason.substring(0, Math.min(reason.length(), WEBSOCKET_MAX_BYTES_FOR_REASON))));
      }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession wsSession, CloseStatus status)
        throws Exception {
      log.info(
          "WebSocket connection closed: {} with status {}", wsSession.getId(), status.toString());
      cleanupSession(wsSession);
    }

    private void cleanupSession(WebSocketSession wsSession) {
      LiveRequestQueue liveRequestQueue =
          (LiveRequestQueue) wsSession.getAttributes().remove(LIVE_REQUEST_QUEUE_ATTR);
      if (liveRequestQueue != null) {
        liveRequestQueue.close(); // Signal end of input to the runner
        log.debug("Called close() on LiveRequestQueue for session {}", wsSession.getId());
      }

      Disposable disposable = (Disposable) wsSession.getAttributes().remove(LIVE_SUBSCRIPTION_ATTR);
      if (disposable != null && !disposable.isDisposed()) {
        disposable.dispose();
      }
      log.debug("Cleaned up resources for WebSocket session {}", wsSession.getId());
    }
  }

  /**
   * Main entry point for the Spring Boot application.
   *
   * @param args Command line arguments.
   */
  public static void main(String[] args) {
    System.setProperty(
        "org.apache.tomcat.websocket.DEFAULT_BUFFER_SIZE", String.valueOf(10 * 1024 * 1024));
    SpringApplication.run(AdkWebServer.class, args);
    log.info("AdkWebServer application started successfully.");
  }
}

/*
 * 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.
 */

// LoadArtifactsTool.java
package com.google.adk.tools;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.google.adk.JsonBaseModel;
import com.google.adk.models.LlmRequest;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.genai.types.Content;
import com.google.genai.types.FunctionDeclaration;
import com.google.genai.types.FunctionResponse;
import com.google.genai.types.Part;
import com.google.genai.types.Schema;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Observable;
import io.reactivex.rxjava3.core.Single;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;

/**
 * A tool that loads artifacts and adds them to the session.
 *
 * <p>This tool informs the model about available artifacts and provides their content when
 * requested by the model through a function call.
 */
public final class LoadArtifactsTool extends BaseTool {
  public LoadArtifactsTool() {
    super("load_artifacts", "Loads the artifacts and adds them to the session.");
  }

  @Override
  public Optional<FunctionDeclaration> declaration() {
    return Optional.of(
        FunctionDeclaration.builder()
            .name(this.name())
            .description(this.description())
            .parameters(
                Schema.builder()
                    .type("OBJECT")
                    .properties(
                        ImmutableMap.of(
                            "artifact_names",
                            Schema.builder()
                                .type("ARRAY")
                                .items(Schema.builder().type("STRING").build())
                                .build()))
                    .build())
            .build());
  }

  @Override
  public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContext toolContext) {
    @SuppressWarnings("unchecked")
    List<String> artifactNames =
        (List<String>) args.getOrDefault("artifact_names", ImmutableList.of());
    return Single.just(ImmutableMap.of("artifact_names", artifactNames));
  }

  @Override
  public Completable processLlmRequest(
      LlmRequest.Builder llmRequestBuilder, ToolContext toolContext) {
    return super.processLlmRequest(llmRequestBuilder, toolContext)
        .andThen(appendArtifactsToLlmRequest(llmRequestBuilder, toolContext));
  }

  public Completable appendArtifactsToLlmRequest(
      LlmRequest.Builder llmRequestBuilder, ToolContext toolContext) {

    return toolContext
        .listArtifacts()
        .flatMapCompletable(
            artifactNamesList -> {
              if (artifactNamesList.isEmpty()) {
                return Completable.complete();
              }

              appendInitialInstructions(llmRequestBuilder, artifactNamesList);

              return processLoadArtifactsFunctionCall(llmRequestBuilder, toolContext);
            });
  }

  private void appendInitialInstructions(
      LlmRequest.Builder llmRequestBuilder, List<String> artifactNamesList) {
    try {
      String instructions =
          String.format(
              "You have a list of artifacts:\n  %s\n\nWhen the user asks questions about"
                  + " any of the artifacts, you should call the `load_artifacts` function"
                  + " to load the artifact. Do not generate any text other than the"
                  + " function call.",
              JsonBaseModel.getMapper().writeValueAsString(artifactNamesList));
      llmRequestBuilder.appendInstructions(ImmutableList.of(instructions));
    } catch (JsonProcessingException e) {
      throw new IllegalStateException("Failed to serialize artifact names to JSON", e);
    }
  }

  private Completable processLoadArtifactsFunctionCall(
      LlmRequest.Builder llmRequestBuilder, ToolContext toolContext) {

    LlmRequest currentRequestState = llmRequestBuilder.build();
    List<Content> currentContents = currentRequestState.contents();

    if (currentContents.isEmpty()) {
      return Completable.complete();
    }

    return Iterables.getLast(currentContents)
        .parts()
        .filter(partsList -> !partsList.isEmpty())
        .flatMap(partsList -> partsList.get(0).functionResponse())
        .filter(fr -> Objects.equals(fr.name().orElse(null), "load_artifacts"))
        .flatMap(FunctionResponse::response)
        .flatMap(responseMap -> Optional.ofNullable(responseMap.get("artifact_names")))
        .filter(obj -> obj instanceof List)
        .map(obj -> (List<?>) obj)
        .filter(list -> !list.isEmpty())
        .map(
            artifactNamesRaw -> {
              @SuppressWarnings("unchecked")
              List<String> artifactNamesToLoad = (List<String>) artifactNamesRaw;

              return Observable.fromIterable(artifactNamesToLoad)
                  .flatMapCompletable(
                      artifactName ->
                          loadAndAppendIndividualArtifact(
                              llmRequestBuilder, toolContext, artifactName));
            })
        .orElse(Completable.complete());
  }

  private Completable loadAndAppendIndividualArtifact(
      LlmRequest.Builder llmRequestBuilder, ToolContext toolContext, String artifactName) {

    return toolContext
        .loadArtifact(artifactName, Optional.empty())
        .flatMapCompletable(
            actualArtifact ->
                Completable.fromAction(
                    () ->
                        appendArtifactToLlmRequest(
                            llmRequestBuilder,
                            "Artifact " + artifactName + " is:",
                            actualArtifact)));
  }

  private void appendArtifactToLlmRequest(
      LlmRequest.Builder llmRequestBuilder, String prefix, Part artifact) {
    llmRequestBuilder.contents(
        ImmutableList.<Content>builder()
            .addAll(llmRequestBuilder.build().contents())
            .add(Content.fromParts(Part.fromText(prefix), artifact))
            .build());
  }
}

/*
 * 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.agents;

import static com.google.common.base.Strings.isNullOrEmpty;

import com.google.adk.events.Event;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import io.reactivex.rxjava3.core.Flowable;
import java.util.ArrayList;
import java.util.List;

/**
 * A shell agent that runs its sub-agents in parallel in isolated manner.
 *
 *
This approach is beneficial for scenarios requiring multiple perspectives or attempts on a
 * single task, such as running different algorithms simultaneously or generating multiple responses
 * for review by a subsequent evaluation agent.
 */
public class ParallelAgent extends BaseAgent {

  /**
   * Constructor for ParallelAgent.
   *
   * @param name The agent's name.
   * @param description The agent's description.
   * @param subAgents The list of sub-agents to run sequentially.
   * @param beforeAgentCallback Optional callback before the agent runs.
   * @param afterAgentCallback Optional callback after the agent runs.
   */
  private ParallelAgent(
      String name,
      String description,
      List<? extends BaseAgent> subAgents,
      List<Callbacks.BeforeAgentCallback> beforeAgentCallback,
      List<Callbacks.AfterAgentCallback> afterAgentCallback) {

    super(name, description, subAgents, beforeAgentCallback, afterAgentCallback);
  }

  /** Builder for {@link ParallelAgent}. */
  public static class Builder {
    private String name;
    private String description;
    private List<? extends BaseAgent> subAgents;
    private ImmutableList<Callbacks.BeforeAgentCallback> beforeAgentCallback;
    private ImmutableList<Callbacks.AfterAgentCallback> afterAgentCallback;

    @CanIgnoreReturnValue
    public Builder name(String name) {
      this.name = name;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder description(String description) {
      this.description = description;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder subAgents(List<? extends BaseAgent> subAgents) {
      this.subAgents = subAgents;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder subAgents(BaseAgent... subAgents) {
      this.subAgents = ImmutableList.copyOf(subAgents);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeAgentCallback(Callbacks.BeforeAgentCallback beforeAgentCallback) {
      this.beforeAgentCallback = ImmutableList.of(beforeAgentCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeAgentCallback(
        List<Callbacks.BeforeAgentCallbackBase> beforeAgentCallback) {
      this.beforeAgentCallback = CallbackUtil.getBeforeAgentCallbacks(beforeAgentCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterAgentCallback(Callbacks.AfterAgentCallback afterAgentCallback) {
      this.afterAgentCallback = ImmutableList.of(afterAgentCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterAgentCallback(List<Callbacks.AfterAgentCallbackBase> afterAgentCallback) {
      this.afterAgentCallback = CallbackUtil.getAfterAgentCallbacks(afterAgentCallback);
      return this;
    }

    public ParallelAgent build() {
      return new ParallelAgent(
          name, description, subAgents, beforeAgentCallback, afterAgentCallback);
    }
  }

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

  /**
   * Sets the branch for the current agent in the invocation context.
   *
   *
Appends the agent name to the current branch, or sets it if undefined.
   *
   * @param currentAgent Current agent.
   * @param invocationContext Invocation context to update.
   */
  private static void setBranchForCurrentAgent(
      BaseAgent currentAgent, InvocationContext invocationContext) {
    String branch = invocationContext.branch().orElse(null);
    if (isNullOrEmpty(branch)) {
      invocationContext.branch(currentAgent.name());
    } else {
      invocationContext.branch(branch + "." + currentAgent.name());
    }
  }

  /**
   * Runs sub-agents in parallel and emits their events.
   *
   *
Sets the branch and merges event streams from all sub-agents.
   *
   * @param invocationContext Invocation context.
   * @return Flowable emitting events from all sub-agents.
   */
  @Override
  protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
    setBranchForCurrentAgent(this, invocationContext);

    List<? extends BaseAgent> currentSubAgents = subAgents();
    if (currentSubAgents == null || currentSubAgents.isEmpty()) {
      return Flowable.empty();
    }

    List<Flowable<Event>> agentFlowables = new ArrayList<>();
    for (BaseAgent subAgent : currentSubAgents) {
      agentFlowables.add(subAgent.runAsync(invocationContext));
    }
    return Flowable.merge(agentFlowables);
  }

  /**
   * Not supported for ParallelAgent.
   *
   * @param invocationContext Invocation context.
   * @return Flowable that always throws UnsupportedOperationException.
   */
  @Override
  protected Flowable<Event> runLiveImpl(InvocationContext invocationContext) {
    return Flowable.error(
        new UnsupportedOperationException("runLive is not defined for ParallelAgent yet."));
  }
}

/*
 * 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 static com.google.common.base.Strings.nullToEmpty;
import static java.util.concurrent.TimeUnit.SECONDS;
import static java.util.stream.Collectors.toCollection;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.adk.JsonBaseModel;
import com.google.adk.events.Event;
import com.google.adk.events.EventActions;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.genai.types.Content;
import com.google.genai.types.FinishReason;
import com.google.genai.types.GroundingMetadata;
import com.google.genai.types.HttpOptions;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import javax.annotation.Nullable;
import okhttp3.ResponseBody;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Connects to the managed Vertex AI Session Service. */
/** TODO: Use the genai HttpApiClient and ApiResponse methods once they are public. */
public final class VertexAiSessionService implements BaseSessionService {
  private static final int MAX_RETRY_ATTEMPTS = 5;
  private static final ObjectMapper objectMapper = JsonBaseModel.getMapper();
  private static final Logger logger = LoggerFactory.getLogger(VertexAiSessionService.class);

  private final HttpApiClient apiClient;

  /**
   * Creates a new instance of the Vertex AI Session Service with a custom ApiClient for testing.
   */
  public VertexAiSessionService(String project, String location, HttpApiClient apiClient) {
    this.apiClient = apiClient;
  }

  /** Creates a session service with default configuration. */
  public VertexAiSessionService() {
    this.apiClient =
        new HttpApiClient(Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty());
  }

  /** Creates a session service with specified project, location, credentials, and HTTP options. */
  public VertexAiSessionService(
      String project,
      String location,
      Optional<GoogleCredentials> credentials,
      Optional<HttpOptions> httpOptions) {
    this.apiClient =
        new HttpApiClient(Optional.of(project), Optional.of(location), credentials, httpOptions);
  }

  /**
   * Parses the JSON response body from the given API response.
   *
   * @throws UncheckedIOException if parsing fails.
   */
  private static JsonNode getJsonResponse(ApiResponse apiResponse) {
    try {
      ResponseBody responseBody = apiResponse.getResponseBody();
      return objectMapper.readTree(responseBody.string());
    } catch (IOException e) {
      throw new UncheckedIOException(e);
    }
  }

  @Override
  public Single<Session> createSession(
      String appName,
      String userId,
      @Nullable ConcurrentMap<String, Object> state,
      @Nullable String sessionId) {

    String reasoningEngineId = parseReasoningEngineId(appName);
    ConcurrentHashMap<String, Object> sessionJsonMap = new ConcurrentHashMap<>();
    sessionJsonMap.put("userId", userId);
    if (state != null) {
      sessionJsonMap.put("sessionState", state);
    }

    ApiResponse apiResponse;
    try {
      apiResponse =
          apiClient.request(
              "POST",
              "reasoningEngines/" + reasoningEngineId + "/sessions",
              objectMapper.writeValueAsString(sessionJsonMap));
    } catch (IOException e) {
      throw new UncheckedIOException(e);
    }

    logger.debug("Create Session response {}", apiResponse.getResponseBody());
    String sessionName = "";
    String operationId = "";
    String sessId = nullToEmpty(sessionId);
    if (apiResponse.getResponseBody() != null) {
      JsonNode jsonResponse = getJsonResponse(apiResponse);
      sessionName = jsonResponse.get("name").asText();
      List<String> parts = Splitter.on('/').splitToList(sessionName);
      sessId = parts.get(parts.size() - 3);
      operationId = Iterables.getLast(parts);
    }
    for (int i = 0; i < MAX_RETRY_ATTEMPTS; i++) {
      ApiResponse lroResponse = apiClient.request("GET", "operations/" + operationId, "");
      JsonNode jsonResponse = getJsonResponse(lroResponse);
      if (jsonResponse.get("done") != null) {
        break;
      }
      try {
        SECONDS.sleep(1);
      } catch (InterruptedException e) {
        logger.warn("Error during sleep", e);
      }
    }

    ApiResponse getSessionApiResponse =
        apiClient.request(
            "GET", "reasoningEngines/" + reasoningEngineId + "/sessions/" + sessId, "");
    JsonNode getSessionResponseMap = getJsonResponse(getSessionApiResponse);
    Instant updateTimestamp = Instant.parse(getSessionResponseMap.get("updateTime").asText());
    ConcurrentMap<String, Object> sessionState = null;
    if (getSessionResponseMap != null && getSessionResponseMap.has("sessionState")) {
      JsonNode sessionStateNode = getSessionResponseMap.get("sessionState");
      if (sessionStateNode != null) {
        sessionState =
            objectMapper.convertValue(
                sessionStateNode, new TypeReference<ConcurrentMap<String, Object>>() {});
      }
    }
    return Single.just(
        Session.builder(sessId)
            .appName(appName)
            .userId(userId)
            .lastUpdateTime(updateTimestamp)
            .state(sessionState == null ? new ConcurrentHashMap<>() : sessionState)
            .build());
  }

  @Override
  public Single<ListSessionsResponse> listSessions(String appName, String userId) {
    String reasoningEngineId = parseReasoningEngineId(appName);

    ApiResponse apiResponse =
        apiClient.request(
            "GET",
            "reasoningEngines/" + reasoningEngineId + "/sessions?filter=user_id=" + userId,
            "");

    // Handles empty response case
    if (apiResponse.getResponseBody() == null) {
      return Single.just(ListSessionsResponse.builder().build());
    }

    JsonNode listSessionsResponseMap = getJsonResponse(apiResponse);
    List<Map<String, Object>> apiSessions =
        objectMapper.convertValue(
            listSessionsResponseMap.get("sessions"),
            new TypeReference<List<Map<String, Object>>>() {});

    List<Session> sessions = new ArrayList<>();
    for (Map<String, Object> apiSession : apiSessions) {
      String sessionId =
          Iterables.getLast(Splitter.on('/').splitToList((String) apiSession.get("name")));
      Instant updateTimestamp = Instant.parse((String) apiSession.get("updateTime"));
      Session session =
          Session.builder(sessionId)
              .appName(appName)
              .userId(userId)
              .state(new ConcurrentHashMap<>())
              .lastUpdateTime(updateTimestamp)
              .build();
      sessions.add(session);
    }
    return Single.just(ListSessionsResponse.builder().sessions(sessions).build());
  }

  @Override
  public Single<ListEventsResponse> listEvents(String appName, String userId, String sessionId) {
    String reasoningEngineId = parseReasoningEngineId(appName);
    ApiResponse apiResponse =
        apiClient.request(
            "GET",
            "reasoningEngines/" + reasoningEngineId + "/sessions/" + sessionId + "/events",
            "");

    logger.debug("List events response {}", apiResponse);

    if (apiResponse.getResponseBody() == null) {
      return Single.just(ListEventsResponse.builder().build());
    }

    JsonNode sessionEventsNode = getJsonResponse(apiResponse).get("sessionEvents");
    if (sessionEventsNode == null || sessionEventsNode.isEmpty()) {
      return Single.just(ListEventsResponse.builder().events(new ArrayList<>()).build());
    }
    return Single.just(
        ListEventsResponse.builder()
            .events(
                objectMapper
                    .convertValue(
                        sessionEventsNode,
                        new TypeReference<List<ConcurrentMap<String, Object>>>() {})
                    .stream()
                    .map(event -> fromApiEvent(event))
                    .collect(toCollection(ArrayList::new)))
            .build());
  }

  @Override
  public Maybe<Session> getSession(
      String appName, String userId, String sessionId, Optional<GetSessionConfig> config) {
    String reasoningEngineId = parseReasoningEngineId(appName);
    ApiResponse apiResponse =
        apiClient.request(
            "GET", "reasoningEngines/" + reasoningEngineId + "/sessions/" + sessionId, "");
    JsonNode getSessionResponseMap = getJsonResponse(apiResponse);

    if (getSessionResponseMap == null) {
      return Maybe.empty();
    }

    String sessId =
        Optional.ofNullable(getSessionResponseMap.get("name"))
            .map(name -> Iterables.getLast(Splitter.on('/').splitToList(name.asText())))
            .orElse(sessionId);
    Instant updateTimestamp =
        Optional.ofNullable(getSessionResponseMap.get("updateTime"))
            .map(updateTime -> Instant.parse(updateTime.asText()))
            .orElse(null);

    ConcurrentMap<String, Object> sessionState = new ConcurrentHashMap<>();
    if (getSessionResponseMap != null && getSessionResponseMap.has("sessionState")) {
      sessionState.putAll(
          objectMapper.convertValue(
              getSessionResponseMap.get("sessionState"),
              new TypeReference<ConcurrentMap<String, Object>>() {}));
    }

    return listEvents(appName, userId, sessionId)
        .map(
            response -> {
              Session.Builder sessionBuilder =
                  Session.builder(sessId)
                      .appName(appName)
                      .userId(userId)
                      .lastUpdateTime(updateTimestamp)
                      .state(sessionState);
              List<Event> events = response.events();
              if (events.isEmpty()) {
                return sessionBuilder.build();
              }
              events =
                  events.stream()
                      .filter(
                          event ->
                              updateTimestamp == null
                                  || Instant.ofEpochMilli(event.timestamp())
                                      .isBefore(updateTimestamp))
                      .sorted(Comparator.comparing(Event::timestamp))
                      .collect(toCollection(ArrayList::new));

              if (config.isPresent()) {
                if (config.get().numRecentEvents().isPresent()) {
                  int numRecentEvents = config.get().numRecentEvents().get();
                  if (events.size() > numRecentEvents) {
                    events = events.subList(events.size() - numRecentEvents, events.size());
                  }
                } else if (config.get().afterTimestamp().isPresent()) {
                  Instant afterTimestamp = config.get().afterTimestamp().get();
                  int i = events.size() - 1;
                  while (i >= 0) {
                    if (Instant.ofEpochMilli(events.get(i).timestamp()).isBefore(afterTimestamp)) {
                      break;
                    }
                    i -= 1;
                  }
                  if (i >= 0) {
                    events = events.subList(i, events.size());
                  }
                }
              }
              return sessionBuilder.events(events).build();
            })
        .toMaybe();
  }

  @Override
  public Completable deleteSession(String appName, String userId, String sessionId) {
    String reasoningEngineId = parseReasoningEngineId(appName);
    ApiResponse unused =
        apiClient.request(
            "DELETE", "reasoningEngines/" + reasoningEngineId + "/sessions/" + sessionId, "");
    return Completable.complete();
  }

  @Override
  public Single<Event> appendEvent(Session session, Event event) {
    BaseSessionService.super.appendEvent(session, event);

    String reasoningEngineId = parseReasoningEngineId(session.appName());
    ApiResponse response =
        apiClient.request(
            "POST",
            "reasoningEngines/" + reasoningEngineId + "/sessions/" + session.id() + ":appendEvent",
            convertEventToJson(event));
    // TODO(b/414263934)): Improve error handling for appendEvent.
    try {
      if (response.getResponseBody().string().contains("com.google.genai.errors.ClientException")) {
        logger.warn("Failed to append event: ", event);
      }
    } catch (IOException e) {
      throw new UncheckedIOException(e);
    }

    response.close();
    return Single.just(event);
  }

  /**
   * Converts an {@link Event} to its JSON string representation for API transmission.
   *
   * @return JSON string of the event.
   * @throws UncheckedIOException if serialization fails.
   */
  static String convertEventToJson(Event event) {
    Map<String, Object> metadataJson = new HashMap<>();
    metadataJson.put("partial", event.partial());
    metadataJson.put("turnComplete", event.turnComplete());
    metadataJson.put("interrupted", event.interrupted());
    metadataJson.put("branch", event.branch().orElse(null));
    metadataJson.put(
        "long_running_tool_ids",
        event.longRunningToolIds() != null ? event.longRunningToolIds().orElse(null) : null);
    if (event.groundingMetadata() != null) {
      metadataJson.put("grounding_metadata", event.groundingMetadata());
    }

    Map<String, Object> eventJson = new HashMap<>();
    eventJson.put("author", event.author());
    eventJson.put("invocationId", event.invocationId());
    eventJson.put(
        "timestamp",
        new HashMap<>(
            ImmutableMap.of(
                "seconds",
                event.timestamp() / 1000,
                "nanos",
                (event.timestamp() % 1000) * 1000000)));
    if (event.errorCode().isPresent()) {
      eventJson.put("errorCode", event.errorCode());
    }
    if (event.errorMessage().isPresent()) {
      eventJson.put("errorMessage", event.errorMessage());
    }
    eventJson.put("eventMetadata", metadataJson);

    if (event.actions() != null) {
      Map<String, Object> actionsJson = new HashMap<>();
      actionsJson.put("skipSummarization", event.actions().skipSummarization());
      actionsJson.put("stateDelta", event.actions().stateDelta());
      actionsJson.put("artifactDelta", event.actions().artifactDelta());
      actionsJson.put("transferAgent", event.actions().transferToAgent());
      actionsJson.put("escalate", event.actions().escalate());
      actionsJson.put("requestedAuthConfigs", event.actions().requestedAuthConfigs());
      eventJson.put("actions", actionsJson);
    }
    if (event.content().isPresent()) {
      eventJson.put("content", SessionUtils.encodeContent(event.content().get()));
    }
    if (event.errorCode().isPresent()) {
      eventJson.put("errorCode", event.errorCode().get());
    }
    if (event.errorMessage().isPresent()) {
      eventJson.put("errorMessage", event.errorMessage().get());
    }
    try {
      return objectMapper.writeValueAsString(eventJson);
    } catch (JsonProcessingException e) {
      throw new UncheckedIOException(e);
    }
  }

  /**
   * Converts a raw value to a {@link Content} object.
   *
   * @return parsed {@link Content}, or {@code null} if conversion fails.
   */
  @Nullable
  @SuppressWarnings("unchecked")
  private static Content convertMapToContent(Object rawContentValue) {
    if (rawContentValue == null) {
      return null;
    }

    if (rawContentValue instanceof Map) {
      Map<String, Object> contentMap = (Map<String, Object>) rawContentValue;
      try {
        return objectMapper.convertValue(contentMap, Content.class);
      } catch (IllegalArgumentException e) {
        logger.warn("Error converting Map to Content", e);
        return null;
      }
    } else {
      logger.warn(
          "Unexpected type for 'content' in apiEvent: {}", rawContentValue.getClass().getName());
      return null;
    }
  }

  /**
   * Extracts the reasoning engine ID from the given app name or full resource name.
   *
   * @return reasoning engine ID.
   * @throws IllegalArgumentException if format is invalid.
   */
  static String parseReasoningEngineId(String appName) {
    if (appName.matches("\\d+")) {
      return appName;
    }

    Matcher matcher = APP_NAME_PATTERN.matcher(appName);

    if (!matcher.matches()) {
      throw new IllegalArgumentException(
          "App name "
              + appName
              + " is not valid. It should either be the full"
              + " ReasoningEngine resource name, or the reasoning engine id.");
    }

    return matcher.group(matcher.groupCount());
  }

  /**
   * Converts raw API event data into an {@link Event} object.
   *
   * @return parsed {@link Event}.
   */
  @SuppressWarnings("unchecked")
  static Event fromApiEvent(Map<String, Object> apiEvent) {
    EventActions eventActions = new EventActions();
    if (apiEvent.get("actions") != null) {
      Map<String, Object> actionsMap = (Map<String, Object>) apiEvent.get("actions");
      eventActions.setSkipSummarization(
          Optional.ofNullable(actionsMap.get("skipSummarization")).map(value -> (Boolean) value));
      eventActions.setStateDelta(
          actionsMap.get("stateDelta") != null
              ? new ConcurrentHashMap<>((Map<String, Object>) actionsMap.get("stateDelta"))
              : new ConcurrentHashMap<>());
      eventActions.setArtifactDelta(
          actionsMap.get("artifactDelta") != null
              ? new ConcurrentHashMap<>((Map<String, Part>) actionsMap.get("artifactDelta"))
              : new ConcurrentHashMap<>());
      eventActions.setTransferToAgent(
          actionsMap.get("transferAgent") != null
              ? (String) actionsMap.get("transferAgent")
              : null);
      eventActions.setEscalate(
          Optional.ofNullable(actionsMap.get("escalate")).map(value -> (Boolean) value));
      eventActions.setRequestedAuthConfigs(
          Optional.ofNullable(actionsMap.get("requestedAuthConfigs"))
              .map(VertexAiSessionService::asConcurrentMapOfConcurrentMaps)
              .orElse(new ConcurrentHashMap<>()));
    }

    Event event =
        Event.builder()
            .id((String) Iterables.getLast(Splitter.on('/').split(apiEvent.get("name").toString())))
            .invocationId((String) apiEvent.get("invocationId"))
            .author((String) apiEvent.get("author"))
            .actions(eventActions)
            .content(
                Optional.ofNullable(apiEvent.get("content"))
                    .map(VertexAiSessionService::convertMapToContent)
                    .map(SessionUtils::decodeContent)
                    .orElse(null))
            .timestamp(convertToInstant(apiEvent.get("timestamp")).toEpochMilli())
            .errorCode(
                Optional.ofNullable(apiEvent.get("errorCode"))
                    .map(value -> new FinishReason((String) value)))
            .errorMessage(
                Optional.ofNullable(apiEvent.get("errorMessage")).map(value -> (String) value))
            .branch(Optional.ofNullable(apiEvent.get("branch")).map(value -> (String) value))
            .build();
    // TODO(b/414263934): Add Event branch and grounding metadata for python parity.
    if (apiEvent.get("eventMetadata") != null) {
      Map<String, Object> eventMetadata = (Map<String, Object>) apiEvent.get("eventMetadata");
      List<String> longRunningToolIdsList = (List<String>) eventMetadata.get("longRunningToolIds");

      GroundingMetadata groundingMetadata = null;
      Object rawGroundingMetadata = eventMetadata.get("groundingMetadata");
      if (rawGroundingMetadata != null) {
        groundingMetadata =
            objectMapper.convertValue(rawGroundingMetadata, GroundingMetadata.class);
      }

      event =
          event.toBuilder()
              .partial(Optional.ofNullable((Boolean) eventMetadata.get("partial")).orElse(false))
              .turnComplete(
                  Optional.ofNullable((Boolean) eventMetadata.get("turnComplete")).orElse(false))
              .interrupted(
                  Optional.ofNullable((Boolean) eventMetadata.get("interrupted")).orElse(false))
              .branch(Optional.ofNullable((String) eventMetadata.get("branch")))
              .groundingMetadata(groundingMetadata)
              .longRunningToolIds(
                  longRunningToolIdsList != null ? new HashSet<>(longRunningToolIdsList) : null)
              .build();
    }
    return event;
  }

  /**
   * Converts a timestamp from a Map or String into an {@link Instant}.
   *
   * @param timestampObj map with "seconds"/"nanos" or an ISO string.
   * @return parsed {@link Instant}.
   */
  private static Instant convertToInstant(Object timestampObj) {
    if (timestampObj instanceof Map<?, ?> timestampMap) {
      return Instant.ofEpochSecond(
          ((Number) timestampMap.get("seconds")).longValue(),
          ((Number) timestampMap.get("nanos")).longValue());
    } else if (timestampObj != null) {
      return Instant.parse(timestampObj.toString());
    } else {
      throw new IllegalArgumentException("Timestamp not found in apiEvent");
    }
  }

  /**
   * Converts a nested map into a {@link ConcurrentMap} of {@link ConcurrentMap}s.
   *
   * @return thread-safe nested map.
   */
  @SuppressWarnings("unchecked")
  private static ConcurrentMap<String, ConcurrentMap<String, Object>>
      asConcurrentMapOfConcurrentMaps(Object value) {
    return ((Map<String, Map<String, Object>>) value)
        .entrySet().stream()
            .collect(
                ConcurrentHashMap::new,
                (map, entry) -> map.put(entry.getKey(), new ConcurrentHashMap<>(entry.getValue())),
                ConcurrentHashMap::putAll);
  }

  /** Regex for parsing full ReasoningEngine resource names. */
  private static final Pattern APP_NAME_PATTERN =
      Pattern.compile(
          "^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\\d+)$");
}

/*
 * 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.flows.llmflows;

import com.google.adk.agents.InvocationContext;
import com.google.adk.events.Event;
import com.google.adk.models.LlmResponse;
import com.google.auto.value.AutoValue;
import io.reactivex.rxjava3.core.Single;
import java.util.Optional;

public interface ResponseProcessor {

  @AutoValue
  public abstract static class ResponseProcessingResult {
    public abstract LlmResponse updatedResponse();

    public abstract Iterable<Event> events();

    public abstract Optional<String> transferToAgent();

    public static ResponseProcessingResult create(
        LlmResponse updatedResponse, Iterable<Event> events, Optional<String> transferToAgent) {
      return new AutoValue_ResponseProcessor_ResponseProcessingResult(
          updatedResponse, events, transferToAgent);
    }
  }

  /**
   * Process the LLM response as part of the post-processing stage.
   *
   * @param context the invocation context.
   * @param response the LLM response to process.
   * @return a list of events generated during processing (if any).
   */
  Single<ResponseProcessingResult> processResponse(InvocationContext context, LlmResponse response);
}

/*
 * 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;

import static com.google.common.collect.ImmutableList.toImmutableList;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.google.adk.agents.InvocationContext;
import com.google.adk.events.Event;
import com.google.adk.models.LlmRequest;
import com.google.adk.models.LlmResponse;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.Content;
import com.google.genai.types.Part;
import io.opentelemetry.api.GlobalOpenTelemetry;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Tracer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Utility class for capturing and reporting telemetry data within the ADK. This class provides
 * methods to trace various aspects of the agent's execution, including tool calls, tool responses,
 * LLM interactions, and data handling. It leverages OpenTelemetry for tracing and logging for
 * detailed information. These traces can then be exported through the ADK Dev Server UI.
 */
public class Telemetry {

  private static final Logger log = LoggerFactory.getLogger(Telemetry.class);
  private static final Tracer tracer = GlobalOpenTelemetry.getTracer("gcp.vertex.agent");

  private Telemetry() {}

  /**
   * Traces tool call arguments.
   *
   * @param args The arguments to the tool call.
   */
  public static void traceToolCall(Map<String, Object> args) {
    Span span = Span.current();
    if (span == null || !span.getSpanContext().isValid()) {
      log.trace("traceToolCall: No valid span in current context.");
      return;
    }

    span.setAttribute("gen_ai.system", "gcp.vertex.agent");
    try {
      span.setAttribute(
          "gcp.vertex.agent.tool_call_args", JsonBaseModel.getMapper().writeValueAsString(args));
    } catch (JsonProcessingException e) {
      log.warn("traceToolCall: Failed to serialize tool call args to JSON", e);
    }
  }

  /**
   * Traces tool response event.
   *
   * @param invocationContext The invocation context for the current agent run.
   * @param eventId The ID of the event.
   * @param functionResponseEvent The function response event.
   */
  public static void traceToolResponse(
      InvocationContext invocationContext, String eventId, Event functionResponseEvent) {
    Span span = Span.current();
    if (span == null || !span.getSpanContext().isValid()) {
      log.trace("traceToolResponse: No valid span in current context.");
      return;
    }

    span.setAttribute("gen_ai.system", "gcp.vertex.agent");
    span.setAttribute("gcp.vertex.agent.invocation_id", invocationContext.invocationId());
    span.setAttribute("gcp.vertex.agent.event_id", eventId);
    span.setAttribute("gcp.vertex.agent.tool_response", functionResponseEvent.toJson());

    // Setting empty llm request and response (as the AdkDevServer UI expects these)
    span.setAttribute("gcp.vertex.agent.llm_request", "{}");
    span.setAttribute("gcp.vertex.agent.llm_response", "{}");
    if (invocationContext.session() != null && invocationContext.session().id() != null) {
      span.setAttribute("gcp.vertex.agent.session_id", invocationContext.session().id());
    }
  }

  /**
   * Builds a dictionary representation of the LLM request for tracing. {@code GenerationConfig} is
   * included as a whole. For other fields like {@code Content}, parts that cannot be easily
   * serialized or are not needed for the trace (e.g., inlineData) are excluded.
   *
   * @param llmRequest The LlmRequest object.
   * @return A Map representation of the LLM request for tracing.
   */
  private static Map<String, Object> buildLlmRequestForTrace(LlmRequest llmRequest) {
    Map<String, Object> result = new HashMap<>();
    result.put("model", llmRequest.model().orElse(null));
    llmRequest.config().ifPresent(config -> result.put("config", config));

    List<Content> contentsList = new ArrayList<>();
    for (Content content : llmRequest.contents()) {
      ImmutableList<Part> filteredParts =
          content.parts().orElse(ImmutableList.of()).stream()
              .filter(part -> part.inlineData().isEmpty())
              .collect(toImmutableList());

      Content.Builder contentBuilder = Content.builder();
      content.role().ifPresent(contentBuilder::role);
      contentBuilder.parts(filteredParts);
      contentsList.add(contentBuilder.build());
    }
    result.put("contents", contentsList);
    return result;
  }

  /**
   * Traces a call to the LLM.
   *
   * @param invocationContext The invocation context.
   * @param eventId The ID of the event associated with this LLM call/response.
   * @param llmRequest The LLM request object.
   * @param llmResponse The LLM response object.
   */
  public static void traceCallLlm(
      InvocationContext invocationContext,
      String eventId,
      LlmRequest llmRequest,
      LlmResponse llmResponse) {
    Span span = Span.current();
    if (span == null || !span.getSpanContext().isValid()) {
      log.trace("traceCallLlm: No valid span in current context.");
      return;
    }

    span.setAttribute("gen_ai.system", "gcp.vertex.agent");
    llmRequest.model().ifPresent(modelName -> span.setAttribute("gen_ai.request.model", modelName));
    span.setAttribute("gcp.vertex.agent.invocation_id", invocationContext.invocationId());
    span.setAttribute("gcp.vertex.agent.event_id", eventId);

    if (invocationContext.session() != null && invocationContext.session().id() != null) {
      span.setAttribute("gcp.vertex.agent.session_id", invocationContext.session().id());
    } else {
      log.trace(
          "traceCallLlm: InvocationContext session or session ID is null, cannot set"
              + " gcp.vertex.agent.session_id");
    }

    try {
      span.setAttribute(
          "gcp.vertex.agent.llm_request",
          JsonBaseModel.getMapper().writeValueAsString(buildLlmRequestForTrace(llmRequest)));
      span.setAttribute("gcp.vertex.agent.llm_response", llmResponse.toJson());
    } catch (JsonProcessingException e) {
      log.warn("traceCallLlm: Failed to serialize LlmRequest or LlmResponse to JSON", e);
    }
  }

  /**
   * Traces the sending of data (history or new content) to the agent/model.
   *
   * @param invocationContext The invocation context.
   * @param eventId The ID of the event, if applicable.
   * @param data A list of content objects being sent.
   */
  public static void traceSendData(
      InvocationContext invocationContext, String eventId, List<Content> data) {
    Span span = Span.current();
    if (span == null || !span.getSpanContext().isValid()) {
      log.trace("traceSendData: No valid span in current context.");
      return;
    }

    span.setAttribute("gcp.vertex.agent.invocation_id", invocationContext.invocationId());
    if (eventId != null && !eventId.isEmpty()) {
      span.setAttribute("gcp.vertex.agent.event_id", eventId);
    }

    if (invocationContext.session() != null && invocationContext.session().id() != null) {
      span.setAttribute("gcp.vertex.agent.session_id", invocationContext.session().id());
    }

    try {
      List<Map<String, Object>> dataList = new ArrayList<>();
      if (data != null) {
        for (Content content : data) {
          if (content != null) {
            dataList.add(
                JsonBaseModel.getMapper()
                    .convertValue(content, new TypeReference<Map<String, Object>>() {}));
          }
        }
      }
      span.setAttribute("gcp.vertex.agent.data", JsonBaseModel.toJsonString(dataList));
    } catch (IllegalStateException e) {
      log.warn("traceSendData: Failed to serialize data to JSON", e);
    }
  }

  /**
   * Gets the tracer.
   *
   * @return The tracer.
   */
  public static Tracer getTracer() {
    return tracer;
  }
}

/*
 * 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.models.langchain4j;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.adk.models.BaseLlm;
import com.google.adk.models.BaseLlmConnection;
import com.google.adk.models.LlmRequest;
import com.google.adk.models.LlmResponse;
import com.google.genai.types.Blob;
import com.google.genai.types.Content;
import com.google.genai.types.FunctionCall;
import com.google.genai.types.FunctionCallingConfigMode;
import com.google.genai.types.FunctionDeclaration;
import com.google.genai.types.FunctionResponse;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.Part;
import com.google.genai.types.Schema;
import com.google.genai.types.ToolConfig;
import com.google.genai.types.Type;
import dev.langchain4j.Experimental;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.data.audio.Audio;
import dev.langchain4j.data.image.Image;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.AudioContent;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ImageContent;
import dev.langchain4j.data.message.PdfFileContent;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.TextContent;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.message.VideoContent;
import dev.langchain4j.data.pdf.PdfFile;
import dev.langchain4j.data.video.Video;
import dev.langchain4j.exception.UnsupportedFeatureException;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ToolChoice;
import dev.langchain4j.model.chat.request.json.JsonArraySchema;
import dev.langchain4j.model.chat.request.json.JsonBooleanSchema;
import dev.langchain4j.model.chat.request.json.JsonIntegerSchema;
import dev.langchain4j.model.chat.request.json.JsonNumberSchema;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.model.chat.request.json.JsonSchemaElement;
import dev.langchain4j.model.chat.request.json.JsonStringSchema;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import io.reactivex.rxjava3.core.BackpressureStrategy;
import io.reactivex.rxjava3.core.Flowable;
import java.util.ArrayList;
import java.util.Base64;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.UUID;

@Experimental
public class LangChain4j extends BaseLlm {

  private static final TypeReference<Map<String, Object>> MAP_TYPE_REFERENCE =
      new TypeReference<>() {};

  private final ChatModel chatModel;
  private final StreamingChatModel streamingChatModel;
  private final ObjectMapper objectMapper;

  public LangChain4j(ChatModel chatModel) {
    super(
        Objects.requireNonNull(
            chatModel.defaultRequestParameters().modelName(), "chat model name cannot be null"));
    this.chatModel = Objects.requireNonNull(chatModel, "chatModel cannot be null");
    this.streamingChatModel = null;
    this.objectMapper = new ObjectMapper();
  }

  public LangChain4j(ChatModel chatModel, String modelName) {
    super(Objects.requireNonNull(modelName, "chat model name cannot be null"));
    this.chatModel = Objects.requireNonNull(chatModel, "chatModel cannot be null");
    this.streamingChatModel = null;
    this.objectMapper = new ObjectMapper();
  }

  public LangChain4j(StreamingChatModel streamingChatModel) {
    super(
        Objects.requireNonNull(
            streamingChatModel.defaultRequestParameters().modelName(),
            "streaming chat model name cannot be null"));
    this.chatModel = null;
    this.streamingChatModel =
        Objects.requireNonNull(streamingChatModel, "streamingChatModel cannot be null");
    this.objectMapper = new ObjectMapper();
  }

  public LangChain4j(StreamingChatModel streamingChatModel, String modelName) {
    super(Objects.requireNonNull(modelName, "streaming chat model name cannot be null"));
    this.chatModel = null;
    this.streamingChatModel =
        Objects.requireNonNull(streamingChatModel, "streamingChatModel cannot be null");
    this.objectMapper = new ObjectMapper();
  }

  public LangChain4j(ChatModel chatModel, StreamingChatModel streamingChatModel, String modelName) {
    super(Objects.requireNonNull(modelName, "model name cannot be null"));
    this.chatModel = Objects.requireNonNull(chatModel, "chatModel cannot be null");
    this.streamingChatModel =
        Objects.requireNonNull(streamingChatModel, "streamingChatModel cannot be null");
    this.objectMapper = new ObjectMapper();
  }

  @Override
  public Flowable<LlmResponse> generateContent(LlmRequest llmRequest, boolean stream) {
    if (stream) {
      if (this.streamingChatModel == null) {
        return Flowable.error(new IllegalStateException("StreamingChatModel is not configured"));
      }

      ChatRequest chatRequest = toChatRequest(llmRequest);

      return Flowable.create(
          emitter -> {
            streamingChatModel.chat(
                chatRequest,
                new StreamingChatResponseHandler() {
                  @Override
                  public void onPartialResponse(String s) {
                    emitter.onNext(
                        LlmResponse.builder().content(Content.fromParts(Part.fromText(s))).build());
                  }

                  @Override
                  public void onCompleteResponse(ChatResponse chatResponse) {
                    if (chatResponse.aiMessage().hasToolExecutionRequests()) {
                      AiMessage aiMessage = chatResponse.aiMessage();
                      toParts(aiMessage).stream()
                          .map(Part::functionCall)
                          .forEach(
                              functionCall -> {
                                functionCall.ifPresent(
                                    function -> {
                                      emitter.onNext(
                                          LlmResponse.builder()
                                              .content(
                                                  Content.fromParts(
                                                      Part.fromFunctionCall(
                                                          function.name().orElse(""),
                                                          function.args().orElse(Map.of()))))
                                              .build());
                                    });
                              });
                    }
                    emitter.onComplete();
                  }

                  @Override
                  public void onError(Throwable throwable) {
                    emitter.onError(throwable);
                  }
                });
          },
          BackpressureStrategy.BUFFER);
    } else {
      if (this.chatModel == null) {
        return Flowable.error(new IllegalStateException("ChatModel is not configured"));
      }

      ChatRequest chatRequest = toChatRequest(llmRequest);
      ChatResponse chatResponse = chatModel.chat(chatRequest);
      LlmResponse llmResponse = toLlmResponse(chatResponse);

      return Flowable.just(llmResponse);
    }
  }

  private ChatRequest toChatRequest(LlmRequest llmRequest) {
    ChatRequest.Builder requestBuilder = ChatRequest.builder();

    List<ToolSpecification> toolSpecifications = toToolSpecifications(llmRequest);
    requestBuilder.toolSpecifications(toolSpecifications);

    if (llmRequest.config().isPresent()) {
      GenerateContentConfig generateContentConfig = llmRequest.config().get();

      generateContentConfig
          .temperature()
          .ifPresent(temp -> requestBuilder.temperature(temp.doubleValue()));
      generateContentConfig.topP().ifPresent(topP -> requestBuilder.topP(topP.doubleValue()));
      generateContentConfig.topK().ifPresent(topK -> requestBuilder.topK(topK.intValue()));
      generateContentConfig.maxOutputTokens().ifPresent(requestBuilder::maxOutputTokens);
      generateContentConfig.stopSequences().ifPresent(requestBuilder::stopSequences);
      generateContentConfig
          .frequencyPenalty()
          .ifPresent(freqPenalty -> requestBuilder.frequencyPenalty(freqPenalty.doubleValue()));
      generateContentConfig
          .presencePenalty()
          .ifPresent(presPenalty -> requestBuilder.presencePenalty(presPenalty.doubleValue()));

      if (generateContentConfig.toolConfig().isPresent()) {
        ToolConfig toolConfig = generateContentConfig.toolConfig().get();
        toolConfig
            .functionCallingConfig()
            .ifPresent(
                functionCallingConfig -> {
                  functionCallingConfig
                      .mode()
                      .ifPresent(
                          functionMode -> {
                            if (functionMode
                                .knownEnum()
                                .equals(FunctionCallingConfigMode.Known.AUTO)) {
                              requestBuilder.toolChoice(ToolChoice.AUTO);
                            } else if (functionMode
                                .knownEnum()
                                .equals(FunctionCallingConfigMode.Known.ANY)) {
                              // TODO check if it's the correct
                              // mapping
                              requestBuilder.toolChoice(ToolChoice.REQUIRED);
                              functionCallingConfig
                                  .allowedFunctionNames()
                                  .ifPresent(
                                      allowedFunctionNames -> {
                                        requestBuilder.toolSpecifications(
                                            toolSpecifications.stream()
                                                .filter(
                                                    toolSpecification ->
                                                        allowedFunctionNames.contains(
                                                            toolSpecification.name()))
                                                .toList());
                                      });
                            } else if (functionMode
                                .knownEnum()
                                .equals(FunctionCallingConfigMode.Known.NONE)) {
                              requestBuilder.toolSpecifications(List.of());
                            }
                          });
                });
        toolConfig
            .retrievalConfig()
            .ifPresent(
                retrievalConfig -> {
                  // TODO? It exposes Latitude / Longitude, what to do with this?
                });
      }
    }

    return requestBuilder.messages(toMessages(llmRequest)).build();
  }

  private List<ChatMessage> toMessages(LlmRequest llmRequest) {
    List<ChatMessage> messages =
        new ArrayList<>(
            llmRequest.getSystemInstructions().stream().map(SystemMessage::from).toList());
    llmRequest.contents().forEach(content -> messages.addAll(toChatMessage(content)));
    return messages;
  }

  private List<ChatMessage> toChatMessage(Content content) {
    String role = content.role().orElseThrow().toLowerCase();
    return switch (role) {
      case "user" -> toUserOrToolResultMessage(content);
      case "model", "assistant" -> List.of(toAiMessage(content));
      default -> throw new IllegalStateException("Unexpected role: " + role);
    };
  }

  private List<ChatMessage> toUserOrToolResultMessage(Content content) {
    List<ToolExecutionResultMessage> toolExecutionResultMessages = new ArrayList<>();
    List<ToolExecutionRequest> toolExecutionRequests = new ArrayList<>();

    List<dev.langchain4j.data.message.Content> lc4jContents = new ArrayList<>();

    for (Part part : content.parts().orElse(List.of())) {
      if (part.text().isPresent()) {
        lc4jContents.add(TextContent.from(part.text().get()));
      } else if (part.functionResponse().isPresent()) {
        FunctionResponse functionResponse = part.functionResponse().get();
        toolExecutionResultMessages.add(
            ToolExecutionResultMessage.from(
                functionResponse.id().orElseThrow(),
                functionResponse.name().orElseThrow(),
                toJson(functionResponse.response().orElseThrow())));
      } else if (part.functionCall().isPresent()) {
        FunctionCall functionCall = part.functionCall().get();
        toolExecutionRequests.add(
            ToolExecutionRequest.builder()
                .id(functionCall.id().orElseThrow())
                .name(functionCall.name().orElseThrow())
                .arguments(toJson(functionCall.args().orElse(Map.of())))
                .build());
      } else if (part.inlineData().isPresent()) {
        Blob blob = part.inlineData().get();

        if (blob.mimeType().isEmpty() || blob.data().isEmpty()) {
          throw new IllegalArgumentException("Mime type and data required");
        }

        byte[] bytes = blob.data().get();
        String mimeType = blob.mimeType().get();

        Base64.Encoder encoder = Base64.getEncoder();

        dev.langchain4j.data.message.Content lc4jContent = null;

        if (mimeType.startsWith("audio/")) {
          lc4jContent =
              AudioContent.from(
                  Audio.builder()
                      .base64Data(encoder.encodeToString(bytes))
                      .mimeType(mimeType)
                      .build());
        } else if (mimeType.startsWith("video/")) {
          lc4jContent =
              VideoContent.from(
                  Video.builder()
                      .base64Data(encoder.encodeToString(bytes))
                      .mimeType(mimeType)
                      .build());
        } else if (mimeType.startsWith("image/")) {
          lc4jContent =
              ImageContent.from(
                  Image.builder()
                      .base64Data(encoder.encodeToString(bytes))
                      .mimeType(mimeType)
                      .build());
        } else if (mimeType.startsWith("application/pdf")) {
          lc4jContent =
              PdfFileContent.from(
                  PdfFile.builder()
                      .base64Data(encoder.encodeToString(bytes))
                      .mimeType(mimeType)
                      .build());
        } else if (mimeType.startsWith("text/")
            || mimeType.equals("application/json")
            || mimeType.endsWith("+json")
            || mimeType.endsWith("+xml")) {
          // TODO are there missing text based mime types?
          // TODO should we assume UTF_8?
          lc4jContents.add(
              TextContent.from(new String(bytes, java.nio.charset.StandardCharsets.UTF_8)));
        }

        if (lc4jContent != null) {
          lc4jContents.add(lc4jContent);
        } else {
          throw new IllegalArgumentException("Unknown or unhandled mime type: " + mimeType);
        }
      } else {
        throw new IllegalStateException(
            "Text, media or functionCall is expected, but was: " + part);
      }
    }

    if (!toolExecutionResultMessages.isEmpty()) {
      return new ArrayList<ChatMessage>(toolExecutionResultMessages);
    } else if (!toolExecutionRequests.isEmpty()) {
      return toolExecutionRequests.stream()
          .map(AiMessage::aiMessage)
          .map(msg -> (ChatMessage) msg)
          .toList();
    } else {
      return List.of(UserMessage.from(lc4jContents));
    }
  }

  private AiMessage toAiMessage(Content content) {
    List<String> texts = new ArrayList<>();
    List<ToolExecutionRequest> toolExecutionRequests = new ArrayList<>();

    content
        .parts()
        .orElse(List.of())
        .forEach(
            part -> {
              if (part.text().isPresent()) {
                texts.add(part.text().get());
              } else if (part.functionCall().isPresent()) {
                FunctionCall functionCall = part.functionCall().get();
                ToolExecutionRequest toolExecutionRequest =
                    ToolExecutionRequest.builder()
                        .id(functionCall.id().orElseThrow())
                        .name(functionCall.name().orElseThrow())
                        .arguments(toJson(functionCall.args().orElseThrow()))
                        .build();
                toolExecutionRequests.add(toolExecutionRequest);
              } else {
                throw new IllegalStateException(
                    "Either text or functionCall is expected, but was: " + part);
              }
            });

    return AiMessage.builder()
        .text(String.join("\n", texts))
        .toolExecutionRequests(toolExecutionRequests)
        .build();
  }

  private String toJson(Object object) {
    try {
      return objectMapper.writeValueAsString(object);
    } catch (JsonProcessingException e) {
      throw new RuntimeException(e);
    }
  }

  private List<ToolSpecification> toToolSpecifications(LlmRequest llmRequest) {
    List<ToolSpecification> toolSpecifications = new ArrayList<>();

    llmRequest
        .tools()
        .values()
        .forEach(
            baseTool -> {
              if (baseTool.declaration().isPresent()) {
                FunctionDeclaration functionDeclaration = baseTool.declaration().get();
                if (functionDeclaration.parameters().isPresent()) {
                  Schema schema = functionDeclaration.parameters().get();
                  ToolSpecification toolSpecification =
                      ToolSpecification.builder()
                          .name(baseTool.name())
                          .description(baseTool.description())
                          .parameters(toParameters(schema))
                          .build();
                  toolSpecifications.add(toolSpecification);
                } else {
                  // TODO exception or something else?
                  throw new IllegalStateException("Tool lacking parameters: " + baseTool);
                }
              } else {
                // TODO exception or something else?
                throw new IllegalStateException("Tool lacking declaration: " + baseTool);
              }
            });

    return toolSpecifications;
  }

  private JsonObjectSchema toParameters(Schema schema) {
    if (schema.type().isPresent() && schema.type().get().knownEnum().equals(Type.Known.OBJECT)) {
      return JsonObjectSchema.builder()
          .addProperties(toProperties(schema))
          .required(schema.required().orElse(List.of()))
          .build();
    } else {
      throw new UnsupportedOperationException(
          "LangChain4jLlm does not support schema of type: " + schema.type());
    }
  }

  private Map<String, JsonSchemaElement> toProperties(Schema schema) {
    Map<String, Schema> properties = schema.properties().orElse(Map.of());
    Map<String, JsonSchemaElement> result = new HashMap<>();
    properties.forEach((k, v) -> result.put(k, toJsonSchemaElement(v)));
    return result;
  }

  private JsonSchemaElement toJsonSchemaElement(Schema schema) {
    if (schema != null && schema.type().isPresent()) {
      Type type = schema.type().get();
      return switch (type.knownEnum()) {
        case STRING ->
            JsonStringSchema.builder().description(schema.description().orElse(null)).build();
        case NUMBER ->
            JsonNumberSchema.builder().description(schema.description().orElse(null)).build();
        case INTEGER ->
            JsonIntegerSchema.builder().description(schema.description().orElse(null)).build();
        case BOOLEAN ->
            JsonBooleanSchema.builder().description(schema.description().orElse(null)).build();
        case ARRAY ->
            JsonArraySchema.builder()
                .description(schema.description().orElse(null))
                .items(toJsonSchemaElement(schema.items().orElseThrow()))
                .build();
        case OBJECT -> toParameters(schema);
        case TYPE_UNSPECIFIED ->
            throw new UnsupportedFeatureException(
                "LangChain4jLlm does not support schema of type: " + type);
      };
    } else {
      throw new IllegalArgumentException("Schema type cannot be null or absent");
    }
  }

  private LlmResponse toLlmResponse(ChatResponse chatResponse) {
    Content content =
        Content.builder().role("model").parts(toParts(chatResponse.aiMessage())).build();

    return LlmResponse.builder().content(content).build();
  }

  private List<Part> toParts(AiMessage aiMessage) {
    if (aiMessage.hasToolExecutionRequests()) {
      List<Part> parts = new ArrayList<>();
      aiMessage
          .toolExecutionRequests()
          .forEach(
              toolExecutionRequest -> {
                FunctionCall functionCall =
                    FunctionCall.builder()
                        .id(
                            toolExecutionRequest.id() != null
                                ? toolExecutionRequest.id()
                                : UUID.randomUUID().toString())
                        .name(toolExecutionRequest.name())
                        .args(toArgs(toolExecutionRequest))
                        .build();
                Part part = Part.builder().functionCall(functionCall).build();
                parts.add(part);
              });
      return parts;
    } else {
      Part part = Part.builder().text(aiMessage.text()).build();
      return List.of(part);
    }
  }

  private Map<String, Object> toArgs(ToolExecutionRequest toolExecutionRequest) {
    try {
      return objectMapper.readValue(toolExecutionRequest.arguments(), MAP_TYPE_REFERENCE);
    } catch (JsonProcessingException e) {
      throw new RuntimeException(e);
    }
  }

  @Override
  public BaseLlmConnection connect(LlmRequest llmRequest) {
    throw new UnsupportedOperationException(
        "Live connection is not supported for LangChain4j models.");
  }
}

/*
 * 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.agents;

import com.google.adk.events.Event;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import io.reactivex.rxjava3.core.Flowable;
import java.util.List;
import java.util.Optional;

/**
 * An agent that runs its sub-agents sequentially in a loop.
 *
 * <p>The loop continues until a sub-agent escalates, or until the maximum number of iterations is
 * reached (if specified).
 */
public class LoopAgent extends BaseAgent {

  private final Optional<Integer> maxIterations;

  /**
   * Constructor for LoopAgent.
   *
   * @param name The agent's name.
   * @param description The agent's description.
   * @param subAgents The list of sub-agents to run in the loop.
   * @param maxIterations Optional termination condition: maximum number of loop iterations.
   * @param beforeAgentCallback Optional callback before the agent runs.
   * @param afterAgentCallback Optional callback after the agent runs.
   */
  private LoopAgent(
      String name,
      String description,
      List<? extends BaseAgent> subAgents,
      Optional<Integer> maxIterations,
      List<Callbacks.BeforeAgentCallback> beforeAgentCallback,
      List<Callbacks.AfterAgentCallback> afterAgentCallback) {

    super(name, description, subAgents, beforeAgentCallback, afterAgentCallback);
    this.maxIterations = maxIterations;
  }

  /** Builder for {@link LoopAgent}. */
  public static class Builder {
    private String name;
    private String description;
    private List<? extends BaseAgent> subAgents;
    private Optional<Integer> maxIterations = Optional.empty();
    private ImmutableList<Callbacks.BeforeAgentCallback> beforeAgentCallback;
    private ImmutableList<Callbacks.AfterAgentCallback> afterAgentCallback;

    @CanIgnoreReturnValue
    public Builder name(String name) {
      this.name = name;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder description(String description) {
      this.description = description;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder subAgents(List<? extends BaseAgent> subAgents) {
      this.subAgents = subAgents;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder subAgents(BaseAgent... subAgents) {
      this.subAgents = ImmutableList.copyOf(subAgents);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder maxIterations(int maxIterations) {
      this.maxIterations = Optional.of(maxIterations);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder maxIterations(Optional<Integer> maxIterations) {
      this.maxIterations = maxIterations;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeAgentCallback(Callbacks.BeforeAgentCallback beforeAgentCallback) {
      this.beforeAgentCallback = ImmutableList.of(beforeAgentCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeAgentCallback(
        List<Callbacks.BeforeAgentCallbackBase> beforeAgentCallback) {
      this.beforeAgentCallback = CallbackUtil.getBeforeAgentCallbacks(beforeAgentCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterAgentCallback(Callbacks.AfterAgentCallback afterAgentCallback) {
      this.afterAgentCallback = ImmutableList.of(afterAgentCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterAgentCallback(List<Callbacks.AfterAgentCallbackBase> afterAgentCallback) {
      this.afterAgentCallback = CallbackUtil.getAfterAgentCallbacks(afterAgentCallback);
      return this;
    }

    public LoopAgent build() {
      // TODO(b/410859954): Add validation for required fields like name.
      return new LoopAgent(
          name, description, subAgents, maxIterations, beforeAgentCallback, afterAgentCallback);
    }
  }

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

  @Override
  protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
    List<? extends BaseAgent> subAgents = subAgents();
    if (subAgents == null || subAgents.isEmpty()) {
      return Flowable.empty();
    }

    return Flowable.fromIterable(subAgents)
        .concatMap(subAgent -> subAgent.runAsync(invocationContext))
        .repeat(maxIterations.orElse(Integer.MAX_VALUE))
        .takeUntil(LoopAgent::hasEscalateAction);
  }

  @Override
  protected Flowable<Event> runLiveImpl(InvocationContext invocationContext) {
    return Flowable.error(
        new UnsupportedOperationException("runLive is not defined for LoopAgent yet."));
  }

  private static boolean hasEscalateAction(Event event) {
    return event.actions().escalate().orElse(false);
  }
}

/*
 * 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.agents;

import com.google.adk.events.Event;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import io.reactivex.rxjava3.core.Flowable;
import java.util.List;

/** An agent that runs its sub-agents sequentially. */
public class SequentialAgent extends BaseAgent {

  /**
   * Constructor for SequentialAgent.
   *
   * @param name The agent's name.
   * @param description The agent's description.
   * @param subAgents The list of sub-agents to run sequentially.
   * @param beforeAgentCallback Optional callback before the agent runs.
   * @param afterAgentCallback Optional callback after the agent runs.
   */
  private SequentialAgent(
      String name,
      String description,
      List<? extends BaseAgent> subAgents,
      List<Callbacks.BeforeAgentCallback> beforeAgentCallback,
      List<Callbacks.AfterAgentCallback> afterAgentCallback) {

    super(name, description, subAgents, beforeAgentCallback, afterAgentCallback);
  }

  /** Builder for {@link SequentialAgent}. */
  public static class Builder {
    private String name;
    private String description;
    private List<? extends BaseAgent> subAgents;
    private ImmutableList<Callbacks.BeforeAgentCallback> beforeAgentCallback;
    private ImmutableList<Callbacks.AfterAgentCallback> afterAgentCallback;

    @CanIgnoreReturnValue
    public Builder name(String name) {
      this.name = name;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder description(String description) {
      this.description = description;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder subAgents(List<? extends BaseAgent> subAgents) {
      this.subAgents = subAgents;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder subAgents(BaseAgent... subAgents) {
      this.subAgents = ImmutableList.copyOf(subAgents);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeAgentCallback(Callbacks.BeforeAgentCallback beforeAgentCallback) {
      this.beforeAgentCallback = ImmutableList.of(beforeAgentCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder beforeAgentCallback(
        List<Callbacks.BeforeAgentCallbackBase> beforeAgentCallback) {
      this.beforeAgentCallback = CallbackUtil.getBeforeAgentCallbacks(beforeAgentCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterAgentCallback(Callbacks.AfterAgentCallback afterAgentCallback) {
      this.afterAgentCallback = ImmutableList.of(afterAgentCallback);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder afterAgentCallback(List<Callbacks.AfterAgentCallbackBase> afterAgentCallback) {
      this.afterAgentCallback = CallbackUtil.getAfterAgentCallbacks(afterAgentCallback);
      return this;
    }

    public SequentialAgent build() {
      // TODO(b/410859954): Add validation for required fields like name.
      return new SequentialAgent(
          name, description, subAgents, beforeAgentCallback, afterAgentCallback);
    }
  }

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

  /**
   * Runs sub-agents sequentially.
   *
   * @param invocationContext Invocation context.
   * @return Flowable emitting events from sub-agents.
   */
  @Override
  protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
    return Flowable.fromIterable(subAgents())
        .concatMap(subAgent -> subAgent.runAsync(invocationContext));
  }

  /**
   * Runs sub-agents sequentially in live mode.
   *
   * @param invocationContext Invocation context.
   * @return Flowable emitting events from sub-agents in live mode.
   */
  @Override
  protected Flowable<Event> runLiveImpl(InvocationContext invocationContext) {
    return Flowable.fromIterable(subAgents())
        .concatMap(subAgent -> subAgent.runLive(invocationContext));
  }
}

/*
 * 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.agents;

import com.google.adk.artifacts.BaseArtifactService;
import com.google.adk.exceptions.LlmCallsLimitExceededException;
import com.google.adk.sessions.BaseSessionService;
import com.google.adk.sessions.Session;
import com.google.genai.types.Content;
import java.util.Objects;
import java.util.Optional;
import java.util.UUID;
import javax.annotation.Nullable;

/** The context for an agent invocation. */
public class InvocationContext {

  private final BaseSessionService sessionService;
  private final BaseArtifactService artifactService;
  private final Optional<LiveRequestQueue> liveRequestQueue;

  private Optional<String> branch;
  private final String invocationId;
  private BaseAgent agent;
  private final Session session;

  private final Optional<Content> userContent;
  private final RunConfig runConfig;
  private boolean endInvocation;
  private final InvocationCostManager invocationCostManager = new InvocationCostManager();

  private InvocationContext(
      BaseSessionService sessionService,
      BaseArtifactService artifactService,
      Optional<LiveRequestQueue> liveRequestQueue,
      Optional<String> branch,
      String invocationId,
      BaseAgent agent,
      Session session,
      Optional<Content> userContent,
      RunConfig runConfig,
      boolean endInvocation) {
    this.sessionService = sessionService;
    this.artifactService = artifactService;
    this.liveRequestQueue = liveRequestQueue;
    this.branch = branch;
    this.invocationId = invocationId;
    this.agent = agent;
    this.session = session;
    this.userContent = userContent;
    this.runConfig = runConfig;
    this.endInvocation = endInvocation;
  }

  public static InvocationContext create(
      BaseSessionService sessionService,
      BaseArtifactService artifactService,
      String invocationId,
      BaseAgent agent,
      Session session,
      Content userContent,
      RunConfig runConfig) {
    return new InvocationContext(
        sessionService,
        artifactService,
        Optional.empty(),
        /* branch= */ Optional.empty(),
        invocationId,
        agent,
        session,
        Optional.ofNullable(userContent),
        runConfig,
        false);
  }

  public static InvocationContext create(
      BaseSessionService sessionService,
      BaseArtifactService artifactService,
      BaseAgent agent,
      Session session,
      LiveRequestQueue liveRequestQueue,
      RunConfig runConfig) {
    return new InvocationContext(
        sessionService,
        artifactService,
        Optional.ofNullable(liveRequestQueue),
        /* branch= */ Optional.empty(),
        InvocationContext.newInvocationContextId(),
        agent,
        session,
        Optional.empty(),
        runConfig,
        false);
  }

  public static InvocationContext copyOf(InvocationContext other) {
    return new InvocationContext(
        other.sessionService,
        other.artifactService,
        other.liveRequestQueue,
        other.branch,
        other.invocationId,
        other.agent,
        other.session,
        other.userContent,
        other.runConfig,
        other.endInvocation);
  }

  public BaseSessionService sessionService() {
    return sessionService;
  }

  public BaseArtifactService artifactService() {
    return artifactService;
  }

  public Optional<LiveRequestQueue> liveRequestQueue() {
    return liveRequestQueue;
  }

  public String invocationId() {
    return invocationId;
  }

  public void branch(@Nullable String branch) {
    this.branch = Optional.ofNullable(branch);
  }

  public Optional<String> branch() {
    return branch;
  }

  public BaseAgent agent() {
    return agent;
  }

  public void agent(BaseAgent agent) {
    this.agent = agent;
  }

  public Session session() {
    return session;
  }

  public Optional<Content> userContent() {
    return userContent;
  }

  public RunConfig runConfig() {
    return runConfig;
  }

  public boolean endInvocation() {
    return endInvocation;
  }

  public void setEndInvocation(boolean endInvocation) {
    this.endInvocation = endInvocation;
  }

  public String appName() {
    return session.appName();
  }

  public String userId() {
    return session.userId();
  }

  public static String newInvocationContextId() {
    return "e-" + UUID.randomUUID();
  }

  public void incrementLlmCallsCount() throws LlmCallsLimitExceededException {
    this.invocationCostManager.incrementAndEnforceLlmCallsLimit(this.runConfig);
  }

  private static class InvocationCostManager {
    private int numberOfLlmCalls = 0;

    public void incrementAndEnforceLlmCallsLimit(RunConfig runConfig)
        throws LlmCallsLimitExceededException {
      this.numberOfLlmCalls++;

      if (runConfig != null
          && runConfig.maxLlmCalls() > 0
          && this.numberOfLlmCalls > runConfig.maxLlmCalls()) {
        throw new LlmCallsLimitExceededException(
            "Max number of llm calls limit of " + runConfig.maxLlmCalls() + " exceeded");
      }
    }
  }

  @Override
  public boolean equals(Object o) {
    if (this == o) {
      return true;
    }
    if (!(o instanceof InvocationContext that)) {
      return false;
    }
    return endInvocation == that.endInvocation
        && Objects.equals(sessionService, that.sessionService)
        && Objects.equals(artifactService, that.artifactService)
        && Objects.equals(liveRequestQueue, that.liveRequestQueue)
        && Objects.equals(branch, that.branch)
        && Objects.equals(invocationId, that.invocationId)
        && Objects.equals(agent, that.agent)
        && Objects.equals(session, that.session)
        && Objects.equals(userContent, that.userContent)
        && Objects.equals(runConfig, that.runConfig);
  }

  @Override
  public int hashCode() {
    return Objects.hash(
        sessionService,
        artifactService,
        liveRequestQueue,
        branch,
        invocationId,
        agent,
        session,
        userContent,
        runConfig,
        endInvocation);
  }
}

/*
 * 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.tools;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.genai.types.FunctionDeclaration;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.lang.reflect.Parameter;
import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import javax.annotation.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** FunctionTool implements a customized function calling tool. */
public class FunctionTool extends BaseTool {
  private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
  private static final Logger logger = LoggerFactory.getLogger(FunctionTool.class);

  @Nullable private final Object instance;
  private final Method func;
  private final FunctionDeclaration funcDeclaration;

  public static FunctionTool create(Object instance, Method func) {
    if (!areParametersAnnotatedWithSchema(func) && wasCompiledWithDefaultParameterNames(func)) {
      logger.error(
          "Functions used in tools must have their parameters annotated with @Schema or at least"
              + " the code must be compiled with the -parameters flag as a fallback. Your function"
              + " tool will likely not work as expected and exit at runtime.");
    }
    if (!Modifier.isStatic(func.getModifiers()) && !func.getDeclaringClass().isInstance(instance)) {
      throw new IllegalArgumentException(
          String.format(
              "The instance provided is not an instance of the declaring class of the method."
                  + " Expected: %s, Actual: %s",
              func.getDeclaringClass().getName(), instance.getClass().getName()));
    }
    return new FunctionTool(instance, func, /* isLongRunning= */ false);
  }

  public static FunctionTool create(Method func) {
    if (!areParametersAnnotatedWithSchema(func) && wasCompiledWithDefaultParameterNames(func)) {
      logger.error(
          "Functions used in tools must have their parameters annotated with @Schema or at least"
              + " the code must be compiled with the -parameters flag as a fallback. Your function"
              + " tool will likely not work as expected and exit at runtime.");
    }
    if (!Modifier.isStatic(func.getModifiers())) {
      throw new IllegalArgumentException("The method provided must be static.");
    }
    return new FunctionTool(null, func, /* isLongRunning= */ false);
  }

  public static FunctionTool create(Class<?> cls, String methodName) {
    for (Method method : cls.getMethods()) {
      if (method.getName().equals(methodName) && Modifier.isStatic(method.getModifiers())) {
        return create(null, method);
      }
    }
    throw new IllegalArgumentException(
        String.format("Static method %s not found in class %s.", methodName, cls.getName()));
  }

  public static FunctionTool create(Object instance, String methodName) {
    Class<?> cls = instance.getClass();
    for (Method method : cls.getMethods()) {
      if (method.getName().equals(methodName) && !Modifier.isStatic(method.getModifiers())) {
        return create(instance, method);
      }
    }
    throw new IllegalArgumentException(
        String.format("Instance method %s not found in class %s.", methodName, cls.getName()));
  }

  private static boolean areParametersAnnotatedWithSchema(Method func) {
    for (Parameter parameter : func.getParameters()) {
      if (!parameter.isAnnotationPresent(Annotations.Schema.class)
          || parameter.getAnnotation(Annotations.Schema.class).name().isEmpty()) {
        return false;
      }
    }
    return true;
  }

  // Rough check to see if the code wasn't compiled with the -parameters flag.
  private static boolean wasCompiledWithDefaultParameterNames(Method func) {
    for (Parameter parameter : func.getParameters()) {
      String parameterName = parameter.getName();
      if (!parameterName.matches("arg\\d+")) {
        return false;
      }
    }
    return true;
  }

  protected FunctionTool(@Nullable Object instance, Method func, boolean isLongRunning) {
    super(
        func.isAnnotationPresent(Annotations.Schema.class)
                && !func.getAnnotation(Annotations.Schema.class).name().isEmpty()
            ? func.getAnnotation(Annotations.Schema.class).name()
            : func.getName(),
        func.isAnnotationPresent(Annotations.Schema.class)
            ? func.getAnnotation(Annotations.Schema.class).description()
            : "",
        isLongRunning);
    boolean isStatic = Modifier.isStatic(func.getModifiers());
    if (isStatic && instance != null) {
      throw new IllegalArgumentException("Static function tool must not have an instance.");
    } else if (!isStatic && instance == null) {
      throw new IllegalArgumentException("Instance function tool must have an instance.");
    }

    this.instance = instance;
    this.func = func;
    this.funcDeclaration =
        FunctionCallingUtils.buildFunctionDeclaration(this.func, ImmutableList.of("toolContext"));
  }

  @Override
  public Optional<FunctionDeclaration> declaration() {
    return Optional.of(this.funcDeclaration);
  }

  @Override
  public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContext toolContext) {
    try {
      return this.call(args, toolContext).defaultIfEmpty(ImmutableMap.of());
    } catch (Exception e) {
      e.printStackTrace();
      return Single.just(ImmutableMap.of());
    }
  }

  @SuppressWarnings("unchecked") // For tool parameter type casting.
  private Maybe<Map<String, Object>> call(Map<String, Object> args, ToolContext toolContext)
      throws IllegalAccessException, InvocationTargetException {
    Parameter[] parameters = func.getParameters();
    Object[] arguments = new Object[parameters.length];
    for (int i = 0; i < parameters.length; i++) {
      String paramName =
          parameters[i].isAnnotationPresent(Annotations.Schema.class)
                  && !parameters[i].getAnnotation(Annotations.Schema.class).name().isEmpty()
              ? parameters[i].getAnnotation(Annotations.Schema.class).name()
              : parameters[i].getName();
      if (paramName.equals("toolContext")) {
        arguments[i] = toolContext;
        continue;
      }
      if (!args.containsKey(paramName)) {
        throw new IllegalArgumentException(
            String.format(
                "The parameter '%s' was not found in the arguments provided by the model.",
                paramName));
      }
      Class<?> paramType = parameters[i].getType();
      Object argValue = args.get(paramName);
      if (paramType.equals(List.class)) {
        if (argValue instanceof List) {
          Type type =
              ((ParameterizedType) parameters[i].getParameterizedType())
                  .getActualTypeArguments()[0];
          arguments[i] = createList((List<Object>) argValue, (Class) type);
          continue;
        }
      } else if (argValue instanceof Map) {
        arguments[i] = OBJECT_MAPPER.convertValue(argValue, paramType);
        continue;
      }
      arguments[i] = castValue(argValue, paramType);
    }
    Object result = func.invoke(instance, arguments);
    if (result == null) {
      return Maybe.empty();
    } else if (result instanceof Maybe) {
      return (Maybe<Map<String, Object>>) result;
    } else if (result instanceof Single) {
      return ((Single<Map<String, Object>>) result).toMaybe();
    } else {
      return Maybe.just((Map<String, Object>) result);
    }
  }

  private static List<Object> createList(List<Object> values, Class<?> type) {
    List<Object> list = new ArrayList<>();
    // List of parameterized type is not supported.
    if (type == null) {
      return list;
    }
    Class<?> cls = type;
    for (Object value : values) {
      if (cls == Integer.class
          || cls == Long.class
          || cls == Double.class
          || cls == Float.class
          || cls == Boolean.class
          || cls == String.class) {
        list.add(castValue(value, cls));
      } else {
        list.add(OBJECT_MAPPER.convertValue(value, type));
      }
    }
    return list;
  }

  private static Object castValue(Object value, Class<?> type) {
    if (type.equals(Integer.class) || type.equals(int.class)) {
      if (value instanceof Integer) {
        return value;
      }
    }
    if (type.equals(Long.class) || type.equals(long.class)) {
      if (value instanceof Long || value instanceof Integer) {
        return value;
      }
    } else if (type.equals(Double.class) || type.equals(double.class)) {
      if (value instanceof Double d) {
        return d.doubleValue();
      }
      if (value instanceof Float f) {
        return f.doubleValue();
      }
      if (value instanceof Integer i) {
        return i.doubleValue();
      }
      if (value instanceof Long l) {
        return l.doubleValue();
      }
    } else if (type.equals(Float.class) || type.equals(float.class)) {
      if (value instanceof Double d) {
        return d.floatValue();
      }
      if (value instanceof Float f) {
        return f.floatValue();
      }
      if (value instanceof Integer i) {
        return i.floatValue();
      }
      if (value instanceof Long l) {
        return l.floatValue();
      }
    } else if (type.equals(Boolean.class) || type.equals(boolean.class)) {
      if (value instanceof Boolean) {
        return value;
      }
    } else if (type.equals(String.class)) {
      if (value instanceof String) {
        return value;
      }
    }
    return OBJECT_MAPPER.convertValue(value, type);
  }
}

/*
 * 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.models;

import static com.google.common.base.StandardSystemProperty.JAVA_VERSION;
import static com.google.common.collect.ImmutableList.toImmutableList;

import com.google.adk.Version;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.genai.Client;
import com.google.genai.ResponseStream;
import com.google.genai.types.Blob;
import com.google.genai.types.Candidate;
import com.google.genai.types.Content;
import com.google.genai.types.FileData;
import com.google.genai.types.FinishReason;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.GenerateContentResponse;
import com.google.genai.types.HttpOptions;
import com.google.genai.types.LiveConnectConfig;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Flowable;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Stream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Represents the Gemini Generative AI model.
 *
 * <p>This class provides methods for interacting with the Gemini model, including standard
 * request-response generation and establishing persistent bidirectional connections.
 */
public class Gemini extends BaseLlm {

  private static final Logger logger = LoggerFactory.getLogger(Gemini.class);
  private static final ImmutableMap<String, String> TRACKING_HEADERS;

  static {
    String frameworkLabel = "google-adk/" + Version.JAVA_ADK_VERSION;
    String languageLabel = "gl-java/" + JAVA_VERSION.value();
    String versionHeaderValue = String.format("%s %s", frameworkLabel, languageLabel);

    TRACKING_HEADERS =
        ImmutableMap.of(
            "x-goog-api-client", versionHeaderValue,
            "user-agent", versionHeaderValue);
  }

  private final Client apiClient;

  private static final String CONTINUE_OUTPUT_MESSAGE =
      "Continue output. DO NOT look at this line. ONLY look at the content before this line and"
          + " system instruction.";

  /**
   * Constructs a new Gemini instance.
   *
   * @param modelName The name of the Gemini model to use (e.g., "gemini-2.0-flash").
   * @param apiClient The genai {@link com.google.genai.Client} instance for making API calls.
   */
  public Gemini(String modelName, Client apiClient) {
    super(modelName);
    this.apiClient = Objects.requireNonNull(apiClient, "apiClient cannot be null");
  }

  /**
   * Constructs a new Gemini instance with a Google Gemini API key.
   *
   * @param modelName The name of the Gemini model to use (e.g., "gemini-2.0-flash").
   * @param apiKey The Google Gemini API key.
   */
  public Gemini(String modelName, String apiKey) {
    super(modelName);
    Objects.requireNonNull(apiKey, "apiKey cannot be null");
    this.apiClient =
        Client.builder()
            .apiKey(apiKey)
            .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
            .build();
  }

  /**
   * Constructs a new Gemini instance with a Google Gemini API key.
   *
   * @param modelName The name of the Gemini model to use (e.g., "gemini-2.0-flash").
   * @param vertexCredentials The Vertex AI credentials to access the Gemini model.
   */
  public Gemini(String modelName, VertexCredentials vertexCredentials) {
    super(modelName);
    Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null");
    Client.Builder apiClientBuilder =
        Client.builder().httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build());
    vertexCredentials.project().ifPresent(apiClientBuilder::project);
    vertexCredentials.location().ifPresent(apiClientBuilder::location);
    vertexCredentials.credentials().ifPresent(apiClientBuilder::credentials);
    this.apiClient = apiClientBuilder.build();
  }

  /**
   * Returns a new Builder instance for constructing Gemini objects. Note that when building a
   * Gemini object, at least one of apiKey, vertexCredentials, or an explicit apiClient must be set.
   * If multiple are set, the explicit apiClient will take precedence.
   *
   * @return A new {@link Builder}.
   */
  public static Builder builder() {
    return new Builder();
  }

  /** Builder for {@link Gemini}. */
  public static class Builder {
    private String modelName;
    private Client apiClient;
    private String apiKey;
    private VertexCredentials vertexCredentials;

    private Builder() {}

    /**
     * Sets the name of the Gemini model to use.
     *
     * @param modelName The model name (e.g., "gemini-2.0-flash").
     * @return This builder.
     */
    @CanIgnoreReturnValue
    public Builder modelName(String modelName) {
      this.modelName = modelName;
      return this;
    }

    /**
     * Sets the explicit {@link com.google.genai.Client} instance for making API calls. If this is
     * set, apiKey and vertexCredentials will be ignored.
     *
     * @param apiClient The client instance.
     * @return This builder.
     */
    @CanIgnoreReturnValue
    public Builder apiClient(Client apiClient) {
      this.apiClient = apiClient;
      return this;
    }

    /**
     * Sets the Google Gemini API key. If {@link #apiClient(Client)} is also set, the explicit
     * client will take precedence. If {@link #vertexCredentials(VertexCredentials)} is also set,
     * this apiKey will take precedence.
     *
     * @param apiKey The API key.
     * @return This builder.
     */
    @CanIgnoreReturnValue
    public Builder apiKey(String apiKey) {
      this.apiKey = apiKey;
      return this;
    }

    /**
     * Sets the Vertex AI credentials. If {@link #apiClient(Client)} or {@link #apiKey(String)} are
     * also set, they will take precedence over these credentials.
     *
     * @param vertexCredentials The Vertex AI credentials.
     * @return This builder.
     */
    @CanIgnoreReturnValue
    public Builder vertexCredentials(VertexCredentials vertexCredentials) {
      this.vertexCredentials = vertexCredentials;
      return this;
    }

    /**
     * Builds the {@link Gemini} instance.
     *
     * @return A new {@link Gemini} instance.
     * @throws NullPointerException if modelName is null.
     */
    public Gemini build() {
      Objects.requireNonNull(modelName, "modelName must be set.");

      if (apiClient != null) {
        return new Gemini(modelName, apiClient);
      } else if (apiKey != null) {
        return new Gemini(modelName, apiKey);
      } else if (vertexCredentials != null) {
        return new Gemini(modelName, vertexCredentials);
      } else {
        return new Gemini(
            modelName,
            Client.builder()
                .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
                .build());
      }
    }
  }

  /**
   * Sanitizes the request to ensure it is compatible with the configured API backend. Required as
   * there are some parameters that if included in the request will raise a runtime error if sent to
   * the wrong backend (e.g. image names when the backend isn't Vertex AI).
   *
   * @param llmRequest The request to sanitize.
   * @return The sanitized request.
   */
  private LlmRequest sanitizeRequest(LlmRequest llmRequest) {
    if (apiClient.vertexAI()) {
      return llmRequest;
    }
    LlmRequest.Builder requestBuilder = llmRequest.toBuilder();

    // Using API key from Google AI Studio to call model doesn't support labels.
    llmRequest
        .config()
        .ifPresent(
            config -> {
              if (config.labels().isPresent()) {
                requestBuilder.config(config.toBuilder().labels(null).build());
              }
            });

    if (llmRequest.contents().isEmpty()) {
      return requestBuilder.build();
    }

    // This backend does not support the display_name parameter for file uploads,
    // so it must be removed to prevent request failures.
    ImmutableList<Content> updatedContents =
        llmRequest.contents().stream()
            .map(
                content -> {
                  if (content.parts().isEmpty() || content.parts().get().isEmpty()) {
                    return content;
                  }

                  ImmutableList<Part> updatedParts =
                      content.parts().get().stream()
                          .map(
                              part -> {
                                Part.Builder partBuilder = part.toBuilder();
                                if (part.inlineData().flatMap(Blob::displayName).isPresent()) {
                                  Blob blob = part.inlineData().get();
                                  Blob.Builder newBlobBuilder = Blob.builder();
                                  blob.data().ifPresent(newBlobBuilder::data);
                                  blob.mimeType().ifPresent(newBlobBuilder::mimeType);
                                  partBuilder.inlineData(newBlobBuilder.build());
                                }
                                if (part.fileData().flatMap(FileData::displayName).isPresent()) {
                                  FileData fileData = part.fileData().get();
                                  FileData.Builder newFileDataBuilder = FileData.builder();
                                  fileData.fileUri().ifPresent(newFileDataBuilder::fileUri);
                                  fileData.mimeType().ifPresent(newFileDataBuilder::mimeType);
                                  partBuilder.fileData(newFileDataBuilder.build());
                                }
                                return partBuilder.build();
                              })
                          .collect(toImmutableList());

                  return content.toBuilder().parts(updatedParts).build();
                })
            .collect(toImmutableList());
    return requestBuilder.contents(updatedContents).build();
  }

  @Override
  public Flowable<LlmResponse> generateContent(LlmRequest llmRequest, boolean stream) {
    llmRequest = sanitizeRequest(llmRequest);
    List<Content> contents = llmRequest.contents();
    // Last content must be from the user, otherwise the model won't respond.
    if (contents.isEmpty() || !Iterables.getLast(contents).role().orElse("").equals("user")) {
      Content userContent = Content.fromParts(Part.fromText(CONTINUE_OUTPUT_MESSAGE));
      contents =
          Stream.concat(contents.stream(), Stream.of(userContent)).collect(toImmutableList());
    }

    List<Content> finalContents = stripThoughts(contents);
    GenerateContentConfig config = llmRequest.config().orElse(null);
    String effectiveModelName = llmRequest.model().orElse(model());

    logger.trace("Request Contents: {}", finalContents);
    logger.trace("Request Config: {}", config);

    if (stream) {
      logger.debug("Sending streaming generateContent request to model {}", effectiveModelName);
      CompletableFuture<ResponseStream<GenerateContentResponse>> streamFuture =
          apiClient.async.models.generateContentStream(effectiveModelName, finalContents, config);

      return Flowable.defer(
          () -> {
            final StringBuilder accumulatedText = new StringBuilder();
            // Array to bypass final local variable reassignment in lambda.
            final GenerateContentResponse[] lastRawResponseHolder = {null};

            return Flowable.fromFuture(streamFuture)
                .flatMapIterable(iterable -> iterable)
                .concatMap(
                    rawResponse -> {
                      lastRawResponseHolder[0] = rawResponse;
                      logger.trace("Raw streaming response: {}", rawResponse);

                      List<LlmResponse> responsesToEmit = new ArrayList<>();
                      LlmResponse currentProcessedLlmResponse = LlmResponse.create(rawResponse);
                      String currentTextChunk = getTextFromLlmResponse(currentProcessedLlmResponse);

                      if (!currentTextChunk.isEmpty()) {
                        accumulatedText.append(currentTextChunk);
                        LlmResponse partialResponse =
                            currentProcessedLlmResponse.toBuilder().partial(true).build();
                        responsesToEmit.add(partialResponse);
                      } else {
                        if (accumulatedText.length() > 0
                            && shouldEmitAccumulatedText(currentProcessedLlmResponse)) {
                          LlmResponse aggregatedTextResponse =
                              LlmResponse.builder()
                                  .content(
                                      Content.builder()
                                          .parts(
                                              ImmutableList.of(
                                                  Part.builder()
                                                      .text(accumulatedText.toString())
                                                      .build()))
                                          .build())
                                  .build();
                          responsesToEmit.add(aggregatedTextResponse);
                          accumulatedText.setLength(0);
                        }
                        responsesToEmit.add(currentProcessedLlmResponse);
                      }
                      logger.debug("Responses to emit: {}", responsesToEmit);
                      return Flowable.fromIterable(responsesToEmit);
                    })
                .concatWith(
                    Flowable.defer(
                        () -> {
                          if (accumulatedText.length() > 0 && lastRawResponseHolder[0] != null) {
                            GenerateContentResponse finalRawResp = lastRawResponseHolder[0];
                            boolean isStop =
                                finalRawResp
                                    .candidates()
                                    .flatMap(
                                        candidates ->
                                            candidates.isEmpty()
                                                ? Optional.empty()
                                                : Optional.of(candidates.get(0)))
                                    .flatMap(Candidate::finishReason)
                                    .map(
                                        finishReason ->
                                            finishReason.equals(
                                                new FinishReason(FinishReason.Known.STOP)))
                                    .orElse(false);

                            if (isStop) {
                              LlmResponse finalAggregatedTextResponse =
                                  LlmResponse.builder()
                                      .content(
                                          Content.builder()
                                              .parts(
                                                  ImmutableList.of(
                                                      Part.builder()
                                                          .text(accumulatedText.toString())
                                                          .build()))
                                              .build())
                                      .build();
                              return Flowable.just(finalAggregatedTextResponse);
                            }
                          }
                          return Flowable.empty();
                        }));
          });
    } else {
      logger.debug("Sending generateContent request to model {}", effectiveModelName);
      return Flowable.fromFuture(
          apiClient
              .async
              .models
              .generateContent(effectiveModelName, finalContents, config)
              .thenApplyAsync(LlmResponse::create));
    }
  }

  /**
   * Extracts text content from the first part of an LlmResponse, if available.
   *
   * @param llmResponse The LlmResponse to extract text from.
   * @return The text content, or an empty string if not found.
   */
  private String getTextFromLlmResponse(LlmResponse llmResponse) {
    return llmResponse
        .content()
        .flatMap(Content::parts)
        .filter(parts -> !parts.isEmpty())
        .map(parts -> parts.get(0))
        .flatMap(Part::text)
        .orElse("");
  }

  /**
   * Determines if accumulated text should be emitted based on the current LlmResponse. We flush if
   * current response is not a text continuation (e.g., no content, no parts, or the first part is
   * not inline_data, meaning it's something else or just empty, thereby warranting a flush of
   * preceding text).
   *
   * @param currentLlmResponse The current LlmResponse being processed.
   * @return True if accumulated text should be emitted, false otherwise.
   */
  private boolean shouldEmitAccumulatedText(LlmResponse currentLlmResponse) {
    Optional<Content> contentOpt = currentLlmResponse.content();
    if (contentOpt.isEmpty()) {
      return true;
    }

    Optional<List<Part>> partsOpt = contentOpt.get().parts();
    if (partsOpt.isEmpty() || partsOpt.get().isEmpty()) {
      return true;
    }

    // If content and parts are present, and parts list is not empty, we want to yield accumulated
    // text only if `text` is present AND (`not llm_response.content` OR `not
    // llm_response.content.parts` OR `not llm_response.content.parts[0].inline_data`)
    // This means we flush if the first part does NOT have inline_data.
    // If it *has* inline_data, the condition below is false,
    // and we would not flush based on this specific sub-condition.
    Part firstPart = partsOpt.get().get(0);
    return firstPart.inlineData().isEmpty();
  }

  @Override
  public BaseLlmConnection connect(LlmRequest llmRequest) {
    llmRequest = sanitizeRequest(llmRequest);
    logger.debug("Establishing Gemini connection.");
    LiveConnectConfig liveConnectConfig = llmRequest.liveConnectConfig();
    String effectiveModelName = llmRequest.model().orElse(model());

    logger.debug("Connecting to model {}", effectiveModelName);
    logger.trace("Connection Config: {}", liveConnectConfig);

    return new GeminiLlmConnection(apiClient, effectiveModelName, liveConnectConfig);
  }

  /** Removes any `Part` that contains only a `thought` from the content list. */
  private List<Content> stripThoughts(List<Content> originalContents) {
    List<Content> updatedContents = new ArrayList<>();
    for (Content content : originalContents) {
      ImmutableList<Part> nonThoughtParts =
          content.parts().orElse(ImmutableList.of()).stream()
              // Keep if thought is not present OR if thought is present but false
              .filter(part -> part.thought().map(isThought -> !isThought).orElse(true))
              .collect(toImmutableList());
      updatedContents.add(content.toBuilder().parts(nonThoughtParts).build());
    }
    return updatedContents;
  }
}

/*
 * 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.agents;

import com.google.adk.models.LlmRequest;
import com.google.adk.models.LlmResponse;
import com.google.adk.tools.BaseTool;
import com.google.adk.tools.ToolContext;
import com.google.genai.types.Content;
import io.reactivex.rxjava3.core.Maybe;
import java.util.Map;
import java.util.Optional;

/** Functional interfaces for agent lifecycle callbacks. */
public final class Callbacks {

  interface BeforeModelCallbackBase {}

  @FunctionalInterface
  public interface BeforeModelCallback extends BeforeModelCallbackBase {
    /**
     * Async callback before LLM invocation.
     *
     * @param callbackContext Callback context.
     * @param llmRequest LLM request.
     * @return response override, or empty to continue.
     */
    Maybe<LlmResponse> call(CallbackContext callbackContext, LlmRequest llmRequest);
  }

  /**
   * Helper interface to allow for sync beforeModelCallback. The function is wrapped into an async
   * one before being processed further.
   */
  @FunctionalInterface
  public interface BeforeModelCallbackSync extends BeforeModelCallbackBase {
    Optional<LlmResponse> call(CallbackContext callbackContext, LlmRequest llmRequest);
  }

  interface AfterModelCallbackBase {}

  @FunctionalInterface
  public interface AfterModelCallback extends AfterModelCallbackBase {
    /**
     * Async callback after LLM response.
     *
     * @param callbackContext Callback context.
     * @param llmResponse LLM response.
     * @return modified response, or empty to keep original.
     */
    Maybe<LlmResponse> call(CallbackContext callbackContext, LlmResponse llmResponse);
  }

  /**
   * Helper interface to allow for sync afterModelCallback. The function is wrapped into an async
   * one before being processed further.
   */
  @FunctionalInterface
  public interface AfterModelCallbackSync extends AfterModelCallbackBase {
    Optional<LlmResponse> call(CallbackContext callbackContext, LlmResponse llmResponse);
  }

  interface BeforeAgentCallbackBase {}

  @FunctionalInterface
  public interface BeforeAgentCallback extends BeforeAgentCallbackBase {
    /**
     * Async callback before agent runs.
     *
     * @param callbackContext Callback context.
     * @return content override, or empty to continue.
     */
    Maybe<Content> call(CallbackContext callbackContext);
  }

  /**
   * Helper interface to allow for sync beforeAgentCallback. The function is wrapped into an async
   * one before being processed further.
   */
  @FunctionalInterface
  public interface BeforeAgentCallbackSync extends BeforeAgentCallbackBase {
    Optional<Content> call(CallbackContext callbackContext);
  }

  interface AfterAgentCallbackBase {}

  @FunctionalInterface
  public interface AfterAgentCallback extends AfterAgentCallbackBase {
    /**
     * Async callback after agent runs.
     *
     * @param callbackContext Callback context.
     * @return modified content, or empty to keep original.
     */
    Maybe<Content> call(CallbackContext callbackContext);
  }

  /**
   * Helper interface to allow for sync afterAgentCallback. The function is wrapped into an async
   * one before being processed further.
   */
  @FunctionalInterface
  public interface AfterAgentCallbackSync extends AfterAgentCallbackBase {
    Optional<Content> call(CallbackContext callbackContext);
  }

  interface BeforeToolCallbackBase {}

  @FunctionalInterface
  public interface BeforeToolCallback extends BeforeToolCallbackBase {
    /**
     * Async callback before tool runs.
     *
     * @param invocationContext Invocation context.
     * @param baseTool Tool instance.
     * @param input Tool input arguments.
     * @param toolContext Tool context.
     * @return override result, or empty to continue.
     */
    Maybe<Map<String, Object>> call(
        InvocationContext invocationContext,
        BaseTool baseTool,
        Map<String, Object> input,
        ToolContext toolContext);
  }

  /**
   * Helper interface to allow for sync beforeToolCallback. The function is wrapped into an async
   * one before being processed further.
   */
  @FunctionalInterface
  public interface BeforeToolCallbackSync extends BeforeToolCallbackBase {
    Optional<Map<String, Object>> call(
        InvocationContext invocationContext,
        BaseTool baseTool,
        Map<String, Object> input,
        ToolContext toolContext);
  }

  interface AfterToolCallbackBase {}

  @FunctionalInterface
  public interface AfterToolCallback extends AfterToolCallbackBase {
    /**
     * Async callback after tool runs.
     *
     * @param invocationContext Invocation context.
     * @param baseTool Tool instance.
     * @param input Tool input arguments.
     * @param toolContext Tool context.
     * @param response Raw tool response.
     * @return processed result, or empty to keep original.
     */
    Maybe<Map<String, Object>> call(
        InvocationContext invocationContext,
        BaseTool baseTool,
        Map<String, Object> input,
        ToolContext toolContext,
        Object response);
  }

  /**
   * Helper interface to allow for sync afterToolCallback. The function is wrapped into an async one
   * before being processed further.
   */
  @FunctionalInterface
  public interface AfterToolCallbackSync extends AfterToolCallbackBase {
    Optional<Map<String, Object>> call(
        InvocationContext invocationContext,
        BaseTool baseTool,
        Map<String, Object> input,
        ToolContext toolContext,
        Object response);
  }

  private Callbacks() {}
}

/*
 * 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.agents;

import com.google.adk.agents.Callbacks.AfterAgentCallback;
import com.google.adk.agents.Callbacks.AfterAgentCallbackBase;
import com.google.adk.agents.Callbacks.AfterAgentCallbackSync;
import com.google.adk.agents.Callbacks.BeforeAgentCallback;
import com.google.adk.agents.Callbacks.BeforeAgentCallbackBase;
import com.google.adk.agents.Callbacks.BeforeAgentCallbackSync;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import io.reactivex.rxjava3.core.Maybe;
import java.util.List;
import org.jspecify.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Utility methods for normalizing agent callbacks. */
public final class CallbackUtil {
  private static final Logger logger = LoggerFactory.getLogger(CallbackUtil.class);

  /**
   * Normalizes before-agent callbacks.
   *
   * @param beforeAgentCallback Callback list (sync or async).
   * @return normalized async callbacks, or null if input is null.
   */
  @CanIgnoreReturnValue
  public static @Nullable ImmutableList<BeforeAgentCallback> getBeforeAgentCallbacks(
      List<BeforeAgentCallbackBase> beforeAgentCallback) {
    if (beforeAgentCallback == null) {
      return null;
    } else if (beforeAgentCallback.isEmpty()) {
      return ImmutableList.of();
    } else {
      ImmutableList.Builder<BeforeAgentCallback> builder = ImmutableList.builder();
      for (BeforeAgentCallbackBase callback : beforeAgentCallback) {
        if (callback instanceof BeforeAgentCallback beforeAgentCallbackInstance) {
          builder.add(beforeAgentCallbackInstance);
        } else if (callback instanceof BeforeAgentCallbackSync beforeAgentCallbackSyncInstance) {
          builder.add(
              (BeforeAgentCallback)
                  (callbackContext) ->
                      Maybe.fromOptional(beforeAgentCallbackSyncInstance.call(callbackContext)));
        } else {
          logger.warn(
              "Invalid beforeAgentCallback callback type: %s. Ignoring this callback.",
              callback.getClass().getName());
        }
      }
      return builder.build();
    }
  }

  /**
   * Normalizes after-agent callbacks.
   *
   * @param afterAgentCallback Callback list (sync or async).
   * @return normalized async callbacks, or null if input is null.
   */
  @CanIgnoreReturnValue
  public static @Nullable ImmutableList<AfterAgentCallback> getAfterAgentCallbacks(
      List<AfterAgentCallbackBase> afterAgentCallback) {
    if (afterAgentCallback == null) {
      return null;
    } else if (afterAgentCallback.isEmpty()) {
      return ImmutableList.of();
    } else {
      ImmutableList.Builder<AfterAgentCallback> builder = ImmutableList.builder();
      for (AfterAgentCallbackBase callback : afterAgentCallback) {
        if (callback instanceof AfterAgentCallback afterAgentCallbackInstance) {
          builder.add(afterAgentCallbackInstance);
        } else if (callback instanceof AfterAgentCallbackSync afterAgentCallbackSyncInstance) {
          builder.add(
              (AfterAgentCallback)
                  (callbackContext) ->
                      Maybe.fromOptional(afterAgentCallbackSyncInstance.call(callbackContext)));
        } else {
          logger.warn(
              "Invalid afterAgentCallback callback type: %s. Ignoring this callback.",
              callback.getClass().getName());
        }
      }
      return builder.build();
    }
  }

  private CallbackUtil() {}
}

/*
 * 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.tools.retrieval;

import static com.google.common.collect.ImmutableList.toImmutableList;

import com.google.adk.models.LlmRequest;
import com.google.adk.tools.ToolContext;
import com.google.cloud.aiplatform.v1.RagContexts;
import com.google.cloud.aiplatform.v1.RagQuery;
import com.google.cloud.aiplatform.v1.RetrieveContextsRequest;
import com.google.cloud.aiplatform.v1.RetrieveContextsRequest.VertexRagStore.RagResource;
import com.google.cloud.aiplatform.v1.RetrieveContextsResponse;
import com.google.cloud.aiplatform.v1.VertexRagServiceClient;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.Retrieval;
import com.google.genai.types.Tool;
import com.google.genai.types.VertexRagStore;
import com.google.genai.types.VertexRagStoreRagResource;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Single;
import java.util.List;
import java.util.Map;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A retrieval tool that fetches context from Vertex AI RAG.
 *
 * <p>This tool allows to retrieve relevant information based on a query using Vertex AI RAG
 * service. It supports configuration of rag resources and a vector distance threshold.
 */
public class VertexAiRagRetrieval extends BaseRetrievalTool {
  private static final Logger logger = LoggerFactory.getLogger(VertexAiRagRetrieval.class);
  private final VertexRagServiceClient vertexRagServiceClient;
  private final String parent;
  private final List<RagResource> ragResources;
  private final Double vectorDistanceThreshold;
  private final VertexRagStore vertexRagStore;
  private final RetrieveContextsRequest.VertexRagStore apiVertexRagStore;

  public VertexAiRagRetrieval(
      @Nonnull String name,
      @Nonnull String description,
      @Nonnull VertexRagServiceClient vertexRagServiceClient,
      @Nonnull String parent,
      @Nullable List<RagResource> ragResources,
      @Nullable Double vectorDistanceThreshold) {
    super(name, description);
    this.vertexRagServiceClient = vertexRagServiceClient;
    this.parent = parent;
    this.ragResources = ragResources;
    this.vectorDistanceThreshold = vectorDistanceThreshold;

    // For Gemini 2
    VertexRagStore.Builder vertexRagStoreBuilder = VertexRagStore.builder();
    if (this.ragResources != null) {
      vertexRagStoreBuilder.ragResources(
          this.ragResources.stream()
              .map(
                  ragResource ->
                      VertexRagStoreRagResource.builder()
                          .ragCorpus(ragResource.getRagCorpus())
                          .build())
              .collect(toImmutableList()));
    }
    if (this.vectorDistanceThreshold != null) {
      vertexRagStoreBuilder.vectorDistanceThreshold(this.vectorDistanceThreshold);
    }
    this.vertexRagStore = vertexRagStoreBuilder.build();

    // For runAsync
    RetrieveContextsRequest.VertexRagStore.Builder apiVertexRagStoreBuilder =
        RetrieveContextsRequest.VertexRagStore.newBuilder();
    if (this.ragResources != null) {
      apiVertexRagStoreBuilder.addAllRagResources(this.ragResources);
    }
    if (this.vectorDistanceThreshold != null) {
      apiVertexRagStoreBuilder.setVectorDistanceThreshold(this.vectorDistanceThreshold);
    }
    this.apiVertexRagStore = apiVertexRagStoreBuilder.build();
  }

  @Override
  @CanIgnoreReturnValue
  public Completable processLlmRequest(
      LlmRequest.Builder llmRequestBuilder, ToolContext toolContext) {
    LlmRequest llmRequest = llmRequestBuilder.build();
    // Use Gemini built-in Vertex AI RAG tool for Gemini 2 models or when using Vertex AI API Model
    boolean useVertexAi = Boolean.parseBoolean(System.getenv("GOOGLE_GENAI_USE_VERTEXAI"));
    if (useVertexAi
        && (llmRequest.model().isPresent() && llmRequest.model().get().startsWith("gemini-2"))) {
      GenerateContentConfig config =
          llmRequest.config().orElse(GenerateContentConfig.builder().build());
      ImmutableList.Builder<Tool> toolsBuilder = ImmutableList.builder();
      if (config.tools().isPresent()) {
        toolsBuilder.addAll(config.tools().get());
      }
      toolsBuilder.add(
          Tool.builder()
              .retrieval(Retrieval.builder().vertexRagStore(this.vertexRagStore).build())
              .build());
      logger.info(
          "Using Gemini built-in Vertex AI RAG tool for model: {}", llmRequest.model().get());
      llmRequestBuilder.config(config.toBuilder().tools(toolsBuilder.build()).build());
      return Completable.complete();
    } else {
      // Add the function declaration to the tools
      return super.processLlmRequest(llmRequestBuilder, toolContext);
    }
  }

  @Override
  public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContext toolContext) {
    String query = (String) args.get("query");
    logger.info("VertexAiRagRetrieval.runAsync called with query: {}", query);
    return Single.fromCallable(
        () -> {
          logger.info("Retrieving context for query: {}", query);
          RetrieveContextsRequest retrieveContextsRequest =
              RetrieveContextsRequest.newBuilder()
                  .setParent(this.parent)
                  .setQuery(RagQuery.newBuilder().setText(query))
                  .setVertexRagStore(this.apiVertexRagStore)
                  .build();
          logger.info("Request to VertexRagService: {}", retrieveContextsRequest);
          RetrieveContextsResponse response =
              this.vertexRagServiceClient.retrieveContexts(retrieveContextsRequest);
          logger.info("Response from VertexRagService: {}", response);
          if (response.getContexts().getContextsList().isEmpty()) {
            logger.warn("No matching result found for query: {}", query);
            return ImmutableMap.of(
                "response",
                String.format(
                    "No matching result found with the config: resources: %s", this.ragResources));
          } else {
            logger.info(
                "Found {} matching results for query: {}",
                response.getContexts().getContextsCount(),
                query);
            ImmutableList<String> contexts =
                response.getContexts().getContextsList().stream()
                    .map(RagContexts.Context::getText)
                    .collect(toImmutableList());
            logger.info("Returning contexts: {}", contexts);
            return ImmutableMap.of("response", contexts);
          }
        });
  }
}

/*
 * 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.events;

import static com.google.common.collect.ImmutableList.toImmutableList;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
import com.google.adk.JsonBaseModel;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.genai.types.Content;
import com.google.genai.types.FinishReason;
import com.google.genai.types.FunctionCall;
import com.google.genai.types.FunctionResponse;
import com.google.genai.types.GroundingMetadata;
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 javax.annotation.Nullable;

// TODO - b/413761119 update Agent.java when resolved.
/** Represents an event in a session. */
@JsonDeserialize(builder = Event.Builder.class)
public class Event extends JsonBaseModel {

  private String id;
  private String invocationId;
  private String author;
  private Optional<Content> content = Optional.empty();
  private EventActions actions;
  private Optional<Set<String>> longRunningToolIds = Optional.empty();
  private Optional<Boolean> partial = Optional.empty();
  private Optional<Boolean> turnComplete = Optional.empty();
  private Optional<FinishReason> errorCode = Optional.empty();
  private Optional<String> errorMessage = Optional.empty();
  private Optional<Boolean> interrupted = Optional.empty();
  private Optional<String> branch = Optional.empty();
  private Optional<GroundingMetadata> groundingMetadata = Optional.empty();
  private long timestamp;

  private Event() {}

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

  /** The event id. */
  @JsonProperty("id")
  public String id() {
    return id;
  }

  public void setId(String id) {
    this.id = id;
  }

  /** Id of the invocation that this event belongs to. */
  @JsonProperty("invocationId")
  public String invocationId() {
    return invocationId;
  }

  public void setInvocationId(String invocationId) {
    this.invocationId = invocationId;
  }

  /** The author of the event, it could be the name of the agent or "user" literal. */
  @JsonProperty("author")
  public String author() {
    return author;
  }

  public void setAuthor(String author) {
    this.author = author;
  }

  @JsonProperty("content")
  public Optional<Content> content() {
    return content;
  }

  public void setContent(Optional<Content> content) {
    this.content = content;
  }

  @JsonProperty("actions")
  public EventActions actions() {
    return actions;
  }

  public void setActions(EventActions actions) {
    this.actions = actions;
  }

  /**
   * Set of ids of the long running function calls. Agent client will know from this field about
   * which function call is long running.
   */
  @JsonProperty("longRunningToolIds")
  public Optional<Set<String>> longRunningToolIds() {
    return longRunningToolIds;
  }

  public void setLongRunningToolIds(Optional<Set<String>> longRunningToolIds) {
    this.longRunningToolIds = longRunningToolIds;
  }

  /**
   * partial is true for incomplete chunks from the LLM streaming response. The last chunk's partial
   * is False.
   */
  @JsonProperty("partial")
  public Optional<Boolean> partial() {
    return partial;
  }

  public void setPartial(Optional<Boolean> partial) {
    this.partial = partial;
  }

  @JsonProperty("turnComplete")
  public Optional<Boolean> turnComplete() {
    return turnComplete;
  }

  public void setTurnComplete(Optional<Boolean> turnComplete) {
    this.turnComplete = turnComplete;
  }

  @JsonProperty("errorCode")
  public Optional<FinishReason> errorCode() {
    return errorCode;
  }

  public void setErrorCode(Optional<FinishReason> errorCode) {
    this.errorCode = errorCode;
  }

  @JsonProperty("errorMessage")
  public Optional<String> errorMessage() {
    return errorMessage;
  }

  public void setErrorMessage(Optional<String> errorMessage) {
    this.errorMessage = errorMessage;
  }

  @JsonProperty("interrupted")
  public Optional<Boolean> interrupted() {
    return interrupted;
  }

  public void setInterrupted(Optional<Boolean> interrupted) {
    this.interrupted = interrupted;
  }

  /**
   * The branch of the event. The format is like agent_1.agent_2.agent_3, where agent_1 is the
   * parent of agent_2, and agent_2 is the parent of agent_3. Branch is used when multiple sub-agent
   * shouldn't see their peer agents' conversation history.
   */
  @JsonProperty("branch")
  public Optional<String> branch() {
    return branch;
  }

  /**
   * Sets the branch for this event.
   *
   *
Format: agentA.agentB.agentC — shows hierarchy of nested agents.
   *
   * @param branch Branch identifier.
   */
  public void branch(@Nullable String branch) {
    this.branch = Optional.ofNullable(branch);
  }

  public void branch(Optional<String> branch) {
    this.branch = branch;
  }

  /** The grounding metadata of the event. */
  @JsonProperty("groundingMetadata")
  public Optional<GroundingMetadata> groundingMetadata() {
    return groundingMetadata;
  }

  public void setGroundingMetadata(Optional<GroundingMetadata> groundingMetadata) {
    this.groundingMetadata = groundingMetadata;
  }

  /** The timestamp of the event. */
  @JsonProperty("timestamp")
  public long timestamp() {
    return timestamp;
  }

  public void setTimestamp(long timestamp) {
    this.timestamp = timestamp;
  }

  /** Returns all function calls from this event. */
  @JsonIgnore
  public final ImmutableList<FunctionCall> functionCalls() {
    return content().flatMap(Content::parts).stream()
        .flatMap(List::stream)
        .flatMap(part -> part.functionCall().stream())
        .collect(toImmutableList());
  }

  /** Returns all function responses from this event. */
  @JsonIgnore
  public final ImmutableList<FunctionResponse> functionResponses() {
    return content().flatMap(Content::parts).stream()
        .flatMap(List::stream)
        .flatMap(part -> part.functionResponse().stream())
        .collect(toImmutableList());
  }

  /** Returns true if this is a final response. */
  @JsonIgnore
  public final boolean finalResponse() {
    if (actions().skipSummarization().orElse(false)
        || (longRunningToolIds().isPresent() && !longRunningToolIds().get().isEmpty())) {
      return true;
    }
    return functionCalls().isEmpty() && functionResponses().isEmpty() && !partial().orElse(false);
  }

  /**
   * Converts the event content into a readable string.
   *
   *
Includes text, function calls, and responses.
   *
   * @return Stringified content.
   */
  public final String stringifyContent() {
    StringBuilder sb = new StringBuilder();
    content().flatMap(Content::parts).stream()
        .flatMap(List::stream)
        .forEach(
            part -> {
              part.text().ifPresent(sb::append);
              part.functionCall()
                  .ifPresent(functionCall -> sb.append("Function Call: ").append(functionCall));
              part.functionResponse()
                  .ifPresent(
                      functionResponse ->
                          sb.append("Function Response: ").append(functionResponse));
            });
    return sb.toString();
  }

  /** Builder for {@link Event}. */
  public static class Builder {

    private String id;
    private String invocationId;
    private String author;
    private Optional<Content> content = Optional.empty();
    private EventActions actions;
    private Optional<Set<String>> longRunningToolIds = Optional.empty();
    private Optional<Boolean> partial = Optional.empty();
    private Optional<Boolean> turnComplete = Optional.empty();
    private Optional<FinishReason> errorCode = Optional.empty();
    private Optional<String> errorMessage = Optional.empty();
    private Optional<Boolean> interrupted = Optional.empty();
    private Optional<String> branch = Optional.empty();
    private Optional<GroundingMetadata> groundingMetadata = Optional.empty();
    private Optional<Long> timestamp = Optional.empty();

    @JsonCreator
    private static Builder create() {
      return new Builder();
    }

    @CanIgnoreReturnValue
    @JsonProperty("id")
    public Builder id(String value) {
      this.id = value;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("invocationId")
    public Builder invocationId(String value) {
      this.invocationId = value;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("author")
    public Builder author(String value) {
      this.author = value;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("content")
    public Builder content(@Nullable Content value) {
      this.content = Optional.ofNullable(value);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder content(Optional<Content> value) {
      this.content = value;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("actions")
    public Builder actions(EventActions value) {
      this.actions = value;
      return this;
    }

    Optional<EventActions> actions() {
      return Optional.ofNullable(actions);
    }

    @CanIgnoreReturnValue
    @JsonProperty("longRunningToolIds")
    public Builder longRunningToolIds(@Nullable Set<String> value) {
      this.longRunningToolIds = Optional.ofNullable(value);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder longRunningToolIds(Optional<Set<String>> value) {
      this.longRunningToolIds = value;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("partial")
    public Builder partial(@Nullable Boolean value) {
      this.partial = Optional.ofNullable(value);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder partial(Optional<Boolean> value) {
      this.partial = value;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("turnComplete")
    public Builder turnComplete(@Nullable Boolean value) {
      this.turnComplete = Optional.ofNullable(value);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder turnComplete(Optional<Boolean> value) {
      this.turnComplete = value;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("errorCode")
    public Builder errorCode(@Nullable FinishReason value) {
      this.errorCode = Optional.ofNullable(value);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder errorCode(Optional<FinishReason> value) {
      this.errorCode = value;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("errorMessage")
    public Builder errorMessage(@Nullable String value) {
      this.errorMessage = Optional.ofNullable(value);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder errorMessage(Optional<String> value) {
      this.errorMessage = value;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("interrupted")
    public Builder interrupted(@Nullable Boolean value) {
      this.interrupted = Optional.ofNullable(value);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder interrupted(Optional<Boolean> value) {
      this.interrupted = value;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("timestamp")
    public Builder timestamp(long value) {
      this.timestamp = Optional.of(value);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder timestamp(Optional<Long> value) {
      this.timestamp = value;
      return this;
    }

    // Getter for builder's timestamp, used in build()
    Optional<Long> timestamp() {
      return timestamp;
    }

    @CanIgnoreReturnValue
    @JsonProperty("branch")
    public Builder branch(@Nullable String value) {
      this.branch = Optional.ofNullable(value);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder branch(Optional<String> value) {
      this.branch = value;
      return this;
    }

    // Getter for builder's branch, used in build()
    Optional<String> branch() {
      return branch;
    }

    @CanIgnoreReturnValue
    @JsonProperty("groundingMetadata")
    public Builder groundingMetadata(@Nullable GroundingMetadata value) {
      this.groundingMetadata = Optional.ofNullable(value);
      return this;
    }

    @CanIgnoreReturnValue
    public Builder groundingMetadata(Optional<GroundingMetadata> value) {
      this.groundingMetadata = value;
      return this;
    }

    Optional<GroundingMetadata> groundingMetadata() {
      return groundingMetadata;
    }

    public Event build() {
      Event event = new Event();
      event.setId(id);
      event.setInvocationId(invocationId);
      event.setAuthor(author);
      event.setContent(content);
      event.setLongRunningToolIds(longRunningToolIds);
      event.setPartial(partial);
      event.setTurnComplete(turnComplete);
      event.setErrorCode(errorCode);
      event.setErrorMessage(errorMessage);
      event.setInterrupted(interrupted);
      event.branch(branch);
      event.setGroundingMetadata(groundingMetadata);

      event.setActions(actions().orElse(EventActions.builder().build()));
      event.setTimestamp(timestamp().orElse(Instant.now().toEpochMilli()));
      return event;
    }
  }

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

  /** Parses an event from a JSON string. */
  public static Event fromJson(String json) {
    return fromJsonString(json, Event.class);
  }

  /** Creates a builder pre-filled with this event's values. */
  public Builder toBuilder() {
    Builder builder =
        new Builder()
            .id(this.id)
            .invocationId(this.invocationId)
            .author(this.author)
            .content(this.content)
            .actions(this.actions)
            .longRunningToolIds(this.longRunningToolIds)
            .partial(this.partial)
            .turnComplete(this.turnComplete)
            .errorCode(this.errorCode)
            .errorMessage(this.errorMessage)
            .interrupted(this.interrupted)
            .branch(this.branch)
            .groundingMetadata(this.groundingMetadata);
    if (this.timestamp != 0) {
      builder.timestamp(this.timestamp);
    }
    return builder;
  }

  @Override
  public boolean equals(Object obj) {
    if (this == obj) {
      return true;
    }
    if (!(obj instanceof Event other)) {
      return false;
    }
    return timestamp == other.timestamp
        && Objects.equals(id, other.id)
        && Objects.equals(invocationId, other.invocationId)
        && Objects.equals(author, other.author)
        && Objects.equals(content, other.content)
        && Objects.equals(actions, other.actions)
        && Objects.equals(longRunningToolIds, other.longRunningToolIds)
        && Objects.equals(partial, other.partial)
        && Objects.equals(turnComplete, other.turnComplete)
        && Objects.equals(errorCode, other.errorCode)
        && Objects.equals(errorMessage, other.errorMessage)
        && Objects.equals(interrupted, other.interrupted)
        && Objects.equals(branch, other.branch)
        && Objects.equals(groundingMetadata, other.groundingMetadata);
  }

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

  @Override
  public int hashCode() {
    return Objects.hash(
        id,
        invocationId,
        author,
        content,
        actions,
        longRunningToolIds,
        partial,
        turnComplete,
        errorCode,
        errorMessage,
        interrupted,
        branch,
        groundingMetadata,
        timestamp);
  }
}

/*
 * 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.utils;

import com.google.adk.agents.InvocationContext;
import com.google.adk.sessions.Session;
import com.google.adk.sessions.State;
import com.google.common.collect.ImmutableSet;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.regex.MatchResult;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/** Utility methods for handling instruction templates. */
public final class InstructionUtils {

  private static final Pattern INSTRUCTION_PLACEHOLDER_PATTERN =
      Pattern.compile("\\{+[^\\{\\}]*\\}+");

  private InstructionUtils() {}

  /**
   * Populates placeholders in an instruction template string with values from the session state or
   * loaded artifacts.
   *
   *
Placeholder Syntax:
   *
   *
Placeholders are enclosed by one or more curly braces at the start and end, e.g., {@code
   * {key}} or {@code {{key}}}. The core {@code key} is extracted from whatever is between the
   * innermost pair of braces after trimming whitespace and possibly removing the {@code ?} which
   * denotes optionality (e.g. {@code {key?}}). The {@code key} itself must not contain curly
   * braces. For typical usage, a single pair of braces like {@code {my_variable}} is standard.
   *
   *
The extracted {@code key} determines the source and name of the value:
   *
   *

   *
Session State Variables: The {@code key} (e.g., {@code "variable_name"} or {@code
   *       "prefix:variable_name"}) refers to a variable in session state.
   *

   *
Simple name: {@code {variable_name}}. The {@code variable_name} part must be a
   *             valid identifier as per {@link #isValidStateName(String)}. Invalid names will
   *             result in the placeholder being returned as is.
   *
Prefixed name: {@code {prefix:variable_name}}. Valid prefixes are: {@value
   *             com.google.adk.sessions.State#APP_PREFIX}, {@value
   *             com.google.adk.sessions.State#USER_PREFIX}, and {@value
   *             com.google.adk.sessions.State#TEMP_PREFIX} The part of the name following the
   *             prefix must also be a valid identifier. Invalid prefixes will result in the
   *             placeholder being returned as is.
   *

   *
Artifacts: The {@code key} starts with "{@code artifact.}" (e.g., {@code
   *       "artifact.file_name"}).
   *
Optional Placeholders: A {@code key} can be marked as optional by appending a
   *       question mark {@code ?} at its very end, inside the braces.
   *

   *
Example: {@code {optional_variable?}}, {@code {{artifact.optional_file.txt?}}}
   *
If an optional placeholder cannot be resolved (e.g., variable not found, artifact
   *             not found), it is replaced with an empty string.
   *

   *

   *
   * Example Usage:
   *
   *
{@code
   * InvocationContext context = ...; // Assume this is initialized with session and artifact service
   * Session session = context.session();
   *
   * session.state().put("user:name", "Alice");
   *
   * context.artifactService().saveArtifact(
   *     session.appName(), session.userId(), session.id(), "knowledge.txt", Part.fromText("Origins of the universe: At first, there was-"));
   *
   * String template = "You are {user:name}'s assistant. Answer questions based on your knowledge. Your knowledge: {artifact.knowledge.txt}." +
   *                   " Your extra knowledge: {artifact.missing_artifact.txt?}";
   *
   * Single<String> populatedStringSingle = InstructionUtils.injectSessionState(context, template);
   * populatedStringSingle.subscribe(
   *     result -> System.out.println(result),
   *     // Expected: "You are Alice's assistant. Answer questions based on your knowledge. Your knowledge: Origins of the universe: At first, there was-. Your extra knowledge: "
   *     error -> System.err.println("Error populating template: " + error.getMessage())
   * );
   * }

   *
   * @param context The invocation context providing access to session state and artifact services.
   * @param template The instruction template string containing placeholders to be populated.
   * @return A {@link Single} that will emit the populated instruction string upon successful
   *     resolution of all non-optional placeholders. Emits the original template if it is empty or
   *     contains no placeholders that are processed.
   * @throws NullPointerException if the template or context is null.
   * @throws IllegalArgumentException if a non-optional variable or artifact is not found.
   */
  public static Single<String> injectSessionState(InvocationContext context, String template) {
    if (template == null) {
      return Single.error(new NullPointerException("template cannot be null"));
    }
    if (context == null) {
      return Single.error(new NullPointerException("context cannot be null"));
    }
    Matcher matcher = INSTRUCTION_PLACEHOLDER_PATTERN.matcher(template);
    List<Single<String>> parts = new ArrayList<>();
    int lastEnd = 0;

    while (matcher.find()) {
      if (matcher.start() > lastEnd) {
        parts.add(Single.just(template.substring(lastEnd, matcher.start())));
      }
      MatchResult matchResult = matcher.toMatchResult();
      parts.add(resolveMatchAsync(context, matchResult));
      lastEnd = matcher.end();
    }
    if (lastEnd < template.length()) {
      parts.add(Single.just(template.substring(lastEnd)));
    }

    if (parts.isEmpty()) {
      return Single.just(template);
    }

    return Single.zip(
        parts,
        objects -> {
          StringBuilder sb = new StringBuilder();
          for (Object obj : objects) {
            sb.append(obj);
          }
          return sb.toString();
        });
  }

  private static Single<String> resolveMatchAsync(InvocationContext context, MatchResult match) {
    String placeholder = match.group();
    String varNameFromPlaceholder =
        placeholder.replaceAll("^\\{+", "").replaceAll("\\}+$", "").trim();

    final boolean optional;
    final String cleanVarName;
    if (varNameFromPlaceholder.endsWith("?")) {
      optional = true;
      cleanVarName = varNameFromPlaceholder.substring(0, varNameFromPlaceholder.length() - 1);
    } else {
      optional = false;
      cleanVarName = varNameFromPlaceholder;
    }

    if (cleanVarName.startsWith("artifact.")) {
      final String artifactName = cleanVarName.substring("artifact.".length());
      Session session = context.session();

      Maybe<Part> artifactMaybe =
          context
              .artifactService()
              .loadArtifact(
                  session.appName(),
                  session.userId(),
                  session.id(),
                  artifactName,
                  Optional.empty());

      return artifactMaybe
          .map(Part::toJson)
          .switchIfEmpty(
              Single.defer(
                  () -> {
                    if (optional) {
                      return Single.just("");
                    } else {
                      return Single.error(
                          new IllegalArgumentException(
                              String.format("Artifact %s not found.", artifactName)));
                    }
                  }));

    } else if (!isValidStateName(cleanVarName)) {
      return Single.just(placeholder);
    } else if (context.session().state().containsKey(cleanVarName)) {
      Object value = context.session().state().get(cleanVarName);
      return Single.just(String.valueOf(value));
    } else if (optional) {
      return Single.just("");
    } else {
      return Single.error(
          new IllegalArgumentException(
              String.format("Context variable not found: `%s`.", cleanVarName)));
    }
  }

  /**
   * Checks if a given string is a valid state variable name.
   *
   *
A valid state variable name must either:
   *
   *

   *
Be a valid identifier (as defined by {@link Character#isJavaIdentifierStart(int)} and
   *       {@link Character#isJavaIdentifierPart(int)}).
   *
Start with a valid prefix ({@value com.google.adk.sessions.State#APP_PREFIX}, {@value
   *       com.google.adk.sessions.State#USER_PREFIX}, or {@value
   *       com.google.adk.sessions.State#TEMP_PREFIX}) followed by a valid identifier.
   *

   *
   * @param varName The string to check.
   * @return True if the string is a valid state variable name, false otherwise.
   */
  private static boolean isValidStateName(String varName) {
    if (varName.isEmpty()) {
      return false;
    }
    String[] parts = varName.split(":", 2);
    if (parts.length == 1) {
      return isValidIdentifier(parts[0]);
    }

    if (parts.length == 2) {
      String prefixPart = parts[0] + ":";
      ImmutableSet<String> validPrefixes =
          ImmutableSet.of(State.APP_PREFIX, State.USER_PREFIX, State.TEMP_PREFIX);
      if (validPrefixes.contains(prefixPart)) {
        return isValidIdentifier(parts[1]);
      }
    }
    return false;
  }

  private static boolean isValidIdentifier(String s) {
    if (s.isEmpty()) {
      return false;
    }
    if (!Character.isJavaIdentifierStart(s.charAt(0))) {
      return false;
    }
    for (int i = 1; i < s.length(); i++) {
      if (!Character.isJavaIdentifierPart(s.charAt(i))) {
        return false;
      }
    }
    return true;
  }
}

/*
 * 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.tools;

import static com.google.common.collect.ImmutableList.toImmutableList;

import com.google.adk.models.LlmRequest;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.genai.types.FunctionDeclaration;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.Tool;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Single;
import java.util.Map;
import java.util.Optional;
import javax.annotation.Nonnull;
import org.jspecify.annotations.Nullable;

/** The base class for all ADK tools. */
public abstract class BaseTool {
  private final String name;
  private final String description;
  private final boolean isLongRunning;

  protected BaseTool(@Nonnull String name, @Nonnull String description) {
    this(name, description, /* isLongRunning= */ false);
  }

  protected BaseTool(@Nonnull String name, @Nonnull String description, boolean isLongRunning) {
    this.name = name;
    this.description = description;
    this.isLongRunning = isLongRunning;
  }

  public String name() {
    return name;
  }

  public String description() {
    return description;
  }

  public boolean longRunning() {
    return isLongRunning;
  }

  /** Gets the {@link FunctionDeclaration} representation of this tool. */
  public Optional<FunctionDeclaration> declaration() {
    return Optional.empty();
  }

  /** Calls a tool. */
  public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContext toolContext) {
    throw new UnsupportedOperationException("This method is not implemented.");
  }

  /**
   * Processes the outgoing {@link LlmRequest.Builder}.
   *
   * <p>This implementation adds the current tool's {@link #declaration()} to the {@link
   * GenerateContentConfig} within the builder. If a tool with function declarations already exists,
   * the current tool's declaration is merged into it. Otherwise, a new tool definition with the
   * current tool's declaration is created. The current tool itself is also added to the builder's
   * internal list of tools. Override this method for processing the outgoing request.
   */
  @CanIgnoreReturnValue
  public Completable processLlmRequest(
      LlmRequest.Builder llmRequestBuilder, ToolContext toolContext) {
    if (declaration().isEmpty()) {
      return Completable.complete();
    }

    llmRequestBuilder.appendTools(ImmutableList.of(this));

    LlmRequest llmRequest = llmRequestBuilder.build();
    ImmutableList<Tool> toolsWithoutFunctionDeclarations =
        findToolsWithoutFunctionDeclarations(llmRequest);
    Tool toolWithFunctionDeclarations = findToolWithFunctionDeclarations(llmRequest);
    // If LlmRequest GenerateContentConfig already has a function calling tool,
    // merge the function declarations.
    // Otherwise, add a new tool definition with function calling declaration..
    if (toolWithFunctionDeclarations == null) {
      toolWithFunctionDeclarations =
          Tool.builder().functionDeclarations(ImmutableList.of(declaration().get())).build();
    } else {
      toolWithFunctionDeclarations =
          toolWithFunctionDeclarations.toBuilder()
              .functionDeclarations(
                  ImmutableList.<FunctionDeclaration>builder()
                      .addAll(
                          toolWithFunctionDeclarations
                              .functionDeclarations()
                              .orElse(ImmutableList.of()))
                      .add(declaration().get())
                      .build())
              .build();
    }
    // Patch the GenerateContentConfig with the new tool definition.
    GenerateContentConfig generateContentConfig =
        llmRequest
            .config()
            .map(GenerateContentConfig::toBuilder)
            .orElse(GenerateContentConfig.builder())
            .tools(
                new ImmutableList.Builder<Tool>()
                    .addAll(toolsWithoutFunctionDeclarations)
                    .add(toolWithFunctionDeclarations)
                    .build())
            .build();
    llmRequestBuilder.config(generateContentConfig);
    return Completable.complete();
  }

  /**
   * Finds a tool in GenerateContentConfig that has function calling declarations, or returns null
   * otherwise.
   */
  private static @Nullable Tool findToolWithFunctionDeclarations(LlmRequest llmRequest) {
    return llmRequest
        .config()
        .flatMap(config -> config.tools())
        .flatMap(
            tools -> tools.stream().filter(t -> t.functionDeclarations().isPresent()).findFirst())
        .orElse(null);
  }

  /** Finds all tools in GenerateContentConfig that do not have function calling declarations. */
  private static ImmutableList<Tool> findToolsWithoutFunctionDeclarations(LlmRequest llmRequest) {
    return llmRequest
        .config()
        .flatMap(config -> config.tools())
        .map(
            tools ->
                tools.stream()
                    .filter(t -> t.functionDeclarations().isEmpty())
                    .collect(toImmutableList()))
        .orElse(ImmutableList.of());
  }
}

/*
 * 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.models;

import com.anthropic.client.AnthropicClient;
import com.anthropic.models.messages.ContentBlock;
import com.anthropic.models.messages.ContentBlockParam;
import com.anthropic.models.messages.Message;
import com.anthropic.models.messages.MessageCreateParams;
import com.anthropic.models.messages.MessageParam;
import com.anthropic.models.messages.MessageParam.Role;
import com.anthropic.models.messages.TextBlockParam;
import com.anthropic.models.messages.Tool;
import com.anthropic.models.messages.ToolChoice;
import com.anthropic.models.messages.ToolChoiceAuto;
import com.anthropic.models.messages.ToolResultBlockParam;
import com.anthropic.models.messages.ToolUnion;
import com.anthropic.models.messages.ToolUseBlockParam;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.datatype.jdk8.Jdk8Module;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.genai.types.Content;
import com.google.genai.types.FunctionCall;
import com.google.genai.types.FunctionDeclaration;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Flowable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Represents the Claude Generative AI model by Anthropic.
 *
 * <p>This class provides methods for interacting with Claude models. Streaming and live connections
 * are not currently supported for Claude.
 */
public class Claude extends BaseLlm {

  private static final Logger logger = LoggerFactory.getLogger(Claude.class);
  private static final int MAX_TOKEN = 8192;
  private final AnthropicClient anthropicClient;

  /**
   * Constructs a new Claude instance.
   *
   * @param modelName The name of the Claude model to use (e.g., "claude-3-opus-20240229").
   * @param anthropicClient The Anthropic API client instance.
   */
  public Claude(String modelName, AnthropicClient anthropicClient) {
    super(modelName);
    this.anthropicClient = anthropicClient;
  }

  @Override
  public Flowable<LlmResponse> generateContent(LlmRequest llmRequest, boolean stream) {
    // TODO: Switch to streaming API.
    List<MessageParam> messages =
        llmRequest.contents().stream()
            .map(this::contentToAnthropicMessageParam)
            .collect(Collectors.toList());

    List<ToolUnion> tools = ImmutableList.of();
    if (llmRequest.config().isPresent()
        && llmRequest.config().get().tools().isPresent()
        && !llmRequest.config().get().tools().get().isEmpty()
        && llmRequest.config().get().tools().get().get(0).functionDeclarations().isPresent()) {
      tools =
          llmRequest.config().get().tools().get().get(0).functionDeclarations().get().stream()
              .map(this::functionDeclarationToAnthropicTool)
              .map(tool -> ToolUnion.ofTool(tool))
              .collect(Collectors.toList());
    }

    ToolChoice toolChoice =
        llmRequest.tools().isEmpty()
            ? null
            : ToolChoice.ofAuto(ToolChoiceAuto.builder().disableParallelToolUse(true).build());

    String systemText = "";
    Optional<GenerateContentConfig> configOpt = llmRequest.config();
    if (configOpt.isPresent()) {
      Optional<Content> systemInstructionOpt = configOpt.get().systemInstruction();
      if (systemInstructionOpt.isPresent()) {
        String extractedSystemText =
            systemInstructionOpt.get().parts().orElse(ImmutableList.of()).stream()
                .filter(p -> p.text().isPresent())
                .map(p -> p.text().get())
                .collect(Collectors.joining("\n"));
        if (!extractedSystemText.isEmpty()) {
          systemText = extractedSystemText;
        }
      }
    }

    var message =
        this.anthropicClient
            .messages()
            .create(
                MessageCreateParams.builder()
                    .model(llmRequest.model().orElse(model()))
                    .system(systemText)
                    .messages(messages)
                    .tools(tools)
                    .toolChoice(toolChoice)
                    .maxTokens(MAX_TOKEN)
                    .build());

    logger.debug("Claude response: {}", message);

    return Flowable.just(convertAnthropicResponseToLlmResponse(message));
  }

  private Role toClaudeRole(String role) {
    return role.equals("model") || role.equals("assistant") ? Role.ASSISTANT : Role.USER;
  }

  private MessageParam contentToAnthropicMessageParam(Content content) {
    return MessageParam.builder()
        .role(toClaudeRole(content.role().orElse("")))
        .contentOfBlockParams(
            content.parts().orElse(ImmutableList.of()).stream()
                .map(this::partToAnthropicMessageBlock)
                .filter(Objects::nonNull)
                .collect(Collectors.toList()))
        .build();
  }

  private ContentBlockParam partToAnthropicMessageBlock(Part part) {
    if (part.text().isPresent()) {
      return ContentBlockParam.ofText(TextBlockParam.builder().text(part.text().get()).build());
    } else if (part.functionCall().isPresent()) {
      return ContentBlockParam.ofToolUse(
          ToolUseBlockParam.builder()
              .id(part.functionCall().get().id().orElse(""))
              .name(part.functionCall().get().name().orElseThrow())
              .type(com.anthropic.core.JsonValue.from("tool_use"))
              .input(
                  com.anthropic.core.JsonValue.from(
                      part.functionCall().get().args().orElse(ImmutableMap.of())))
              .build());
    } else if (part.functionResponse().isPresent()) {
      String content = "";
      if (part.functionResponse().get().response().isPresent()
          && part.functionResponse().get().response().get().getOrDefault("result", null) != null) {
        content = part.functionResponse().get().response().get().get("result").toString();
      }
      return ContentBlockParam.ofToolResult(
          ToolResultBlockParam.builder()
              .toolUseId(part.functionResponse().get().id().orElse(""))
              .content(content)
              .isError(false)
              .build());
    }
    throw new UnsupportedOperationException("Not supported yet.");
  }

  private void updateTypeString(Map<String, Object> valueDict) {
    if (valueDict == null) {
      return;
    }
    if (valueDict.containsKey("type")) {
      valueDict.put("type", ((String) valueDict.get("type")).toLowerCase());
    }

    if (valueDict.containsKey("items")) {
      updateTypeString((Map<String, Object>) valueDict.get("items"));

      if (valueDict.get("items") instanceof Map
          && ((Map) valueDict.get("items")).containsKey("properties")) {
        Map<String, Object> properties =
            (Map<String, Object>) ((Map) valueDict.get("items")).get("properties");
        if (properties != null) {
          for (Object value : properties.values()) {
            if (value instanceof Map) {
              updateTypeString((Map<String, Object>) value);
            }
          }
        }
      }
    }
  }

  private Tool functionDeclarationToAnthropicTool(FunctionDeclaration functionDeclaration) {
    Map<String, Map<String, Object>> properties = new HashMap<>();
    if (functionDeclaration.parameters().isPresent()
        && functionDeclaration.parameters().get().properties().isPresent()) {
      functionDeclaration
          .parameters()
          .get()
          .properties()
          .get()
          .forEach(
              (key, schema) -> {
                ObjectMapper objectMapper = new ObjectMapper();
                objectMapper.registerModule(new Jdk8Module());
                Map<String, Object> schemaMap =
                    objectMapper.convertValue(schema, new TypeReference<Map<String, Object>>() {});
                updateTypeString(schemaMap);
                properties.put(key, schemaMap);
              });
    }

    return Tool.builder()
        .name(functionDeclaration.name().orElseThrow())
        .description(functionDeclaration.description().orElse(""))
        .inputSchema(
            Tool.InputSchema.builder()
                .properties(com.anthropic.core.JsonValue.from(properties))
                .build())
        .build();
  }

  private LlmResponse convertAnthropicResponseToLlmResponse(Message message) {
    LlmResponse.Builder responseBuilder = LlmResponse.builder();
    List<Part> parts = new ArrayList<>();

    if (message.content() != null) {
      for (ContentBlock block : message.content()) {
        Part part = anthropicContentBlockToPart(block);
        if (part != null) {
          parts.add(part);
        }
      }
      responseBuilder.content(
          Content.builder().role("model").parts(ImmutableList.copyOf(parts)).build());
    }
    return responseBuilder.build();
  }

  private Part anthropicContentBlockToPart(ContentBlock block) {
    if (block.isText()) {
      return Part.builder().text(block.asText().text()).build();
    } else if (block.isToolUse()) {
      return Part.builder()
          .functionCall(
              FunctionCall.builder()
                  .id(block.asToolUse().id())
                  .name(block.asToolUse().name())
                  .args(
                      block
                          .asToolUse()
                          ._input()
                          .convert(new TypeReference<Map<String, Object>>() {}))
                  .build())
          .build();
    }
    throw new UnsupportedOperationException("Not supported yet.");
  }

  @Override
  public BaseLlmConnection connect(LlmRequest llmRequest) {
    throw new UnsupportedOperationException("Live connection is not supported for Claude models.");
  }
}

package com.google.adk.tools.applicationintegrationtoolset;

import static com.google.common.base.Strings.isNullOrEmpty;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.adk.tools.applicationintegrationtoolset.IntegrationConnectorTool.HttpExecutor;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.io.IOException;
import java.net.URI;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;

/**
 * Utility class for interacting with Google Cloud Application Integration.
 *
 * <p>This class provides methods for retrieving OpenAPI spec for an integration or a connection.
 */
public class IntegrationClient {
  String project;
  String location;
  String integration;
  List<String> triggers;
  String connection;
  Map<String, List<String>> entityOperations;
  List<String> actions;
  private final HttpExecutor httpExecutor;
  public static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

  IntegrationClient(
      String project,
      String location,
      String integration,
      List<String> triggers,
      String connection,
      Map<String, List<String>> entityOperations,
      List<String> actions,
      HttpExecutor httpExecutor) {
    this.project = project;
    this.location = location;
    this.integration = integration;
    this.triggers = triggers;
    this.connection = connection;
    this.entityOperations = entityOperations;
    this.actions = actions;
    this.httpExecutor = httpExecutor;
    if (!isNullOrEmpty(connection)) {
      validate();
    }
  }

  private void validate() {
    // Check if both are null, throw exception

    if (this.entityOperations == null && this.actions == null) {
      throw new IllegalArgumentException(
          "No entity operations or actions provided. Please provide at least one of them.");
    }

    if (this.entityOperations != null) {
      Preconditions.checkArgument(
          !this.entityOperations.isEmpty(), "entityOperations map cannot be empty");
      for (Map.Entry<String, List<String>> entry : this.entityOperations.entrySet()) {
        String key = entry.getKey();
        List<String> value = entry.getValue();
        Preconditions.checkArgument(
            key != null && !key.isEmpty(),
            "Enitity in entityOperations map cannot be null or empty");
        Preconditions.checkArgument(
            value != null, "Operations for entity '%s' cannot be null", key);
        for (String str : value) {
          Preconditions.checkArgument(
              str != null && !str.isEmpty(),
              "Operation for entity '%s' cannot be null or empty",
              key);
        }
      }
    }

    // Validate actions if it's not null
    if (this.actions != null) {
      Preconditions.checkArgument(!this.actions.isEmpty(), "Actions list cannot be empty");
      Preconditions.checkArgument(
          this.actions.stream().allMatch(Objects::nonNull),
          "Actions list cannot contain null values");
      Preconditions.checkArgument(
          this.actions.stream().noneMatch(String::isEmpty),
          "Actions list cannot contain empty strings");
    }
  }

  String generateOpenApiSpec() throws Exception {
    String url =
        String.format(
            "https://%s-integrations.googleapis.com/v1/projects/%s/locations/%s:generateOpenApiSpec",
            this.location, this.project, this.location);

    String jsonRequestBody =
        OBJECT_MAPPER.writeValueAsString(
            ImmutableMap.of(
                "apiTriggerResources",
                ImmutableList.of(
                    ImmutableMap.of(
                        "integrationResource",
                        this.integration,
                        "triggerId",
                        Arrays.asList(this.triggers))),
                "fileFormat",
                "JSON"));
    HttpRequest request =
        HttpRequest.newBuilder()
            .uri(URI.create(url))
            .header("Authorization", "Bearer " + httpExecutor.getToken())
            .header("Content-Type", "application/json")
            .POST(HttpRequest.BodyPublishers.ofString(jsonRequestBody))
            .build();
    HttpResponse<String> response =
        httpExecutor.send(request, HttpResponse.BodyHandlers.ofString());

    if (response.statusCode() < 200 || response.statusCode() >= 300) {
      throw new Exception("Error fetching OpenAPI spec. Status: " + response.statusCode());
    }
    return response.body();
  }

  @SuppressWarnings("unchecked")
  ObjectNode getOpenApiSpecForConnection(String toolName, String toolInstructions)
      throws IOException, InterruptedException {
    final String integrationName = "ExecuteConnection";

    ConnectionsClient connectionsClient = createConnectionsClient();

    ImmutableMap<String, Object> baseSpecMap = ConnectionsClient.getConnectorBaseSpec();
    ObjectNode connectorSpec = OBJECT_MAPPER.valueToTree(baseSpecMap);

    ObjectNode paths = (ObjectNode) connectorSpec.path("paths");
    ObjectNode schemas = (ObjectNode) connectorSpec.path("components").path("schemas");

    if (this.entityOperations != null) {
      for (Map.Entry<String, List<String>> entry : this.entityOperations.entrySet()) {
        String entity = entry.getKey();
        List<String> operations = entry.getValue();

        ConnectionsClient.EntitySchemaAndOperations schemaInfo;
        try {
          schemaInfo = connectionsClient.getEntitySchemaAndOperations(entity);
        } catch (InterruptedException e) {
          Thread.currentThread().interrupt();
          throw new IOException("Operation was interrupted while getting entity schema", e);
        }

        Map<String, Object> schemaMap = schemaInfo.schema;
        List<String> supportedOperations = schemaInfo.operations;

        if (operations == null || operations.isEmpty()) {
          operations = supportedOperations;
        }

        String jsonSchemaAsString = OBJECT_MAPPER.writeValueAsString(schemaMap);
        String entityLower = entity.toLowerCase(Locale.ROOT);

        schemas.set(
            "connectorInputPayload_" + entityLower,
            OBJECT_MAPPER.valueToTree(connectionsClient.connectorPayload(schemaMap)));

        for (String operation : operations) {
          String operationLower = operation.toLowerCase(Locale.ROOT);
          String path =
              String.format(
                  "/v2/projects/%s/locations/%s/integrations/%s:execute?triggerId=api_trigger/%s#%s_%s",
                  this.project,
                  this.location,
                  integrationName,
                  integrationName,
                  operationLower,
                  entityLower);

          switch (operationLower) {
            case "create":
              paths.set(
                  path,
                  OBJECT_MAPPER.valueToTree(
                      ConnectionsClient.createOperation(entityLower, toolName, toolInstructions)));
              schemas.set(
                  "create_" + entityLower + "_Request",
                  OBJECT_MAPPER.valueToTree(ConnectionsClient.createOperationRequest(entityLower)));
              break;
            case "update":
              paths.set(
                  path,
                  OBJECT_MAPPER.valueToTree(
                      ConnectionsClient.updateOperation(entityLower, toolName, toolInstructions)));
              schemas.set(
                  "update_" + entityLower + "_Request",
                  OBJECT_MAPPER.valueToTree(ConnectionsClient.updateOperationRequest(entityLower)));
              break;
            case "delete":
              paths.set(
                  path,
                  OBJECT_MAPPER.valueToTree(
                      ConnectionsClient.deleteOperation(entityLower, toolName, toolInstructions)));
              schemas.set(
                  "delete_" + entityLower + "_Request",
                  OBJECT_MAPPER.valueToTree(ConnectionsClient.deleteOperationRequest()));
              break;
            case "list":
              paths.set(
                  path,
                  OBJECT_MAPPER.valueToTree(
                      ConnectionsClient.listOperation(
                          entityLower, jsonSchemaAsString, toolName, toolInstructions)));
              schemas.set(
                  "list_" + entityLower + "_Request",
                  OBJECT_MAPPER.valueToTree(ConnectionsClient.listOperationRequest()));
              break;
            case "get":
              paths.set(
                  path,
                  OBJECT_MAPPER.valueToTree(
                      ConnectionsClient.getOperation(
                          entityLower, jsonSchemaAsString, toolName, toolInstructions)));
              schemas.set(
                  "get_" + entityLower + "_Request",
                  OBJECT_MAPPER.valueToTree(ConnectionsClient.getOperationRequest()));
              break;
            default:
              throw new IllegalArgumentException(
                  "Invalid operation: " + operation + " for entity: " + entity);
          }
        }
      }
    } else if (this.actions != null) {
      for (String action : this.actions) {
        ObjectNode actionDetails =
            OBJECT_MAPPER.valueToTree(connectionsClient.getActionSchema(action));

        JsonNode inputSchemaNode = actionDetails.path("inputSchema");
        JsonNode outputSchemaNode = actionDetails.path("outputSchema");

        String actionDisplayName = actionDetails.path("displayName").asText("").replace(" ", "");
        String operation = "EXECUTE_ACTION";

        Map<String, Object> inputSchemaMap = OBJECT_MAPPER.treeToValue(inputSchemaNode, Map.class);
        Map<String, Object> outputSchemaMap =
            OBJECT_MAPPER.treeToValue(outputSchemaNode, Map.class);

        if (Objects.equals(action, "ExecuteCustomQuery")) {
          schemas.set(
              actionDisplayName + "_Request",
              OBJECT_MAPPER.valueToTree(ConnectionsClient.executeCustomQueryRequest()));
          operation = "EXECUTE_QUERY";
        } else {
          schemas.set(
              actionDisplayName + "_Request",
              OBJECT_MAPPER.valueToTree(ConnectionsClient.actionRequest(actionDisplayName)));
          schemas.set(
              "connectorInputPayload_" + actionDisplayName,
              OBJECT_MAPPER.valueToTree(connectionsClient.connectorPayload(inputSchemaMap)));
        }

        schemas.set(
            "connectorOutputPayload_" + actionDisplayName,
            OBJECT_MAPPER.valueToTree(connectionsClient.connectorPayload(outputSchemaMap)));
        schemas.set(
            actionDisplayName + "_Response",
            OBJECT_MAPPER.valueToTree(ConnectionsClient.actionResponse(actionDisplayName)));

        String path =
            String.format(
                "/v2/projects/%s/locations/%s/integrations/%s:execute?triggerId=api_trigger/%s#%s",
                this.project, this.location, integrationName, integrationName, action);

        paths.set(
            path,
            OBJECT_MAPPER.valueToTree(
                ConnectionsClient.getActionOperation(
                    action, operation, actionDisplayName, toolName, toolInstructions)));
      }
    } else {
      throw new IllegalArgumentException(
          "No entity operations or actions provided. Please provide at least one of them.");
    }
    return connectorSpec;
  }

  String getOperationIdFromPathUrl(String openApiSchemaString, String pathUrl) throws Exception {
    JsonNode topLevelNode = OBJECT_MAPPER.readTree(openApiSchemaString);
    JsonNode specNode = topLevelNode.path("openApiSpec");
    if (specNode.isMissingNode() || !specNode.isTextual()) {
      throw new IllegalArgumentException(
          "Failed to get OpenApiSpec, please check the project and region for the integration.");
    }
    JsonNode rootNode = OBJECT_MAPPER.readTree(specNode.asText());
    JsonNode paths = rootNode.path("paths");

    Iterator<Map.Entry<String, JsonNode>> pathsFields = paths.fields();
    while (pathsFields.hasNext()) {
      Map.Entry<String, JsonNode> pathEntry = pathsFields.next();
      String currentPath = pathEntry.getKey();
      if (!currentPath.equals(pathUrl)) {
        continue;
      }
      JsonNode pathItem = pathEntry.getValue();

      Iterator<Map.Entry<String, JsonNode>> methods = pathItem.fields();
      while (methods.hasNext()) {
        Map.Entry<String, JsonNode> methodEntry = methods.next();
        JsonNode operationNode = methodEntry.getValue();

        if (operationNode.has("operationId")) {
          return operationNode.path("operationId").asText();
        }
      }
    }
    throw new Exception("Could not find operationId for pathUrl: " + pathUrl);
  }

  ConnectionsClient createConnectionsClient() {
    return new ConnectionsClient(
        this.project, this.location, this.connection, this.httpExecutor, OBJECT_MAPPER);
  }
}

/*
 * 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.tools.mcp;

import static java.util.concurrent.TimeUnit.MILLISECONDS;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.adk.JsonBaseModel;
import com.google.adk.tools.BaseTool;
import com.google.adk.tools.ToolContext;
import com.google.common.collect.ImmutableMap;
import com.google.genai.types.FunctionDeclaration;
import com.google.genai.types.Schema;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.spec.McpSchema.CallToolRequest;
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
import io.modelcontextprotocol.spec.McpSchema.Content;
import io.modelcontextprotocol.spec.McpSchema.JsonSchema;
import io.modelcontextprotocol.spec.McpSchema.TextContent;
import io.modelcontextprotocol.spec.McpSchema.Tool;
import io.reactivex.rxjava3.core.Single;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;

// TODO(b/413489523): Add support for auth. This is a TODO for Python as well.
/**
 * Initializes a MCP tool.
 *
 * <p>This wraps a MCP Tool interface and an active MCP Session. It invokes the MCP Tool through
 * executing the tool from remote MCP Session.
 */
public final class McpTool extends BaseTool {

  Tool mcpTool;
  McpSyncClient mcpSession;
  McpSessionManager mcpSessionManager;
  ObjectMapper objectMapper;

  /**
   * Creates a new McpTool with the default ObjectMapper.
   *
   * @param mcpTool The MCP tool to wrap.
   * @param mcpSession The MCP session to use to call the tool.
   * @param mcpSessionManager The MCP session manager to use to create new sessions.
   * @throws IllegalArgumentException If mcpTool or mcpSession are null.
   */
  public McpTool(Tool mcpTool, McpSyncClient mcpSession, McpSessionManager mcpSessionManager) {
    this(mcpTool, mcpSession, mcpSessionManager, JsonBaseModel.getMapper());
  }

  /**
   * Creates a new McpTool with the default ObjectMapper.
   *
   * @param mcpTool The MCP tool to wrap.
   * @param mcpSession The MCP session to use to call the tool.
   * @param mcpSessionManager The MCP session manager to use to create new sessions.
   * @param objectMapper The ObjectMapper to use to convert JSON schemas.
   * @throws IllegalArgumentException If mcpTool or mcpSession are null.
   */
  public McpTool(
      Tool mcpTool,
      McpSyncClient mcpSession,
      McpSessionManager mcpSessionManager,
      ObjectMapper objectMapper) {
    super(
        mcpTool == null ? "" : mcpTool.name(),
        mcpTool == null ? "" : (mcpTool.description().isEmpty() ? "" : mcpTool.description()));

    if (mcpTool == null) {
      throw new IllegalArgumentException("mcpTool cannot be null");
    }
    if (mcpSession == null) {
      throw new IllegalArgumentException("mcpSession cannot be null");
    }
    if (objectMapper == null) {
      throw new IllegalArgumentException("objectMapper cannot be null");
    }
    this.mcpTool = mcpTool;
    this.mcpSession = mcpSession;
    this.mcpSessionManager = mcpSessionManager;
    this.objectMapper = objectMapper;
  }

  public Schema toGeminiSchema(JsonSchema openApiSchema) {
    return Schema.fromJson(objectMapper.valueToTree(openApiSchema).toString());
  }

  private void reintializeSession() {
    this.mcpSession = this.mcpSessionManager.createSession();
  }

  @Override
  public Optional<FunctionDeclaration> declaration() {
    return Optional.of(
        FunctionDeclaration.builder()
            .name(this.name())
            .description(this.description())
            .parameters(toGeminiSchema(this.mcpTool.inputSchema()))
            .build());
  }

  @Override
  public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContext toolContext) {
    return Single.<Map<String, Object>>fromCallable(
            () -> {
              CallToolResult callResult =
                  mcpSession.callTool(new CallToolRequest(this.name(), ImmutableMap.copyOf(args)));

              if (callResult == null) {
                return ImmutableMap.of("error", "MCP framework error: CallToolResult was null");
              }

              List<Content> contents = callResult.content();
              Boolean isToolError = callResult.isError();

              if (isToolError != null && isToolError) {
                String errorMessage = "Tool execution failed.";
                if (contents != null
                    && !contents.isEmpty()
                    && contents.get(0) instanceof TextContent) {
                  TextContent textContent = (TextContent) contents.get(0);
                  if (textContent.text() != null && !textContent.text().isEmpty()) {
                    errorMessage += " Details: " + textContent.text();
                  }
                }
                return ImmutableMap.of("error", errorMessage);
              }

              if (contents == null || contents.isEmpty()) {
                return ImmutableMap.of();
              }

              List<String> textOutputs = new ArrayList<>();
              for (Content content : contents) {
                if (content instanceof TextContent textContent) {
                  if (textContent.text() != null) {
                    textOutputs.add(textContent.text());
                  }
                }
              }

              if (textOutputs.isEmpty()) {
                return ImmutableMap.of(
                    "error",
                    "Tool '" + this.name() + "' returned content that is not TextContent.",
                    "content_details",
                    contents.toString());
              }

              List<Map<String, Object>> resultMaps = new ArrayList<>();
              for (String textOutput : textOutputs) {
                try {
                  resultMaps.add(
                      objectMapper.readValue(
                          textOutput, new TypeReference<Map<String, Object>>() {}));
                } catch (JsonProcessingException e) {
                  resultMaps.add(ImmutableMap.of("text", textOutput));
                }
              }
              return ImmutableMap.of("text_output", resultMaps);
            })
        .retryWhen(
            errors ->
                errors
                    .delay(100, MILLISECONDS)
                    .take(3)
                    .doOnNext(
                        error -> {
                          System.err.println("Retrying callTool due to: " + error);
                          reintializeSession();
                        }));
  }
}

/*
 * 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.models;

import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableMap.toImmutableMap;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
import com.google.adk.JsonBaseModel;
import com.google.adk.tools.BaseTool;
import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.genai.types.Content;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.LiveConnectConfig;
import com.google.genai.types.Part;
import com.google.genai.types.Schema;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Stream;

/** Represents a request to be sent to the LLM. */
@AutoValue
@JsonDeserialize(builder = LlmRequest.Builder.class)
public abstract class LlmRequest extends JsonBaseModel {

  /**
   * Returns the name of the LLM model to be used. If not set, the default model of the LLM class
   * will be used.
   *
   * @return An optional string representing the model name.
   */
  @JsonProperty("model")
  public abstract Optional<String> model();

  /**
   * Returns the list of content sent to the LLM.
   *
   * @return A list of {@link Content} objects.
   */
  @JsonProperty("contents")
  public abstract List<Content> contents();

  /**
   * Returns the configuration for content generation.
   *
   * @return An optional {@link GenerateContentConfig} object containing the generation settings.
   */
  @JsonProperty("config")
  public abstract Optional<GenerateContentConfig> config();

  /**
   * Returns the configuration for live connections. Populated using the RunConfig in the
   * InvocationContext.
   *
   * @return An optional {@link LiveConnectConfig} object containing the live connection settings.
   */
  @JsonProperty("liveConnectConfig")
  public abstract LiveConnectConfig liveConnectConfig();

  /**
   * Returns a map of tools available to the LLM.
   *
   * @return A map where keys are tool names and values are {@link BaseTool} instances.
   */
  @JsonIgnore
  public abstract Map<String, BaseTool> tools();

  /** returns the first system instruction text from the request if present. */
  @JsonIgnore
  public Optional<String> getFirstSystemInstruction() {
    return this.config()
        .flatMap(GenerateContentConfig::systemInstruction)
        .flatMap(content -> content.parts().flatMap(partList -> partList.stream().findFirst()))
        .flatMap(Part::text);
  }

  /** Returns all system instruction texts from the request as an immutable list. */
  @JsonIgnore
  public ImmutableList<String> getSystemInstructions() {
    return config()
        .flatMap(GenerateContentConfig::systemInstruction)
        .flatMap(Content::parts)
        .map(
            partList ->
                partList.stream()
                    .map(Part::text)
                    .flatMap(Optional::stream)
                    .collect(toImmutableList()))
        .orElse(ImmutableList.of());
  }

  public static Builder builder() {
    return new AutoValue_LlmRequest.Builder()
        .tools(ImmutableMap.of())
        .contents(ImmutableList.of())
        .liveConnectConfig(LiveConnectConfig.builder().build());
  }

  public abstract Builder toBuilder();

  /** Builder for constructing {@link LlmRequest} instances. */
  @AutoValue.Builder
  public abstract static class Builder {

    @JsonCreator
    private static Builder create() {
      return builder();
    }

    @CanIgnoreReturnValue
    @JsonProperty("model")
    public abstract Builder model(String model);

    @CanIgnoreReturnValue
    @JsonProperty("contents")
    public abstract Builder contents(List<Content> contents);

    @CanIgnoreReturnValue
    @JsonProperty("config")
    public abstract Builder config(GenerateContentConfig config);

    public abstract Optional<GenerateContentConfig> config();

    @CanIgnoreReturnValue
    @JsonProperty("liveConnectConfig")
    public abstract Builder liveConnectConfig(LiveConnectConfig liveConnectConfig);

    abstract LiveConnectConfig liveConnectConfig();

    @CanIgnoreReturnValue
    abstract Builder tools(Map<String, BaseTool> tools);

    abstract Map<String, BaseTool> tools();

    @CanIgnoreReturnValue
    public final Builder appendInstructions(List<String> instructions) {
      if (instructions.isEmpty()) {
        return this;
      }
      GenerateContentConfig config = config().orElse(GenerateContentConfig.builder().build());
      ImmutableList.Builder<Part> parts = ImmutableList.builder();
      if (config.systemInstruction().isPresent()) {
        parts.addAll(config.systemInstruction().get().parts().orElse(ImmutableList.of()));
      }
      parts.addAll(
          instructions.stream()
              .map(instruction -> Part.builder().text(instruction).build())
              .collect(toImmutableList()));
      config(
          config.toBuilder()
              .systemInstruction(
                  Content.builder()
                      .parts(parts.build())
                      .role(
                          config
                              .systemInstruction()
                              .map(c -> c.role().orElse("user"))
                              .orElse("user"))
                      .build())
              .build());

      LiveConnectConfig liveConfig = liveConnectConfig();
      ImmutableList.Builder<Part> livePartsBuilder = ImmutableList.builder();

      if (liveConfig.systemInstruction().isPresent()) {
        livePartsBuilder.addAll(
            liveConfig.systemInstruction().get().parts().orElse(ImmutableList.of()));
      }

      livePartsBuilder.addAll(
          instructions.stream()
              .map(instruction -> Part.builder().text(instruction).build())
              .collect(toImmutableList()));

      return liveConnectConfig(
          liveConfig.toBuilder()
              .systemInstruction(
                  Content.builder()
                      .parts(livePartsBuilder.build())
                      .role(
                          liveConfig
                              .systemInstruction()
                              .map(c -> c.role().orElse("user"))
                              .orElse("user"))
                      .build())
              .build());
    }

    @CanIgnoreReturnValue
    public final Builder appendTools(List<BaseTool> tools) {
      if (tools.isEmpty()) {
        return this;
      }
      return tools(
          ImmutableMap.<String, BaseTool>builder()
              .putAll(
                  Stream.concat(tools.stream(), tools().values().stream())
                      .collect(
                          toImmutableMap(
                              BaseTool::name,
                              tool -> tool,
                              (tool1, tool2) -> {
                                throw new IllegalArgumentException(
                                    String.format("Duplicate tool name: %s", tool1.name()));
                              })))
              .buildOrThrow());
    }

    /**
     * Sets the output schema for the LLM response. If set, The output content will always be a JSON
     * string that conforms to the schema.
     */
    @CanIgnoreReturnValue
    public final Builder outputSchema(Schema schema) {
      GenerateContentConfig config = config().orElse(GenerateContentConfig.builder().build());
      return config(
          config.toBuilder().responseSchema(schema).responseMimeType("application/json").build());
    }

    public abstract LlmRequest build();
  }
}

/*
 * 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.models;

import static com.google.common.collect.ImmutableList.toImmutableList;

import com.google.common.collect.ImmutableList;
import com.google.genai.AsyncSession;
import com.google.genai.Client;
import com.google.genai.types.Blob;
import com.google.genai.types.Content;
import com.google.genai.types.FinishReason;
import com.google.genai.types.FunctionCall;
import com.google.genai.types.FunctionResponse;
import com.google.genai.types.LiveConnectConfig;
import com.google.genai.types.LiveSendClientContentParameters;
import com.google.genai.types.LiveSendRealtimeInputParameters;
import com.google.genai.types.LiveSendToolResponseParameters;
import com.google.genai.types.LiveServerContent;
import com.google.genai.types.LiveServerMessage;
import com.google.genai.types.LiveServerToolCall;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.processors.PublishProcessor;
import java.net.SocketException;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.atomic.AtomicBoolean;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Manages a persistent, bidirectional connection to the Gemini model via WebSockets for real-time
 * interaction.
 *
 * <p>This connection allows sending conversation history, individual messages, function responses,
 * and real-time media blobs (like audio chunks) while continuously receiving responses from the
 * model.
 */
public final class GeminiLlmConnection implements BaseLlmConnection {

  private static final Logger logger = LoggerFactory.getLogger(GeminiLlmConnection.class);

  private final Client apiClient;
  private final String modelName;
  private final LiveConnectConfig connectConfig;
  private final CompletableFuture<AsyncSession> sessionFuture;
  private final PublishProcessor<LlmResponse> responseProcessor = PublishProcessor.create();
  private final Flowable<LlmResponse> responseFlowable = responseProcessor.serialize();
  private final AtomicBoolean closed = new AtomicBoolean(false);

  /**
   * Establishes a new connection.
   *
   * @param apiClient The API client for communication.
   * @param modelName The specific Gemini model endpoint (e.g., "gemini-2.0-flash).
   * @param connectConfig Configuration parameters for the live session.
   */
  GeminiLlmConnection(Client apiClient, String modelName, LiveConnectConfig connectConfig) {
    this.apiClient = Objects.requireNonNull(apiClient);
    this.modelName = Objects.requireNonNull(modelName);
    this.connectConfig = Objects.requireNonNull(connectConfig);

    this.sessionFuture =
        this.apiClient
            .async
            .live
            .connect(this.modelName, this.connectConfig)
            .whenCompleteAsync(
                (session, throwable) -> {
                  if (throwable != null) {
                    handleConnectionError(throwable);
                  } else if (session != null) {
                    setupReceiver(session);
                  } else if (!closed.get()) {
                    handleConnectionError(
                        new SocketException("WebSocket connection failed without explicit error."));
                  }
                });
  }

  /** Configures the session to forward incoming messages to the response processor. */
  private void setupReceiver(AsyncSession session) {
    if (closed.get()) {
      closeSessionIgnoringErrors(session);
      return;
    }
    session
        .receive(this::handleServerMessage)
        .exceptionally(
            error -> {
              handleReceiveError(error);
              return null;
            });
  }

  /** Processes messages received from the WebSocket server. */
  private void handleServerMessage(LiveServerMessage message) {
    if (closed.get()) {
      return;
    }

    logger.debug("Received server message: {}", message.toJson());

    Optional<LlmResponse> llmResponse = convertToServerResponse(message);
    llmResponse.ifPresent(responseProcessor::onNext);
  }

  /** Converts a server message into the standardized LlmResponse format. */
  private Optional<LlmResponse> convertToServerResponse(LiveServerMessage message) {
    LlmResponse.Builder builder = LlmResponse.builder();

    if (message.serverContent().isPresent()) {
      LiveServerContent serverContent = message.serverContent().get();
      serverContent.modelTurn().ifPresent(builder::content);
      builder
          .partial(serverContent.turnComplete().map(completed -> !completed).orElse(false))
          .turnComplete(serverContent.turnComplete().orElse(false));
    } else if (message.toolCall().isPresent()) {
      LiveServerToolCall toolCall = message.toolCall().get();
      toolCall
          .functionCalls()
          .ifPresent(
              calls -> {
                for (FunctionCall call : calls) {
                  builder.content(
                      Content.builder()
                          .parts(ImmutableList.of(Part.builder().functionCall(call).build()))
                          .build());
                }
              });
      builder.partial(false).turnComplete(false);
    } else if (message.usageMetadata().isPresent()) {
      logger.debug("Received usage metadata: {}", message.usageMetadata().get());
      return Optional.empty();
    } else if (message.toolCallCancellation().isPresent()) {
      logger.debug("Received tool call cancellation: {}", message.toolCallCancellation().get());
      // TODO: implement proper CFC and thus tool call cancellation handling.
      return Optional.empty();
    } else if (message.setupComplete().isPresent()) {
      logger.debug("Received setup complete.");
      return Optional.empty();
    } else {
      logger.warn("Received unknown or empty server message: {}", message.toJson());
      builder
          .errorCode(new FinishReason("Unknown server message."))
          .errorMessage("Received unknown server message.");
    }

    return Optional.of(builder.build());
  }

  /** Handles errors that occur *during* the initial connection attempt. */
  private void handleConnectionError(Throwable throwable) {
    if (closed.compareAndSet(false, true)) {
      logger.error("WebSocket connection failed", throwable);
      Throwable cause =
          (throwable instanceof CompletionException) ? throwable.getCause() : throwable;
      responseProcessor.onError(cause);
    }
  }

  /** Handles errors reported by the WebSocket client *after* connection (e.g., receive errors). */
  private void handleReceiveError(Throwable throwable) {
    if (closed.compareAndSet(false, true)) {
      logger.error("Error during WebSocket receive operation", throwable);
      responseProcessor.onError(throwable);
      sessionFuture.thenAccept(this::closeSessionIgnoringErrors).exceptionally(err -> null);
    }
  }

  @Override
  public Completable sendHistory(List<Content> history) {
    return sendClientContentInternal(
        LiveSendClientContentParameters.builder().turns(history).build());
  }

  @Override
  public Completable sendContent(Content content) {
    Objects.requireNonNull(content, "content cannot be null");

    Optional<List<FunctionResponse>> functionResponses = extractFunctionResponses(content);

    if (functionResponses.isPresent()) {
      return sendToolResponseInternal(
          LiveSendToolResponseParameters.builder()
              .functionResponses(functionResponses.get())
              .build());
    } else {
      return sendClientContentInternal(
          LiveSendClientContentParameters.builder().turns(ImmutableList.of(content)).build());
    }
  }

  /** Extracts FunctionResponse parts from a Content object if all parts are FunctionResponses. */
  private Optional<List<FunctionResponse>> extractFunctionResponses(Content content) {
    if (content.parts().isEmpty() || content.parts().get().isEmpty()) {
      return Optional.empty();
    }

    ImmutableList<FunctionResponse> responses =
        content.parts().get().stream()
            .map(Part::functionResponse)
            .flatMap(Optional::stream)
            .collect(toImmutableList());

    // Ensure *all* parts were function responses
    if (responses.size() == content.parts().get().size()) {
      return Optional.of(responses);
    } else {
      return Optional.empty();
    }
  }

  @Override
  public Completable sendRealtime(Blob blob) {
    return Completable.fromFuture(
        sessionFuture.thenCompose(
            session ->
                session.sendRealtimeInput(
                    LiveSendRealtimeInputParameters.builder().media(blob).build())));
  }

  /** Helper to send client content parameters. */
  private Completable sendClientContentInternal(LiveSendClientContentParameters parameters) {
    return Completable.fromFuture(
        sessionFuture.thenCompose(session -> session.sendClientContent(parameters)));
  }

  /** Helper to send tool response parameters. */
  private Completable sendToolResponseInternal(LiveSendToolResponseParameters parameters) {
    return Completable.fromFuture(
        sessionFuture.thenCompose(session -> session.sendToolResponse(parameters)));
  }

  @Override
  public Flowable<LlmResponse> receive() {
    return responseFlowable;
  }

  @Override
  public void close() {
    closeInternal(null);
  }

  @Override
  public void close(Throwable throwable) {
    Objects.requireNonNull(throwable, "throwable cannot be null for close");
    closeInternal(throwable);
  }

  /** Internal method to handle closing logic and signal completion/error. */
  private void closeInternal(Throwable throwable) {
    if (closed.compareAndSet(false, true)) {
      logger.debug("Closing GeminiConnection.", throwable);

      if (throwable == null) {
        responseProcessor.onComplete();
      } else {
        responseProcessor.onError(throwable);
      }

      if (sessionFuture.isDone()) {
        sessionFuture.thenAccept(this::closeSessionIgnoringErrors).exceptionally(err -> null);
      } else {
        sessionFuture.cancel(false);
      }
    }
  }

  /** Closes the AsyncSession safely, logging any errors. */
  private void closeSessionIgnoringErrors(AsyncSession session) {
    if (session != null) {
      session
          .close()
          .exceptionally(
              closeError -> {
                logger.warn("Error occurred while closing AsyncSession", closeError);
                return null; // Suppress error during close
              });
    }
  }
}

/*
 * 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.flows.llmflows;

import com.google.common.collect.ImmutableList;
import java.util.Optional;

/** LLM flow with automatic agent transfer support. */
public class AutoFlow extends SingleFlow {

  /** Adds {@link AgentTransfer} to base request processors. */
  private static final ImmutableList<RequestProcessor> REQUEST_PROCESSORS =
      ImmutableList.<RequestProcessor>builder()
          .addAll(SingleFlow.REQUEST_PROCESSORS)
          .add(new AgentTransfer())
          .build();

  /** No additional response processors. */
  private static final ImmutableList<ResponseProcessor> RESPONSE_PROCESSORS = ImmutableList.of();

  public AutoFlow() {
    this(/* maxSteps= */ Optional.empty());
  }

  public AutoFlow(Optional<Integer> maxSteps) {
    super(REQUEST_PROCESSORS, RESPONSE_PROCESSORS, maxSteps);
  }
}

package com.google.adk.tools.applicationintegrationtoolset;

import static com.google.common.base.Strings.isNullOrEmpty;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.adk.agents.ReadonlyContext;
import com.google.adk.tools.BaseTool;
import com.google.adk.tools.BaseToolset;
import com.google.adk.tools.applicationintegrationtoolset.IntegrationConnectorTool.DefaultHttpExecutor;
import com.google.adk.tools.applicationintegrationtoolset.IntegrationConnectorTool.HttpExecutor;
import io.reactivex.rxjava3.core.Flowable;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import org.jspecify.annotations.Nullable;

/** Application Integration Toolset */
public class ApplicationIntegrationToolset implements BaseToolset {
  String project;
  String location;
  @Nullable String integration;
  @Nullable List<String> triggers;
  @Nullable String connection;
  @Nullable Map<String, List<String>> entityOperations;
  @Nullable List<String> actions;
  String serviceAccountJson;
  @Nullable String toolNamePrefix;
  @Nullable String toolInstructions;
  public static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
  HttpExecutor httpExecutor;

  /**
   * ApplicationIntegrationToolset generates tools from a given Application Integration resource.
   *
   *
Example Usage:
   *
   *
integrationTool = new ApplicationIntegrationToolset( project="test-project",
   * location="us-central1", integration="test-integration",
   * triggers=ImmutableList.of("api_trigger/test_trigger", "api_trigger/test_trigger_2",
   * serviceAccountJson="{....}"),connection=null,enitityOperations=null,actions=null,toolNamePrefix="test-integration-tool",toolInstructions="This
   * tool is used to get response from test-integration.");
   *
   *
connectionTool = new ApplicationIntegrationToolset( project="test-project",
   * location="us-central1", integration=null, triggers=null, connection="test-connection",
   * entityOperations=ImmutableMap.of("Entity1", ImmutableList.of("LIST", "GET", "UPDATE")),
   * "Entity2", ImmutableList.of()), actions=ImmutableList.of("ExecuteCustomQuery"),
   * serviceAccountJson="{....}", toolNamePrefix="test-tool", toolInstructions="This tool is used to
   * list, get and update issues in Jira.");
   *
   * @param project The GCP project ID.
   * @param location The GCP location of integration.
   * @param integration The integration name.
   * @param triggers(Optional) The list of trigger ids in the integration.
   * @param connection(Optional) The connection name.
   * @param entityOperations(Optional) The entity operations.
   * @param actions(Optional) The actions.
   * @param serviceAccountJson(Optional) The service account configuration as a dictionary. Required
   *     if not using default service credential. Used for fetching the Application Integration or
   *     Integration Connector resource.
   * @param toolNamePrefix(Optional) The tool name prefix.
   * @param toolInstructions(Optional) The tool instructions.
   */
  public ApplicationIntegrationToolset(
      String project,
      String location,
      String integration,
      List<String> triggers,
      String connection,
      Map<String, List<String>> entityOperations,
      List<String> actions,
      String serviceAccountJson,
      String toolNamePrefix,
      String toolInstructions) {
    this(
        project,
        location,
        integration,
        triggers,
        connection,
        entityOperations,
        actions,
        serviceAccountJson,
        toolNamePrefix,
        toolInstructions,
        new DefaultHttpExecutor().createExecutor(serviceAccountJson));
  }

  ApplicationIntegrationToolset(
      String project,
      String location,
      String integration,
      List<String> triggers,
      String connection,
      Map<String, List<String>> entityOperations,
      List<String> actions,
      String serviceAccountJson,
      String toolNamePrefix,
      String toolInstructions,
      HttpExecutor httpExecutor) {
    this.project = project;
    this.location = location;
    this.integration = integration;
    this.triggers = triggers;
    this.connection = connection;
    this.entityOperations = entityOperations;
    this.actions = actions;
    this.serviceAccountJson = serviceAccountJson;
    this.toolNamePrefix = toolNamePrefix;
    this.toolInstructions = toolInstructions;
    this.httpExecutor = httpExecutor;
  }

  List<String> getPathUrl(String openApiSchemaString) throws Exception {
    List<String> pathUrls = new ArrayList<>();
    JsonNode topLevelNode = OBJECT_MAPPER.readTree(openApiSchemaString);
    JsonNode specNode = topLevelNode.path("openApiSpec");
    if (specNode.isMissingNode() || !specNode.isTextual()) {
      throw new IllegalArgumentException(
          "Failed to get OpenApiSpec, please check the project and region for the integration.");
    }
    JsonNode rootNode = OBJECT_MAPPER.readTree(specNode.asText());
    JsonNode pathsNode = rootNode.path("paths");
    Iterator<Map.Entry<String, JsonNode>> paths = pathsNode.fields();
    while (paths.hasNext()) {
      Map.Entry<String, JsonNode> pathEntry = paths.next();
      String pathUrl = pathEntry.getKey();
      pathUrls.add(pathUrl);
    }
    return pathUrls;
  }

  private List<BaseTool> getAllTools() throws Exception {
    String openApiSchemaString = null;
    List<BaseTool> tools = new ArrayList<>();
    if (!isNullOrEmpty(this.integration)) {
      IntegrationClient integrationClient =
          new IntegrationClient(
              this.project,
              this.location,
              this.integration,
              this.triggers,
              null,
              null,
              null,
              this.httpExecutor);
      openApiSchemaString = integrationClient.generateOpenApiSpec();
      List<String> pathUrls = getPathUrl(openApiSchemaString);
      for (String pathUrl : pathUrls) {
        String toolName = integrationClient.getOperationIdFromPathUrl(openApiSchemaString, pathUrl);
        if (toolName != null) {
          tools.add(
              new IntegrationConnectorTool(
                  openApiSchemaString,
                  pathUrl,
                  toolName,
                  toolInstructions,
                  null,
                  null,
                  null,
                  this.serviceAccountJson,
                  this.httpExecutor));
        }
      }
    } else if (!isNullOrEmpty(this.connection)
        && (this.entityOperations != null || this.actions != null)) {
      IntegrationClient integrationClient =
          new IntegrationClient(
              this.project,
              this.location,
              null,
              null,
              this.connection,
              this.entityOperations,
              this.actions,
              this.httpExecutor);
      ObjectNode parentOpenApiSpec = OBJECT_MAPPER.createObjectNode();
      ObjectNode openApiSpec =
          integrationClient.getOpenApiSpecForConnection(toolNamePrefix, toolInstructions);
      String openApiSpecString = OBJECT_MAPPER.writeValueAsString(openApiSpec);
      parentOpenApiSpec.put("openApiSpec", openApiSpecString);
      openApiSchemaString = OBJECT_MAPPER.writeValueAsString(parentOpenApiSpec);
      List<String> pathUrls = getPathUrl(openApiSchemaString);
      for (String pathUrl : pathUrls) {
        String toolName = integrationClient.getOperationIdFromPathUrl(openApiSchemaString, pathUrl);
        if (!isNullOrEmpty(toolName)) {
          ConnectionsClient connectionsClient =
              new ConnectionsClient(
                  this.project, this.location, this.connection, this.httpExecutor, OBJECT_MAPPER);
          ConnectionsClient.ConnectionDetails connectionDetails =
              connectionsClient.getConnectionDetails();

          tools.add(
              new IntegrationConnectorTool(
                  openApiSchemaString,
                  pathUrl,
                  toolName,
                  "",
                  connectionDetails.name,
                  connectionDetails.serviceName,
                  connectionDetails.host,
                  this.serviceAccountJson,
                  this.httpExecutor));
        }
      }
    } else {
      throw new IllegalArgumentException(
          "Invalid request, Either integration or (connection and"
              + " (entityOperations or actions)) should be provided.");
    }

    return tools;
  }

  @Override
  public Flowable<BaseTool> getTools(@Nullable ReadonlyContext readonlyContext) {
    try {
      List<BaseTool> allTools = getAllTools();
      return Flowable.fromIterable(allTools);
    } catch (Exception e) {
      return Flowable.error(e);
    }
  }

  @Override
  public void close() throws Exception {
    // Nothing to close.
  }
}

/*
 * 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.tools;

import static com.google.common.base.Preconditions.checkArgument;

import com.fasterxml.jackson.databind.BeanDescription;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.introspect.BeanPropertyDefinition;
import com.google.adk.JsonBaseModel;
import com.google.common.base.Strings;
import com.google.genai.types.FunctionDeclaration;
import com.google.genai.types.Schema;
import java.lang.reflect.Method;
import java.lang.reflect.Parameter;
import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/** Utility class for function calling. */
public final class FunctionCallingUtils {

  private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

  public static FunctionDeclaration buildFunctionDeclaration(
      Method func, List<String> ignoreParams) {
    String name =
        func.isAnnotationPresent(Annotations.Schema.class)
                && !func.getAnnotation(Annotations.Schema.class).name().isEmpty()
            ? func.getAnnotation(Annotations.Schema.class).name()
            : func.getName();
    FunctionDeclaration.Builder builder = FunctionDeclaration.builder().name(name);
    if (func.isAnnotationPresent(Annotations.Schema.class)
        && !func.getAnnotation(Annotations.Schema.class).description().isEmpty()) {
      builder.description(func.getAnnotation(Annotations.Schema.class).description());
    }
    List<String> required = new ArrayList<>();
    Map<String, Schema> properties = new LinkedHashMap<>();
    for (Parameter param : func.getParameters()) {
      String paramName =
          param.isAnnotationPresent(Annotations.Schema.class)
                  && !param.getAnnotation(Annotations.Schema.class).name().isEmpty()
              ? param.getAnnotation(Annotations.Schema.class).name()
              : param.getName();
      if (ignoreParams.contains(paramName)) {
        continue;
      }
      required.add(paramName);
      properties.put(paramName, buildSchemaFromParameter(param));
    }
    builder.parameters(
        Schema.builder().required(required).properties(properties).type("OBJECT").build());

    Type returnType = func.getGenericReturnType();
    if (returnType != Void.TYPE) {
      Type realReturnType = returnType;
      if (returnType instanceof ParameterizedType) {
        ParameterizedType parameterizedReturnType = (ParameterizedType) returnType;
        String returnTypeName = ((Class<?>) parameterizedReturnType.getRawType()).getName();
        if (returnTypeName.equals("io.reactivex.rxjava3.core.Maybe")
            || returnTypeName.equals("io.reactivex.rxjava3.core.Single")) {
          returnType = parameterizedReturnType.getActualTypeArguments()[0];
          if (returnType instanceof ParameterizedType) {
            ParameterizedType maybeParameterizedType = (ParameterizedType) returnType;
            returnTypeName = ((Class<?>) maybeParameterizedType.getRawType()).getName();
          }
        }
        if (returnTypeName.equals("java.util.Map")
            || returnTypeName.equals("com.google.common.collect.ImmutableMap")) {
          return builder.response(buildSchemaFromType(returnType)).build();
        }
      }
      throw new IllegalArgumentException(
          "Return type should be Map or Maybe<Map> or Single<Map>, but it was "
              + realReturnType.getTypeName());
    }
    return builder.build();
  }

  static FunctionDeclaration buildFunctionDeclaration(JsonBaseModel func, String description) {
    // Create function declaration through json string.
    String jsonString = func.toJson();
    checkArgument(!Strings.isNullOrEmpty(jsonString), "Input String can't be null or empty.");
    FunctionDeclaration declaration = FunctionDeclaration.fromJson(jsonString);
    declaration = declaration.toBuilder().description(description).build();
    if (declaration.name().isEmpty() || declaration.name().get().isEmpty()) {
      throw new IllegalArgumentException("name field must be present.");
    }
    return declaration;
  }

  private static Schema buildSchemaFromParameter(Parameter param) {
    Schema.Builder builder = Schema.builder();
    if (param.isAnnotationPresent(Annotations.Schema.class)
        && !param.getAnnotation(Annotations.Schema.class).description().isEmpty()) {
      builder.description(param.getAnnotation(Annotations.Schema.class).description());
    }
    switch (param.getType().getName()) {
      case "java.lang.String" -> builder.type("STRING");
      case "boolean", "java.lang.Boolean" -> builder.type("BOOLEAN");
      case "int", "java.lang.Integer" -> builder.type("INTEGER");
      case "double", "java.lang.Double", "float", "java.lang.Float", "long", "java.lang.Long" ->
          builder.type("NUMBER");
      case "java.util.List" ->
          builder
              .type("ARRAY")
              .items(
                  buildSchemaFromType(
                      ((ParameterizedType) param.getParameterizedType())
                          .getActualTypeArguments()[0]));
      case "java.util.Map" -> builder.type("OBJECT");
      default -> {
        BeanDescription beanDescription =
            OBJECT_MAPPER
                .getSerializationConfig()
                .introspect(OBJECT_MAPPER.constructType(param.getType()));
        Map<String, Schema> properties = new LinkedHashMap<>();
        for (BeanPropertyDefinition property : beanDescription.findProperties()) {
          properties.put(property.getName(), buildSchemaFromType(property.getRawPrimaryType()));
        }
        builder.type("OBJECT").properties(properties);
      }
    }
    return builder.build();
  }

  public static Schema buildSchemaFromType(Type type) {
    Schema.Builder builder = Schema.builder();
    if (type instanceof ParameterizedType parameterizedType) {
      switch (((Class<?>) parameterizedType.getRawType()).getName()) {
        case "java.util.List" ->
            builder
                .type("ARRAY")
                .items(buildSchemaFromType(parameterizedType.getActualTypeArguments()[0]));
        case "java.util.Map", "com.google.common.collect.ImmutableMap" -> builder.type("OBJECT");
        default -> throw new IllegalArgumentException("Unsupported generic type: " + type);
      }
    } else if (type instanceof Class<?> clazz) {
      switch (clazz.getName()) {
        case "java.lang.String" -> builder.type("STRING");
        case "boolean", "java.lang.Boolean" -> builder.type("BOOLEAN");
        case "int", "java.lang.Integer" -> builder.type("INTEGER");
        case "double", "java.lang.Double", "float", "java.lang.Float", "long", "java.lang.Long" ->
            builder.type("NUMBER");
        case "java.util.Map", "com.google.common.collect.ImmutableMap" -> builder.type("OBJECT");
        default -> {
          BeanDescription beanDescription =
              OBJECT_MAPPER.getSerializationConfig().introspect(OBJECT_MAPPER.constructType(type));
          Map<String, Schema> properties = new LinkedHashMap<>();
          for (BeanPropertyDefinition property : beanDescription.findProperties()) {
            properties.put(property.getName(), buildSchemaFromType(property.getRawPrimaryType()));
          }
          builder.type("OBJECT").properties(properties);
        }
      }
    }
    return builder.build();
  }

  private FunctionCallingUtils() {}
}

package com.google.adk.tools.applicationintegrationtoolset;

import static com.google.common.base.Strings.isNullOrEmpty;
import static java.nio.charset.StandardCharsets.UTF_8;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.adk.tools.BaseTool;
import com.google.adk.tools.ToolContext;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Streams;
import com.google.genai.types.FunctionDeclaration;
import com.google.genai.types.Schema;
import io.reactivex.rxjava3.core.Single;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import org.jspecify.annotations.Nullable;

/** Application Integration Tool */
public class IntegrationConnectorTool extends BaseTool {

  private final String openApiSpec;
  private final String pathUrl;
  private final HttpExecutor httpExecutor;
  private final String connectionName;
  private final String serviceName;
  private final String host;
  private String entity;
  private String operation;
  private String action;
  private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

  interface HttpExecutor {
    <T> HttpResponse<T> send(HttpRequest request, HttpResponse.BodyHandler<T> responseBodyHandler)
        throws IOException, InterruptedException;

    String getToken() throws IOException;

    public HttpExecutor createExecutor(String serviceAccountJson);
  }

  static class DefaultHttpExecutor implements HttpExecutor {
    private final HttpClient client = HttpClient.newHttpClient();
    private final String serviceAccountJson;

    /** Default constructor for when no service account is specified. */
    DefaultHttpExecutor() {
      this(null);
    }

    /**
     * Constructor that accepts an optional service account JSON string.
     *
     * @param serviceAccountJson The service account key as a JSON string, or null.
     */
    DefaultHttpExecutor(@Nullable String serviceAccountJson) {
      this.serviceAccountJson = serviceAccountJson;
    }

    @Override
    public <T> HttpResponse<T> send(
        HttpRequest request, HttpResponse.BodyHandler<T> responseBodyHandler)
        throws IOException, InterruptedException {
      return client.send(request, responseBodyHandler);
    }

    @Override
    public String getToken() throws IOException {
      GoogleCredentials credentials;

      if (this.serviceAccountJson != null && !this.serviceAccountJson.trim().isEmpty()) {
        try (InputStream is = new ByteArrayInputStream(this.serviceAccountJson.getBytes(UTF_8))) {
          credentials =
              GoogleCredentials.fromStream(is)
                  .createScoped("https://www.googleapis.com/auth/cloud-platform");
        } catch (IOException e) {
          throw new IOException("Failed to load credentials from service_account_json.", e);
        }
      } else {
        try {
          credentials =
              GoogleCredentials.getApplicationDefault()
                  .createScoped("https://www.googleapis.com/auth/cloud-platform");
        } catch (IOException e) {
          throw new IOException(
              "Please provide a service account or configure Application Default Credentials. To"
                  + " set up ADC, see"
                  + " https://cloud.google.com/docs/authentication/external/set-up-adc.",
              e);
        }
      }

      credentials.refreshIfExpired();
      return credentials.getAccessToken().getTokenValue();
    }

    @Override
    public HttpExecutor createExecutor(String serviceAccountJson) {
      if (isNullOrEmpty(serviceAccountJson)) {
        return new DefaultHttpExecutor();
      } else {
        return new DefaultHttpExecutor(serviceAccountJson);
      }
    }
  }

  private static final ImmutableList<String> EXCLUDE_FIELDS =
      ImmutableList.of("connectionName", "serviceName", "host", "entity", "operation", "action");

  private static final ImmutableList<String> OPTIONAL_FIELDS =
      ImmutableList.of("pageSize", "pageToken", "filter", "sortByColumns");

  /** Constructor for Application Integration Tool for integration */
  IntegrationConnectorTool(
      String openApiSpec,
      String pathUrl,
      String toolName,
      String toolDescription,
      String serviceAccountJson) {
    this(
        openApiSpec,
        pathUrl,
        toolName,
        toolDescription,
        null,
        null,
        null,
        serviceAccountJson,
        new DefaultHttpExecutor().createExecutor(serviceAccountJson));
  }

  /**
   * Constructor for Application Integration Tool with connection name, service name, host, entity,
   * operation, and action
   */
  IntegrationConnectorTool(
      String openApiSpec,
      String pathUrl,
      String toolName,
      String toolDescription,
      String connectionName,
      String serviceName,
      String host,
      String serviceAccountJson) {
    this(
        openApiSpec,
        pathUrl,
        toolName,
        toolDescription,
        connectionName,
        serviceName,
        host,
        serviceAccountJson,
        new DefaultHttpExecutor().createExecutor(serviceAccountJson));
  }

  IntegrationConnectorTool(
      String openApiSpec,
      String pathUrl,
      String toolName,
      String toolDescription,
      @Nullable String connectionName,
      @Nullable String serviceName,
      @Nullable String host,
      @Nullable String serviceAccountJson,
      HttpExecutor httpExecutor) {
    super(toolName, toolDescription);
    this.openApiSpec = openApiSpec;
    this.pathUrl = pathUrl;
    this.connectionName = connectionName;
    this.serviceName = serviceName;
    this.host = host;
    this.httpExecutor = httpExecutor;
  }

  Schema toGeminiSchema(String openApiSchema, String operationId) throws Exception {
    String resolvedSchemaString = getResolvedRequestSchemaByOperationId(openApiSchema, operationId);
    return Schema.fromJson(resolvedSchemaString);
  }

  @Override
  public Optional<FunctionDeclaration> declaration() {
    try {
      String operationId = getOperationIdFromPathUrl(openApiSpec, pathUrl);
      Schema parametersSchema = toGeminiSchema(openApiSpec, operationId);
      String operationDescription = getOperationDescription(openApiSpec, operationId);

      FunctionDeclaration declaration =
          FunctionDeclaration.builder()
              .name(operationId)
              .description(operationDescription)
              .parameters(parametersSchema)
              .build();
      return Optional.of(declaration);
    } catch (Exception e) {
      System.err.println("Failed to get OpenAPI spec: " + e.getMessage());
      return Optional.empty();
    }
  }

  @Override
  public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContext toolContext) {
    if (this.connectionName != null) {
      args.put("connectionName", this.connectionName);
      args.put("serviceName", this.serviceName);
      args.put("host", this.host);
      if (!isNullOrEmpty(this.entity)) {
        args.put("entity", this.entity);
      } else if (!isNullOrEmpty(this.action)) {
        args.put("action", this.action);
      }
      if (!isNullOrEmpty(this.operation)) {
        args.put("operation", this.operation);
      }
    }

    return Single.fromCallable(
        () -> {
          try {
            String response = executeIntegration(args);
            return ImmutableMap.of("result", response);
          } catch (Exception e) {
            System.err.println("Failed to execute integration: " + e.getMessage());
            return ImmutableMap.of("error", e.getMessage());
          }
        });
  }

  private String executeIntegration(Map<String, Object> args) throws Exception {
    String url = String.format("https://integrations.googleapis.com%s", this.pathUrl);
    String jsonRequestBody;
    try {
      jsonRequestBody = OBJECT_MAPPER.writeValueAsString(args);
    } catch (IOException e) {
      throw new Exception("Error converting args to JSON: " + e.getMessage(), e);
    }
    HttpRequest request =
        HttpRequest.newBuilder()
            .uri(URI.create(url))
            .header("Authorization", "Bearer " + httpExecutor.getToken())
            .header("Content-Type", "application/json")
            .POST(HttpRequest.BodyPublishers.ofString(jsonRequestBody))
            .build();
    HttpResponse<String> response =
        httpExecutor.send(request, HttpResponse.BodyHandlers.ofString());

    if (response.statusCode() < 200 || response.statusCode() >= 300) {
      throw new Exception(
          "Error executing integration. Status: "
              + response.statusCode()
              + " , Response: "
              + response.body());
    }
    return response.body();
  }

  String getOperationIdFromPathUrl(String openApiSchemaString, String pathUrl) throws Exception {
    JsonNode topLevelNode = OBJECT_MAPPER.readTree(openApiSchemaString);
    JsonNode specNode = topLevelNode.path("openApiSpec");
    if (specNode.isMissingNode() || !specNode.isTextual()) {
      throw new IllegalArgumentException(
          "Failed to get OpenApiSpec, please check the project and region for the integration.");
    }
    JsonNode rootNode = OBJECT_MAPPER.readTree(specNode.asText());
    JsonNode paths = rootNode.path("paths");

    // Iterate through each path in the OpenAPI spec.
    Iterator<Map.Entry<String, JsonNode>> pathsFields = paths.fields();
    while (pathsFields.hasNext()) {
      Map.Entry<String, JsonNode> pathEntry = pathsFields.next();
      String currentPath = pathEntry.getKey();
      if (!currentPath.equals(pathUrl)) {
        continue;
      }
      JsonNode pathItem = pathEntry.getValue();

      Iterator<Map.Entry<String, JsonNode>> methods = pathItem.fields();
      while (methods.hasNext()) {
        Map.Entry<String, JsonNode> methodEntry = methods.next();
        JsonNode operationNode = methodEntry.getValue();
        // Set  values for entity, operation, and action
        this.entity = "";
        this.operation = "";
        this.action = "";
        if (operationNode.has("x-entity")) {
          this.entity = operationNode.path("x-entity").asText();
        } else if (operationNode.has("x-action")) {
          this.action = operationNode.path("x-action").asText();
        }
        if (operationNode.has("x-operation")) {
          this.operation = operationNode.path("x-operation").asText();
        }
        // Get the operationId from the operationNode
        if (operationNode.has("operationId")) {
          return operationNode.path("operationId").asText();
        }
      }
    }
    throw new Exception("Could not find operationId for pathUrl: " + pathUrl);
  }

  private String getResolvedRequestSchemaByOperationId(
      String openApiSchemaString, String operationId) throws Exception {
    JsonNode topLevelNode = OBJECT_MAPPER.readTree(openApiSchemaString);
    JsonNode specNode = topLevelNode.path("openApiSpec");
    if (specNode.isMissingNode() || !specNode.isTextual()) {
      throw new IllegalArgumentException(
          "Failed to get OpenApiSpec, please check the project and region for the integration.");
    }
    JsonNode rootNode = OBJECT_MAPPER.readTree(specNode.asText());
    JsonNode operationNode = findOperationNodeById(rootNode, operationId);
    if (operationNode == null) {
      throw new Exception("Could not find operation with operationId: " + operationId);
    }
    JsonNode requestSchemaNode =
        operationNode.path("requestBody").path("content").path("application/json").path("schema");

    if (requestSchemaNode.isMissingNode()) {
      throw new Exception("Could not find request body schema for operationId: " + operationId);
    }

    JsonNode resolvedSchema = resolveRefs(requestSchemaNode, rootNode);

    if (resolvedSchema.isObject()) {
      ObjectNode schemaObject = (ObjectNode) resolvedSchema;

      // 1. Remove excluded fields from the 'properties' object.
      JsonNode propertiesNode = schemaObject.path("properties");
      if (propertiesNode.isObject()) {
        ObjectNode propertiesObject = (ObjectNode) propertiesNode;
        for (String field : EXCLUDE_FIELDS) {
          propertiesObject.remove(field);
        }
      }

      // 2. Remove optional and excluded fields from the 'required' array.
      JsonNode requiredNode = schemaObject.path("required");
      if (requiredNode.isArray()) {
        // Combine the lists of fields to remove
        List<String> fieldsToRemove =
            Streams.concat(OPTIONAL_FIELDS.stream(), EXCLUDE_FIELDS.stream()).toList();

        // To safely remove items from a list while iterating, we must use an Iterator.
        ArrayNode requiredArray = (ArrayNode) requiredNode;
        Iterator<JsonNode> elements = requiredArray.elements();
        while (elements.hasNext()) {
          JsonNode element = elements.next();
          if (element.isTextual() && fieldsToRemove.contains(element.asText())) {
            // This removes the current element from the underlying array.
            elements.remove();
          }
        }
      }
    }
    return OBJECT_MAPPER.writerWithDefaultPrettyPrinter().writeValueAsString(resolvedSchema);
  }

  private @Nullable JsonNode findOperationNodeById(JsonNode rootNode, String operationId) {
    JsonNode paths = rootNode.path("paths");
    for (JsonNode pathItem : paths) {
      Iterator<Map.Entry<String, JsonNode>> methods = pathItem.fields();
      while (methods.hasNext()) {
        Map.Entry<String, JsonNode> methodEntry = methods.next();
        JsonNode operationNode = methodEntry.getValue();
        if (operationNode.path("operationId").asText().equals(operationId)) {
          return operationNode;
        }
      }
    }
    return null;
  }

  private JsonNode resolveRefs(JsonNode currentNode, JsonNode rootNode) {
    if (currentNode.isObject()) {
      ObjectNode objectNode = (ObjectNode) currentNode;
      if (objectNode.has("$ref")) {
        String refPath = objectNode.get("$ref").asText();
        if (refPath.isEmpty() || !refPath.startsWith("#/")) {
          return objectNode;
        }
        JsonNode referencedNode = rootNode.at(refPath.substring(1));
        if (referencedNode.isMissingNode()) {
          return objectNode;
        }
        return resolveRefs(referencedNode, rootNode);
      } else {
        ObjectNode newObjectNode = OBJECT_MAPPER.createObjectNode();
        Iterator<Map.Entry<String, JsonNode>> fields = currentNode.fields();
        while (fields.hasNext()) {
          Map.Entry<String, JsonNode> field = fields.next();
          newObjectNode.set(field.getKey(), resolveRefs(field.getValue(), rootNode));
        }
        return newObjectNode;
      }
    }
    return currentNode;
  }

  private String getOperationDescription(String openApiSchemaString, String operationId)
      throws Exception {
    JsonNode topLevelNode = OBJECT_MAPPER.readTree(openApiSchemaString);
    JsonNode specNode = topLevelNode.path("openApiSpec");
    if (specNode.isMissingNode() || !specNode.isTextual()) {
      return "";
    }
    JsonNode rootNode = OBJECT_MAPPER.readTree(specNode.asText());
    JsonNode operationNode = findOperationNodeById(rootNode, operationId);
    if (operationNode == null) {
      return "";
    }
    return operationNode.path("summary").asText();
  }
}

/*
 * 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.tools;

import com.google.adk.models.LlmRequest;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.GoogleSearch;
import com.google.genai.types.GoogleSearchRetrieval;
import com.google.genai.types.Tool;
import io.reactivex.rxjava3.core.Completable;
import java.util.List;

/**
 * A built-in tool that is automatically invoked by Gemini 2 models to retrieve search results from
 * Google Search.
 *
 * <p>This tool operates internally within the model and does not require or perform local code
 * execution.
 */
public final class GoogleSearchTool extends BaseTool {

  public GoogleSearchTool() {
    super("google_search", "google_search");
  }

  @Override
  public Completable processLlmRequest(
      LlmRequest.Builder llmRequestBuilder, ToolContext toolContext) {

    GenerateContentConfig.Builder configBuilder =
        llmRequestBuilder
            .build()
            .config()
            .map(GenerateContentConfig::toBuilder)
            .orElse(GenerateContentConfig.builder());

    List<Tool> existingTools = configBuilder.build().tools().orElse(ImmutableList.of());
    ImmutableList.Builder<Tool> updatedToolsBuilder = ImmutableList.builder();
    updatedToolsBuilder.addAll(existingTools);

    String model = llmRequestBuilder.build().model().get();
    if (model != null && model.startsWith("gemini-1")) {
      if (!updatedToolsBuilder.build().isEmpty()) {
        System.out.println(configBuilder.build().tools().get());
        return Completable.error(
            new IllegalArgumentException(
                "Google search tool cannot be used with other tools in Gemini 1.x."));
      }
      updatedToolsBuilder.add(
          Tool.builder().googleSearchRetrieval(GoogleSearchRetrieval.builder().build()).build());
      configBuilder.tools(updatedToolsBuilder.build());
    } else if (model != null && model.startsWith("gemini-2")) {

      updatedToolsBuilder.add(Tool.builder().googleSearch(GoogleSearch.builder().build()).build());
      configBuilder.tools(updatedToolsBuilder.build());
    } else {
      return Completable.error(
          new IllegalArgumentException("Google search tool is not supported for model " + model));
    }

    llmRequestBuilder.config(configBuilder.build());
    return Completable.complete();
  }
}

/*
 * 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
 *
 *      https://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.google.auth.oauth2.GoogleCredentials;
import com.google.common.base.Ascii;
import com.google.common.collect.ImmutableMap;
import com.google.genai.errors.GenAiIOException;
import com.google.genai.types.HttpOptions;
import java.io.IOException;
import java.util.Map;
import java.util.Optional;
import okhttp3.MediaType;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;

/** Base client for the HTTP APIs. */
public class HttpApiClient extends ApiClient {
  public static final MediaType MEDIA_TYPE_APPLICATION_JSON =
      MediaType.parse("application/json; charset=utf-8");

  /** Constructs an ApiClient for Google AI APIs. */
  HttpApiClient(Optional<String> apiKey, Optional<HttpOptions> httpOptions) {
    super(apiKey, httpOptions);
  }

  /** Constructs an ApiClient for Vertex AI APIs. */
  HttpApiClient(
      Optional<String> project,
      Optional<String> location,
      Optional<GoogleCredentials> credentials,
      Optional<HttpOptions> httpOptions) {
    super(project, location, credentials, httpOptions);
  }

  /** Sends a Http request given the http method, path, and request json string. */
  @Override
  public ApiResponse request(String httpMethod, String path, String requestJson) {
    boolean queryBaseModel =
        Ascii.equalsIgnoreCase(httpMethod, "GET") && path.startsWith("publishers/google/models/");
    if (this.vertexAI() && !path.startsWith("projects/") && !queryBaseModel) {
      path =
          String.format("projects/%s/locations/%s/", this.project.get(), this.location.get())
              + path;
    }
    String requestUrl =
        String.format(
            "%s/%s/%s", httpOptions.baseUrl().get(), httpOptions.apiVersion().get(), path);

    Request.Builder requestBuilder = new Request.Builder().url(requestUrl);
    setHeaders(requestBuilder);

    if (Ascii.equalsIgnoreCase(httpMethod, "POST")) {
      requestBuilder.post(RequestBody.create(MEDIA_TYPE_APPLICATION_JSON, requestJson));

    } else if (Ascii.equalsIgnoreCase(httpMethod, "GET")) {
      requestBuilder.get();
    } else if (Ascii.equalsIgnoreCase(httpMethod, "DELETE")) {
      requestBuilder.delete();
    } else {
      throw new IllegalArgumentException("Unsupported HTTP method: " + httpMethod);
    }
    return executeRequest(requestBuilder.build());
  }

  /** Sets the required headers (including auth) on the request object. */
  private void setHeaders(Request.Builder requestBuilder) {
    for (Map.Entry<String, String> header :
        httpOptions.headers().orElse(ImmutableMap.of()).entrySet()) {
      requestBuilder.header(header.getKey(), header.getValue());
    }

    if (apiKey.isPresent()) {
      requestBuilder.header("x-goog-api-key", apiKey.get());
    } else {
      GoogleCredentials cred =
          credentials.orElseThrow(() -> new IllegalStateException("credentials is required"));
      try {
        cred.refreshIfExpired();
      } catch (IOException e) {
        throw new GenAiIOException("Failed to refresh credentials.", e);
      }
      String accessToken;
      try {
        accessToken = cred.getAccessToken().getTokenValue();
      } catch (NullPointerException e) {
        // For test cases where the access token is not available.
        if (e.getMessage()
            .contains(
                "because the return value of"
                    + " \"com.google.auth.oauth2.GoogleCredentials.getAccessToken()\" is null")) {
          accessToken = "";
        } else {
          throw e;
        }
      }
      requestBuilder.header("Authorization", "Bearer " + accessToken);

      if (cred.getQuotaProjectId() != null) {
        requestBuilder.header("x-goog-user-project", cred.getQuotaProjectId());
      }
    }
  }

  /** Executes the given HTTP request. */
  private ApiResponse executeRequest(Request request) {
    try {
      Response response = httpClient.newCall(request).execute();
      return new HttpApiResponse(response);
    } catch (IOException e) {
      throw new GenAiIOException("Failed to execute HTTP request.", e);
    }
  }
}

/*
 * 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.tools.mcp;

import static com.google.common.collect.ImmutableList.toImmutableList;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.adk.JsonBaseModel;
import com.google.adk.agents.ReadonlyContext;
import com.google.adk.tools.BaseTool;
import com.google.adk.tools.BaseToolset;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.ServerParameters;
import io.modelcontextprotocol.spec.McpSchema.ListToolsResult;
import io.reactivex.rxjava3.core.Flowable;
import java.util.Objects;
import java.util.Optional;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Connects to a MCP Server, and retrieves MCP Tools into ADK Tools.
 *
 *
Attributes:
 *
 *

 *
{@code connectionParams}: The connection parameters to the MCP server. Can be either {@code
 *       ServerParameters} or {@code SseServerParameters}.
 *
{@code session}: The MCP session being initialized with the connection.
 *

 */
public class McpToolset implements BaseToolset {
  private static final Logger logger = LoggerFactory.getLogger(McpToolset.class);
  private final McpSessionManager mcpSessionManager;
  private McpSyncClient mcpSession;
  private final ObjectMapper objectMapper;
  private final Optional<Object> toolFilter;

  private static final int MAX_RETRIES = 3;
  private static final long RETRY_DELAY_MILLIS = 100;

  /**
   * Initializes the McpToolset with SSE server parameters.
   *
   * @param connectionParams The SSE connection parameters to the MCP server.
   * @param objectMapper An ObjectMapper instance for parsing schemas.
   * @param toolFilter An Optional containing either a ToolPredicate or a List of tool names.
   */
  public McpToolset(
      SseServerParameters connectionParams,
      ObjectMapper objectMapper,
      Optional<Object> toolFilter) {
    Objects.requireNonNull(connectionParams);
    Objects.requireNonNull(objectMapper);
    this.objectMapper = objectMapper;
    this.mcpSessionManager = new McpSessionManager(connectionParams);
    this.toolFilter = toolFilter;
  }

  /**
   * Initializes the McpToolset with SSE server parameters and no tool filter.
   *
   * @param connectionParams The SSE connection parameters to the MCP server.
   * @param objectMapper An ObjectMapper instance for parsing schemas.
   */
  public McpToolset(SseServerParameters connectionParams, ObjectMapper objectMapper) {
    this(connectionParams, objectMapper, Optional.empty());
  }

  /**
   * Initializes the McpToolset with local server parameters.
   *
   * @param connectionParams The local server connection parameters to the MCP server.
   * @param objectMapper An ObjectMapper instance for parsing schemas.
   * @param toolFilter An Optional containing either a ToolPredicate or a List of tool names.
   */
  public McpToolset(
      ServerParameters connectionParams, ObjectMapper objectMapper, Optional<Object> toolFilter) {
    Objects.requireNonNull(connectionParams);
    Objects.requireNonNull(objectMapper);
    this.objectMapper = objectMapper;
    this.mcpSessionManager = new McpSessionManager(connectionParams);
    this.toolFilter = toolFilter;
  }

  /**
   * Initializes the McpToolset with local server parameters and no tool filter.
   *
   * @param connectionParams The local server connection parameters to the MCP server.
   * @param objectMapper An ObjectMapper instance for parsing schemas.
   */
  public McpToolset(ServerParameters connectionParams, ObjectMapper objectMapper) {
    this(connectionParams, objectMapper, Optional.empty());
  }

  /**
   * Initializes the McpToolset with SSE server parameters, using the ObjectMapper used across the
   * ADK.
   *
   * @param connectionParams The SSE connection parameters to the MCP server.
   * @param toolFilter An Optional containing either a ToolPredicate or a List of tool names.
   */
  public McpToolset(SseServerParameters connectionParams, Optional<Object> toolFilter) {
    this(connectionParams, JsonBaseModel.getMapper(), toolFilter);
  }

  /**
   * Initializes the McpToolset with SSE server parameters, using the ObjectMapper used across the
   * ADK and no tool filter.
   *
   * @param connectionParams The SSE connection parameters to the MCP server.
   */
  public McpToolset(SseServerParameters connectionParams) {
    this(connectionParams, JsonBaseModel.getMapper(), Optional.empty());
  }

  /**
   * Initializes the McpToolset with local server parameters, using the ObjectMapper used across the
   * ADK.
   *
   * @param connectionParams The local server connection parameters to the MCP server.
   * @param toolFilter An Optional containing either a ToolPredicate or a List of tool names.
   */
  public McpToolset(ServerParameters connectionParams, Optional<Object> toolFilter) {
    this(connectionParams, JsonBaseModel.getMapper(), toolFilter);
  }

  /**
   * Initializes the McpToolset with local server parameters, using the ObjectMapper used across the
   * ADK and no tool filter.
   *
   * @param connectionParams The local server connection parameters to the MCP server.
   */
  public McpToolset(ServerParameters connectionParams) {
    this(connectionParams, JsonBaseModel.getMapper(), Optional.empty());
  }

  @Override
  public Flowable<BaseTool> getTools(ReadonlyContext readonlyContext) {
    return Flowable.fromCallable(
            () -> {
              for (int i = 0; i < MAX_RETRIES; i++) {
                try {
                  if (this.mcpSession == null) {
                    logger.info("MCP session is null or closed, initializing (attempt {}).", i + 1);
                    this.mcpSession = this.mcpSessionManager.createSession();
                  }

                  ListToolsResult toolsResponse = this.mcpSession.listTools();
                  return toolsResponse.tools().stream()
                      .map(
                          tool ->
                              new McpTool(
                                  tool, this.mcpSession, this.mcpSessionManager, this.objectMapper))
                      .filter(
                          tool ->
                              isToolSelected(
                                  tool, toolFilter, Optional.ofNullable(readonlyContext)))
                      .collect(toImmutableList());
                } catch (IllegalArgumentException e) {
                  // This could happen if parameters for tool loading are somehow invalid.
                  // This is likely a fatal error and should not be retried.
                  logger.error("Invalid argument encountered during tool loading.", e);
                  throw new McpToolLoadingException(
                      "Invalid argument encountered during tool loading.", e);
                } catch (RuntimeException e) { // Catch any other unexpected runtime exceptions
                  logger.error("Unexpected error during tool loading, retry attempt " + (i + 1), e);
                  if (i < MAX_RETRIES - 1) {
                    // For other general exceptions, we might still want to retry if they are
                    // potentially transient, or if we don't have more specific handling. But it's
                    // better to be specific. For now, we'll treat them as potentially retryable but
                    // log
                    // them at a higher level.
                    try {
                      logger.info(
                          "Reinitializing MCP session before next retry for unexpected error.");
                      this.mcpSession = this.mcpSessionManager.createSession();
                      Thread.sleep(RETRY_DELAY_MILLIS);
                    } catch (InterruptedException ie) {
                      Thread.currentThread().interrupt();
                      logger.error(
                          "Interrupted during retry delay for loadTools (unexpected error).", ie);
                      throw new McpToolLoadingException(
                          "Interrupted during retry delay (unexpected error)", ie);
                    } catch (RuntimeException reinitE) {
                      logger.error(
                          "Failed to reinitialize session during retry (unexpected error).",
                          reinitE);
                      throw new McpInitializationException(
                          "Failed to reinitialize session during tool loading retry (unexpected"
                              + " error).",
                          reinitE);
                    }
                  } else {
                    logger.error(
                        "Failed to load tools after multiple retries due to unexpected error.", e);
                    throw new McpToolLoadingException(
                        "Failed to load tools after multiple retries due to unexpected error.", e);
                  }
                }
              }
              // This line should ideally not be reached if retries are handled correctly or an
              // exception is always thrown.
              throw new IllegalStateException("Unexpected state in getTools retry loop");
            })
        .flatMapIterable(tools -> tools);
  }

  @Override
  public void close() {
    if (this.mcpSession != null) {
      try {
        this.mcpSession.close();
        logger.debug("MCP session closed successfully.");
      } catch (RuntimeException e) {
        logger.error("Failed to close MCP session", e);
        // We don't throw an exception here, as closing is a cleanup operation and
        // failing to close shouldn't prevent the program from continuing (or exiting).
        // However, we log the error for debugging purposes.
      } finally {
        this.mcpSession = null;
      }
    }
  }

  /** Base exception for all errors originating from {@code McpToolset}. */
  public static class McpToolsetException extends RuntimeException {
    public McpToolsetException(String message, Throwable cause) {
      super(message, cause);
    }
  }

  /** Exception thrown when there's an error during MCP session initialization. */
  public static class McpInitializationException extends McpToolsetException {
    public McpInitializationException(String message, Throwable cause) {
      super(message, cause);
    }
  }

  /** Exception thrown when there's an error during loading tools from the MCP server. */
  public static class McpToolLoadingException extends McpToolsetException {
    public McpToolLoadingException(String message, Throwable cause) {
      super(message, cause);
    }
  }
}

/*
 * 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.artifacts;

import static java.util.Collections.max;

import com.google.cloud.storage.Blob;
import com.google.cloud.storage.BlobId;
import com.google.cloud.storage.BlobInfo;
import com.google.cloud.storage.Storage;
import com.google.cloud.storage.Storage.BlobListOption;
import com.google.cloud.storage.StorageException;
import com.google.common.base.Splitter;
import com.google.common.base.VerifyException;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Optional;
import java.util.Set;

/** An artifact service implementation using Google Cloud Storage (GCS). */
public final class GcsArtifactService implements BaseArtifactService {
  private final String bucketName;
  private final Storage storageClient;

  /**
   * Initializes the GcsArtifactService.
   *
   * @param bucketName The name of the GCS bucket to use.
   * @param storageClient The GCS storage client instance.
   */
  public GcsArtifactService(String bucketName, Storage storageClient) {
    this.bucketName = bucketName;
    this.storageClient = storageClient;
  }

  /**
   * Checks if a filename uses the user namespace.
   *
   * @param filename Filename to check.
   * @return true if prefixed with "user:", false otherwise.
   */
  private boolean fileHasUserNamespace(String filename) {
    return filename != null && filename.startsWith("user:");
  }

  /**
   * Constructs the blob prefix for an artifact (excluding version).
   *
   * @param appName Application name.
   * @param userId User ID.
   * @param sessionId Session ID.
   * @param filename Artifact filename.
   * @return prefix string for blob location.
   */
  private String getBlobPrefix(String appName, String userId, String sessionId, String filename) {
    if (fileHasUserNamespace(filename)) {
      return String.format("%s/%s/user/%s/", appName, userId, filename);
    } else {
      return String.format("%s/%s/%s/%s/", appName, userId, sessionId, filename);
    }
  }

  /**
   * Constructs the full blob name for an artifact, including version.
   *
   * @param appName Application name.
   * @param userId User ID.
   * @param sessionId Session ID.
   * @param filename Artifact filename.
   * @param version Artifact version.
   * @return full blob name.
   */
  private String getBlobName(
      String appName, String userId, String sessionId, String filename, int version) {
    return getBlobPrefix(appName, userId, sessionId, filename) + version;
  }

  /**
   * Saves an artifact to GCS and assigns a new version.
   *
   * @param appName Application name.
   * @param userId User ID.
   * @param sessionId Session ID.
   * @param filename Artifact filename.
   * @param artifact Artifact content to save.
   * @return Single with assigned version number.
   */
  @Override
  public Single<Integer> saveArtifact(
      String appName, String userId, String sessionId, String filename, Part artifact) {
    return listVersions(appName, userId, sessionId, filename)
        .map(versions -> versions.isEmpty() ? 0 : max(versions) + 1)
        .map(
            nextVersion -> {
              String blobName = getBlobName(appName, userId, sessionId, filename, nextVersion);
              BlobId blobId = BlobId.of(bucketName, blobName);

              BlobInfo blobInfo =
                  BlobInfo.newBuilder(blobId)
                      .setContentType(artifact.inlineData().get().mimeType().orElse(null))
                      .build();

              try {
                byte[] dataToSave =
                    artifact
                        .inlineData()
                        .get()
                        .data()
                        .orElseThrow(
                            () ->
                                new IllegalArgumentException(
                                    "Saveable artifact data must be non-empty."));
                storageClient.create(blobInfo, dataToSave);
                return nextVersion;
              } catch (StorageException e) {
                throw new VerifyException("Failed to save artifact to GCS", e);
              }
            });
  }

  /**
   * Loads an artifact from GCS.
   *
   * @param appName Application name.
   * @param userId User ID.
   * @param sessionId Session ID.
   * @param filename Artifact filename.
   * @param version Optional version to load. Loads latest if empty.
   * @return Maybe with loaded artifact, or empty if not found.
   */
  @Override
  public Maybe<Part> loadArtifact(
      String appName, String userId, String sessionId, String filename, Optional<Integer> version) {
    return version
        .map(Maybe::just)
        .orElseGet(
            () ->
                listVersions(appName, userId, sessionId, filename)
                    .flatMapMaybe(
                        versions -> versions.isEmpty() ? Maybe.empty() : Maybe.just(max(versions))))
        .flatMap(
            versionToLoad -> {
              String blobName = getBlobName(appName, userId, sessionId, filename, versionToLoad);
              BlobId blobId = BlobId.of(bucketName, blobName);

              try {
                Blob blob = storageClient.get(blobId);
                if (blob == null || !blob.exists()) {
                  return Maybe.empty();
                }
                byte[] data = blob.getContent();
                String mimeType = blob.getContentType();
                return Maybe.just(Part.fromBytes(data, mimeType));
              } catch (StorageException e) {
                return Maybe.empty();
              }
            });
  }

  /**
   * Lists artifact filenames for a user and session.
   *
   * @param appName Application name.
   * @param userId User ID.
   * @param sessionId Session ID.
   * @return Single with sorted list of artifact filenames.
   */
  @Override
  public Single<ListArtifactsResponse> listArtifactKeys(
      String appName, String userId, String sessionId) {
    Set<String> filenames = new HashSet<>();

    // List session-specific files
    String sessionPrefix = String.format("%s/%s/%s/", appName, userId, sessionId);
    try {
      for (Blob blob :
          storageClient.list(bucketName, BlobListOption.prefix(sessionPrefix)).iterateAll()) {
        List<String> parts = Splitter.on('/').splitToList(blob.getName());
        filenames.add(parts.get(3)); // appName/userId/sessionId/filename/version
      }
    } catch (StorageException e) {
      throw new VerifyException("Failed to list session artifacts from GCS", e);
    }

    // List user-namespace files
    String userPrefix = String.format("%s/%s/user/", appName, userId);
    try {
      for (Blob blob :
          storageClient.list(bucketName, BlobListOption.prefix(userPrefix)).iterateAll()) {
        List<String> parts = Splitter.on('/').splitToList(blob.getName());
        filenames.add(parts.get(3)); // appName/userId/user/filename/version
      }
    } catch (StorageException e) {
      throw new VerifyException("Failed to list user artifacts from GCS", e);
    }

    return Single.just(
        ListArtifactsResponse.builder().filenames(ImmutableList.sortedCopyOf(filenames)).build());
  }

  /**
   * Deletes all versions of the specified artifact from GCS.
   *
   * @param appName Application name.
   * @param userId User ID.
   * @param sessionId Session ID.
   * @param filename Artifact filename.
   * @return Completable indicating operation completion.
   */
  @Override
  public Completable deleteArtifact(
      String appName, String userId, String sessionId, String filename) {
    ImmutableList<Integer> versions =
        listVersions(appName, userId, sessionId, filename).blockingGet();
    List<BlobId> blobIdsToDelete = new ArrayList<>();
    for (int version : versions) {
      String blobName = getBlobName(appName, userId, sessionId, filename, version);
      blobIdsToDelete.add(BlobId.of(bucketName, blobName));
    }

    if (!blobIdsToDelete.isEmpty()) {
      try {
        var unused = storageClient.delete(blobIdsToDelete);
      } catch (StorageException e) {
        throw new VerifyException("Failed to delete artifact versions from GCS", e);
      }
    }
    return Completable.complete();
  }

  /**
   * Lists all available versions for a given artifact.
   *
   * @param appName Application name.
   * @param userId User ID.
   * @param sessionId Session ID.
   * @param filename Artifact filename.
   * @return Single with sorted list of version numbers.
   */
  @Override
  public Single<ImmutableList<Integer>> listVersions(
      String appName, String userId, String sessionId, String filename) {
    String prefix = getBlobPrefix(appName, userId, sessionId, filename);
    List<Integer> versions = new ArrayList<>();
    try {
      for (Blob blob : storageClient.list(bucketName, BlobListOption.prefix(prefix)).iterateAll()) {
        String name = blob.getName();
        int versionDelimiterIndex = name.lastIndexOf('/'); // immediately before the version number
        if (versionDelimiterIndex != -1 && versionDelimiterIndex < name.length() - 1) {
          versions.add(Integer.parseInt(name.substring(versionDelimiterIndex + 1)));
        }
      }
      return Single.just(ImmutableList.sortedCopyOf(versions));
    } catch (StorageException e) {
      return Single.just(ImmutableList.of());
    }
  }
}

package com.google.adk.tools.applicationintegrationtoolset;

import static com.google.common.base.Strings.isNullOrEmpty;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.adk.tools.applicationintegrationtoolset.IntegrationConnectorTool.HttpExecutor;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.io.IOException;
import java.net.URI;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * Utility class for interacting with the Google Cloud Connectors API.
 *
 * <p>This class provides methods to fetch connection details, schemas for entities and actions, and
 * to generate OpenAPI specifications for creating tools based on these connections.
 */
public class ConnectionsClient {

  private final String project;
  private final String location;
  private final String connection;
  private static final String CONNECTOR_URL = "https://connectors.googleapis.com";
  private final HttpExecutor httpExecutor;
  private final ObjectMapper objectMapper;

  /** Represents details of a connection. */
  public static class ConnectionDetails {
    public String name;
    public String serviceName;
    public String host;
  }

  /** Represents the schema and available operations for an entity. */
  public static class EntitySchemaAndOperations {
    public Map<String, Object> schema;
    public List<String> operations;
  }

  /** Represents the schema for an action. */
  public static class ActionSchema {
    public Map<String, Object> inputSchema;
    public Map<String, Object> outputSchema;
    public String description;
    public String displayName;
  }

  /**
   * Initializes the ConnectionsClient.
   *
   * @param project The Google Cloud project ID.
   * @param location The Google Cloud location (e.g., us-central1).
   * @param connection The connection name.
   */
  public ConnectionsClient(
      String project,
      String location,
      String connection,
      HttpExecutor httpExecutor,
      ObjectMapper objectMapper) {
    this.project = project;
    this.location = location;
    this.connection = connection;
    this.httpExecutor = httpExecutor;
    this.objectMapper = objectMapper;
  }

  /**
   * Retrieves service details for a given connection.
   *
   * @return A {@link ConnectionDetails} object with the connection's info.
   * @throws IOException If there is an issue with network communication or credentials.
   * @throws InterruptedException If the thread is interrupted during the API call.
   */
  public ConnectionDetails getConnectionDetails() throws IOException, InterruptedException {
    String url =
        String.format(
            "%s/v1/projects/%s/locations/%s/connections/%s?view=BASIC",
            CONNECTOR_URL, project, location, connection);

    HttpResponse<String> response = executeApiCall(url);
    Map<String, Object> connectionData = parseJson(response.body());

    ConnectionDetails details = new ConnectionDetails();
    details.name = (String) connectionData.getOrDefault("name", "");
    details.serviceName = (String) connectionData.getOrDefault("serviceDirectory", "");
    details.host = (String) connectionData.getOrDefault("host", "");
    if (details.host != null && !details.host.isEmpty()) {
      details.serviceName = (String) connectionData.getOrDefault("tlsServiceDirectory", "");
    }
    return details;
  }

  /**
   * Retrieves the JSON schema and available operations for a given entity.
   *
   * @param entity The entity name.
   * @return A {@link EntitySchemaAndOperations} object.
   * @throws IOException If there is an issue with network communication or credentials.
   * @throws InterruptedException If the thread is interrupted during polling.
   */
  @SuppressWarnings("unchecked")
  public EntitySchemaAndOperations getEntitySchemaAndOperations(String entity)
      throws IOException, InterruptedException {
    String url =
        String.format(
            "%s/v1/projects/%s/locations/%s/connections/%s/connectionSchemaMetadata:getEntityType?entityId=%s",
            CONNECTOR_URL, project, location, connection, entity);

    HttpResponse<String> initialResponse = executeApiCall(url);
    String operationId = (String) parseJson(initialResponse.body()).get("name");

    if (isNullOrEmpty(operationId)) {
      throw new IOException("Failed to get operation ID for entity: " + entity);
    }

    Map<String, Object> operationResponse = pollOperation(operationId);
    Map<String, Object> responseData =
        (Map<String, Object>) operationResponse.getOrDefault("response", ImmutableMap.of());

    Map<String, Object> schema =
        (Map<String, Object>) responseData.getOrDefault("jsonSchema", ImmutableMap.of());
    List<String> operations =
        (List<String>) responseData.getOrDefault("operations", ImmutableList.of());
    EntitySchemaAndOperations entitySchemaAndOperations = new EntitySchemaAndOperations();
    entitySchemaAndOperations.schema = schema;
    entitySchemaAndOperations.operations = operations;
    return entitySchemaAndOperations;
  }

  /**
   * Retrieves the input and output JSON schema for a given action.
   *
   * @param action The action name.
   * @return An {@link ActionSchema} object.
   * @throws IOException If there is an issue with network communication or credentials.
   * @throws InterruptedException If the thread is interrupted during polling.
   */
  @SuppressWarnings("unchecked")
  public ActionSchema getActionSchema(String action) throws IOException, InterruptedException {
    String url =
        String.format(
            "%s/v1/projects/%s/locations/%s/connections/%s/connectionSchemaMetadata:getAction?actionId=%s",
            CONNECTOR_URL, project, location, connection, action);

    HttpResponse<String> initialResponse = executeApiCall(url);
    String operationId = (String) parseJson(initialResponse.body()).get("name");

    if (isNullOrEmpty(operationId)) {
      throw new IOException("Failed to get operation ID for action: " + action);
    }

    Map<String, Object> operationResponse = pollOperation(operationId);
    Map<String, Object> responseData =
        (Map<String, Object>) operationResponse.getOrDefault("response", ImmutableMap.of());

    ActionSchema actionSchema = new ActionSchema();
    actionSchema.inputSchema =
        (Map<String, Object>) responseData.getOrDefault("inputJsonSchema", ImmutableMap.of());
    actionSchema.outputSchema =
        (Map<String, Object>) responseData.getOrDefault("outputJsonSchema", ImmutableMap.of());
    actionSchema.description = (String) responseData.getOrDefault("description", "");
    actionSchema.displayName = (String) responseData.getOrDefault("displayName", "");

    return actionSchema;
  }

  private HttpResponse<String> executeApiCall(String url) throws IOException, InterruptedException {
    HttpRequest request =
        HttpRequest.newBuilder()
            .uri(URI.create(url))
            .header("Content-Type", "application/json")
            .header("Authorization", "Bearer " + httpExecutor.getToken())
            .GET()
            .build();

    HttpResponse<String> response =
        httpExecutor.send(request, HttpResponse.BodyHandlers.ofString());

    if (response.statusCode() >= 400) {
      String body = response.body();
      if (response.statusCode() == 400 || response.statusCode() == 404) {
        throw new IllegalArgumentException(
            String.format(
                "Invalid request. Please check the provided values of project(%s), location(%s),"
                    + " connection(%s). Error: %s",
                project, location, connection, body));
      }
      if (response.statusCode() == 401 || response.statusCode() == 403) {
        throw new SecurityException(
            String.format("Permission error (status %d): %s", response.statusCode(), body));
      }
      throw new IOException(
          String.format("API call failed with status %d: %s", response.statusCode(), body));
    }
    return response;
  }

  private Map<String, Object> pollOperation(String operationId)
      throws IOException, InterruptedException {
    boolean operationDone = false;
    Map<String, Object> operationResponse = null;

    while (!operationDone) {
      String getOperationUrl = String.format("%s/v1/%s", CONNECTOR_URL, operationId);
      HttpResponse<String> response = executeApiCall(getOperationUrl);
      operationResponse = parseJson(response.body());

      Object doneObj = operationResponse.get("done");
      if (doneObj instanceof Boolean b) {
        operationDone = b;
      }

      if (!operationDone) {
        Thread.sleep(1000);
      }
    }
    return operationResponse;
  }

  /**
   * Converts a JSON Schema dictionary to an OpenAPI schema dictionary.
   *
   * @param jsonSchema The input JSON schema map.
   * @return The converted OpenAPI schema map.
   */
  public Map<String, Object> convertJsonSchemaToOpenApiSchema(Map<String, Object> jsonSchema) {
    Map<String, Object> openapiSchema = new HashMap<>();

    if (jsonSchema.containsKey("description")) {
      openapiSchema.put("description", jsonSchema.get("description"));
    }

    if (jsonSchema.containsKey("type")) {
      Object type = jsonSchema.get("type");
      if (type instanceof List) {
        List<?> typeList = (List<?>) type;
        if (typeList.contains("null")) {
          openapiSchema.put("nullable", true);
          typeList.stream()
              .filter(t -> t instanceof String && !t.equals("null"))
              .findFirst()
              .ifPresent(t -> openapiSchema.put("type", t));
        } else if (!typeList.isEmpty()) {
          openapiSchema.put("type", typeList.get(0));
        }
      } else {
        openapiSchema.put("type", type);
      }
    }
    if (Objects.equals(openapiSchema.get("type"), "object")
        && jsonSchema.containsKey("properties")) {
      @SuppressWarnings("unchecked")
      Map<String, Map<String, Object>> properties =
          (Map<String, Map<String, Object>>) jsonSchema.get("properties");
      Map<String, Object> convertedProperties = new HashMap<>();
      for (Map.Entry<String, Map<String, Object>> entry : properties.entrySet()) {
        convertedProperties.put(entry.getKey(), convertJsonSchemaToOpenApiSchema(entry.getValue()));
      }
      openapiSchema.put("properties", convertedProperties);
    } else if (Objects.equals(openapiSchema.get("type"), "array")
        && jsonSchema.containsKey("items")) {
      @SuppressWarnings("unchecked")
      Map<String, Object> itemsSchema = (Map<String, Object>) jsonSchema.get("items");
      openapiSchema.put("items", convertJsonSchemaToOpenApiSchema(itemsSchema));
    }

    return openapiSchema;
  }

  public Map<String, Object> connectorPayload(Map<String, Object> jsonSchema) {
    return convertJsonSchemaToOpenApiSchema(jsonSchema);
  }

  private Map<String, Object> parseJson(String json) throws IOException {
    return objectMapper.readValue(json, new TypeReference<>() {});
  }

  public static ImmutableMap<String, Object> getConnectorBaseSpec() {
    return ImmutableMap.ofEntries(
        Map.entry("openapi", "3.0.1"),
        Map.entry(
            "info",
            ImmutableMap.of(
                "title", "ExecuteConnection",
                "description", "This tool can execute a query on connection",
                "version", "4")),
        Map.entry(
            "servers",
            ImmutableList.of(ImmutableMap.of("url", "https://integrations.googleapis.com"))),
        Map.entry(
            "security",
            ImmutableList.of(
                ImmutableMap.of(
                    "google_auth",
                    ImmutableList.of("https://www.googleapis.com/auth/cloud-platform")))),
        Map.entry("paths", ImmutableMap.of()),
        Map.entry(
            "components",
            ImmutableMap.ofEntries(
                Map.entry(
                    "schemas",
                    ImmutableMap.ofEntries(
                        Map.entry(
                            "operation",
                            ImmutableMap.of(
                                "type", "string",
                                "default", "LIST_ENTITIES",
                                "description",
                                    "Operation to execute. Possible values are LIST_ENTITIES,"
                                        + " GET_ENTITY, CREATE_ENTITY, UPDATE_ENTITY, DELETE_ENTITY"
                                        + " in case of entities. EXECUTE_ACTION in case of"
                                        + " actions. and EXECUTE_QUERY in case of custom"
                                        + " queries.")),
                        Map.entry(
                            "entityId",
                            ImmutableMap.of("type", "string", "description", "Name of the entity")),
                        Map.entry("connectorInputPayload", ImmutableMap.of("type", "object")),
                        Map.entry(
                            "filterClause",
                            ImmutableMap.of(
                                "type", "string",
                                "default", "",
                                "description", "WHERE clause in SQL query")),
                        Map.entry(
                            "pageSize",
                            ImmutableMap.of(
                                "type", "integer",
                                "default", 50,
                                "description", "Number of entities to return in the response")),
                        Map.entry(
                            "pageToken",
                            ImmutableMap.of(
                                "type", "string",
                                "default", "",
                                "description", "Page token to return the next page of entities")),
                        Map.entry(
                            "connectionName",
                            ImmutableMap.of(
                                "type", "string",
                                "default", "",
                                "description", "Connection resource name to run the query for")),
                        Map.entry(
                            "serviceName",
                            ImmutableMap.of(
                                "type", "string",
                                "default", "",
                                "description", "Service directory for the connection")),
                        Map.entry(
                            "host",
                            ImmutableMap.of(
                                "type", "string",
                                "default", "",
                                "description", "Host name incase of tls service directory")),
                        Map.entry(
                            "entity",
                            ImmutableMap.of(
                                "type", "string",
                                "default", "Issues",
                                "description", "Entity to run the query for")),
                        Map.entry(
                            "action",
                            ImmutableMap.of(
                                "type", "string",
                                "default", "ExecuteCustomQuery",
                                "description", "Action to run the query for")),
                        Map.entry(
                            "query",
                            ImmutableMap.of(
                                "type", "string",
                                "default", "",
                                "description", "Custom Query to execute on the connection")),
                        Map.entry(
                            "timeout",
                            ImmutableMap.of(
                                "type", "integer",
                                "default", 120,
                                "description", "Timeout in seconds for execution of custom query")),
                        Map.entry(
                            "sortByColumns",
                            ImmutableMap.of(
                                "type",
                                "array",
                                "items",
                                ImmutableMap.of("type", "string"),
                                "default",
                                ImmutableList.of(),
                                "description",
                                "Column to sort the results by")),
                        Map.entry("connectorOutputPayload", ImmutableMap.of("type", "object")),
                        Map.entry("nextPageToken", ImmutableMap.of("type", "string")),
                        Map.entry(
                            "execute-connector_Response",
                            ImmutableMap.of(
                                "required", ImmutableList.of("connectorOutputPayload"),
                                "type", "object",
                                "properties",
                                    ImmutableMap.of(
                                        "connectorOutputPayload",
                                        ImmutableMap.of(
                                            "$ref", "#/components/schemas/connectorOutputPayload"),
                                        "nextPageToken",
                                        ImmutableMap.of(
                                            "$ref", "#/components/schemas/nextPageToken")))))),
                Map.entry(
                    "securitySchemes",
                    ImmutableMap.of(
                        "google_auth",
                        ImmutableMap.of(
                            "type",
                            "oauth2",
                            "flows",
                            ImmutableMap.of(
                                "implicit",
                                ImmutableMap.of(
                                    "authorizationUrl",
                                    "https://accounts.google.com/o/oauth2/auth",
                                    "scopes",
                                    ImmutableMap.of(
                                        "https://www.googleapis.com/auth/cloud-platform",
                                        "Auth for google cloud services")))))))));
  }

  public static ImmutableMap<String, Object> getActionOperation(
      String action,
      String operation,
      String actionDisplayName,
      String toolName,
      String toolInstructions) {
    String description = "Use this tool to execute " + action;
    if (Objects.equals(operation, "EXECUTE_QUERY")) {
      description +=
          " Use pageSize = 50 and timeout = 120 until user specifies a different value"
              + " otherwise. If user provides a query in natural language, convert it to SQL query"
              + " and then execute it using the tool.";
    }

    return ImmutableMap.of(
        "post",
        ImmutableMap.ofEntries(
            Map.entry("summary", actionDisplayName),
            Map.entry("description", description + " " + toolInstructions),
            Map.entry("operationId", toolName + "_" + actionDisplayName),
            Map.entry("x-action", action),
            Map.entry("x-operation", operation),
            Map.entry(
                "requestBody",
                ImmutableMap.of(
                    "content",
                    ImmutableMap.of(
                        "application/json",
                        ImmutableMap.of(
                            "schema",
                            ImmutableMap.of(
                                "$ref",
                                String.format(
                                    "#/components/schemas/%s_Request", actionDisplayName)))))),
            Map.entry(
                "responses",
                ImmutableMap.of(
                    "200",
                    ImmutableMap.of(
                        "description",
                        "Success response",
                        "content",
                        ImmutableMap.of(
                            "application/json",
                            ImmutableMap.of(
                                "schema",
                                ImmutableMap.of(
                                    "$ref",
                                    String.format(
                                        "#/components/schemas/%s_Response",
                                        actionDisplayName)))))))));
  }

  public static ImmutableMap<String, Object> listOperation(
      String entity, String schemaAsString, String toolName, String toolInstructions) {
    return ImmutableMap.of(
        "post",
        ImmutableMap.ofEntries(
            Map.entry("summary", "List " + entity),
            Map.entry(
                "description",
                String.format(
                    "Returns the list of %s data. If the page token was available in the response,"
                        + " let users know there are more records available. Ask if the user wants"
                        + " to fetch the next page of results. When passing filter use the"
                        + " following format: `field_name1='value1' AND field_name2='value2'`. %s",
                    entity, toolInstructions)),
            Map.entry("x-operation", "LIST_ENTITIES"),
            Map.entry("x-entity", entity),
            Map.entry("operationId", toolName + "_list_" + entity),
            Map.entry(
                "requestBody",
                ImmutableMap.of(
                    "content",
                    ImmutableMap.of(
                        "application/json",
                        ImmutableMap.of(
                            "schema",
                            ImmutableMap.of(
                                "$ref", "#/components/schemas/list_" + entity + "_Request"))))),
            Map.entry(
                "responses",
                ImmutableMap.of(
                    "200",
                    ImmutableMap.of(
                        "description",
                        "Success response",
                        "content",
                        ImmutableMap.of(
                            "application/json",
                            ImmutableMap.of(
                                "schema",
                                ImmutableMap.of(
                                    "description",
                                    String.format(
                                        "Returns a list of %s of json schema: %s",
                                        entity, schemaAsString),
                                    "$ref",
                                    "#/components/schemas/execute-connector_Response"))))))));
  }

  public static ImmutableMap<String, Object> getOperation(
      String entity, String schemaAsString, String toolName, String toolInstructions) {
    return ImmutableMap.of(
        "post",
        ImmutableMap.ofEntries(
            Map.entry("summary", "Get " + entity),
            Map.entry(
                "description",
                String.format("Returns the details of the %s. %s", entity, toolInstructions)),
            Map.entry("operationId", toolName + "_get_" + entity),
            Map.entry("x-operation", "GET_ENTITY"),
            Map.entry("x-entity", entity),
            Map.entry(
                "requestBody",
                ImmutableMap.of(
                    "content",
                    ImmutableMap.of(
                        "application/json",
                        ImmutableMap.of(
                            "schema",
                            ImmutableMap.of(
                                "$ref", "#/components/schemas/get_" + entity + "_Request"))))),
            Map.entry(
                "responses",
                ImmutableMap.of(
                    "200",
                    ImmutableMap.of(
                        "description",
                        "Success response",
                        "content",
                        ImmutableMap.of(
                            "application/json",
                            ImmutableMap.of(
                                "schema",
                                ImmutableMap.of(
                                    "description",
                                    String.format(
                                        "Returns %s of json schema: %s", entity, schemaAsString),
                                    "$ref",
                                    "#/components/schemas/execute-connector_Response"))))))));
  }

  public static ImmutableMap<String, Object> createOperation(
      String entity, String toolName, String toolInstructions) {
    return ImmutableMap.of(
        "post",
        ImmutableMap.ofEntries(
            Map.entry("summary", "Creates a new " + entity),
            Map.entry(
                "description", String.format("Creates a new %s. %s", entity, toolInstructions)),
            Map.entry("x-operation", "CREATE_ENTITY"),
            Map.entry("x-entity", entity),
            Map.entry("operationId", toolName + "_create_" + entity),
            Map.entry(
                "requestBody",
                ImmutableMap.of(
                    "content",
                    ImmutableMap.of(
                        "application/json",
                        ImmutableMap.of(
                            "schema",
                            ImmutableMap.of(
                                "$ref", "#/components/schemas/create_" + entity + "_Request"))))),
            Map.entry(
                "responses",
                ImmutableMap.of(
                    "200",
                    ImmutableMap.of(
                        "description",
                        "Success response",
                        "content",
                        ImmutableMap.of(
                            "application/json",
                            ImmutableMap.of(
                                "schema",
                                ImmutableMap.of(
                                    "$ref",
                                    "#/components/schemas/execute-connector_Response"))))))));
  }

  public static ImmutableMap<String, Object> updateOperation(
      String entity, String toolName, String toolInstructions) {
    return ImmutableMap.of(
        "post",
        ImmutableMap.ofEntries(
            Map.entry("summary", "Updates the " + entity),
            Map.entry("description", String.format("Updates the %s. %s", entity, toolInstructions)),
            Map.entry("x-operation", "UPDATE_ENTITY"),
            Map.entry("x-entity", entity),
            Map.entry("operationId", toolName + "_update_" + entity),
            Map.entry(
                "requestBody",
                ImmutableMap.of(
                    "content",
                    ImmutableMap.of(
                        "application/json",
                        ImmutableMap.of(
                            "schema",
                            ImmutableMap.of(
                                "$ref", "#/components/schemas/update_" + entity + "_Request"))))),
            Map.entry(
                "responses",
                ImmutableMap.of(
                    "200",
                    ImmutableMap.of(
                        "description",
                        "Success response",
                        "content",
                        ImmutableMap.of(
                            "application/json",
                            ImmutableMap.of(
                                "schema",
                                ImmutableMap.of(
                                    "$ref",
                                    "#/components/schemas/execute-connector_Response"))))))));
  }

  public static ImmutableMap<String, Object> deleteOperation(
      String entity, String toolName, String toolInstructions) {
    return ImmutableMap.of(
        "post",
        ImmutableMap.ofEntries(
            Map.entry("summary", "Delete the " + entity),
            Map.entry("description", String.format("Deletes the %s. %s", entity, toolInstructions)),
            Map.entry("x-operation", "DELETE_ENTITY"),
            Map.entry("x-entity", entity),
            Map.entry("operationId", toolName + "_delete_" + entity),
            Map.entry(
                "requestBody",
                ImmutableMap.of(
                    "content",
                    ImmutableMap.of(
                        "application/json",
                        ImmutableMap.of(
                            "schema",
                            ImmutableMap.of(
                                "$ref", "#/components/schemas/delete_" + entity + "_Request"))))),
            Map.entry(
                "responses",
                ImmutableMap.of(
                    "200",
                    ImmutableMap.of(
                        "description",
                        "Success response",
                        "content",
                        ImmutableMap.of(
                            "application/json",
                            ImmutableMap.of(
                                "schema",
                                ImmutableMap.of(
                                    "$ref",
                                    "#/components/schemas/execute-connector_Response"))))))));
  }

  public static ImmutableMap<String, Object> createOperationRequest(String entity) {
    return ImmutableMap.of(
        "type",
        "object",
        "required",
        ImmutableList.of(
            "connectorInputPayload",
            "operation",
            "connectionName",
            "serviceName",
            "host",
            "entity"),
        "properties",
        ImmutableMap.ofEntries(
            Map.entry(
                "connectorInputPayload",
                ImmutableMap.of("$ref", "#/components/schemas/connectorInputPayload_" + entity)),
            Map.entry("operation", ImmutableMap.of("$ref", "#/components/schemas/operation")),
            Map.entry(
                "connectionName", ImmutableMap.of("$ref", "#/components/schemas/connectionName")),
            Map.entry("serviceName", ImmutableMap.of("$ref", "#/components/schemas/serviceName")),
            Map.entry("host", ImmutableMap.of("$ref", "#/components/schemas/host")),
            Map.entry("entity", ImmutableMap.of("$ref", "#/components/schemas/entity"))));
  }

  public static ImmutableMap<String, Object> updateOperationRequest(String entity) {
    return ImmutableMap.of(
        "type",
        "object",
        "required",
        ImmutableList.of(
            "connectorInputPayload",
            "entityId",
            "operation",
            "connectionName",
            "serviceName",
            "host",
            "entity"),
        "properties",
        ImmutableMap.ofEntries(
            Map.entry(
                "connectorInputPayload",
                ImmutableMap.of("$ref", "#/components/schemas/connectorInputPayload_" + entity)),
            Map.entry("entityId", ImmutableMap.of("$ref", "#/components/schemas/entityId")),
            Map.entry("operation", ImmutableMap.of("$ref", "#/components/schemas/operation")),
            Map.entry(
                "connectionName", ImmutableMap.of("$ref", "#/components/schemas/connectionName")),
            Map.entry("serviceName", ImmutableMap.of("$ref", "#/components/schemas/serviceName")),
            Map.entry("host", ImmutableMap.of("$ref", "#/components/schemas/host")),
            Map.entry("entity", ImmutableMap.of("$ref", "#/components/schemas/entity")),
            Map.entry(
                "filterClause", ImmutableMap.of("$ref", "#/components/schemas/filterClause"))));
  }

  public static ImmutableMap<String, Object> getOperationRequest() {
    return ImmutableMap.of(
        "type",
        "object",
        "required",
        ImmutableList.of(
            "entityId", "operation", "connectionName", "serviceName", "host", "entity"),
        "properties",
        ImmutableMap.ofEntries(
            Map.entry("entityId", ImmutableMap.of("$ref", "#/components/schemas/entityId")),
            Map.entry("operation", ImmutableMap.of("$ref", "#/components/schemas/operation")),
            Map.entry(
                "connectionName", ImmutableMap.of("$ref", "#/components/schemas/connectionName")),
            Map.entry("serviceName", ImmutableMap.of("$ref", "#/components/schemas/serviceName")),
            Map.entry("host", ImmutableMap.of("$ref", "#/components/schemas/host")),
            Map.entry("entity", ImmutableMap.of("$ref", "#/components/schemas/entity"))));
  }

  public static ImmutableMap<String, Object> deleteOperationRequest() {
    return ImmutableMap.of(
        "type",
        "object",
        "required",
        ImmutableList.of(
            "entityId", "operation", "connectionName", "serviceName", "host", "entity"),
        "properties",
        ImmutableMap.ofEntries(
            Map.entry("entityId", ImmutableMap.of("$ref", "#/components/schemas/entityId")),
            Map.entry("operation", ImmutableMap.of("$ref", "#/components/schemas/operation")),
            Map.entry(
                "connectionName", ImmutableMap.of("$ref", "#/components/schemas/connectionName")),
            Map.entry("serviceName", ImmutableMap.of("$ref", "#/components/schemas/serviceName")),
            Map.entry("host", ImmutableMap.of("$ref", "#/components/schemas/host")),
            Map.entry("entity", ImmutableMap.of("$ref", "#/components/schemas/entity")),
            Map.entry(
                "filterClause", ImmutableMap.of("$ref", "#/components/schemas/filterClause"))));
  }

  public static ImmutableMap<String, Object> listOperationRequest() {
    return ImmutableMap.of(
        "type",
        "object",
        "required",
        ImmutableList.of("operation", "connectionName", "serviceName", "host", "entity"),
        "properties",
        ImmutableMap.ofEntries(
            Map.entry("filterClause", ImmutableMap.of("$ref", "#/components/schemas/filterClause")),
            Map.entry("pageSize", ImmutableMap.of("$ref", "#/components/schemas/pageSize")),
            Map.entry("pageToken", ImmutableMap.of("$ref", "#/components/schemas/pageToken")),
            Map.entry("operation", ImmutableMap.of("$ref", "#/components/schemas/operation")),
            Map.entry(
                "connectionName", ImmutableMap.of("$ref", "#/components/schemas/connectionName")),
            Map.entry("serviceName", ImmutableMap.of("$ref", "#/components/schemas/serviceName")),
            Map.entry("host", ImmutableMap.of("$ref", "#/components/schemas/host")),
            Map.entry("entity", ImmutableMap.of("$ref", "#/components/schemas/entity")),
            Map.entry(
                "sortByColumns", ImmutableMap.of("$ref", "#/components/schemas/sortByColumns"))));
  }

  public static ImmutableMap<String, Object> actionRequest(String action) {
    return ImmutableMap.of(
        "type",
        "object",
        "required",
        ImmutableList.of(
            "operation",
            "connectionName",
            "serviceName",
            "host",
            "action",
            "connectorInputPayload"),
        "properties",
        ImmutableMap.ofEntries(
            Map.entry("operation", ImmutableMap.of("$ref", "#/components/schemas/operation")),
            Map.entry(
                "connectionName", ImmutableMap.of("$ref", "#/components/schemas/connectionName")),
            Map.entry("serviceName", ImmutableMap.of("$ref", "#/components/schemas/serviceName")),
            Map.entry("host", ImmutableMap.of("$ref", "#/components/schemas/host")),
            Map.entry("action", ImmutableMap.of("$ref", "#/components/schemas/action")),
            Map.entry(
                "connectorInputPayload",
                ImmutableMap.of("$ref", "#/components/schemas/connectorInputPayload_" + action))));
  }

  public static ImmutableMap<String, Object> actionResponse(String action) {
    return ImmutableMap.of(
        "type",
        "object",
        "properties",
        ImmutableMap.of(
            "connectorOutputPayload",
            ImmutableMap.of("$ref", "#/components/schemas/connectorOutputPayload_" + action)));
  }

  public static ImmutableMap<String, Object> executeCustomQueryRequest() {
    return ImmutableMap.of(
        "type",
        "object",
        "required",
        ImmutableList.of(
            "operation",
            "connectionName",
            "serviceName",
            "host",
            "action",
            "query",
            "timeout",
            "pageSize"),
        "properties",
        ImmutableMap.ofEntries(
            Map.entry("operation", ImmutableMap.of("$ref", "#/components/schemas/operation")),
            Map.entry(
                "connectionName", ImmutableMap.of("$ref", "#/components/schemas/connectionName")),
            Map.entry("serviceName", ImmutableMap.of("$ref", "#/components/schemas/serviceName")),
            Map.entry("host", ImmutableMap.of("$ref", "#/components/schemas/host")),
            Map.entry("action", ImmutableMap.of("$ref", "#/components/schemas/action")),
            Map.entry("query", ImmutableMap.of("$ref", "#/components/schemas/query")),
            Map.entry("timeout", ImmutableMap.of("$ref", "#/components/schemas/timeout")),
            Map.entry("pageSize", ImmutableMap.of("$ref", "#/components/schemas/pageSize"))));
  }
}

/*
 * 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 static java.util.stream.Collectors.toCollection;

import com.google.adk.events.Event;
import com.google.adk.events.EventActions;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import org.jspecify.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * An in-memory implementation of {@link BaseSessionService} assuming {@link Session} objects are
 * mutable regarding their state map, events list, and last update time.
 *
 *
This implementation stores sessions, user state, and app state directly in memory using
 * concurrent maps for basic thread safety. It is suitable for testing or single-node deployments
 * where persistence is not required.
 *
 *
Note: State merging (app/user state prefixed with {@code _app_} / {@code _user_}) occurs
 * during retrieval operations ({@code getSession}, {@code createSession}).
 */
public final class InMemorySessionService implements BaseSessionService {

  private static final Logger logger = LoggerFactory.getLogger(InMemorySessionService.class);

  // Structure: appName -> userId -> sessionId -> Session
  private final ConcurrentMap<String, ConcurrentMap<String, ConcurrentMap<String, Session>>>
      sessions;
  // Structure: appName -> userId -> stateKey -> stateValue
  private final ConcurrentMap<String, ConcurrentMap<String, ConcurrentMap<String, Object>>>
      userState;
  // Structure: appName -> stateKey -> stateValue
  private final ConcurrentMap<String, ConcurrentMap<String, Object>> appState;

  /** Creates a new instance of the in-memory session service with empty storage. */
  public InMemorySessionService() {
    this.sessions = new ConcurrentHashMap<>();
    this.userState = new ConcurrentHashMap<>();
    this.appState = new ConcurrentHashMap<>();
  }

  @Override
  public Single<Session> createSession(
      String appName,
      String userId,
      @Nullable ConcurrentMap<String, Object> state,
      @Nullable String sessionId) {
    Objects.requireNonNull(appName, "appName cannot be null");
    Objects.requireNonNull(userId, "userId cannot be null");

    String resolvedSessionId =
        Optional.ofNullable(sessionId)
            .map(String::trim)
            .filter(s -> !s.isEmpty())
            .orElseGet(() -> UUID.randomUUID().toString());

    // Ensure state map and events list are mutable for the new session
    ConcurrentMap<String, Object> initialState =
        (state == null) ? new ConcurrentHashMap<>() : new ConcurrentHashMap<>(state);
    List<Event> initialEvents = new ArrayList<>();

    // Assuming Session constructor or setters allow setting these mutable collections
    Session newSession =
        Session.builder(resolvedSessionId)
            .appName(appName)
            .userId(userId)
            .state(initialState)
            .events(initialEvents)
            .lastUpdateTime(Instant.now())
            .build();

    sessions
        .computeIfAbsent(appName, k -> new ConcurrentHashMap<>())
        .computeIfAbsent(userId, k -> new ConcurrentHashMap<>())
        .put(resolvedSessionId, newSession);

    // Create a mutable copy for the return value
    Session returnCopy = copySession(newSession);
    // Merge state into the copy before returning
    return Single.just(mergeWithGlobalState(appName, userId, returnCopy));
  }

  @Override
  public Maybe<Session> getSession(
      String appName, String userId, String sessionId, Optional<GetSessionConfig> configOpt) {
    Objects.requireNonNull(appName, "appName cannot be null");
    Objects.requireNonNull(userId, "userId cannot be null");
    Objects.requireNonNull(sessionId, "sessionId cannot be null");
    Objects.requireNonNull(configOpt, "configOpt cannot be null");

    Session storedSession =
        sessions
            .getOrDefault(appName, new ConcurrentHashMap<>())
            .getOrDefault(userId, new ConcurrentHashMap<>())
            .get(sessionId);

    if (storedSession == null) {
      return Maybe.empty();
    }

    Session sessionCopy = copySession(storedSession);

    // Apply filtering based on config directly to the mutable list in the copy
    GetSessionConfig config = configOpt.orElse(GetSessionConfig.builder().build());
    List<Event> eventsInCopy = sessionCopy.events();

    config
        .numRecentEvents()
        .ifPresent(
            num -> {
              if (!eventsInCopy.isEmpty() && num < eventsInCopy.size()) {
                // Keep the last 'num' events by removing older ones
                // Create sublist view (modifications affect original list)

                List<Event> eventsToRemove = eventsInCopy.subList(0, eventsInCopy.size() - num);
                eventsToRemove.clear(); // Clear the sublist view, modifying eventsInCopy
              }
            });

    // Only apply timestamp filter if numRecentEvents was not applied
    if (config.numRecentEvents().isEmpty() && config.afterTimestamp().isPresent()) {
      Instant threshold = config.afterTimestamp().get();

      eventsInCopy.removeIf(
          event -> getEventTimestampEpochSeconds(event) < threshold.getEpochSecond());
    }

    // Merge state into the potentially filtered copy and return
    return Maybe.just(mergeWithGlobalState(appName, userId, sessionCopy));
  }

  // Helper to get event timestamp as epoch seconds
  private long getEventTimestampEpochSeconds(Event event) {
    return event.timestamp() / 1000L;
  }

  @Override
  public Single<ListSessionsResponse> listSessions(String appName, String userId) {
    Objects.requireNonNull(appName, "appName cannot be null");
    Objects.requireNonNull(userId, "userId cannot be null");

    Map<String, Session> userSessionsMap =
        sessions.getOrDefault(appName, new ConcurrentHashMap<>()).get(userId);

    if (userSessionsMap == null || userSessionsMap.isEmpty()) {
      return Single.just(ListSessionsResponse.builder().build());
    }

    // Create copies with empty events and state for the response
    List<Session> sessionCopies =
        userSessionsMap.values().stream()
            .map(this::copySessionMetadata)
            .collect(toCollection(ArrayList::new));

    return Single.just(ListSessionsResponse.builder().sessions(sessionCopies).build());
  }

  @Override
  public Completable deleteSession(String appName, String userId, String sessionId) {
    Objects.requireNonNull(appName, "appName cannot be null");
    Objects.requireNonNull(userId, "userId cannot be null");
    Objects.requireNonNull(sessionId, "sessionId cannot be null");

    ConcurrentMap<String, Session> userSessionsMap =
        sessions.getOrDefault(appName, new ConcurrentHashMap<>()).get(userId);

    if (userSessionsMap != null) {
      userSessionsMap.remove(sessionId);
    }
    return Completable.complete();
  }

  @Override
  public Single<ListEventsResponse> listEvents(String appName, String userId, String sessionId) {
    Objects.requireNonNull(appName, "appName cannot be null");
    Objects.requireNonNull(userId, "userId cannot be null");
    Objects.requireNonNull(sessionId, "sessionId cannot be null");

    Session storedSession =
        sessions
            .getOrDefault(appName, new ConcurrentHashMap<>())
            .getOrDefault(userId, new ConcurrentHashMap<>())
            .get(sessionId);

    if (storedSession == null) {
      return Single.just(ListEventsResponse.builder().build());
    }

    ImmutableList<Event> eventsCopy = ImmutableList.copyOf(storedSession.events());
    return Single.just(ListEventsResponse.builder().events(eventsCopy).build());
  }

  @CanIgnoreReturnValue
  @Override
  public Single<Event> appendEvent(Session session, Event event) {
    Objects.requireNonNull(session, "session cannot be null");
    Objects.requireNonNull(event, "event cannot be null");
    Objects.requireNonNull(session.appName(), "session.appName cannot be null");
    Objects.requireNonNull(session.userId(), "session.userId cannot be null");
    Objects.requireNonNull(session.id(), "session.id cannot be null");

    String appName = session.appName();
    String userId = session.userId();
    String sessionId = session.id();

    // --- Update User/App State (Same as before) ---
    EventActions actions = event.actions();
    if (actions != null) {
      Map<String, Object> stateDelta = actions.stateDelta();
      if (stateDelta != null && !stateDelta.isEmpty()) {
        stateDelta.forEach(
            (key, value) -> {
              if (key.startsWith(State.APP_PREFIX)) {
                String appStateKey = key.substring(State.APP_PREFIX.length());
                appState
                    .computeIfAbsent(appName, k -> new ConcurrentHashMap<>())
                    .put(appStateKey, value);
              } else if (key.startsWith(State.USER_PREFIX)) {
                String userStateKey = key.substring(State.USER_PREFIX.length());
                userState
                    .computeIfAbsent(appName, k -> new ConcurrentHashMap<>())
                    .computeIfAbsent(userId, k -> new ConcurrentHashMap<>())
                    .put(userStateKey, value);
              }
            });
      }
    }

    BaseSessionService.super.appendEvent(session, event);
    session.lastUpdateTime(getInstantFromEvent(event));

    // --- Update the session stored in this service ---
    sessions
        .getOrDefault(appName, new ConcurrentHashMap<>())
        .getOrDefault(userId, new ConcurrentHashMap<>())
        .put(sessionId, session);

    return Single.just(event);
  }

  /** Converts an event's timestamp to an Instant. Adapt based on actual Event structure. */
  // TODO: have Event.timestamp() return Instant directly
  private Instant getInstantFromEvent(Event event) {
    double epochSeconds = getEventTimestampEpochSeconds(event);
    long seconds = (long) epochSeconds;
    long nanos = (long) ((epochSeconds % 1.0) * 1_000_000_000L);
    return Instant.ofEpochSecond(seconds, nanos);
  }

  /**
   * Creates a shallow copy of the session, but with deep copies of the mutable state map and events
   * list. Assumes Session provides necessary getters and a suitable constructor/setters.
   *
   * @param original The session to copy.
   * @return A new Session instance with copied data, including mutable collections.
   */
  private Session copySession(Session original) {
    return Session.builder(original.id())
        .appName(original.appName())
        .userId(original.userId())
        .state(new ConcurrentHashMap<>(original.state()))
        .events(new ArrayList<>(original.events()))
        .lastUpdateTime(original.lastUpdateTime())
        .build();
  }

  /**
   * Creates a copy of the session containing only metadata fields (ID, appName, userId, timestamp).
   * State and Events are explicitly *not* copied.
   *
   * @param original The session whose metadata to copy.
   * @return A new Session instance with only metadata fields populated.
   */
  private Session copySessionMetadata(Session original) {
    return Session.builder(original.id())
        .appName(original.appName())
        .userId(original.userId())
        .lastUpdateTime(original.lastUpdateTime())
        .build();
  }

  /**
   * Merges the app-specific and user-specific state into the provided *mutable* session's state
   * map.
   *
   * @param appName The application name.
   * @param userId The user ID.
   * @param session The mutable session whose state map will be augmented.
   * @return The same session instance passed in, now with merged state.
   */
  @CanIgnoreReturnValue
  private Session mergeWithGlobalState(String appName, String userId, Session session) {
    Map<String, Object> sessionState = session.state();

    // Merge App State directly into the session's state map
    appState
        .getOrDefault(appName, new ConcurrentHashMap<String, Object>())
        .forEach((key, value) -> sessionState.put(State.APP_PREFIX + key, value));

    userState
        .getOrDefault(appName, new ConcurrentHashMap<>())
        .getOrDefault(userId, new ConcurrentHashMap<>())
        .forEach((key, value) -> sessionState.put(State.USER_PREFIX + key, value));

    return session;
  }
}

/*
 * 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.tools;

import com.google.adk.models.LlmRequest;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.Tool;
import com.google.genai.types.ToolCodeExecution;
import io.reactivex.rxjava3.core.Completable;
import java.util.List;

/**
 * A built-in code execution tool that is automatically invoked by Gemini 2 models.
 *
 * <p>This tool operates internally within the model and does not require or perform local code
 * execution.
 */
public final class BuiltInCodeExecutionTool extends BaseTool {

  public BuiltInCodeExecutionTool() {
    super("code_execution", "code_execution");
  }

  @Override
  public Completable processLlmRequest(
      LlmRequest.Builder llmRequestBuilder, ToolContext toolContext) {

    String model = llmRequestBuilder.build().model().get();
    if (model.isEmpty() || !model.startsWith("gemini-2")) {
      return Completable.error(
          new IllegalArgumentException("Code execution tool is not supported for model " + model));
    }
    GenerateContentConfig.Builder configBuilder =
        llmRequestBuilder
            .build()
            .config()
            .map(GenerateContentConfig::toBuilder)
            .orElse(GenerateContentConfig.builder());

    List<Tool> existingTools = configBuilder.build().tools().orElse(ImmutableList.of());
    ImmutableList.Builder<Tool> updatedToolsBuilder = ImmutableList.builder();
    updatedToolsBuilder
        .addAll(existingTools)
        .add(Tool.builder().codeExecution(ToolCodeExecution.builder().build()).build());
    configBuilder.tools(updatedToolsBuilder.build());
    llmRequestBuilder.config(configBuilder.build());
    return Completable.complete();
  }
}

/*
 * 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.agents;

import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.genai.types.AudioTranscriptionConfig;
import com.google.genai.types.Modality;
import com.google.genai.types.SpeechConfig;
import org.jspecify.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Configuration to modify an agent's LLM's underlying behavior. */
@AutoValue
public abstract class RunConfig {
  private static final Logger logger = LoggerFactory.getLogger(RunConfig.class);

  /** Streaming mode for the runner. Required for BaseAgent.runLive() to work. */
  public enum StreamingMode {
    NONE,
    SSE,
    BIDI
  }

  public abstract @Nullable SpeechConfig speechConfig();

  public abstract ImmutableList<Modality> responseModalities();

  public abstract boolean saveInputBlobsAsArtifacts();

  public abstract StreamingMode streamingMode();

  public abstract @Nullable AudioTranscriptionConfig outputAudioTranscription();

  public abstract int maxLlmCalls();

  public static Builder builder() {
    return new AutoValue_RunConfig.Builder()
        .setSaveInputBlobsAsArtifacts(false)
        .setResponseModalities(ImmutableList.of())
        .setStreamingMode(StreamingMode.NONE)
        .setMaxLlmCalls(500);
  }

  public static Builder builder(RunConfig runConfig) {
    return new AutoValue_RunConfig.Builder()
        .setSaveInputBlobsAsArtifacts(runConfig.saveInputBlobsAsArtifacts())
        .setStreamingMode(runConfig.streamingMode())
        .setMaxLlmCalls(runConfig.maxLlmCalls())
        .setResponseModalities(runConfig.responseModalities())
        .setSpeechConfig(runConfig.speechConfig())
        .setOutputAudioTranscription(runConfig.outputAudioTranscription());
  }

  /** Builder for {@link RunConfig}. */
  @AutoValue.Builder
  public abstract static class Builder {

    @CanIgnoreReturnValue
    public abstract Builder setSpeechConfig(SpeechConfig speechConfig);

    @CanIgnoreReturnValue
    public abstract Builder setResponseModalities(Iterable<Modality> responseModalities);

    @CanIgnoreReturnValue
    public abstract Builder setSaveInputBlobsAsArtifacts(boolean saveInputBlobsAsArtifacts);

    @CanIgnoreReturnValue
    public abstract Builder setStreamingMode(StreamingMode streamingMode);

    @CanIgnoreReturnValue
    public abstract Builder setOutputAudioTranscription(
        AudioTranscriptionConfig outputAudioTranscription);

    @CanIgnoreReturnValue
    public abstract Builder setMaxLlmCalls(int maxLlmCalls);

    abstract RunConfig autoBuild();

    public RunConfig build() {
      RunConfig runConfig = autoBuild();
      if (runConfig.maxLlmCalls() < 0) {
        logger.warn(
            "maxLlmCalls is negative. This will result in no enforcement on total"
                + " number of llm calls that will be made for a run. This may not be ideal, as this"
                + " could result in a never ending communication between the model and the agent in"
                + " certain cases.");
      }
      return runConfig;
    }
  }
}

/*
 * 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.tools;

import com.google.adk.agents.CallbackContext;
import com.google.adk.agents.InvocationContext;
import com.google.adk.artifacts.ListArtifactsResponse;
import com.google.adk.events.EventActions;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import io.reactivex.rxjava3.core.Single;
import java.util.List;
import java.util.Optional;

/** ToolContext object provides a structured context for executing tools or functions. */
public class ToolContext extends CallbackContext {
  private Optional<String> functionCallId = Optional.empty();

  private ToolContext(
      InvocationContext invocationContext,
      EventActions eventActions,
      Optional<String> functionCallId) {
    super(invocationContext, eventActions);
    this.functionCallId = functionCallId;
  }

  public EventActions actions() {
    return this.eventActions;
  }

  public void setActions(EventActions actions) {
    this.eventActions = actions;
  }

  public Optional<String> functionCallId() {
    return functionCallId;
  }

  public void functionCallId(String functionCallId) {
    this.functionCallId = Optional.ofNullable(functionCallId);
  }

  @SuppressWarnings("unused")
  private void requestCredential() {
    // TODO: b/414678311 - Implement credential request logic. Make this public.
    throw new UnsupportedOperationException("Credential request not implemented yet.");
  }

  @SuppressWarnings("unused")
  private void getAuthResponse() {
    // TODO: b/414678311 - Implement auth response retrieval logic. Make this public.
    throw new UnsupportedOperationException("Auth response retrieval not implemented yet.");
  }

  @SuppressWarnings("unused")
  private void searchMemory() {
    // TODO: b/414680316 - Implement search memory logic. Make this public.
    throw new UnsupportedOperationException("Search memory not implemented yet.");
  }

  /** Lists the filenames of the artifacts attached to the current session. */
  public Single<List<String>> listArtifacts() {
    if (invocationContext.artifactService() == null) {
      throw new IllegalStateException("Artifact service is not initialized.");
    }
    return invocationContext
        .artifactService()
        .listArtifactKeys(
            invocationContext.session().appName(),
            invocationContext.session().userId(),
            invocationContext.session().id())
        .map(ListArtifactsResponse::filenames);
  }

  public static Builder builder(InvocationContext invocationContext) {
    return new Builder(invocationContext);
  }

  public Builder toBuilder() {
    return new Builder(invocationContext)
        .actions(eventActions)
        .functionCallId(functionCallId.orElse(null));
  }

  /** Builder for {@link ToolContext}. */
  public static final class Builder {
    private final InvocationContext invocationContext;
    private EventActions eventActions = EventActions.builder().build(); // Default empty actions
    private Optional<String> functionCallId = Optional.empty();

    private Builder(InvocationContext invocationContext) {
      this.invocationContext = invocationContext;
    }

    @CanIgnoreReturnValue
    public Builder actions(EventActions actions) {
      this.eventActions = actions;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder functionCallId(String functionCallId) {
      this.functionCallId = Optional.ofNullable(functionCallId);
      return this;
    }

    public ToolContext build() {
      return new ToolContext(invocationContext, eventActions, functionCallId);
    }
  }
}

/*
 * 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.memory;

import static com.google.common.collect.ImmutableList.toImmutableList;

import com.google.adk.events.Event;
import com.google.adk.sessions.Session;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Single;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * An in-memory memory service for prototyping purposes only.
 *
 * <p>Uses keyword matching instead of semantic search.
 */
public final class InMemoryMemoryService implements BaseMemoryService {

  // Pattern to extract words, matching the Python version.
  private static final Pattern WORD_PATTERN = Pattern.compile("[A-Za-z]+");

  /** Keys are "app_name/user_id", values are maps of "session_id" to a list of events. */
  private final Map<String, Map<String, List<Event>>> sessionEvents;

  public InMemoryMemoryService() {
    this.sessionEvents = new ConcurrentHashMap<>();
  }

  private static String userKey(String appName, String userId) {
    return appName + "/" + userId;
  }

  @Override
  public Completable addSessionToMemory(Session session) {
    return Completable.fromAction(
        () -> {
          String key = userKey(session.appName(), session.userId());
          Map<String, List<Event>> userSessions =
              sessionEvents.computeIfAbsent(key, k -> new ConcurrentHashMap<>());
          ImmutableList<Event> nonEmptyEvents =
              session.events().stream()
                  .filter(
                      event ->
                          event.content().isPresent()
                              && event.content().get().parts().isPresent()
                              && !event.content().get().parts().get().isEmpty())
                  .collect(toImmutableList());
          userSessions.put(session.id(), nonEmptyEvents);
        });
  }

  @Override
  public Single<SearchMemoryResponse> searchMemory(String appName, String userId, String query) {
    return Single.fromCallable(
        () -> {
          String key = userKey(appName, userId);

          if (!sessionEvents.containsKey(key)) {
            return SearchMemoryResponse.builder().build();
          }

          Map<String, List<Event>> userSessions = sessionEvents.get(key);

          ImmutableSet<String> wordsInQuery =
              ImmutableSet.copyOf(query.toLowerCase(Locale.ROOT).split("\\s+"));

          List<MemoryEntry> matchingMemories = new ArrayList<>();

          for (List<Event> eventsInSession : userSessions.values()) {
            for (Event event : eventsInSession) {
              if (event.content().isEmpty() || event.content().get().parts().isEmpty()) {
                continue;
              }

              Set<String> wordsInEvent = new HashSet<>();
              for (Part part : event.content().get().parts().get()) {
                if (!Strings.isNullOrEmpty(part.text().get())) {
                  Matcher matcher = WORD_PATTERN.matcher(part.text().get());
                  while (matcher.find()) {
                    wordsInEvent.add(matcher.group().toLowerCase(Locale.ROOT));
                  }
                }
              }

              if (wordsInEvent.isEmpty()) {
                continue;
              }

              if (!Collections.disjoint(wordsInQuery, wordsInEvent)) {
                MemoryEntry memory =
                    MemoryEntry.builder()
                        .setContent(event.content().get())
                        .setAuthor(event.author())
                        .setTimestamp(formatTimestamp(event.timestamp()))
                        .build();
                matchingMemories.add(memory);
              }
            }
          }

          return SearchMemoryResponse.builder()
              .setMemories(ImmutableList.copyOf(matchingMemories))
              .build();
        });
  }

  private String formatTimestamp(long timestamp) {
    return Instant.ofEpochSecond(timestamp).toString();
  }
}

/*
 * 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.flows.llmflows;

import com.google.adk.agents.InvocationContext;
import com.google.adk.events.Event;
import com.google.adk.models.LlmRequest;
import com.google.auto.value.AutoValue;
import io.reactivex.rxjava3.core.Single;

public interface RequestProcessor {

  @AutoValue
  public abstract static class RequestProcessingResult {
    public abstract LlmRequest updatedRequest();

    public abstract Iterable<Event> events();

    public static RequestProcessingResult create(
        LlmRequest updatedRequest, Iterable<Event> events) {
      return new AutoValue_RequestProcessor_RequestProcessingResult(updatedRequest, events);
    }
  }

  /**
   * Process the LLM request as part of the pre-processing stage.
   *
   * @param context the invocation context.
   * @param request the LLM request to process.
   * @return a list of events generated during processing (if any).
   */
  Single<RequestProcessingResult> processRequest(InvocationContext context, LlmRequest request);
}

/*
 * 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
 *
 *      https://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 static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.StandardSystemProperty.JAVA_VERSION;

import com.google.auth.oauth2.GoogleCredentials;
import com.google.common.base.Ascii;
import com.google.common.collect.ImmutableMap;
import com.google.genai.errors.GenAiIOException;
import com.google.genai.types.HttpOptions;
import java.io.IOException;
import java.time.Duration;
import java.util.Map;
import java.util.Optional;
import okhttp3.OkHttpClient;
import org.jspecify.annotations.Nullable;

/** Interface for an API client which issues HTTP requests to the GenAI APIs. */
abstract class ApiClient {
  OkHttpClient httpClient;
  // For Google AI APIs
  final Optional<String> apiKey;
  // For Vertex AI APIs
  final Optional<String> project;
  final Optional<String> location;
  final Optional<GoogleCredentials> credentials;
  HttpOptions httpOptions;
  final boolean vertexAI;

  /** Constructs an ApiClient for Google AI APIs. */
  ApiClient(Optional<String> apiKey, Optional<HttpOptions> customHttpOptions) {
    checkNotNull(apiKey, "API Key cannot be null");
    checkNotNull(customHttpOptions, "customHttpOptions cannot be null");

    try {
      this.apiKey = Optional.of(apiKey.orElse(System.getenv("GOOGLE_API_KEY")));
    } catch (NullPointerException e) {
      throw new IllegalArgumentException(
          "API key must either be provided or set in the environment variable" + " GOOGLE_API_KEY.",
          e);
    }

    this.project = Optional.empty();
    this.location = Optional.empty();
    this.credentials = Optional.empty();
    this.vertexAI = false;

    this.httpOptions = defaultHttpOptions(/* vertexAI= */ false, this.location);

    if (customHttpOptions.isPresent()) {
      applyHttpOptions(customHttpOptions.get());
    }

    this.httpClient = createHttpClient(httpOptions.timeout());
  }

  ApiClient(
      Optional<String> project,
      Optional<String> location,
      Optional<GoogleCredentials> credentials,
      Optional<HttpOptions> customHttpOptions) {
    checkNotNull(project, "project cannot be null");
    checkNotNull(location, "location cannot be null");
    checkNotNull(credentials, "credentials cannot be null");
    checkNotNull(customHttpOptions, "customHttpOptions cannot be null");

    try {
      this.project = Optional.of(project.orElse(System.getenv("GOOGLE_CLOUD_PROJECT")));
    } catch (NullPointerException e) {
      throw new IllegalArgumentException(
          "Project must either be provided or set in the environment variable"
              + " GOOGLE_CLOUD_PROJECT.",
          e);
    }
    if (this.project.get().isEmpty()) {
      throw new IllegalArgumentException("Project must not be empty.");
    }

    try {
      this.location = Optional.of(location.orElse(System.getenv("GOOGLE_CLOUD_LOCATION")));
    } catch (NullPointerException e) {
      throw new IllegalArgumentException(
          "Location must either be provided or set in the environment variable"
              + " GOOGLE_CLOUD_LOCATION.",
          e);
    }
    if (this.location.get().isEmpty()) {
      throw new IllegalArgumentException("Location must not be empty.");
    }

    this.credentials = Optional.of(credentials.orElseGet(this::defaultCredentials));

    this.httpOptions = defaultHttpOptions(/* vertexAI= */ true, this.location);

    if (customHttpOptions.isPresent()) {
      applyHttpOptions(customHttpOptions.get());
    }
    this.apiKey = Optional.empty();
    this.vertexAI = true;
    this.httpClient = createHttpClient(httpOptions.timeout());
  }

  private OkHttpClient createHttpClient(Optional<Integer> timeout) {
    OkHttpClient.Builder builder = new OkHttpClient().newBuilder();
    if (timeout.isPresent()) {
      builder.connectTimeout(Duration.ofMillis(timeout.get()));
    }
    return builder.build();
  }

  /** Sends a Http request given the http method, path, and request json string. */
  public abstract ApiResponse request(String httpMethod, String path, String requestJson);

  /** Returns the library version. */
  static String libraryVersion() {
    // TODO: Automate revisions to the SDK library version.
    String libraryLabel = "google-genai-sdk/0.1.0";
    String languageLabel = "gl-java/" + JAVA_VERSION.value();
    return libraryLabel + " " + languageLabel;
  }

  /** Returns whether the client is using Vertex AI APIs. */
  public boolean vertexAI() {
    return vertexAI;
  }

  /** Returns the project ID for Vertex AI APIs. */
  public @Nullable String project() {
    return project.orElse(null);
  }

  /** Returns the location for Vertex AI APIs. */
  public @Nullable String location() {
    return location.orElse(null);
  }

  /** Returns the API key for Google AI APIs. */
  public @Nullable String apiKey() {
    return apiKey.orElse(null);
  }

  /** Returns the HttpClient for API calls. */
  OkHttpClient httpClient() {
    return httpClient;
  }

  private Optional<Map<String, String>> getTimeoutHeader(HttpOptions httpOptionsToApply) {
    if (httpOptionsToApply.timeout().isPresent()) {
      int timeoutInSeconds = (int) Math.ceil((double) httpOptionsToApply.timeout().get() / 1000.0);
      // TODO(b/329147724): Document the usage of X-Server-Timeout header.
      return Optional.of(ImmutableMap.of("X-Server-Timeout", Integer.toString(timeoutInSeconds)));
    }
    return Optional.empty();
  }

  private void applyHttpOptions(HttpOptions httpOptionsToApply) {
    HttpOptions.Builder mergedHttpOptionsBuilder = this.httpOptions.toBuilder();
    if (httpOptionsToApply.baseUrl().isPresent()) {
      mergedHttpOptionsBuilder.baseUrl(httpOptionsToApply.baseUrl().get());
    }
    if (httpOptionsToApply.apiVersion().isPresent()) {
      mergedHttpOptionsBuilder.apiVersion(httpOptionsToApply.apiVersion().get());
    }
    if (httpOptionsToApply.timeout().isPresent()) {
      mergedHttpOptionsBuilder.timeout(httpOptionsToApply.timeout().get());
    }
    if (httpOptionsToApply.headers().isPresent()) {
      ImmutableMap<String, String> mergedHeaders =
          ImmutableMap.<String, String>builder()
              .putAll(httpOptionsToApply.headers().orElse(ImmutableMap.of()))
              .putAll(this.httpOptions.headers().orElse(ImmutableMap.of()))
              .putAll(getTimeoutHeader(httpOptionsToApply).orElse(ImmutableMap.of()))
              .buildOrThrow();
      mergedHttpOptionsBuilder.headers(mergedHeaders);
    }
    this.httpOptions = mergedHttpOptionsBuilder.build();
  }

  static HttpOptions defaultHttpOptions(boolean vertexAI, Optional<String> location) {
    ImmutableMap.Builder<String, String> defaultHeaders = ImmutableMap.builder();
    defaultHeaders
        .put("Content-Type", "application/json")
        .put("user-agent", libraryVersion())
        .put("x-goog-api-client", libraryVersion());

    HttpOptions.Builder defaultHttpOptionsBuilder =
        HttpOptions.builder().headers(defaultHeaders.buildOrThrow());

    if (vertexAI && location.isPresent()) {
      defaultHttpOptionsBuilder
          .baseUrl(
              Ascii.equalsIgnoreCase(location.get(), "global")
                  ? "https://aiplatform.googleapis.com"
                  : String.format("https://%s-aiplatform.googleapis.com", location.get()))
          .apiVersion("v1beta1");
    } else if (vertexAI && location.isEmpty()) {
      throw new IllegalArgumentException("Location must be provided for Vertex AI APIs.");
    } else {
      defaultHttpOptionsBuilder
          .baseUrl("https://generativelanguage.googleapis.com")
          .apiVersion("v1beta");
    }
    return defaultHttpOptionsBuilder.build();
  }

  GoogleCredentials defaultCredentials() {
    try {
      return GoogleCredentials.getApplicationDefault()
          .createScoped("https://www.googleapis.com/auth/cloud-platform");
    } catch (IOException e) {
      throw new GenAiIOException(
          "Failed to get application default credentials, please explicitly provide credentials.",
          e);
    }
  }
}

package com.google.adk.tools;

import com.google.adk.agents.ReadonlyContext;
import io.reactivex.rxjava3.core.Flowable;
import java.util.List;
import java.util.Optional;

/** Base interface for toolsets. */
public interface BaseToolset extends AutoCloseable {

  /**
   * Return all tools in the toolset based on the provided context.
   *
   * @param readonlyContext Context used to filter tools available to the agent.
   * @return A Single emitting a list of tools available under the specified context.
   */
  Flowable<BaseTool> getTools(ReadonlyContext readonlyContext);

  /**
   * Performs cleanup and releases resources held by the toolset.
   *
   * <p>NOTE: This method is invoked, for example, at the end of an agent server's lifecycle or when
   * the toolset is no longer needed. Implementations should ensure that any open connections,
   * files, or other managed resources are properly released to prevent leaks.
   */
  @Override
  void close() throws Exception;

  /**
   * Helper method to be used by implementers that returns true if the given tool is in the provided
   * list of tools of if testing against the given ToolPredicate returns true (otherwise false).
   *
   * @param tool The tool to check.
   * @param toolFilter An Optional containing either a ToolPredicate or a List of tool names.
   * @param readonlyContext The current context.
   * @return true if the tool is selected.
   */
  default boolean isToolSelected(
      BaseTool tool, Optional<Object> toolFilter, Optional<ReadonlyContext> readonlyContext) {
    if (toolFilter.isEmpty()) {
      return true;
    }
    Object filter = toolFilter.get();
    if (filter instanceof ToolPredicate toolPredicate) {
      return toolPredicate.test(tool, readonlyContext);
    }
    if (filter instanceof List) {
      @SuppressWarnings("unchecked")
      List<String> toolNames = (List<String>) filter;
      return toolNames.contains(tool.name());
    }
    return false;
  }
}

/*
 * 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.agents;

import com.google.adk.events.EventActions;
import com.google.adk.sessions.State;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Maybe;
import java.util.Optional;

/** The context of various callbacks for an agent invocation. */
public class CallbackContext extends ReadonlyContext {

  protected EventActions eventActions;
  private final State state;

  /**
   * Initializes callback context.
   *
   * @param invocationContext Current invocation context.
   * @param eventActions Callback event actions.
   */
  public CallbackContext(InvocationContext invocationContext, EventActions eventActions) {
    super(invocationContext);
    this.eventActions = eventActions != null ? eventActions : EventActions.builder().build();
    this.state = new State(invocationContext.session().state(), this.eventActions.stateDelta());
  }

  /** Returns the delta-aware state of the current callback. */
  @Override
  public State state() {
    return state;
  }

  /** Returns the EventActions associated with this context. */
  public EventActions eventActions() {
    return eventActions;
  }

  /**
   * Loads an artifact from the artifact service associated with the current session.
   *
   * @param filename Artifact file name.
   * @param version Artifact version (optional).
   * @return loaded part, or empty if not found.
   * @throws IllegalStateException if the artifact service is not initialized.
   */
  public Maybe<Part> loadArtifact(String filename, Optional<Integer> version) {
    if (invocationContext.artifactService() == null) {
      throw new IllegalStateException("Artifact service is not initialized.");
    }
    return invocationContext
        .artifactService()
        .loadArtifact(
            invocationContext.appName(),
            invocationContext.userId(),
            invocationContext.session().id(),
            filename,
            version);
  }

  /**
   * Saves an artifact and records it as a delta for the current session.
   *
   * @param filename Artifact file name.
   * @param artifact Artifact content to save.
   * @throws IllegalStateException if the artifact service is not initialized.
   */
  public void saveArtifact(String filename, Part artifact) {
    if (invocationContext.artifactService() == null) {
      throw new IllegalStateException("Artifact service is not initialized.");
    }
    var unused =
        invocationContext
            .artifactService()
            .saveArtifact(
                invocationContext.appName(),
                invocationContext.userId(),
                invocationContext.session().id(),
                filename,
                artifact);
    this.eventActions.artifactDelta().put(filename, artifact);
  }
}

/*
 * 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;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.google.genai.types.Schema;
import com.google.genai.types.Type;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

/** Utility class for validating schemas. */
public final class SchemaUtils {

  private SchemaUtils() {} // Private constructor for utility class

  /**
   * Matches a value against a schema type.
   *
   * @param value The value to match.
   * @param schema The schema to match against.
   * @param isInput Whether the value is an input or output.
   * @return True if the value matches the schema type, false otherwise.
   * @throws IllegalArgumentException If the schema type is not supported.
   */
  @SuppressWarnings("unchecked") // For tool parameter type casting.
  private static Boolean matchType(Object value, Schema schema, Boolean isInput) {
    // Based on types from https://cloud.google.com/vertex-ai/docs/reference/rest/v1/Schema
    Type.Known type = schema.type().get().knownEnum();
    switch (type) {
      case STRING:
        return value instanceof String;
      case INTEGER:
        return value instanceof Integer;
      case BOOLEAN:
        return value instanceof Boolean;
      case NUMBER:
        return value instanceof Number;
      case ARRAY:
        if (value instanceof List) {
          for (Object element : (List<?>) value) {
            if (!matchType(element, schema.items().get(), isInput)) {
              return false;
            }
          }
          return true;
        }
        return false;
      case OBJECT:
        if (value instanceof Map) {
          validateMapOnSchema((Map<String, Object>) value, schema, isInput);
          return true;
        } else {
          return false;
        }
      case TYPE_UNSPECIFIED:
        throw new IllegalArgumentException(
            "Unsupported type: " + type + " is not a Open API data type.");
      default:
        // This category includes NULL, which is not supported.
        break;
    }
    return false;
  }

  /**
   * Validates a map against a schema.
   *
   * @param args The map to validate.
   * @param schema The schema to validate against.
   * @param isInput Whether the map is an input or output.
   * @throws IllegalArgumentException If the map does not match the schema.
   */
  public static void validateMapOnSchema(Map<String, Object> args, Schema schema, Boolean isInput) {
    Map<String, Schema> properties = schema.properties().get();
    for (Entry<String, Object> arg : args.entrySet()) {
      // Check if the argument is in the schema.
      if (!properties.containsKey(arg.getKey())) {
        if (isInput) {
          throw new IllegalArgumentException(
              "Input arg: " + arg.getKey() + " does not match agent input schema: " + schema);
        } else {
          throw new IllegalArgumentException(
              "Output arg: " + arg.getKey() + " does not match agent output schema: " + schema);
        }
      }
      // Check if the argument type matches the schema type.
      if (!matchType(arg.getValue(), properties.get(arg.getKey()), isInput)) {
        if (isInput) {
          throw new IllegalArgumentException(
              "Input arg: " + arg.getKey() + " does not match agent input schema: " + schema);
        } else {
          throw new IllegalArgumentException(
              "Output arg: " + arg.getKey() + " does not match agent output schema: " + schema);
        }
      }
    }
    // Check if all required arguments are present.
    if (schema.required().isPresent()) {
      for (String required : schema.required().get()) {
        if (!args.containsKey(required)) {
          if (isInput) {
            throw new IllegalArgumentException("Input args does not contain required " + required);
          } else {
            throw new IllegalArgumentException("Output args does not contain required " + required);
          }
        }
      }
    }
  }

  /**
   * Validates an output string against a schema.
   *
   * @param output The output string to validate.
   * @param schema The schema to validate against.
   * @return The output map.
   * @throws IllegalArgumentException If the output string does not match the schema.
   * @throws JsonProcessingException If the output string cannot be parsed.
   */
  @SuppressWarnings("unchecked") // For tool parameter type casting.
  public static Map<String, Object> validateOutputSchema(String output, Schema schema)
      throws JsonProcessingException {
    Map<String, Object> outputMap = JsonBaseModel.getMapper().readValue(output, HashMap.class);
    validateMapOnSchema(outputMap, schema, false);
    return outputMap;
  }
}

/*
 * 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.runner;

import com.google.adk.agents.BaseAgent;
import com.google.adk.artifacts.InMemoryArtifactService;
import com.google.adk.sessions.InMemorySessionService;

/** The class for the in-memory GenAi runner, using in-memory artifact and session services. */
public class InMemoryRunner extends Runner {

  public InMemoryRunner(BaseAgent agent) {
    // TODO: Change the default appName to InMemoryRunner to align with adk python.
    // Check the dev UI in case we break something there.
    this(agent, /* appName= */ agent.name());
  }

  public InMemoryRunner(BaseAgent agent, String appName) {
    super(agent, appName, new InMemoryArtifactService(), new InMemorySessionService());
  }
}

/*
 * 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.tools;

import java.lang.reflect.Method;

/** A function tool that returns the result asynchronously. */
public class LongRunningFunctionTool extends FunctionTool {

  public static LongRunningFunctionTool create(Method func) {
    return new LongRunningFunctionTool(func);
  }

  public static LongRunningFunctionTool create(Class<?> cls, String methodName) {
    for (Method method : cls.getMethods()) {
      if (method.getName().equals(methodName)) {
        return create(method);
      }
    }
    throw new IllegalArgumentException(
        String.format("Method %s not found in class %s.", methodName, cls.getName()));
  }

  public static LongRunningFunctionTool create(Object instance, String methodName) {
    Class<?> cls = instance.getClass();
    for (Method method : cls.getMethods()) {
      if (method.getName().equals(methodName)) {
        return new LongRunningFunctionTool(instance, method);
      }
    }
    throw new IllegalArgumentException(
        String.format("Method %s not found in class %s.", methodName, cls.getName()));
  }

  private LongRunningFunctionTool(Method func) {
    super(null, func, /* isLongRunning= */ true);
  }

  private LongRunningFunctionTool(Object instance, Method func) {
    super(instance, func, /* isLongRunning= */ true);
  }
}

/*
 * 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.google.adk.events.Event;
import com.google.adk.events.EventActions;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ConcurrentMap;
import javax.annotation.Nullable;

/**
 * Defines the contract for managing {@link Session}s and their associated {@link Event}s. Provides
 * methods for creating, retrieving, listing, and deleting sessions, as well as listing and
 * appending events to a session. Implementations of this interface handle the underlying storage
 * and retrieval logic.
 */
public interface BaseSessionService {

  /**
   * Creates a new session with the specified parameters.
   *
   * @param appName The name of the application associated with the session.
   * @param userId The identifier for the user associated with the session.
   * @param state An optional map representing the initial state of the session. Can be null or
   *     empty.
   * @param sessionId An optional client-provided identifier for the session. If empty or null, the
   *     service should generate a unique ID.
   * @return The newly created {@link Session} instance.
   * @throws SessionException if creation fails.
   */
  Single<Session> createSession(
      String appName,
      String userId,
      @Nullable ConcurrentMap<String, Object> state,
      @Nullable String sessionId);

  /**
   * 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.
   *
   *
This is a shortcut for {@link #createSession(String, String, Map, String)} with null state
   * and a null session ID.
   *
   * @param appName The name of the application associated with the session.
   * @param userId The identifier for the user associated with the session.
   * @return The newly created {@link Session} instance.
   * @throws SessionException if creation fails.
   */
  default Single<Session> createSession(String appName, String userId) {
    return createSession(appName, userId, null, null);
  }

  /**
   * Retrieves a specific session, optionally filtering the events included.
   *
   * @param appName The name of the application.
   * @param userId The identifier of the user.
   * @param sessionId The unique identifier of the session to retrieve.
   * @param config Optional configuration to filter the events returned within the session (e.g.,
   *     limit number of recent events, filter by timestamp). If empty, default retrieval behavior
   *     is used (potentially all events or a service-defined limit).
   * @return An {@link Optional} containing the {@link Session} if found, otherwise {@link
   *     Optional#empty()}.
   * @throws SessionException for retrieval errors other than not found.
   */
  Maybe<Session> getSession(
      String appName, String userId, String sessionId, Optional<GetSessionConfig> config);

  /**
   * Lists sessions associated with a specific application and user.
   *
   *
The {@link Session} objects in the response typically contain only metadata (like ID,
   * creation time) and not the full event list or state to optimize performance.
   *
   * @param appName The name of the application.
   * @param userId The identifier of the user whose sessions are to be listed.
   * @return A {@link ListSessionsResponse} containing a list of matching sessions.
   * @throws SessionException if listing fails.
   */
  Single<ListSessionsResponse> listSessions(String appName, String userId);

  /**
   * Deletes a specific session.
   *
   * @param appName The name of the application.
   * @param userId The identifier of the user.
   * @param sessionId The unique identifier of the session to delete.
   * @throws SessionNotFoundException if the session doesn't exist.
   * @throws SessionException for other deletion errors.
   */
  Completable deleteSession(String appName, String userId, String sessionId);

  /**
   * Lists the events within a specific session. Supports pagination via the response object.
   *
   * @param appName The name of the application.
   * @param userId The identifier of the user.
   * @param sessionId The unique identifier of the session whose events are to be listed.
   * @return A {@link ListEventsResponse} containing a list of events and an optional token for
   *     retrieving the next page.
   * @throws SessionNotFoundException if the session doesn't exist.
   * @throws SessionException for other listing errors.
   */
  Single<ListEventsResponse> listEvents(String appName, String userId, String sessionId);

  /**
   * 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.
   *
   * @param session The session object to close.
   */
  default Completable closeSession(Session session) {
    // Default implementation does nothing.
    // TODO: Determine whether we want to finalize the session here.
    return Completable.complete();
  }

  /**
   * Appends an event to an in-memory session object and updates the session's state based on the
   * event's state delta, if applicable.
   *
   *
This method primarily modifies the passed {@code session} object in memory. Persisting these
   * changes typically requires a separate call to an update/save method provided by the specific
   * service implementation, or might happen implicitly depending on the implementation's design.
   *
   *
If the event is marked as partial (e.g., {@code event.isPartial() == true}), it is returned
   * directly without modifying the session state or event list. State delta keys starting with
   * {@link State#TEMP_PREFIX} are ignored during state updates.
   *
   * @param session The {@link Session} object to which the event should be appended (will be
   *     mutated).
   * @param event The {@link Event} to append.
   * @return The appended {@link Event} instance (or the original event if it was partial).
   * @throws NullPointerException if session or event is null.
   */
  @CanIgnoreReturnValue
  default Single<Event> appendEvent(Session session, Event event) {
    Objects.requireNonNull(session, "session cannot be null");
    Objects.requireNonNull(event, "event cannot be null");

    // If the event indicates it's partial or incomplete, don't process it yet.
    if (event.partial().orElse(false)) {
      return Single.just(event);
    }

    EventActions actions = event.actions();
    if (actions != null) {
      ConcurrentMap<String, Object> stateDelta = actions.stateDelta();
      if (stateDelta != null && !stateDelta.isEmpty()) {
        ConcurrentMap<String, Object> sessionState = session.state();
        if (sessionState != null) {
          stateDelta.forEach(
              (key, value) -> {
                if (!key.startsWith(State.TEMP_PREFIX)) {
                  sessionState.put(key, value);
                }
              });
        }
      }
    }

    List<Event> sessionEvents = session.events();
    if (sessionEvents != null) {
      sessionEvents.add(event);
    }

    return Single.just(event);
  }
}

/*
 * 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.artifacts;

import static com.google.common.collect.ImmutableList.toImmutableList;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Streams;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.IntStream;

/** An in-memory implementation of the {@link BaseArtifactService}. */
public final class InMemoryArtifactService implements BaseArtifactService {
  private final Map<String, Map<String, Map<String, Map<String, List<Part>>>>> artifacts;

  public InMemoryArtifactService() {
    this.artifacts = new HashMap<>();
  }

  /**
   * Saves an artifact in memory and assigns a new version.
   *
   * @return Single with assigned version number.
   */
  @Override
  public Single<Integer> saveArtifact(
      String appName, String userId, String sessionId, String filename, Part artifact) {
    List<Part> versions =
        artifacts
            .computeIfAbsent(appName, k -> new HashMap<>())
            .computeIfAbsent(userId, k -> new HashMap<>())
            .computeIfAbsent(sessionId, k -> new HashMap<>())
            .computeIfAbsent(filename, k -> new ArrayList<>());
    versions.add(artifact);
    return Single.just(versions.size() - 1);
  }

  /**
   * Loads an artifact by version or latest.
   *
   * @return Maybe with the artifact, or empty if not found.
   */
  @Override
  public Maybe<Part> loadArtifact(
      String appName, String userId, String sessionId, String filename, Optional<Integer> version) {
    List<Part> versions =
        artifacts
            .getOrDefault(appName, new HashMap<>())
            .getOrDefault(userId, new HashMap<>())
            .getOrDefault(sessionId, new HashMap<>())
            .getOrDefault(filename, new ArrayList<>());

    if (versions.isEmpty()) {
      return Maybe.empty();
    }
    if (version.isPresent()) {
      int v = version.get();
      if (v >= 0 && v < versions.size()) {
        return Maybe.just(versions.get(v));
      } else {
        return Maybe.empty();
      }
    } else {
      return Maybe.fromOptional(Streams.findLast(versions.stream()));
    }
  }

  /**
   * Lists filenames of stored artifacts for the session.
   *
   * @return Single with list of artifact filenames.
   */
  @Override
  public Single<ListArtifactsResponse> listArtifactKeys(
      String appName, String userId, String sessionId) {
    return Single.just(
        ListArtifactsResponse.builder()
            .filenames(
                ImmutableList.copyOf(
                    artifacts
                        .getOrDefault(appName, new HashMap<>())
                        .getOrDefault(userId, new HashMap<>())
                        .getOrDefault(sessionId, new HashMap<>())
                        .keySet()))
            .build());
  }

  /**
   * Deletes all versions of the given artifact.
   *
   * @return Completable indicating completion.
   */
  @Override
  public Completable deleteArtifact(
      String appName, String userId, String sessionId, String filename) {
    artifacts
        .getOrDefault(appName, new HashMap<>())
        .getOrDefault(userId, new HashMap<>())
        .getOrDefault(sessionId, new HashMap<>())
        .remove(filename);
    return Completable.complete();
  }

  /**
   * Lists all versions of the specified artifact.
   *
   * @return Single with list of version numbers.
   */
  @Override
  public Single<ImmutableList<Integer>> listVersions(
      String appName, String userId, String sessionId, String filename) {
    int size =
        artifacts
            .getOrDefault(appName, new HashMap<>())
            .getOrDefault(userId, new HashMap<>())
            .getOrDefault(sessionId, new HashMap<>())
            .getOrDefault(filename, new ArrayList<>())
            .size();
    if (size == 0) {
      return Single.just(ImmutableList.of());
    }
    return Single.just(IntStream.range(0, size).boxed().collect(toImmutableList()));
  }
}

/*
 * 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.models;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
import com.fasterxml.jackson.databind.annotation.JsonPOJOBuilder;
import com.google.adk.JsonBaseModel;
import com.google.auto.value.AutoValue;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.genai.types.Candidate;
import com.google.genai.types.Content;
import com.google.genai.types.FinishReason;
import com.google.genai.types.GenerateContentResponse;
import com.google.genai.types.GenerateContentResponsePromptFeedback;
import com.google.genai.types.GroundingMetadata;
import java.util.List;
import java.util.Optional;
import javax.annotation.Nullable;

/** Represents a response received from the LLM. */
@AutoValue
@JsonDeserialize(builder = LlmResponse.Builder.class)
public abstract class LlmResponse extends JsonBaseModel {

  LlmResponse() {}

  /**
   * Returns the content of the first candidate in the response, if available.
   *
   * @return An {@link Content} of the first {@link Candidate} in the {@link
   *     GenerateContentResponse} if the response contains at least one candidate., or an empty
   *     optional if no candidates are present in the response.
   */
  @JsonProperty("content")
  public abstract Optional<Content> content();

  /**
   * Returns the grounding metadata of the first candidate in the response, if available.
   *
   * @return An {@link Optional} containing {@link GroundingMetadata} or empty.
   */
  @JsonProperty("groundingMetadata")
  public abstract Optional<GroundingMetadata> groundingMetadata();

  /**
   * Indicates whether the text content is part of a unfinished text stream.
   *
   *
Only used for streaming mode and when the content is plain text.
   */
  @JsonProperty("partial")
  public abstract Optional<Boolean> partial();

  /**
   * Indicates whether the response from the model is complete.
   *
   *
Only used for streaming mode.
   */
  @JsonProperty("turnComplete")
  public abstract Optional<Boolean> turnComplete();

  /** Error code if the response is an error. Code varies by model. */
  @JsonProperty("errorCode")
  public abstract Optional<FinishReason> errorCode();

  /** Error message if the response is an error. */
  @JsonProperty("errorMessage")
  public abstract Optional<String> errorMessage();

  /**
   * Indicates that LLM was interrupted when generating the content. Usually it's due to user
   * interruption during a bidi streaming.
   */
  @JsonProperty("interrupted")
  public abstract Optional<Boolean> interrupted();

  public abstract Builder toBuilder();

  /** Builder for constructing {@link LlmResponse} instances. */
  @AutoValue.Builder
  @JsonPOJOBuilder(buildMethodName = "build", withPrefix = "")
  public abstract static class Builder {

    @JsonCreator
    static LlmResponse.Builder jacksonBuilder() {
      return LlmResponse.builder();
    }

    @JsonProperty("content")
    public abstract Builder content(Content content);

    @JsonProperty("interrupted")
    public abstract Builder interrupted(@Nullable Boolean interrupted);

    public abstract Builder interrupted(Optional<Boolean> interrupted);

    @JsonProperty("groundingMetadata")
    public abstract Builder groundingMetadata(@Nullable GroundingMetadata groundingMetadata);

    public abstract Builder groundingMetadata(Optional<GroundingMetadata> groundingMetadata);

    @JsonProperty("partial")
    public abstract Builder partial(@Nullable Boolean partial);

    public abstract Builder partial(Optional<Boolean> partial);

    @JsonProperty("turnComplete")
    public abstract Builder turnComplete(@Nullable Boolean turnComplete);

    public abstract Builder turnComplete(Optional<Boolean> turnComplete);

    @JsonProperty("errorCode")
    public abstract Builder errorCode(@Nullable FinishReason errorCode);

    public abstract Builder errorCode(Optional<FinishReason> errorCode);

    @JsonProperty("errorMessage")
    public abstract Builder errorMessage(@Nullable String errorMessage);

    public abstract Builder errorMessage(Optional<String> errorMessage);

    @CanIgnoreReturnValue
    public final Builder response(GenerateContentResponse response) {
      Optional<List<Candidate>> candidatesOpt = response.candidates();
      if (candidatesOpt.isPresent() && !candidatesOpt.get().isEmpty()) {
        Candidate candidate = candidatesOpt.get().get(0);
        if (candidate.content().isPresent()) {
          this.content(candidate.content().get());
          this.groundingMetadata(candidate.groundingMetadata());
        } else {
          candidate.finishReason().ifPresent(this::errorCode);
          candidate.finishMessage().ifPresent(this::errorMessage);
        }
      } else {
        Optional<GenerateContentResponsePromptFeedback> promptFeedbackOpt =
            response.promptFeedback();
        if (promptFeedbackOpt.isPresent()) {
          GenerateContentResponsePromptFeedback promptFeedback = promptFeedbackOpt.get();
          promptFeedback
              .blockReason()
              .ifPresent(reason -> this.errorCode(new FinishReason(reason.toString())));
          promptFeedback.blockReasonMessage().ifPresent(this::errorMessage);
        } else {
          this.errorCode(new FinishReason("Unknown error."));
          this.errorMessage("Unknown error.");
        }
      }
      return this;
    }

    abstract LlmResponse autoBuild();

    public LlmResponse build() {
      return autoBuild();
    }
  }

  public static Builder builder() {
    return new AutoValue_LlmResponse.Builder();
  }

  public static LlmResponse create(List<Candidate> candidates) {
    GenerateContentResponse response =
        GenerateContentResponse.builder().candidates(candidates).build();
    return builder().response(response).build();
  }

  public static LlmResponse create(GenerateContentResponse response) {
    return builder().response(response).build();
  }
}

/*
 * 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.models;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/** Central registry for managing Large Language Model (LLM) instances. */
public final class LlmRegistry {

  /** A thread-safe cache mapping model names to LLM instances. */
  private static final Map<String, BaseLlm> instances = new ConcurrentHashMap<>();

  /** The factory interface for creating LLM instances. */
  @FunctionalInterface
  public interface LlmFactory {
    BaseLlm create(String modelName);
  }

  /** Map of model name patterns regex to factories. */
  private static final Map<String, LlmFactory> llmFactories = new ConcurrentHashMap<>();

  /** Registers default LLM factories, e.g. for Gemini models. */
  static {
    registerLlm("gemini-.*", modelName -> Gemini.builder().modelName(modelName).build());
  }

  /**
   * Registers a factory for model names matching the given regex pattern.
   *
   * @param modelNamePattern Regex pattern for matching model names.
   * @param factory Factory to create LLM instances.
   */
  public static void registerLlm(String modelNamePattern, LlmFactory factory) {
    llmFactories.put(modelNamePattern, factory);
  }

  /**
   * Returns an LLM instance for the given model name, using a cached or new factory-created
   * instance.
   *
   * @param modelName Model name to look up.
   * @return Matching {@link BaseLlm} instance.
   * @throws IllegalArgumentException If no factory matches the model name.
   */
  public static BaseLlm getLlm(String modelName) {
    return instances.computeIfAbsent(modelName, LlmRegistry::createLlm);
  }

  /**
   * Creates a {@link BaseLlm} by matching the model name against registered factories.
   *
   * @param modelName Model name to match.
   * @return A new {@link BaseLlm} instance.
   * @throws IllegalArgumentException If no factory matches the model name.
   */
  private static BaseLlm createLlm(String modelName) {
    for (Map.Entry<String, LlmFactory> entry : llmFactories.entrySet()) {
      if (modelName.matches(entry.getKey())) {
        return entry.getValue().create(modelName);
      }
    }
    throw new IllegalArgumentException("Unsupported model: " + modelName);
  }

  /**
   * Registers an LLM factory for testing purposes. Clears cached instances matching the given
   * pattern to ensure test isolation.
   *
   * @param modelNamePattern Regex pattern for matching model names.
   * @param factory The {@link LlmFactory} to register.
   */
  static void registerTestLlm(String modelNamePattern, LlmFactory factory) {
    llmFactories.put(modelNamePattern, factory);
    // Clear any cached instances that match this pattern to ensure test isolation.
    instances.keySet().removeIf(modelName -> modelName.matches(modelNamePattern));
  }

  private LlmRegistry() {}
}

/*
 * 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.flows.llmflows;

import com.google.common.collect.ImmutableList;
import java.util.List;
import java.util.Optional;

/** Basic LLM flow with fixed request processors and no response post-processing. */
public class SingleFlow extends BaseLlmFlow {

  protected static final ImmutableList<RequestProcessor> REQUEST_PROCESSORS =
      ImmutableList.of(
          new Basic(), new Instructions(), new Identity(), new Contents(), new Examples());

  protected static final ImmutableList<ResponseProcessor> RESPONSE_PROCESSORS = ImmutableList.of();

  public SingleFlow() {
    this(/* maxSteps= */ Optional.empty());
  }

  public SingleFlow(Optional<Integer> maxSteps) {
    this(REQUEST_PROCESSORS, RESPONSE_PROCESSORS, maxSteps);
  }

  protected SingleFlow(
      List<RequestProcessor> requestProcessors,
      List<ResponseProcessor> responseProcessors,
      Optional<Integer> maxSteps) {
    super(requestProcessors, responseProcessors, maxSteps);
  }
}

/*
 * 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.google.common.collect.ImmutableList;
import com.google.genai.types.Blob;
import com.google.genai.types.Content;
import com.google.genai.types.Part;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.Optional;

/** Utility functions for session service. */
public final class SessionUtils {

  public SessionUtils() {}

  /** Base64-encodes inline blobs in content. */
  public static Content encodeContent(Content content) {
    List<Part> encodedParts = new ArrayList<>();
    for (Part part : content.parts().orElse(ImmutableList.of())) {
      boolean isInlineDataPresent = false;
      if (part.inlineData() != null) {
        Optional<Blob> inlineDataOptional = part.inlineData();
        if (inlineDataOptional.isPresent()) {
          Blob inlineDataBlob = inlineDataOptional.get();
          Optional<byte[]> dataOptional = inlineDataBlob.data();
          if (dataOptional.isPresent()) {
            byte[] dataBytes = dataOptional.get();
            byte[] encodedData = Base64.getEncoder().encode(dataBytes);
            encodedParts.add(
                part.toBuilder().inlineData(Blob.builder().data(encodedData).build()).build());
            isInlineDataPresent = true;
          }
        }
      }
      if (!isInlineDataPresent) {
        encodedParts.add(part);
      }
    }
    return toContent(encodedParts, content.role());
  }

  /** Decodes Base64-encoded inline blobs in content. */
  public static Content decodeContent(Content content) {
    List<Part> decodedParts = new ArrayList<>();
    for (Part part : content.parts().orElse(ImmutableList.of())) {
      boolean isInlineDataPresent = false;
      if (part.inlineData() != null) {
        Optional<Blob> inlineDataOptional = part.inlineData();
        if (inlineDataOptional.isPresent()) {
          Blob inlineDataBlob = inlineDataOptional.get();
          Optional<byte[]> dataOptional = inlineDataBlob.data();
          if (dataOptional.isPresent()) {
            byte[] dataBytes = dataOptional.get();
            byte[] decodedData = Base64.getDecoder().decode(dataBytes);
            decodedParts.add(
                part.toBuilder().inlineData(Blob.builder().data(decodedData).build()).build());
            isInlineDataPresent = true;
          }
        }
      }
      if (!isInlineDataPresent) {
        decodedParts.add(part);
      }
    }
    return toContent(decodedParts, content.role());
  }

  /** Builds content from parts and optional role. */
  private static Content toContent(List<Part> parts, Optional<String> role) {
    Content.Builder contentBuilder = Content.builder().parts(parts);
    role.ifPresent(contentBuilder::role);
    return contentBuilder.build();
  }
}

/*
 * 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.agents;

import io.reactivex.rxjava3.core.Single;
import java.util.function.Function;

/**
 * Represents an instruction that can be provided to an agent to guide its behavior.
 *
 *
In the instructions, you should describe concisely what the agent will do, when it should
 * defer to other agents/tools, and how it should respond to the user.
 *
 *
Templating is supported using placeholders like {@code {variable_name}} or {@code
 * {artifact.artifact_name}}. These are replaced with values from the agent's session state or
 * loaded artifacts, respectively. For example, an instruction like {@code "Translate the following
 * text to {language}: {user_query}"} would substitute {@code {language}} and {@code {user_query}}
 * with their corresponding values from the session state.
 *
 *
Instructions can also be dynamically constructed using {@link Instruction.Provider}. This
 * allows for more complex logic where the instruction text is generated based on the current {@link
 * ReadonlyContext}. Additionally, an instruction could be built to include specific information
 * based on based on some external factors fetched during the Provider call like the current time,
 * the result of some API call, etc.
 */
public sealed interface Instruction permits Instruction.Static, Instruction.Provider {
  /** Plain instruction directly provided to the agent. */
  record Static(String instruction) implements Instruction {}

  /** Returns an instruction dynamically constructed from the given context. */
  record Provider(Function<ReadonlyContext, Single<String>> getInstruction)
      implements Instruction {}
}

/*
 * 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.agents;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
import com.fasterxml.jackson.databind.annotation.JsonPOJOBuilder;
import com.google.adk.JsonBaseModel;
import com.google.auto.value.AutoValue;
import com.google.common.base.Preconditions;
import com.google.genai.types.Blob;
import com.google.genai.types.Content;
import java.util.Optional;
import javax.annotation.Nullable;

/** Represents a request to be sent to a live connection to the LLM model. */
@AutoValue
@JsonDeserialize(builder = LiveRequest.Builder.class)
public abstract class LiveRequest extends JsonBaseModel {

  LiveRequest() {}

  /**
   * Returns the content of the request.
   *
   *
If set, send the content to the model in turn-by-turn mode.
   *
   * @return An optional {@link Content} object containing the content of the request.
   */
  @JsonProperty("content")
  public abstract Optional<Content> content();

  /**
   * Returns the blob of the request.
   *
   *
If set, send the blob to the model in realtime mode.
   *
   * @return An optional {@link Blob} object containing the blob of the request.
   */
  @JsonProperty("blob")
  public abstract Optional<Blob> blob();

  /**
   * Returns whether the connection should be closed.
   *
   *
If set to true, the connection will be closed after the request is sent.
   *
   * @return A boolean indicating whether the connection should be closed.
   */
  @JsonProperty("close")
  public abstract Optional<Boolean> close();

  /** Extracts boolean value from the close field or returns false if unset. */
  public boolean shouldClose() {
    return close().orElse(false);
  }

  /** Builder for constructing {@link LiveRequest} instances. */
  @AutoValue.Builder
  @JsonPOJOBuilder(buildMethodName = "build", withPrefix = "")
  public abstract static class Builder {
    @JsonProperty("content")
    public abstract Builder content(@Nullable Content content);

    public abstract Builder content(Optional<Content> content);

    @JsonProperty("blob")
    public abstract Builder blob(@Nullable Blob blob);

    public abstract Builder blob(Optional<Blob> blob);

    @JsonProperty("close")
    public abstract Builder close(@Nullable Boolean close);

    public abstract Builder close(Optional<Boolean> close);

    abstract LiveRequest autoBuild();

    public final LiveRequest build() {
      LiveRequest request = autoBuild();
      Preconditions.checkState(
          request.content().isPresent()
              || request.blob().isPresent()
              || request.close().isPresent(),
          "One of content, blob, or close must be set");
      return request;
    }
  }

  public static Builder builder() {
    return new AutoValue_LiveRequest.Builder().close(false);
  }

  public abstract Builder toBuilder();

  /** Deserializes a Json string to a {@link LiveRequest} object. */
  public static LiveRequest fromJsonString(String json) {
    return JsonBaseModel.fromJsonString(json, LiveRequest.class);
  }

  @JsonCreator
  static LiveRequest.Builder jacksonBuilder() {
    return LiveRequest.builder();
  }
}

/*
 * 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.agents;

import com.google.adk.events.Event;
import com.google.genai.types.Content;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;

/** Provides read-only access to the context of an agent run. */
public class ReadonlyContext {

  protected final InvocationContext invocationContext;
  private List<Event> eventsView;
  private Map<String, Object> stateView;

  public ReadonlyContext(InvocationContext invocationContext) {
    this.invocationContext = invocationContext;
  }

  /** Returns the user content that initiated this invocation. */
  public Optional<Content> userContent() {
    return invocationContext.userContent();
  }

  /** Returns the ID of the current invocation. */
  public String invocationId() {
    return invocationContext.invocationId();
  }

  /** Returns the branch of the current invocation, if present. */
  public Optional<String> branch() {
    return invocationContext.branch();
  }

  /** Returns the name of the agent currently running. */
  public String agentName() {
    return invocationContext.agent().name();
  }

  /** Returns the session ID. */
  public String sessionId() {
    return invocationContext.session().id();
  }

  /**
   * Returns an unmodifiable view of the events of the session.
   *
   *
Warning: This is a live view, not a snapshot.
   */
  public List<Event> events() {
    if (eventsView == null) {
      eventsView = Collections.unmodifiableList(invocationContext.session().events());
    }
    return eventsView;
  }

  /**
   * Returns an unmodifiable view of the state of the session.
   *
   *
Warning: This is a live view, not a snapshot.
   */
  public Map<String, Object> state() {
    if (stateView == null) {
      stateView = Collections.unmodifiableMap(invocationContext.session().state());
    }
    return stateView;
  }
}

/*
 * 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.agents;

import com.google.genai.types.Blob;
import com.google.genai.types.Content;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.processors.FlowableProcessor;
import io.reactivex.rxjava3.processors.MulticastProcessor;

/** A queue of live requests to be sent to the model. */
public final class LiveRequestQueue {
  private final FlowableProcessor<LiveRequest> processor;

  public LiveRequestQueue() {
    MulticastProcessor<LiveRequest> processor = MulticastProcessor.<LiveRequest>create();
    processor.start();
    this.processor = processor.toSerialized();
  }

  public void close() {
    processor.onNext(LiveRequest.builder().close(true).build());
    processor.onComplete();
  }

  public void content(Content content) {
    processor.onNext(LiveRequest.builder().content(content).build());
  }

  public void realtime(Blob blob) {
    processor.onNext(LiveRequest.builder().blob(blob).build());
  }

  public void send(LiveRequest request) {
    processor.onNext(request);
    if (request.shouldClose()) {
      processor.onComplete();
    }
  }

  public Flowable<LiveRequest> get() {
    return processor;
  }
}

/*
 * 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;

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.PropertyNamingStrategies;
import com.fasterxml.jackson.datatype.jdk8.Jdk8Module;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;

/** The base class for the types that needs JSON serialization/deserialization capability. */
public abstract class JsonBaseModel {

  private static final ObjectMapper objectMapper = new ObjectMapper();

  static {
    objectMapper
        .setSerializationInclusion(JsonInclude.Include.NON_ABSENT)
        .setPropertyNamingStrategy(PropertyNamingStrategies.LOWER_CAMEL_CASE)
        .registerModule(new Jdk8Module())
        .registerModule(new JavaTimeModule()) // TODO: echo sec module replace, locale
        .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
  }

  /** Serializes an object to a Json string. */
  protected static String toJsonString(Object object) {
    try {
      return objectMapper.writeValueAsString(object);
    } catch (JsonProcessingException e) {
      throw new IllegalStateException(e);
    }
  }

  public static ObjectMapper getMapper() {
    return JsonBaseModel.objectMapper;
  }

  public String toJson() {
    return toJsonString(this);
  }

  /** Serializes an object to a JsonNode. */
  protected static JsonNode toJsonNode(Object object) {
    return objectMapper.valueToTree(object);
  }

  /** Deserializes a Json string to an object of the given type. */
  public static <T extends JsonBaseModel> T fromJsonString(String jsonString, Class<T> clazz) {
    try {
      return objectMapper.readValue(jsonString, clazz);
    } catch (JsonProcessingException e) {
      throw new IllegalStateException(e);
    }
  }

  /** Deserializes a JsonNode to an object of the given type. */
  public static <T extends JsonBaseModel> T fromJsonNode(JsonNode jsonNode, Class<T> clazz) {
    try {
      return objectMapper.treeToValue(jsonNode, clazz);
    } catch (JsonProcessingException e) {
      throw new IllegalStateException(e);
    }
  }
}

/*
 * 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.fasterxml.jackson.databind.annotation.JsonDeserialize;
import com.google.adk.JsonBaseModel;
import com.google.adk.events.Event;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/** A {@link Session} object that encapsulates the {@link State} and {@link Event}s of a session. */
@JsonDeserialize(builder = Session.Builder.class)
public final class Session extends JsonBaseModel {
  private final String id;

  private final String appName;

  private final String userId;

  private final State state;

  private final List<Event> events;

  private Instant lastUpdateTime;

  public static Builder builder(String id) {
    return new Builder(id);
  }

  /** Builder for {@link Session}. */
  public static final class Builder {
    private String id;
    private String appName;
    private String userId;
    private State state = new State(new ConcurrentHashMap<>());
    private List<Event> events = new ArrayList<>();
    private Instant lastUpdateTime = Instant.EPOCH;

    public Builder(String id) {
      this.id = id;
    }

    @JsonCreator
    private Builder() {}

    @CanIgnoreReturnValue
    @JsonProperty("id")
    public Builder id(String id) {
      this.id = id;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder state(State state) {
      this.state = state;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("state")
    public Builder state(ConcurrentMap<String, Object> state) {
      this.state = new State(state);
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("appName")
    public Builder appName(String appName) {
      this.appName = appName;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("userId")
    public Builder userId(String userId) {
      this.userId = userId;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("events")
    public Builder events(List<Event> events) {
      this.events = events;
      return this;
    }

    @CanIgnoreReturnValue
    public Builder lastUpdateTime(Instant lastUpdateTime) {
      this.lastUpdateTime = lastUpdateTime;
      return this;
    }

    @CanIgnoreReturnValue
    @JsonProperty("lastUpdateTime")
    public Builder lastUpdateTimeSeconds(double seconds) {
      long secs = (long) seconds;
      // Convert fractional part to nanoseconds
      long nanos = (long) ((seconds - secs) * Duration.ofSeconds(1).toNanos());
      this.lastUpdateTime = Instant.ofEpochSecond(secs, nanos);
      return this;
    }

    public Session build() {
      if (id == null) {
        throw new IllegalStateException("Session id is null");
      }
      return new Session(appName, userId, id, state, events, lastUpdateTime);
    }
  }

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

  @JsonProperty("state")
  public ConcurrentMap<String, Object> state() {
    return state;
  }

  @JsonProperty("events")
  public List<Event> events() {
    return events;
  }

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

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

  public void lastUpdateTime(Instant lastUpdateTime) {
    this.lastUpdateTime = lastUpdateTime;
  }

  public Instant lastUpdateTime() {
    return lastUpdateTime;
  }

  @JsonProperty("lastUpdateTime")
  public double getLastUpdateTimeAsDouble() {
    if (lastUpdateTime == null) {
      return 0.0;
    }
    long seconds = lastUpdateTime.getEpochSecond();
    int nanos = lastUpdateTime.getNano();
    return seconds + nanos / (double) Duration.ofSeconds(1).toNanos();
  }

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

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

  private Session(
      String appName,
      String userId,
      String id,
      State state,
      List<Event> events,
      Instant lastUpdateTime) {
    this.id = id;
    this.appName = appName;
    this.userId = userId;
    this.state = state;
    this.events = events;
    this.lastUpdateTime = lastUpdateTime;
  }
}

/*
 * 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.tools.retrieval;

import com.google.adk.tools.BaseTool;
import com.google.genai.types.FunctionDeclaration;
import com.google.genai.types.Schema;
import java.util.Collections;
import java.util.Optional;

/** Base class for retrieval tools. */
public abstract class BaseRetrievalTool extends BaseTool {
  public BaseRetrievalTool(String name, String description) {
    super(name, description);
  }

  public BaseRetrievalTool(String name, String description, boolean isLongRunning) {
    super(name, description, isLongRunning);
  }

  @Override
  public Optional<FunctionDeclaration> declaration() {
    Schema querySchema =
        Schema.builder().type("STRING").description("The query to retrieve.").build();
    Schema parametersSchema =
        Schema.builder()
            .type("OBJECT")
            .properties(Collections.singletonMap("query", querySchema))
            .build();

    return Optional.of(
        FunctionDeclaration.builder()
            .name(this.name())
            .description(this.description())
            .parameters(parametersSchema)
            .build());
  }
}

package com.google.adk.tools.mcp;

import com.google.common.collect.ImmutableMap;
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
import io.modelcontextprotocol.client.transport.ServerParameters;
import io.modelcontextprotocol.client.transport.StdioClientTransport;
import io.modelcontextprotocol.spec.McpClientTransport;
import java.util.Collection;
import java.util.Optional;

/**
 * The default builder for creating MCP client transports. Supports StdioClientTransport based on
 * {@link ServerParameters} and the standard HttpClientSseClientTransport based on {@link
 * SseServerParameters}.
 */
public class DefaultMcpTransportBuilder implements McpTransportBuilder {

  @Override
  public McpClientTransport build(Object connectionParams) {
    if (connectionParams instanceof ServerParameters serverParameters) {
      return new StdioClientTransport(serverParameters);
    } else if (connectionParams instanceof SseServerParameters sseServerParams) {
      return HttpClientSseClientTransport.builder(sseServerParams.url())
          .sseEndpoint("sse")
          .customizeRequest(
              builder ->
                  Optional.ofNullable(sseServerParams.headers())
                      .map(ImmutableMap::entrySet)
                      .stream()
                      .flatMap(Collection::stream)
                      .forEach(
                          entry ->
                              builder.header(
                                  entry.getKey(),
                                  Optional.ofNullable(entry.getValue())
                                      .map(Object::toString)
                                      .orElse(""))))
          .build();
    } else {
      throw new IllegalArgumentException(
          "DefaultMcpTransportBuilder supports only ServerParameters or SseServerParameters, but"
              + " got "
              + connectionParams.getClass().getName());
    }
  }
}

/*
 * 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.tools.mcp;

import com.google.adk.tools.BaseTool;
import com.google.genai.types.FunctionDeclaration;
import com.google.genai.types.Schema;
import io.modelcontextprotocol.spec.McpSchema;
import java.util.Optional;

/** Utility class for converting between different representations of MCP tools. */
public final class ConversionUtils {

  public McpSchema.Tool adkToMcpToolType(BaseTool tool) {
    Optional<FunctionDeclaration> toolDeclaration = tool.declaration();
    if (toolDeclaration.isEmpty()) {
      return new McpSchema.Tool(tool.name(), tool.description(), "");
    }
    Schema geminiSchema = toolDeclaration.get().parameters().get();
    return new McpSchema.Tool(tool.name(), tool.description(), geminiSchema.toJson());
  }

  private ConversionUtils() {}
}

/*
 * 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.tools.mcp;

import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.spec.McpClientTransport;
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
import java.time.Duration;
import java.util.Optional;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Manages MCP client sessions.
 *
 * <p>This class provides methods for creating and initializing MCP client sessions, handling
 * different connection parameters and transport builders.
 */
// TODO(b/413489523): Implement this class.
public class McpSessionManager {

  private final Object connectionParams; // ServerParameters or SseServerParameters
  private static final Logger logger = LoggerFactory.getLogger(McpSessionManager.class);

  public McpSessionManager(Object connectionParams) {
    this.connectionParams = connectionParams;
  }

  public McpSyncClient createSession() {
    return initializeSession(this.connectionParams);
  }

  public static McpSyncClient initializeSession(Object connectionParams) {
    return initializeSession(connectionParams, new DefaultMcpTransportBuilder());
  }

  public static McpSyncClient initializeSession(
      Object connectionParams, McpTransportBuilder transportBuilder) {
    Duration initializationTimeout = null;
    Duration requestTimeout = null;
    McpClientTransport transport = transportBuilder.build(connectionParams);
    if (connectionParams instanceof SseServerParameters sseServerParams) {
      initializationTimeout = sseServerParams.timeout();
      requestTimeout = sseServerParams.sseReadTimeout();
    }
    McpSyncClient client =
        McpClient.sync(transport)
            .initializationTimeout(
                Optional.ofNullable(initializationTimeout).orElse(Duration.ofSeconds(10)))
            .requestTimeout(Optional.ofNullable(requestTimeout).orElse(Duration.ofSeconds(10)))
            .capabilities(ClientCapabilities.builder().build())
            .build();
    InitializeResult initResult = client.initialize();
    logger.debug("Initialize Client Result: {}", initResult);
    return client;
  }
}

/*
 * 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.models;

import io.reactivex.rxjava3.core.Flowable;

/**
 * Abstract base class for Large Language Models (LLMs).
 *
 * <p>Provides a common interface for interacting with different LLMs.
 */
public abstract class BaseLlm {

  /** The name of the LLM model, e.g. gemini-1.5-flash or gemini-1.5-flash-001. */
  private final String model;

  public BaseLlm(String model) {
    this.model = model;
  }

  /**
   * Returns the name of the LLM model.
   *
   * @return The name of the LLM model.
   */
  public String model() {
    return model;
  }

  /**
   * Generates one content from the given LLM request and tools.
   *
   * @param llmRequest The LLM request containing the input prompt and parameters.
   * @param stream A boolean flag indicating whether to stream the response.
   * @return A Flowable of LlmResponses. For non-streaming calls, it will only yield one
   *     LlmResponse. For streaming calls, it may yield more than one LlmResponse, but all yielded
   *     LlmResponses should be treated as one content by merging their parts.
   */
  public abstract Flowable<LlmResponse> generateContent(LlmRequest llmRequest, boolean stream);

  /** Creates a live connection to the LLM. */
  public abstract BaseLlmConnection connect(LlmRequest llmRequest);
}

/*
 * 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.flows;

import com.google.adk.agents.InvocationContext;
import com.google.adk.events.Event;
import io.reactivex.rxjava3.core.Flowable;

/** Interface for the execution flows to run a group of agents. */
public interface BaseFlow {

  /**
   * Run this flow.
   *
   *
To implement this method, the flow should follow the below requirements:
   *
   *

   *
1. `session` should be treated as immutable, DO NOT change it.
   *
2. The caller who trigger the flow is responsible for updating the session as the events
   *       being generated. The subclass implementation will assume session is updated after each
   *       yield event statement.
   *
3. A flow may spawn sub-agent flows depending on the agent definition.
   *

   */
  Flowable<Event> run(InvocationContext invocationContext);

  default Flowable<Event> runLive(InvocationContext invocationContext) {
    throw new UnsupportedOperationException("Not implemented");
  }
}

/*
 * 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.events;

import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.function.Supplier;

/** Iterable stream of {@link Event} objects. */
public class EventStream implements Iterable<Event> {

  private final Supplier<Event> eventSupplier;

  /** Constructs a new event stream. */
  public EventStream(Supplier<Event> eventSupplier) {
    this.eventSupplier = eventSupplier;
  }

  /** Returns an iterator that fetches events lazily. */
  @Override
  public Iterator<Event> iterator() {
    return new EventIterator();
  }

  /** Iterator that returns events from the supplier until it returns {@code null}. */
  private class EventIterator implements Iterator<Event> {
    private Event nextEvent = null;
    private boolean finished = false;

    /** Returns {@code true} if another event is available. */
    @Override
    public boolean hasNext() {
      if (finished) {
        return false;
      }
      if (nextEvent == null) {
        nextEvent = eventSupplier.get();
        finished = (nextEvent == null);
      }
      return !finished;
    }

    /**
     * Returns the next event.
     *
     * @throws NoSuchElementException if no more events are available.
     */
    @Override
    public Event next() {
      if (!hasNext()) {
        throw new NoSuchElementException("No more events.");
      }
      Event currentEvent = nextEvent;
      nextEvent = null;
      return currentEvent;
    }
  }
}

/*
 * 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.tools;

/** Exits the loop. */
public final class ExitLoopTool {
  public static void exitLoop(ToolContext toolContext) {
    // Exits the loop.
    // Call this function only when you are instructed to do so.
    toolContext.setActions(toolContext.actions().toBuilder().escalate(true).build());
  }

  private ExitLoopTool() {}
}

/*
 * 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 java.util.Collection;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/** A {@link State} object that also keeps track of the changes to the state. */
@SuppressWarnings("ShouldNotSubclass")
public final class State implements ConcurrentMap<String, Object> {

  public static final String APP_PREFIX = "app:";
  public static final String USER_PREFIX = "user:";
  public static final String TEMP_PREFIX = "temp:";

  // Sentinel object to mark removed entries in the delta map
  private static final Object REMOVED = new Object();

  private final ConcurrentMap<String, Object> state;
  private final ConcurrentMap<String, Object> delta;

  public State(ConcurrentMap<String, Object> state) {
    this(state, new ConcurrentHashMap<>());
  }

  public State(ConcurrentMap<String, Object> state, ConcurrentMap<String, Object> delta) {
    this.state = Objects.requireNonNull(state);
    this.delta = delta;
  }

  @Override
  public void clear() {
    state.clear();
  }

  @Override
  public boolean containsKey(Object key) {
    return state.containsKey(key);
  }

  @Override
  public boolean containsValue(Object value) {
    return state.containsValue(value);
  }

  @Override
  public Set<Entry<String, Object>> entrySet() {
    return state.entrySet();
  }

  @Override
  public boolean equals(Object o) {
    if (o == this) {
      return true;
    }
    if (!(o instanceof State other)) {
      return false;
    }
    return state.equals(other.state);
  }

  @Override
  public Object get(Object key) {
    return state.get(key);
  }

  @Override
  public int hashCode() {
    return state.hashCode();
  }

  @Override
  public boolean isEmpty() {
    return state.isEmpty();
  }

  @Override
  public Set<String> keySet() {
    return state.keySet();
  }

  @Override
  public Object put(String key, Object value) {
    Object oldValue = state.put(key, value);
    delta.put(key, value);
    return oldValue;
  }

  @Override
  public Object putIfAbsent(String key, Object value) {
    Object existingValue = state.putIfAbsent(key, value);
    if (existingValue == null) {
      delta.put(key, value);
    }
    return existingValue;
  }

  @Override
  public void putAll(Map<? extends String, ? extends Object> m) {
    state.putAll(m);
    delta.putAll(m);
  }

  @Override
  public Object remove(Object key) {
    if (state.containsKey(key)) {
      delta.put((String) key, REMOVED);
    }
    return state.remove(key);
  }

  @Override
  public boolean remove(Object key, Object value) {
    boolean removed = state.remove(key, value);
    if (removed) {
      delta.put((String) key, REMOVED);
    }
    return removed;
  }

  @Override
  public boolean replace(String key, Object oldValue, Object newValue) {
    boolean replaced = state.replace(key, oldValue, newValue);
    if (replaced) {
      delta.put(key, newValue);
    }
    return replaced;
  }

  @Override
  public Object replace(String key, Object value) {
    Object oldValue = state.replace(key, value);
    if (oldValue != null) {
      delta.put(key, value);
    }
    return oldValue;
  }

  @Override
  public int size() {
    return state.size();
  }

  @Override
  public Collection<Object> values() {
    return state.values();
  }

  public boolean hasDelta() {
    return !delta.isEmpty();
  }
}

/*
 * 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
 *
 *      https://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.network;

import com.google.genai.errors.ApiException;
import com.google.genai.errors.ClientException;
import com.google.genai.errors.ServerException;
import java.io.IOException;
import okhttp3.Response;
import okhttp3.ResponseBody;

/** Wraps a real HTTP response to expose the methods needed by the GenAI SDK. */
public final class HttpApiResponse extends ApiResponse {

  private final Response response;

  /** Constructs a HttpApiResponse instance with the response. */
  public HttpApiResponse(Response response) {
    throwFromResponse(response);
    this.response = response;
  }

  /** Returns the ResponseBody from the response. */
  @Override
  public ResponseBody getEntity() {
    return response.body();
  }

  /**
   * Throws an ApiException from the response if the response is not a OK status. This method is
   * adapted to work with OkHttp's Response.
   *
   * @param response The response from the API call.
   */
  private void throwFromResponse(Response response) {
    if (response.isSuccessful()) {
      return;
    }
    int code = response.code();
    String status = response.message();
    String message = getErrorMessageFromResponse(response);
    if (code >= 400 && code < 500) { // Client errors.
      throw new ClientException(code, status, message);
    } else if (code >= 500 && code < 600) { // Server errors.
      throw new ServerException(code, status, message);
    } else {
      throw new ApiException(code, status, message);
    }
  }

  private String getErrorMessageFromResponse(Response response) {
    try {
      if (response.body() != null) {
        return response.body().string(); // Consume only once.
      }
    } catch (IOException e) {
      // Log the error.  Return a generic message to avoid masking original exception.
      return "Error reading response body";
    }
    return "";
  }

  /** Closes the Http response. */
  @Override
  public void close() {
    response.close();
  }
}

/*
 * 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.flows.llmflows.audio;

import com.google.cloud.speech.v1.RecognitionAudio;
import com.google.cloud.speech.v1.RecognitionConfig;
import com.google.cloud.speech.v1.RecognizeResponse;
import com.google.cloud.speech.v1.SpeechClient;
import java.io.IOException;

/** Implementation of SpeechClientInterface using Vertex AI SpeechClient. */
public class VertexSpeechClient implements SpeechClientInterface {

  private final SpeechClient speechClient;

  /**
   * Constructs a VertexSpeechClient, initializing the underlying Google Cloud SpeechClient.
   *
   * @throws IOException if SpeechClient creation fails.
   */
  public VertexSpeechClient() throws IOException {
    this.speechClient = SpeechClient.create();
  }

  /**
   * Performs synchronous speech recognition on the given audio input.
   *
   * @param config Recognition configuration (e.g., language, encoding).
   * @param audio Audio data to recognize.
   * @return The recognition result.
   */
  @Override
  public RecognizeResponse recognize(RecognitionConfig config, RecognitionAudio audio) {
    // The original SpeechClient.recognize doesn't declare checked exceptions other than what might
    // be runtime. The interface declares Exception to be more general for other implementations.
    return speechClient.recognize(config, audio);
  }

  @Override
  public void close() throws Exception {
    if (speechClient != null) {
      speechClient.close();
    }
  }
}

/*
 * 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.models;

import com.google.auto.value.AutoValue;
import java.util.Optional;

/** Represents a model by name or instance. */
@AutoValue
public abstract class Model {

  public abstract Optional<String> modelName();

  public abstract Optional<BaseLlm> model();

  public static Builder builder() {
    return new AutoValue_Model.Builder();
  }

  public abstract Builder toBuilder();

  /** Builder for {@link Model}. */
  @AutoValue.Builder
  public abstract static class Builder {

    public abstract Builder modelName(String modelName);

    public abstract Builder model(BaseLlm model);

    public abstract Model build();
  }
}

package com.google.adk.tools;

import com.google.adk.agents.ReadonlyContext;
import java.util.Optional;

/**
 * Functional interface to decide whether a tool should be exposed to the LLM based on the current
 * context.
 */
@FunctionalInterface
public interface ToolPredicate {
  /**
   * Decides if the given tool is selected.
   *
   * @param tool The tool to check.
   * @param readonlyContext The current context.
   * @return true if the tool should be selected, false otherwise.
   */
  boolean test(BaseTool tool, Optional<ReadonlyContext> readonlyContext);
}

/*
 * 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.web.config;

import java.util.List;
import org.springframework.boot.context.properties.ConfigurationProperties;

/**
 * Properties for configuring CORS in ADK Web. This class is used to load CORS settings from
 * application properties.
 */
@ConfigurationProperties(prefix = "adk.web.cors")
public record AdkWebCorsProperties(
    String mapping,
    List<String> origins,
    List<String> methods,
    List<String> headers,
    boolean allowCredentials,
    long maxAge) {

  public AdkWebCorsProperties {
    mapping = mapping != null ? mapping : "/**";
    origins = origins != null && !origins.isEmpty() ? origins : List.of();
    methods =
        methods != null && !methods.isEmpty()
            ? methods
            : List.of("GET", "POST", "PUT", "DELETE", "OPTIONS");
    headers = headers != null && !headers.isEmpty() ? headers : List.of("*");
    maxAge = maxAge > 0 ? maxAge : 3600;
  }
}

/*
 * 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.memory;

import com.google.auto.value.AutoValue;
import com.google.genai.types.Content;
import java.time.Instant;
import javax.annotation.Nullable;

/** Represents one memory entry. */
@AutoValue
public abstract class MemoryEntry {

  /** Returns the main content of the memory. */
  public abstract Content content();

  /** Returns the author of the memory, or null if not set. */
  @Nullable
  public abstract String author();

  /**
   * Returns the timestamp when the original content of this memory happened, or null if not set.
   *
   *
This string will be forwarded to LLM. Preferred format is ISO 8601 format
   */
  @Nullable
  public abstract String timestamp();

  /** Returns a new builder for creating a {@link MemoryEntry}. */
  public static Builder builder() {
    return new AutoValue_MemoryEntry.Builder();
  }

  /**
   * Creates a new builder with a copy of this entry's values.
   *
   * @return a new {@link Builder} instance.
   */
  public abstract Builder toBuilder();

  /** Builder for {@link MemoryEntry}. */
  @AutoValue.Builder
  public abstract static class Builder {

    /**
     * Sets the main content of the memory.
     *
     *
This is a required field.
     */
    public abstract Builder setContent(Content content);

    /** Sets the author of the memory. */
    public abstract Builder setAuthor(@Nullable String author);

    /** Sets the timestamp when the original content of this memory happened. */
    public abstract Builder setTimestamp(@Nullable String timestamp);

    /**
     * A convenience method to set the timestamp from an {@link Instant} object, formatted as an ISO
     * 8601 string.
     *
     * @param instant The timestamp as an Instant object.
     */
    public Builder setTimestamp(Instant instant) {
      return setTimestamp(instant.toString());
    }

    /** Builds the immutable {@link MemoryEntry} object. */
    public abstract MemoryEntry build();
  }
}

/*
 * 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.memory;

import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableList;
import java.util.List;

/** Represents the response from a memory search. */
@AutoValue
public abstract class SearchMemoryResponse {

  /** Returns a list of memory entries that relate to the search query. */
  public abstract ImmutableList<MemoryEntry> memories();

  /** Creates a new builder for {@link SearchMemoryResponse}. */
  public static Builder builder() {
    return new AutoValue_SearchMemoryResponse.Builder().setMemories(ImmutableList.of());
  }

  /** Builder for {@link SearchMemoryResponse}. */
  @AutoValue.Builder
  public abstract static class Builder {

    abstract Builder setMemories(ImmutableList<MemoryEntry> memories);

    /** Sets the list of memory entries using a list. */
    public Builder setMemories(List<MemoryEntry> memories) {
      return setMemories(ImmutableList.copyOf(memories));
    }

    /** Builds the immutable {@link SearchMemoryResponse} object. */
    public abstract SearchMemoryResponse build();
  }
}

/*
 * 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 static com.google.common.collect.ImmutableList.toImmutableList;

import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableList;
import java.util.List;

/** Response for listing sessions. */
@AutoValue
public abstract class ListSessionsResponse {

  public abstract ImmutableList<Session> sessions();

  public List<String> sessionIds() {
    return sessions().stream().map(Session::id).collect(toImmutableList());
  }

  /** Builder for {@link ListSessionsResponse}. */
  @AutoValue.Builder
  public abstract static class Builder {

    public abstract Builder sessions(List<Session> sessions);

    public abstract ListSessionsResponse build();
  }

  public static Builder builder() {
    return new AutoValue_ListSessionsResponse.Builder().sessions(ImmutableList.of());
  }
}

/*
 * 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.web.config;

import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.stereotype.Component;

/** Properties for loading agents. */
@Component
@ConfigurationProperties(prefix = "adk.agents")
public class AgentLoadingProperties {
  private String sourceDir = "src/main/java";
  private String compileClasspath;

  public String getSourceDir() {
    return sourceDir;
  }

  public void setSourceDir(String sourceDir) {
    this.sourceDir = sourceDir;
  }

  public String getCompileClasspath() {
    return compileClasspath;
  }

  public void setCompileClasspath(String compileClasspath) {
    this.compileClasspath = compileClasspath;
  }
}

/*
 * 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.models;

import com.google.genai.types.Blob;
import com.google.genai.types.Content;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Flowable;
import java.util.List;

/** The base class for a live model connection. */
public interface BaseLlmConnection {

  /**
   * Sends the conversation history to the model.
   *
   *
You call this method right after setting up the model connection. The model will respond if
   * the last content is from user, otherwise it will wait for new user input before responding.
   */
  Completable sendHistory(List<Content> history);

  /**
   * Sends a user content to the model.
   *
   *
The model will respond immediately upon receiving the content. If you send function
   * responses, all parts in the content should be function responses.
   */
  Completable sendContent(Content content);

  /**
   * Sends a chunk of audio or a frame of video to the model in realtime.
   *
   *
The model may not respond immediately upon receiving the blob. It will do voice activity
   * detection and decide when to respond.
   */
  Completable sendRealtime(Blob blob);

  /** Receives the model responses. */
  Flowable<LlmResponse> receive();

  /** Closes the connection. */
  void close();

  /** Closes the connection with an error. */
  void close(Throwable throwable);
}

/*
 * 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.tools.mcp;

import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableMap;
import java.time.Duration;
import java.util.Map;
import javax.annotation.Nullable;

/** Parameters for establishing a MCP Server-Sent Events (SSE) connection. */
@AutoValue
public abstract class SseServerParameters {

  /** The URL of the SSE server. */
  public abstract String url();

  /** Optional headers to include in the SSE connection request. */
  @Nullable
  public abstract ImmutableMap<String, Object> headers();

  /** The timeout for the initial connection attempt. */
  public abstract Duration timeout();

  /** The timeout for reading data from the SSE stream. */
  public abstract Duration sseReadTimeout();

  /** Creates a new builder for {@link SseServerParameters}. */
  public static Builder builder() {
    return new AutoValue_SseServerParameters.Builder()
        .timeout(Duration.ofSeconds(5))
        .sseReadTimeout(Duration.ofMinutes(5));
  }

  /** Builder for {@link SseServerParameters}. */
  @AutoValue.Builder
  public abstract static class Builder {
    /** Sets the URL of the SSE server. */
    public abstract Builder url(String url);

    /** Sets the headers for the SSE connection request. */
    public abstract Builder headers(@Nullable Map<String, Object> headers);

    /** Sets the timeout for the initial connection attempt. */
    public abstract Builder timeout(Duration timeout);

    /** Sets the timeout for reading data from the SSE stream. */
    public abstract Builder sseReadTimeout(Duration sseReadTimeout);

    /** Builds a new {@link SseServerParameters} instance. */
    public abstract SseServerParameters build();
  }
}

/*
 * 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.tools;

import static java.lang.annotation.ElementType.METHOD;
import static java.lang.annotation.ElementType.PARAMETER;

import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/** Annotations for tools. */
public final class Annotations {

  /** The annotation for binding the 'Schema' input. */
  @Target({METHOD, PARAMETER})
  @Retention(RetentionPolicy.RUNTIME)
  public @interface Schema {
    String name() default "";

    String description() default "";
  }

  private Annotations() {}
}

/*
 * 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.google.adk.events.Event;
import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableList;
import java.util.List;
import java.util.Optional;

/** Response for listing events. */
@AutoValue
public abstract class ListEventsResponse {

  public abstract ImmutableList<Event> events();

  public abstract Optional<String> nextPageToken();

  /** Builder for {@link ListEventsResponse}. */
  @AutoValue.Builder
  public abstract static class Builder {

    public abstract Builder events(List<Event> events);

    public abstract Builder nextPageToken(String nextPageToken);

    public abstract ListEventsResponse build();
  }

  public static Builder builder() {
    return new AutoValue_ListEventsResponse.Builder().events(ImmutableList.of());
  }
}

/*
 * 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.artifacts;

import com.google.common.collect.ImmutableList;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.util.Optional;

/** Base interface for artifact services. */
public interface BaseArtifactService {

  /**
   * Saves an artifact.
   *
   * @param appName the app name
   * @param userId the user ID
   * @param sessionId the session ID
   * @param filename the filename
   * @param artifact the artifact
   * @return the revision ID (version) of the saved artifact.
   */
  Single<Integer> saveArtifact(
      String appName, String userId, String sessionId, String filename, Part artifact);

  /**
   * Gets an artifact.
   *
   * @param appName the app name
   * @param userId the user ID
   * @param sessionId the session ID
   * @param filename the filename
   * @param version Optional version number. If null, loads the latest version.
   * @return the artifact or empty if not found
   */
  Maybe<Part> loadArtifact(
      String appName, String userId, String sessionId, String filename, Optional<Integer> version);

  /**
   * Lists all the artifact filenames within a session.
   *
   * @param appName the app name
   * @param userId the user ID
   * @param sessionId the session ID
   * @return the list artifact response containing filenames
   */
  Single<ListArtifactsResponse> listArtifactKeys(String appName, String userId, String sessionId);

  /**
   * Deletes an artifact.
   *
   * @param appName the app name
   * @param userId the user ID
   * @param sessionId the session ID
   * @param filename the filename
   */
  Completable deleteArtifact(String appName, String userId, String sessionId, String filename);

  /**
   * Lists all the versions (as revision IDs) of an artifact.
   *
   * @param appName the app name
   * @param userId the user ID
   * @param sessionId the session ID
   * @param filename the artifact filename
   * @return A list of integer version numbers.
   */
  Single<ImmutableList<Integer>> listVersions(
      String appName, String userId, String sessionId, String filename);
}

/*
 * 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.utils;

import com.google.common.collect.ImmutableMap;
import java.util.concurrent.ConcurrentHashMap;

/** Utility class for creating ConcurrentHashMaps. */
public final class Pairs {

  private Pairs() {}

  /**
   * Returns a new, empty {@code ConcurrentHashMap}.
   *
   * @param <K> the {@code ConcurrentHashMap}'s key type
   * @param <V> the {@code ConcurrentHashMap}'s value type
   * @return an empty {@code ConcurrentHashMap}
   */
  public static <K, V> ConcurrentHashMap<K, V> of() {
    return new ConcurrentHashMap<>();
  }

  /**
   * Returns a new {@code ConcurrentHashMap} containing a single mapping.
   *
   * @param <K> the {@code ConcurrentHashMap}'s key type
   * @param <V> the {@code ConcurrentHashMap}'s value type
   * @param k1 the mapping's key
   * @param v1 the mapping's value
   * @return a {@code ConcurrentHashMap} containing the specified mapping
   * @throws NullPointerException if the key or the value is {@code null}
   */
  public static <K, V> ConcurrentHashMap<K, V> of(K k1, V v1) {
    return new ConcurrentHashMap<>(ImmutableMap.of(k1, v1));
  }

  /**
   * Returns a new {@code ConcurrentHashMap} containing two mappings. This method leverages {@code
   * java.util.Map.of} for initial validation.
   *
   * @param <K> the {@code ConcurrentHashMap}'s key type
   * @param <V> the {@code ConcurrentHashMap}'s value type
   * @param k1 the first mapping's key
   * @param v1 the first mapping's value
   * @param k2 the second mapping's key
   * @param v2 the second mapping's value
   * @return a {@code ConcurrentHashMap} containing the specified mappings
   */
  public static <K, V> ConcurrentHashMap<K, V> of(K k1, V v1, K k2, V v2) {
    return new ConcurrentHashMap<>(ImmutableMap.of(k1, v1, k2, v2));
  }

  /**
   * Returns a new {@code ConcurrentHashMap} containing three mappings. This method leverages {@code
   * java.util.Map.of} for initial validation.
   *
   * @param <K> the {@code ConcurrentHashMap}'s key type
   * @param <V> the {@code ConcurrentHashMap}'s value type
   * @param k1 the first mapping's key
   * @param v1 the first mapping's value
   * @param k2 the second mapping's key
   * @param v2 the second mapping's value
   * @param k3 the third mapping's key
   * @param v3 the third mapping's value
   * @return a {@code ConcurrentHashMap} containing the specified mappings
   */
  public static <K, V> ConcurrentHashMap<K, V> of(K k1, V v1, K k2, V v2, K k3, V v3) {
    return new ConcurrentHashMap<>(ImmutableMap.of(k1, v1, k2, v2, k3, v3));
  }

  /**
   * Returns a new {@code ConcurrentHashMap} containing four mappings. This method leverages {@code
   * java.util.Map.of} for initial validation.
   *
   * @param <K> the {@code ConcurrentHashMap}'s key type
   * @param <V> the {@code ConcurrentHashMap}'s value type
   * @param k1 the first mapping's key
   * @param v1 the first mapping's value
   * @param k2 the second mapping's key
   * @param v2 the second mapping's value
   * @param k3 the third mapping's key
   * @param v3 the third mapping's value
   * @param k4 the fourth mapping's key
   * @param v4 the fourth mapping's value
   * @return a {@code ConcurrentHashMap} containing the specified mappings
   */
  public static <K, V> ConcurrentHashMap<K, V> of(K k1, V v1, K k2, V v2, K k3, V v3, K k4, V v4) {
    return new ConcurrentHashMap<>(ImmutableMap.of(k1, v1, k2, v2, k3, v3, k4, v4));
  }

  /**
   * Returns a new {@code ConcurrentHashMap} containing five mappings. This method leverages {@code
   * java.util.Map.of} for initial validation.
   *
   * @param <K> the {@code ConcurrentHashMap}'s key type
   * @param <V> the {@code ConcurrentHashMap}'s value type
   * @param k1 the first mapping's key
   * @param v1 the first mapping's value
   * @param k2 the second mapping's key
   * @param v2 the second mapping's value
   * @param k3 the third mapping's key
   * @param v3 the third mapping's value
   * @param k4 the fourth mapping's key
   * @param v4 the fourth mapping's value
   * @param k5 the fifth mapping's key
   * @param v5 the fifth mapping's value
   * @return a {@code ConcurrentHashMap} containing the specified mappings
   */
  public static <K, V> ConcurrentHashMap<K, V> of(
      K k1, V v1, K k2, V v2, K k3, V v3, K k4, V v4, K k5, V v5) {
    return new ConcurrentHashMap<>(ImmutableMap.of(k1, v1, k2, v2, k3, v3, k4, v4, k5, v5));
  }

  /**
   * Returns a new {@code ConcurrentHashMap} containing six mappings. This method leverages {@code
   * java.util.Map.of} for initial validation.
   *
   * @param <K> the {@code ConcurrentHashMap}'s key type
   * @param <V> the {@code ConcurrentHashMap}'s value type
   * @param k1 the first mapping's key
   * @param v1 the first mapping's value
   * @param k2 the second mapping's key
   * @param v2 the second mapping's value
   * @param k3 the third mapping's key
   * @param v3 the third mapping's value
   * @param k4 the fourth mapping's key
   * @param v4 the fourth mapping's value
   * @param k5 the fifth mapping's key
   * @param v5 the fifth mapping's value
   * @param k6 the sixth mapping's key
   * @param v6 the sixth mapping's value
   * @return a {@code ConcurrentHashMap} containing the specified mappings
   * @throws IllegalArgumentException if there are any duplicate keys (behavior inherited from
   *     Map.of)
   * @throws NullPointerException if any key or value is {@code null} (behavior inherited from
   *     Map.of)
   */
  public static <K, V> ConcurrentHashMap<K, V> of(
      K k1, V v1, K k2, V v2, K k3, V v3, K k4, V v4, K k5, V v5, K k6, V v6) {
    return new ConcurrentHashMap<>(ImmutableMap.of(k1, v1, k2, v2, k3, v3, k4, v4, k5, v5, k6, v6));
  }

  /**
   * Returns a new {@code ConcurrentHashMap} containing seven mappings. This method leverages {@code
   * java.util.Map.of} for initial validation.
   *
   * @param <K> the {@code ConcurrentHashMap}'s key type
   * @param <V> the {@code ConcurrentHashMap}'s value type
   * @param k1 the first mapping's key
   * @param v1 the first mapping's value
   * @param k2 the second mapping's key
   * @param v2 the second mapping's value
   * @param k3 the third mapping's key
   * @param v3 the third mapping's value
   * @param k4 the fourth mapping's key
   * @param v4 the fourth mapping's value
   * @param k5 the fifth mapping's key
   * @param v5 the fifth mapping's value
   * @param k6 the sixth mapping's key
   * @param v6 the sixth mapping's value
   * @param k7 the seventh mapping's key
   * @param v7 the seventh mapping's value
   * @return a {@code ConcurrentHashMap} containing the specified mappings
   */
  public static <K, V> ConcurrentHashMap<K, V> of(
      K k1, V v1, K k2, V v2, K k3, V v3, K k4, V v4, K k5, V v5, K k6, V v6, K k7, V v7) {
    return new ConcurrentHashMap<>(
        ImmutableMap.of(k1, v1, k2, v2, k3, v3, k4, v4, k5, v5, k6, v6, k7, v7));
  }

  /**
   * Returns a new {@code ConcurrentHashMap} containing eight mappings. This method leverages {@code
   * java.util.Map.of} for initial validation.
   *
   * @param <K> the {@code ConcurrentHashMap}'s key type
   * @param <V> the {@code ConcurrentHashMap}'s value type
   * @param k1 the first mapping's key
   * @param v1 the first mapping's value
   * @param k2 the second mapping's key
   * @param v2 the second mapping's value
   * @param k3 the third mapping's key
   * @param v3 the third mapping's value
   * @param k4 the fourth mapping's key
   * @param v4 the fourth mapping's value
   * @param k5 the fifth mapping's key
   * @param v5 the fifth mapping's value
   * @param k6 the sixth mapping's key
   * @param v6 the sixth mapping's value
   * @param k7 the seventh mapping's key
   * @param v7 the seventh mapping's value
   * @param k8 the eighth mapping's key
   * @param v8 the eighth mapping's value
   * @return a {@code ConcurrentHashMap} containing the specified mappings
   */
  public static <K, V> ConcurrentHashMap<K, V> of(
      K k1,
      V v1,
      K k2,
      V v2,
      K k3,
      V v3,
      K k4,
      V v4,
      K k5,
      V v5,
      K k6,
      V v6,
      K k7,
      V v7,
      K k8,
      V v8) {
    return new ConcurrentHashMap<>(
        ImmutableMap.of(k1, v1, k2, v2, k3, v3, k4, v4, k5, v5, k6, v6, k7, v7, k8, v8));
  }

  /**
   * Returns a new {@code ConcurrentHashMap} containing nine mappings. This method leverages {@code
   * java.util.Map.of} for initial validation.
   *
   * @param <K> the {@code ConcurrentHashMap}'s key type
   * @param <V> the {@code ConcurrentHashMap}'s value type
   * @param k1 the first mapping's key
   * @param v1 the first mapping's value
   * @param k2 the second mapping's key
   * @param v2 the second mapping's value
   * @param k3 the third mapping's key
   * @param v3 the third mapping's value
   * @param k4 the fourth mapping's key
   * @param v4 the fourth mapping's value
   * @param k5 the fifth mapping's key
   * @param v5 the fifth mapping's value
   * @param k6 the sixth mapping's key
   * @param v6 the sixth mapping's value
   * @param k7 the seventh mapping's key
   * @param v7 the seventh mapping's value
   * @param k8 the eighth mapping's key
   * @param v8 the eighth mapping's value
   * @param k9 the ninth mapping's key
   * @param v9 the ninth mapping's value
   * @return a {@code ConcurrentHashMap} containing the specified mappings
   */
  public static <K, V> ConcurrentHashMap<K, V> of(
      K k1,
      V v1,
      K k2,
      V v2,
      K k3,
      V v3,
      K k4,
      V v4,
      K k5,
      V v5,
      K k6,
      V v6,
      K k7,
      V v7,
      K k8,
      V v8,
      K k9,
      V v9) {
    return new ConcurrentHashMap<>(
        ImmutableMap.of(k1, v1, k2, v2, k3, v3, k4, v4, k5, v5, k6, v6, k7, v7, k8, v8, k9, v9));
  }

  /**
   * Returns a new {@code ConcurrentHashMap} containing ten mappings. This method leverages {@code
   * java.util.Map.of} for initial validation.
   *
   * @param <K> the {@code ConcurrentHashMap}'s key type
   * @param <V> the {@code ConcurrentHashMap}'s value type
   * @param k1 the first mapping's key
   * @param v1 the first mapping's value
   * @param k2 the second mapping's key
   * @param v2 the second mapping's value
   * @param k3 the third mapping's key
   * @param v3 the third mapping's value
   * @param k4 the fourth mapping's key
   * @param v4 the fourth mapping's value
   * @param k5 the fifth mapping's key
   * @param v5 the fifth mapping's value
   * @param k6 the sixth mapping's key
   * @param v6 the sixth mapping's value
   * @param k7 the seventh mapping's key
   * @param v7 the seventh mapping's value
   * @param k8 the eighth mapping's key
   * @param v8 the eighth mapping's value
   * @param k9 the ninth mapping's key
   * @param v9 the ninth mapping's value
   * @param k10 the tenth mapping's key
   * @param v10 the tenth mapping's value
   * @return a {@code ConcurrentHashMap} containing the specified mappings
   */
  public static <K, V> ConcurrentHashMap<K, V> of(
      K k1,
      V v1,
      K k2,
      V v2,
      K k3,
      V v3,
      K k4,
      V v4,
      K k5,
      V v5,
      K k6,
      V v6,
      K k7,
      V v7,
      K k8,
      V v8,
      K k9,
      V v9,
      K k10,
      V v10) {
    return new ConcurrentHashMap<>(
        ImmutableMap.of(
            k1, v1, k2, v2, k3, v3, k4, v4, k5, v5, k6, v6, k7, v7, k8, v8, k9, v9, k10, v10));
  }
}

/*
 * 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.memory;

import com.google.adk.sessions.Session;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Single;

/**
 * Base contract for memory services.
 *
 *
The service provides functionalities to ingest sessions into memory so that the memory can be
 * used for user queries.
 */
public interface BaseMemoryService {

  /**
   * Adds a session to the memory service.
   *
   *
A session may be added multiple times during its lifetime.
   *
   * @param session The session to add.
   */
  Completable addSessionToMemory(Session session);

  /**
   * Searches for sessions that match the query asynchronously.
   *
   * @param appName The name of the application.
   * @param userId The id of the user.
   * @param query The query to search for.
   * @return A {@link SearchMemoryResponse} containing the matching memories.
   */
  Single<SearchMemoryResponse> searchMemory(String appName, String userId, String query);
}

/*
 * 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.artifacts;

import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.Part;
import java.util.List;

/** Response for listing artifact versions. */
@AutoValue
public abstract class ListArtifactVersionsResponse {

  public abstract ImmutableList<Part> versions();

  /** Builder for {@link ListArtifactVersionsResponse}. */
  @AutoValue.Builder
  public abstract static class Builder {
    public abstract Builder versions(List<Part> versions);

    public abstract ListArtifactVersionsResponse build();
  }

  public static Builder builder() {
    return new AutoValue_ListArtifactVersionsResponse.Builder();
  }
}

/*
 * 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.artifacts;

import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableList;
import java.util.List;

/** Response for listing artifacts. */
@AutoValue
public abstract class ListArtifactsResponse {

  public abstract ImmutableList<String> filenames();

  /** Builder for {@link ListArtifactsResponse}. */
  @AutoValue.Builder
  public abstract static class Builder {
    public abstract Builder filenames(List<String> filenames);

    public abstract ListArtifactsResponse build();
  }

  public static Builder builder() {
    return new AutoValue_ListArtifactsResponse.Builder();
  }
}

/*
 * 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.exceptions;

/** An error indicating that the limit for calls to the LLM has been exceeded. */
public final class LlmCallsLimitExceededException extends Exception {

  public LlmCallsLimitExceededException(String message) {
    super(message);
  }
}

/*
 * 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.web.config;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.cors.CorsConfiguration;
import org.springframework.web.cors.CorsConfigurationSource;
import org.springframework.web.cors.UrlBasedCorsConfigurationSource;
import org.springframework.web.filter.CorsFilter;

/**
 * Configuration class for setting up Cross-Origin Resource Sharing (CORS) in the ADK Web
 * application. This class defines beans for configuring CORS settings based on properties defined
 * in {@link AdkWebCorsProperties}.
 *
 *
CORS allows the application to handle requests from different origins, enabling secure
 * communication between the frontend and backend services.
 *
 *
Beans provided:
 *
 *

 *
{@link CorsConfigurationSource}: Configures CORS settings such as allowed origins, methods,
 *       headers, credentials, and max age.
 *
{@link CorsFilter}: Applies the CORS configuration to incoming requests.
 *

 */
@Configuration
public class AdkWebCorsConfig {

  @Bean
  public CorsConfigurationSource corsConfigurationSource(AdkWebCorsProperties corsProperties) {
    CorsConfiguration configuration = new CorsConfiguration();

    configuration.setAllowedOrigins(corsProperties.origins());
    configuration.setAllowedMethods(corsProperties.methods());
    configuration.setAllowedHeaders(corsProperties.headers());
    configuration.setAllowCredentials(corsProperties.allowCredentials());
    configuration.setMaxAge(corsProperties.maxAge());

    UrlBasedCorsConfigurationSource source = new UrlBasedCorsConfigurationSource();
    source.registerCorsConfiguration(corsProperties.mapping(), configuration);

    return source;
  }

  @Bean
  public CorsFilter corsFilter(CorsConfigurationSource corsConfigurationSource) {
    return new CorsFilter(corsConfigurationSource);
  }
}

/*
 * 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
 *
 *      https://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 okhttp3.Response;
import okhttp3.ResponseBody;

/** Wraps a real HTTP response to expose the methods needed by the GenAI SDK. */
public final class HttpApiResponse extends ApiResponse {

  private final Response response;

  /** Constructs a HttpApiResponse instance with the response. */
  public HttpApiResponse(Response response) {
    this.response = response;
  }

  /** Returns the HttpEntity from the response. */
  @Override
  public ResponseBody getResponseBody() {
    return response.body();
  }

  /** Closes the Http response. */
  @Override
  public void close() {
    response.close();
  }
}

/*
 * 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.models;

import com.google.auth.oauth2.GoogleCredentials;
import com.google.auto.value.AutoValue;
import java.util.Optional;
import javax.annotation.Nullable;

/** Credentials for accessing Gemini models through Vertex. */
@AutoValue
public abstract class VertexCredentials {

  public abstract Optional<String> project();

  public abstract Optional<String> location();

  public abstract Optional<GoogleCredentials> credentials();

  public static Builder builder() {
    return new AutoValue_VertexCredentials.Builder();
  }

  /** Builder for {@link VertexCredentials}. */
  @AutoValue.Builder
  public abstract static class Builder {
    public abstract Builder setProject(Optional<String> value);

    public abstract Builder setProject(@Nullable String value);

    public abstract Builder setLocation(Optional<String> value);

    public abstract Builder setLocation(@Nullable String value);

    public abstract Builder setCredentials(Optional<GoogleCredentials> value);

    public abstract Builder setCredentials(@Nullable GoogleCredentials value);

    public abstract VertexCredentials build();
  }
}

/*
 * 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.google.auto.value.AutoValue;
import java.time.Instant;
import java.util.Optional;

/** Configuration for getting a session. */
@AutoValue
public abstract class GetSessionConfig {

  public abstract Optional<Integer> numRecentEvents();

  public abstract Optional<Instant> afterTimestamp();

  /** Builder for {@link GetSessionConfig}. */
  @AutoValue.Builder
  public abstract static class Builder {

    public abstract Builder numRecentEvents(int numRecentEvents);

    public abstract Builder afterTimestamp(Instant afterTimestamp);

    public abstract GetSessionConfig build();
  }

  public static Builder builder() {
    return new AutoValue_GetSessionConfig.Builder();
  }
}

/*
 * 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.utils;

import com.google.common.collect.Iterables;

/** Frequently used code snippets for collections. */
public final class CollectionUtils {

  /**
   * Checks if the given iterable is null or empty.
   *
   * @param iterable the iterable to check
   * @return true if the iterable is null or empty, false otherwise
   */
  public static <T> boolean isNullOrEmpty(Iterable<T> iterable) {
    return iterable == null || Iterables.isEmpty(iterable);
  }

  private CollectionUtils() {}
}

/*
 * 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.flows.llmflows.audio;

import com.google.cloud.speech.v1.RecognitionAudio;
import com.google.cloud.speech.v1.RecognitionConfig;
import com.google.cloud.speech.v1.RecognizeResponse;

/**
 * Interface for a speech-to-text client. Allows for different implementations (e.g., Cloud, Mocks).
 */
public interface SpeechClientInterface extends AutoCloseable {

  /**
   * Performs synchronous speech recognition.
   *
   * @param config The recognition configuration.
   * @param audio The audio data to transcribe.
   * @return The recognition response.
   * @throws Exception if an error occurs during recognition.
   */
  RecognizeResponse recognize(RecognitionConfig config, RecognitionAudio audio) throws Exception;

  /**
   * Closes the client and releases any resources.
   *
   * @throws Exception if an error occurs during closing.
   */
  @Override
  void close() throws Exception;
}

/*
 * 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;

/** Represents a general error that occurred during session management operations. */
public class SessionException extends RuntimeException {

  public SessionException(String message) {
    super(message);
  }

  public SessionException(String message, Throwable cause) {
    super(message, cause);
  }

  public SessionException(Throwable cause) {
    super(cause);
  }
}

/*
 * 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
 *
 *      https://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 okhttp3.ResponseBody;

/** The API response contains a response to a call to the GenAI APIs. */
public abstract class ApiResponse implements AutoCloseable {
  /** Gets the HttpEntity. */
  public abstract ResponseBody getResponseBody();

  @Override
  public abstract void close();
}

/*
 * 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
 *
 *      https://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.network;

import okhttp3.ResponseBody;

/** The API response contains a response to a call to the GenAI APIs. */
public abstract class ApiResponse implements AutoCloseable {
  /** Gets the ResponseBody. */
  public abstract ResponseBody getEntity();

  @Override
  public abstract void close();
}

/*
 * 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;

/**
 * Tracks the current ADK version. Useful for tracking headers. Kept as a string literal to avoid
 * coupling with the build system.
 */
public final class Version {
  // Don't touch this, release-please should keep it up to date.
  public static final String JAVA_ADK_VERSION = "0.2.1-SNAPSHOT";

  private Version() {}
}

/*
 * 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;

/** Indicates that a requested session could not be found. */
public class SessionNotFoundException extends SessionException {

  public SessionNotFoundException(String message) {
    super(message);
  }

  public SessionNotFoundException(String message, Throwable cause) {
    super(message, cause);
  }
}

package com.google.adk.tools.mcp;

import io.modelcontextprotocol.spec.McpClientTransport;

/**
 * Interface for building McpClientTransport instances. Implementations of this interface are
 * responsible for constructing concrete McpClientTransport objects based on the provided connection
 * parameters.
 */
public interface McpTransportBuilder {
  /**
   * Builds an McpClientTransport based on the provided connection parameters.
   *
   * @param connectionParams The parameters required to configure the transport. The type of this
   *     object determines the type of transport built.
   * @return An instance of McpClientTransport.
   * @throws IllegalArgumentException if the connectionParams are not supported or invalid.
   */
  McpClientTransport build(Object connectionParams);
}

public final
class
AgentTransfer
implements RequestProcessor
{
@Override
  public Single<RequestProcessor.RequestProcessingResult> processRequest (
      InvocationContext context, LlmRequest request){}
private String buildTargetAgentsInfo (BaseAgent targetAgent){}
private String buildTargetAgentsInstructions (LlmAgent agent, List<BaseAgent> transferTargets){}
private List<BaseAgent> getTransferTargets (LlmAgent agent){}
public static void transferToAgent (String agentName, ToolContext toolContext){}
}

This code segment functions to: Facilitate dynamic routing between interconnected agents by evaluating hierarchical relationships and capability descriptions, then provide conditional delegation mechanisms based on suitability assessments while maintaining request processing continuity.

```java
[TASK_BEGIN]
public final class AgentTransfer implements RequestProcessor {
    @Override
    public Single<RequestProcessor.RequestProcessingResult> processRequest(
            InvocationContext context, LlmRequest request) {
        if (request.params().containsKey(EventActions.TRANSFER_AGENT_FLAG)) {
            LlmAgent currentAgent = (LlmAgent) context.currentAgent();
            List<BaseAgent> transferTargets = getTransferTargets(currentAgent);

            String agentsInfo = buildTargetAgentsInfo(transferTargets);
            String instructions = buildTargetAgentsInstructions(currentAgent, transferTargets);

            LlmRequest newRequest = request.toBuilder()
                    .clearParams()
                    .putParam("target_agents_info", agentsInfo)
                    .putParam("llm_instructions", instructions)
                    .build();

            return Single.just(RequestProcessor.RequestProcessingResult.replaceRequest(newRequest));
        }
        return Single.just(RequestProcessor.RequestProcessingResult.proceed());
    }

    private String buildTargetAgentsInfo(List<BaseAgent> transferTargets) {
        StringBuilder sb = new StringBuilder("Available agents:\n");
        for (BaseAgent agent : transferTargets) {
            sb.append("- ").append(agent.name())
              .append(": ").append(agent.description()).append("\n");
        }
        return sb.toString();
    }

    private String buildTargetAgentsInstructions(LlmAgent agent, List<BaseAgent> transferTargets) {
        return String.format(agent.getTransferInstructionsTemplate(),
                buildTargetAgentsInfo(transferTargets));
    }

    private List<BaseAgent> getTransferTargets(LlmAgent agent) {
        return ImmutableList.copyOf(agent.getTransferTargets());
    }

    public static void transferToAgent(String agentName, ToolContext toolContext) {
        toolContext.setTargetAgent(agentName);
    }
}
[TASK_END]
```

public final class AgentTransfer implements RequestProcessor {
    @Override
    public Single<RequestProcessor.RequestProcessingResult> processRequest(
            InvocationContext context, LlmRequest request) {
        if (request.params().containsKey(EventActions.TRANSFER_AGENT_FLAG)) {
            LlmAgent currentAgent = (LlmAgent) context.currentAgent();
            List<BaseAgent> transferTargets = getTransferTargets(currentAgent);
            String agentsInfo = buildTargetAgentsInfo(transferTargets);
            String instructions = buildTargetAgentsInstructions(currentAgent, transferTargets);
            LlmRequest newRequest = request.toBuilder()
                    .clearParams()
                    .putParam("target_agents_info", agentsInfo)
                    .putParam("llm_instructions", instructions)
                    .build();
            return Single.just(RequestProcessor.RequestProcessingResult.replaceRequest(newRequest));
        }
        return Single.just(RequestProcessor.RequestProcessingResult.proceed());
    }
    private String buildTargetAgentsInfo(List<BaseAgent> transferTargets) {
        StringBuilder sb = new StringBuilder("Available agents:\n");
        for (BaseAgent agent : transferTargets) {
            sb.append("- ").append(agent.name())
              .append(": ").append(agent.description()).append("\n");
        }
        return sb.toString();
    }
    private String buildTargetAgentsInstructions(LlmAgent agent, List<BaseAgent> transferTargets) {
        return String.format(agent.getTransferInstructionsTemplate(),
                buildTargetAgentsInfo(transferTargets));
    }
    private List<BaseAgent> getTransferTargets(LlmAgent agent) {
        return ImmutableList.copyOf(agent.getTransferTargets());
    }
    public static void transferToAgent(String agentName, ToolContext toolContext) {
        toolContext.setTargetAgent(agentName);
    }
}