Modeling the agent session

package akka.ask.agent.domain;

import java.time.Instant;
import akka.javasdk.annotations.TypeName;

public sealed interface SessionEvent {

  @TypeName("user-message-added")
  public record UserMessageAdded(
      String userId,
      String sessionId,
      String query,
      int tokensUsed,
      Instant timeStamp)
      implements SessionEvent {
  }

  @TypeName("ai-message-added")
  public record AiMessageAdded(
      String userId,
      String sessionId,
      String response, (1)
      int tokensUsed,
      Instant timeStamp)
      implements SessionEvent {
  }
}

javaCopied!

@ComponentId("session-entity")
public class SessionEntity extends EventSourcedEntity<SessionEntity.State, SessionEvent> {

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

  public record Exchange(String userId, (1)
      String sessionId,
      String userQuestion,
      int queryTokensCount,
      String assistantResponse,
      int responseTokensCount) {
  }

  enum MessageType {
    AI,
    USER
  }

  public record Message(String content, MessageType type) { (2)
  }

  public record State(List<Message> messages, int totalTokenUsage) {
    public static State empty() {
      return new State(new ArrayList<>(), 0);
    }

    public State add(Message content) {
      messages.add(content);
      return new State(messages, totalTokenUsage);
    }

    public State addTokenUsage(int usage) {
      return new State(messages, totalTokenUsage + usage);
    }
  }

  public record Messages(List<Message> messages) {
  }

  public Effect<Done> addExchange(Exchange exchange) {

    var now = Instant.now();

    var userEvt = new SessionEvent.UserMessageAdded(
        exchange.userId,
        exchange.sessionId,
        exchange.userQuestion,
        exchange.queryTokensCount,
        now);

    var assistantEvt = new SessionEvent.AiMessageAdded(
        exchange.userId,
        exchange.sessionId,
        exchange.assistantResponse,
        exchange.responseTokensCount, (3)
        now);

    return effects()
        .persist(userEvt, assistantEvt) (4)
        .thenReply(__ -> Done.getInstance());
  }

  public Effect<Messages> getHistory() {
    logger.debug("Getting history from {}", commandContext().entityId());
    return effects().reply(new Messages(currentState().messages));
  }

  @Override
  public State emptyState() {
    return State.empty();
  }

  @Override
  public State applyEvent(SessionEvent event) {
    return switch (event) {
      case SessionEvent.UserMessageAdded msg ->
        currentState()
            .add(new Message(msg.query(), USER))
            .addTokenUsage(msg.tokensUsed());

      case SessionEvent.AiMessageAdded msg ->
        currentState()
            .add(new Message(msg.response(), AI))
            .addTokenUsage(msg.tokensUsed());
    };
  }
}

javaCopied!

@ComponentId("view_chat_log")
public class ConversationHistoryView extends View {

  public record ConversationHistory(List<Session> sessions) {
  }

  public record Message(String message,
      String origin, long timestamp) { (1)
  }

  public record Session(String userId,
      String sessionId, long creationDate, List<Message> messages) {
    public Session add(Message message) {
      messages.add(message);
      return this;
    }
  }

  @Query("SELECT collect(*) as sessions FROM view_chat_log " +
      "WHERE userId = :userId ORDER by creationDate DESC")
  public QueryEffect<ConversationHistory> getSessionsByUser(String userId) { (2)
    return queryResult();
  }

  @Consume.FromEventSourcedEntity(SessionEntity.class)
  public static class ChatMessageUpdater extends TableUpdater<Session> {

    public Effect<Session> onEvent(SessionEvent event) {
      return switch (event) {
        case SessionEvent.AiMessageAdded added -> aiMessage(added);
        case SessionEvent.UserMessageAdded added -> userMessage(added);
      };
    }

    private Effect<Session> aiMessage(SessionEvent.AiMessageAdded added) {
      Message newMessage = new Message(added.response(), "ai", added.timeStamp().toEpochMilli());
      var rowState = rowStateOrNew(added.userId(), added.sessionId());
      return effects().updateRow(rowState.add(newMessage));
    }

    private Effect<Session> userMessage(SessionEvent.UserMessageAdded added) {
      Message newMessage = new Message(added.query(), "user", added.timeStamp().toEpochMilli());
      var rowState = rowStateOrNew(added.userId(), added.sessionId());
      return effects().updateRow(rowState.add(newMessage));
    }

    private Session rowStateOrNew(String userId, String sessionId) { (3)
      if (rowState() != null)
        return rowState();
      else
        return new Session(
            userId,
            sessionId,
            Instant.now().toEpochMilli(),
            new ArrayList<>());
    }
  }
}

javaCopied!