Implementing agents

@ComponentId("activity-agent")
public class ActivityAgentWithTemplate extends Agent {

  public Effect<String> query(String message) {
    return effects()
      .systemMessageFromTemplate("activity-agent-prompt") (1)
      .userMessage(message) //
      .thenReply();
  }
}

@HttpEndpoint("/activity-prompts")
public class ActivityPromptEndpoint {

  private final ComponentClient componentClient;

  public ActivityPromptEndpoint(ComponentClient componentClient) {
    this.componentClient = componentClient;
  }

  @Put
  public HttpResponse update(String prompt) {
    componentClient
      .forEventSourcedEntity("activity-agent-prompt") (1)
      .method(PromptTemplate::update) (2)
      .invoke(prompt);

    return HttpResponses.ok();
  }

  @Get
  public String get() {
    return componentClient
      .forEventSourcedEntity("activity-agent-prompt") (1)
      .method(PromptTemplate::get) (3)
      .invoke();
  }
}

akka.javasdk.agent.memory {
  enabled = true
  limited-window {
    max-size = 156KiB # max history size before oldest message start being removed
  }
}

public Effect<String> ask(String question) {
  return effects()
    .memory(MemoryProvider.limitedWindow().readLast(5))
    .systemMessage("You are a helpful...")
    .userMessage(question)
    .thenReply();
}

public Effect<String> ask(String question) {
  return effects()
    .memory(MemoryProvider.none())
    .systemMessage("You are a helpful...")
    .userMessage(question)
    .thenReply();
}

@ComponentId("session-memory-consumer")
@Consume.FromEventSourcedEntity(SessionMemoryEntity.class)
public class SessionMemoryConsumer extends Consumer {

  private final Logger logger = LoggerFactory.getLogger(getClass());


  public Effect onSessionMemoryEvent(SessionMemoryEntity.Event event) {
    var sessionId = messageContext().eventSubject().get();

    switch (event) {
      case SessionMemoryEntity.Event.UserMessageAdded userMsg -> logger.info(
        "User message added to session {}: {}",
        sessionId,
        userMsg.message()
      );
      // ...

      default -> logger.debug("Unhandled session memory event: {}", event);
    }

    return effects().done();
  }
}

@ComponentId("compaction-agent")
public class CompactionAgent extends Agent {

  private static final String SYSTEM_MESSAGE =
    """
    You can compact an interaction history with an LLM. From the given
    USER, TOOL_CALL_RESPONSE and AI messages you create one single
    user message and one single ai message.

    The interaction history starts with USER: followed by the user message.
    For each user message there is a corresponding response for AI that starts with AI:
    Keep the original style of user question and AI answer in the summary.

    Note that AI messages may contain TOOL_CALL_REQUEST(S) and
    be followed by TOOL_CALL_RESPONSE(S).
    Make sure to keep this information in the generated ai message.
    Do not keep it as structured tool calls, but make sure to extract
    the relevant context.

    Your response should follow a strict json schema as defined bellow.
    {
      "userMessage": "<the user message summary>",
      "aiMessage: "<the AI message summary>",
    }

    Do not include any explanations or text outside of the JSON structure.
    """.stripIndent(); (1)

  public record Result(String userMessage, String aiMessage) {}

  private final ComponentClient componentClient;

  public CompactionAgent(ComponentClient componentClient) {
    this.componentClient = componentClient;
  }

  public Effect<Result> summarizeSessionHistory(SessionHistory history) { (2)
    String concatenatedMessages = history
      .messages()
      .stream()
      .map(msg -> {
        return switch (msg) {
          case SessionMessage.UserMessage userMsg -> "\n\nUSER:\n" + userMsg.text(); (3)
          case SessionMessage.AiMessage aiMessage -> {
            var aiText = "\n\nAI:\n" + aiMessage.text();
            yield aiMessage
              .toolCallRequests()
              .stream()
              .reduce(
                aiText,
                // if there are tool requests, also append them to the aiText
                (acc, req) ->
                  acc +
                  "\n\tTOOL_CALL_REQUEST: id=" +
                  req.id() +
                  ", name=" +
                  req.name() +
                  ", args=" +
                  req.arguments() +
                  " \n",
                String::concat
              );
          }
          case SessionMessage.ToolCallResponse toolRes -> "\n\nTOOL_CALL_RESPONSE:\n" +
          toolRes.text();
        };
      })
      .collect(Collectors.joining()); (3)

    return effects()
      .memory(MemoryProvider.none()) (4)
      .model(
        ModelProvider.openAi()
          .withModelName("gpt-4o-mini")
          .withApiKey(System.getenv("OPENAI_API_KEY"))
          .withMaxTokens(1000)
      )
      .systemMessage(SYSTEM_MESSAGE)
      .userMessage(concatenatedMessages)
      .responseAs(Result.class)
      .thenReply();
  }
}

@ComponentId("session-memory-consumer")
@Consume.FromEventSourcedEntity(SessionMemoryEntity.class)
public class SessionMemoryConsumer extends Consumer {

  private final Logger logger = LoggerFactory.getLogger(getClass());

  private final ComponentClient componentClient;

  public SessionMemoryConsumer(ComponentClient componentClient) {
    this.componentClient = componentClient;
  }


  public Effect onSessionMemoryEvent(SessionMemoryEntity.Event event) {
    var sessionId = messageContext().eventSubject().get();

    switch (event) {
      case SessionMemoryEntity.Event.UserMessageAdded userMsg -> logger.info(
        "User message added to session {}: {}",
        sessionId,
        userMsg.message()
      );
      // ...
      case SessionMemoryEntity.Event.AiMessageAdded aiMsg -> {
        if (aiMsg.historySizeInBytes() > 100000) { (1)
          var history = componentClient
            .forEventSourcedEntity(sessionId)
            .method(SessionMemoryEntity::getHistory) (2)
            .invoke(new SessionMemoryEntity.GetHistoryCmd(Optional.empty()));

          var summary = componentClient
            .forAgent()
            .inSession(sessionId)
            .method(CompactionAgent::summarizeSessionHistory) (3)
            .invoke(history);

          var now = Instant.now();
          componentClient
            .forEventSourcedEntity(sessionId)
            .method(SessionMemoryEntity::compactHistory) (4)
            .invoke(
              new SessionMemoryEntity.CompactionCmd(
                new SessionMessage.UserMessage(now, summary.userMessage(), ""),
                new SessionMessage.AiMessage(now, summary.aiMessage(), ""),
                history.sequenceNumber() (5)
              )
            );
        }
      }

      default -> logger.debug("Unhandled session memory event: {}", event);
    }

    return effects().done();
  }
}

public Effect<String> query(String message) {
  return effects()
    .systemMessage(SYSTEM_MESSAGE)
    .userMessage(message)
    .onFailure(exception -> {
      // Handle different types of exceptions with appropriate fallback responses
      return switch (exception) { (1)
        case RateLimitException exc -> "Rate limit exceeded, try again later"; (2)
        case ModelTimeoutException exc -> "Request timeout, service is delayed";
        case ToolCallExecutionException exc -> "Tool error: " + exc.getToolName();
        default -> "Unexpected error occurred"; (3)
      };
    })
    .thenReply();

public Effect<AgentResponse> query(String message) {
  return effects()
    .systemMessage(SYSTEM_MESSAGE)
    .tools(WeatherService.class) (1)
    .userMessage(message)
    .responseAs(AgentResponse.class)
    .thenReply();
}

public Effect<AgentResponse> query(String message) {
  return effects()
    .systemMessage(SYSTEM_MESSAGE)
    .mcpTools(
      RemoteMcpTools.fromService("weather-service"), (1)
      RemoteMcpTools.fromServer("https://weather.example.com/mcp") (2)
        .addClientHeader(Authorization.oauth2(System.getenv("WEATHER_API_TOKEN"))) (3)
        .withAllowedToolNames(Set.of("get_weather")) (4)
    )
    .userMessage(message)
    .responseAs(AgentResponse.class)
    .thenReply();
}

# Increase the limit to allow more tool calls
akka.javasdk.agent.max-tool-call-steps = 150

@ComponentId("agent-team")
public class AgentTeamWorkflow extends Workflow<AgentTeamWorkflow.State> {

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

  public record State(String userQuery, String weatherForecast, String answer) {
    State withWeatherForecast(String f) {
      return new State(userQuery, f, answer);
    }

    State withAnswer(String a) {
      return new State(userQuery, weatherForecast, a);
    }
  }

  private final ComponentClient componentClient;

  public AgentTeamWorkflow(ComponentClient componentClient) {
    this.componentClient = componentClient;
  }

  public Effect<Done> start(String query) {
    return effects()
      .updateState(new State(query, "", ""))
      .transitionTo("weather") (1)
      .thenReply(Done.getInstance());
  }

  public Effect<String> getAnswer() {
    if (currentState() == null || currentState().answer.isEmpty()) {
      return effects()
        .error(
          "Workflow '" + commandContext().workflowId() + "' not started, or not completed"
        );
    } else {
      return effects().reply(currentState().answer);
    }
  }

  @Override
  public WorkflowDef<State> definition() {
    return workflow()
      .addStep(askWeather())
      .addStep(suggestActivities())
      .addStep(error())
      .defaultStepRecoverStrategy(maxRetries(2).failoverTo("error"));
  }

  private Step askWeather() { (2)
    return step("weather")
      .call(
        () ->
          componentClient
            .forAgent()
            .inSession(sessionId())
            .method(WeatherAgent::query)
            .invoke(currentState().userQuery)
      )
      .andThen(String.class, forecast -> {
        logger.info("Weather forecast: {}", forecast);

        return effects()
          .updateState(currentState().withWeatherForecast(forecast)) (3)
          .transitionTo("activities");
      })
      .timeout(Duration.ofSeconds(60));
  }

  private Step suggestActivities() {
    return step("activities")
      .call(() -> {
        String request = (4)
          currentState().userQuery + "\nWeather forecast: " + currentState().weatherForecast;

        return componentClient
          .forAgent()
          .inSession(sessionId())
          .method(ActivityAgent::query)
          .invoke(request);
      })
      .andThen(String.class, suggestion -> {
        logger.info("Activities: {}", suggestion);

        return effects()
          .updateState(currentState().withAnswer(suggestion)) (5)
          .end();
      })
      .timeout(Duration.ofSeconds(60));
  }

  private Step error() {
    return step("error").call(() -> null).andThen(() -> effects().end());
  }

  private String sessionId() {
    // the workflow corresponds to the session
    return commandContext().workflowId();
  }
}

@ComponentId("selector-agent")
@AgentDescription(
  name = "Selector Agent",
  description = """
    An agent that analyses the user request and selects useful agents for
    answering the request.
  """
)
public class SelectorAgent extends Agent {

  private final String systemMessage;

  public SelectorAgent(AgentRegistry agentsRegistry) { (1)
    var agents = agentsRegistry.agentsWithRole("worker"); (2)

    this.systemMessage = """
      Your job is to analyse the user request and select the agents that should be
      used to answer the user. In order to do that, you will receive a list of
      available agents. Each agent has an id, a name and a description of its capabilities.

      For example, a user may be asking to book a trip. If you see that there is a
      weather agent, a city trip agent and a hotel booking agent, you should select
      those agents to complete the task. Note that this is just an example. The list
      of available agents may vary, so you need to use reasoning to dissect the original
      user request and using the list of available agents,
      decide which agents must be selected.

      You don't need to come up with an execution order. Your task is to
      analyze user's request and select the agents.

      Your response should follow a strict json schema as defined bellow.
      It should contain a single field 'agents'. The field agents must be array of strings
      containing the agent's IDs. If none of the existing agents are suitable for executing
      the task, you return an empty array.

       {
         "agents": []
       }

      Do not include any explanations or text outside of the JSON structure.

      You can find the list of existing agents below (in JSON format):
      Also important, use the agent id to identify the agents.
      %s
    """.stripIndent()
      .formatted(JsonSupport.encodeToString(agents)); (3)
  }

  public Effect<AgentSelection> selectAgents(String message) {
    return effects()
      .systemMessage(systemMessage)
      .userMessage(message)
      .responseAs(AgentSelection.class)
      .thenReply();
  }
}

@ComponentId("weather-agent")
@AgentDescription(
  name = "Weather Agent",
  description = """
    An agent that provides weather information. It can provide current weather,
    forecasts, and other related information.
  """,
  role = "worker"
)
public class WeatherAgent extends Agent {

@ComponentId("activity-agent")
@AgentDescription(
  name = "Activity Agent",
  description = """
    An agent that suggests activities in the real world. Like for example,
    a team building activity, sports, an indoor or outdoor game,
    board games, a city trip, etc.
  """,
  role = "worker"
)
public class ActivityAgent extends Agent {

public record AgentSelection(List<String> agents) {}

@ComponentId("planner-agent")
@AgentDescription(
  name = "Planner",
  description = """
  An agent that analyzes the user request and available agents to plan the tasks
  to produce a suitable answer.
  """
)
public class PlannerAgent extends Agent {

  public record Request(String message, AgentSelection agentSelection) {}

  private final AgentRegistry agentsRegistry;

  public PlannerAgent(AgentRegistry agentsRegistry) {
    this.agentsRegistry = agentsRegistry;
  }

  private String buildSystemMessage(AgentSelection agentSelection) {
    var agents = agentSelection.agents().stream().map(agentsRegistry::agentInfo).toList(); (1)
    return """
      Your job is to analyse the user request and the list of agents and devise the
      best order in which the agents should be called in order to produce a
      suitable answer to the user.

      You can find the list of exiting agents below (in JSON format):
      %s

      Note that each agent has a description of its capabilities.
      Given the user request, you must define the right ordering.

      Moreover, you must generate a concise request to be sent to each agent.
      This agent request is of course based on the user original request,
      but is tailored to the specific agent. Each individual agent should not
      receive requests or any text that is not related with its domain of expertise.

      Your response should follow a strict json schema as defined bellow.
       {
         "steps": [
            {
              "agentId": "<the id of the agent>",
              "query: "<agent tailored query>",
            }
         ]
       }

      The '<the id of the agent>' should be filled with the agent id.
      The '<agent tailored query>' should contain the agent tailored message.
      The order of the items inside the "steps" array should be the order of execution.

      Do not include any explanations or text outside of the JSON structure.

    """.stripIndent()
      // note: here we are not using the full list of agents, but a pre-selection
      .formatted(JsonSupport.encodeToString(agents)); (2)
  }

  public Effect<Plan> createPlan(Request request) {
    if (request.agentSelection.agents().size() == 1) {
      // no need to call an LLM to make a plan where selection has a single agent
      var step = new PlanStep(request.agentSelection.agents().getFirst(), request.message());
      return effects().reply(new Plan(List.of(step)));
    } else {
      return effects()
        .systemMessage(buildSystemMessage(request.agentSelection))
        .userMessage(request.message())
        .responseAs(Plan.class)
        .thenReply();
    }
  }
}

@ComponentId("agent-team")
public class AgentTeamWorkflow extends Workflow<AgentTeamWorkflow.State> { (1)

  public record Request(String userId, String message) {}

  @Override
  public WorkflowDef<State> definition() {
    return workflow()
      .defaultStepRecoverStrategy(maxRetries(1).failoverTo(INTERRUPT))
      .defaultStepTimeout(Duration.of(30, SECONDS))
      .addStep(selectAgentsStep()) (2)
      .addStep(planStep())
      .addStep(executePlanStep())
      .addStep(summarizeStep())
      .addStep(interruptStep());
  }

  public Effect<Done> start(Request request) {
    if (currentState() == null) {
      return effects()
        .updateState(State.init(request.userId(), request.message()))
        .transitionTo(SELECT_AGENTS) (3)
        .thenReply(Done.getInstance());
    } else {
      return effects()
        .error("Workflow '" + commandContext().workflowId() + "' already started");
    }
  }

  private static final String SELECT_AGENTS = "select-agents";

  private Step selectAgentsStep() {
    return step(SELECT_AGENTS)
      .call(
        () ->
          componentClient
            .forAgent()
            .inSession(sessionId())
            .method(SelectorAgent::selectAgents)
            .invoke(currentState().userQuery)
      ) (4)
      .andThen(AgentSelection.class, selection -> {
        logger.info("Selected agents: {}", selection.agents());
        if (selection.agents().isEmpty()) {
          var newState = currentState()
            .withFinalAnswer(
              "Couldn't find any agent(s) able to respond to the original query."
            )
            .failed();
          return effects().updateState(newState).end(); // terminate workflow
        } else {
          return effects().transitionTo(CREATE_PLAN, selection); (5)
        }
      });
  }

  private static final String CREATE_PLAN = "create-plan";

  private Step planStep() {
    return step(CREATE_PLAN)
      .call(AgentSelection.class, agentSelection -> {
        logger.info(
          "Calling planner with: '{}' / {}",
          currentState().userQuery,
          agentSelection.agents()
        );

        return componentClient
          .forAgent()
          .inSession(sessionId())
          .method(PlannerAgent::createPlan)
          .invoke(new PlannerAgent.Request(currentState().userQuery, agentSelection)); (6)
      })
      .andThen(Plan.class, plan -> {
        logger.info("Execution plan: {}", plan);
        return effects()
          .updateState(currentState().withPlan(plan))
          .transitionTo(EXECUTE_PLAN); (7)
      });
  }

  private static final String EXECUTE_PLAN = "execute-plan";

  private Step executePlanStep() {
    return step(EXECUTE_PLAN)
      .call(() -> {
        var stepPlan = currentState().nextStepPlan(); (8)
        logger.info(
          "Executing plan step (agent:{}), asking {}",
          stepPlan.agentId(),
          stepPlan.query()
        );
        var agentResponse = callAgent(stepPlan.agentId(), stepPlan.query()); (9)
        if (agentResponse.startsWith("ERROR")) {
          throw new RuntimeException(
            "Agent '" + stepPlan.agentId() + "' responded with error: " + agentResponse
          );
        } else {
          logger.info("Response from [agent:{}]: '{}'", stepPlan.agentId(), agentResponse);
          return agentResponse;
        }
      })
      .andThen(String.class, answer -> {
        var newState = currentState().addAgentResponse(answer);

        if (newState.hasMoreSteps()) {
          logger.info("Still {} steps to execute.", newState.plan().steps().size());
          return effects().updateState(newState).transitionTo(EXECUTE_PLAN); (10)
        } else {
          logger.info("No further steps to execute.");
          return effects().updateState(newState).transitionTo(SUMMARIZE);
        }
      });
  }

  private String callAgent(String agentId, String query) {
    // We know the id of the agent to call, but not the agent class.
    // Could be WeatherAgent or ActivityAgent.
    // We can still invoke the agent based on its id, given that we know that it
    // takes a AgentRequest parameter and returns String.
    var request = new AgentRequest(currentState().userId(), query);
    DynamicMethodRef<AgentRequest, String> call = componentClient
      .forAgent()
      .inSession(sessionId())
      .dynamicCall(agentId); (9)
    return call.invoke(request);
  }

}

private String callAgent(String agentId, String query) {
  // We know the id of the agent to call, but not the agent class.
  // Could be WeatherAgent or ActivityAgent.
  // We can still invoke the agent based on its id, given that we know that it
  // takes a AgentRequest parameter and returns String.
  var request = new AgentRequest(currentState().userId(), query);
  DynamicMethodRef<AgentRequest, String> call = componentClient
    .forAgent()
    .inSession(sessionId())
    .dynamicCall(agentId); (9)
  return call.invoke(request);
}

public class AgentTeamWorkflowTest extends TestKitSupport { (1)

  private final TestModelProvider selectorModel = new TestModelProvider(); (2)
  private final TestModelProvider plannerModel = new TestModelProvider();
  private final TestModelProvider activitiesModel = new TestModelProvider();
  private final TestModelProvider weatherModel = new TestModelProvider();
  private final TestModelProvider summaryModel = new TestModelProvider();

  @Override
  protected TestKit.Settings testKitSettings() {
    return TestKit.Settings.DEFAULT.withAdditionalConfig(
      "akka.javasdk.agent.openai.api-key = n/a"
    )
      .withModelProvider(SelectorAgent.class, selectorModel) (3)
      .withModelProvider(PlannerAgent.class, plannerModel)
      .withModelProvider(ActivityAgent.class, activitiesModel)
      .withModelProvider(WeatherAgent.class, weatherModel)
      .withModelProvider(SummarizerAgent.class, summaryModel);
  }

  @Test
  public void test() {
    var selection = new AgentSelection(List.of("activity-agent", "weather-agent"));
    selectorModel.fixedResponse(JsonSupport.encodeToString(selection)); (4)

    var weatherQuery = "What is the current weather in Stockholm?";
    var activityQuery =
      "Suggest activities to do in Stockholm considering the current weather.";
    var plan = new Plan(
      List.of(
        new PlanStep("weather-agent", weatherQuery),
        new PlanStep("activity-agent", activityQuery)
      )
    );
    plannerModel.fixedResponse(JsonSupport.encodeToString(plan));

    weatherModel
      .whenMessage(req -> req.equals(weatherQuery)) (5)
      .reply("The weather in Stockholm is sunny.");

    activitiesModel
      .whenMessage(req -> req.equals(activityQuery))
      .reply(
        "You can take a bike tour around Djurgården Park, " +
        "visit the Vasa Museum, explore Gamla Stan (Old Town)..."
      );

    summaryModel.fixedResponse(
      "The weather in Stockholm is sunny, so you can enjoy " +
      "outdoor activities like a bike tour around Djurgården Park, " +
      "visiting the Vasa Museum, exploring Gamla Stan (Old Town)"
    );

    var query = "I am in Stockholm. What should I do? Beware of the weather";

    var sessionId = UUID.randomUUID().toString();
    var request = new AgentTeamWorkflow.Request("alice", query);

    componentClient
      .forWorkflow(sessionId)
      .method(AgentTeamWorkflow::start) (6)
      .invoke(request);

    Awaitility.await()
      .ignoreExceptions()
      .atMost(10, SECONDS)
      .untilAsserted(() -> {
        var answer = componentClient
          .forWorkflow(sessionId)
          .method(AgentTeamWorkflow::getAnswer)
          .invoke();
        assertThat(answer).isNotBlank();
        assertThat(answer).contains("Stockholm");
        assertThat(answer).contains("sunny");
        assertThat(answer).contains("bike tour");
      });
  }
}