エージェントのメモリ設計 - 何を覚え、何を捨てるか

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;

import com.anthropic.client.AnthropicClient;
import com.anthropic.client.okhttp.AnthropicOkHttpClient;
import com.anthropic.core.http.StreamResponse;
import com.anthropic.models.beta.agents.AgentCreateParams;
import com.anthropic.models.beta.agents.BetaManagedAgentsAgentToolset20260401Params;
import com.anthropic.models.beta.agents.BetaManagedAgentsModel;
import com.anthropic.models.beta.agents.BetaManagedAgentsModelConfigParams;
import com.anthropic.models.beta.environments.EnvironmentCreateParams;
import com.anthropic.models.beta.memorystores.MemoryStoreCreateParams;
import com.anthropic.models.beta.memorystores.MemoryStoreListParams;
import com.anthropic.models.beta.memorystores.memories.MemoryCreateParams;
import com.anthropic.models.beta.memorystores.memories.MemoryRetrieveParams;
import com.anthropic.models.beta.sessions.BetaManagedAgentsMemoryStoreResourceParam;
import com.anthropic.models.beta.sessions.SessionCreateParams;
import com.anthropic.models.beta.sessions.events.BetaManagedAgentsStreamSessionEvents;
import com.anthropic.models.beta.sessions.events.BetaManagedAgentsTextBlock;
import com.anthropic.models.beta.sessions.events.BetaManagedAgentsUserMessageEventParams;
import com.anthropic.models.beta.sessions.events.EventSendParams;
import com.google.common.base.Ascii;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.IOException;
import java.io.StringWriter;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Iterator;
import java.util.List;
import tools.jackson.core.JsonGenerator;
import tools.jackson.core.JsonParser;
import tools.jackson.core.JsonToken;
import tools.jackson.core.ObjectReadContext;
import tools.jackson.core.ObjectWriteContext;
import tools.jackson.core.json.JsonFactory;

/**
 * Managed Agents の Memory を使い、共通ポリシーと顧客別トーンを分離して保持するデモ。
 *
 * <p>初回セットアップ: mvn -q compile exec:java -Dexec.args="setup" 通常実行: mvn -q compile exec:java
 */
public final class CustomerToneReplyDemo {

  private static final Path CONFIG_FILE = Path.of(".customer-tone-demo.json");
  private static final JsonFactory JSON = JsonFactory.builder().build();

  private static final String POLICY_STORE_NAME = "store-policy";
  private static final String STYLE_GUIDE_FILE = "style-guide.md";
  private static final String COMMON_INQUIRY = "営業時間を教えてください。";

  private static final ImmutableList<Customer> CUSTOMERS =
      ImmutableList.of(
          new Customer(
              "sato",
              "佐藤さん（カジュアル）",
              """
              [問い合わせ]
              こんにちは〜！この前のホーローケトル、母がすごく喜んでました。
              同じシリーズのポットってあります？（佐藤）

              [店主の返信]
              佐藤さん、こんにちは！わあ、それは嬉しいです。
              お母さまにもよろしくお伝えください！
              同じシリーズのポット、ありますよ〜。
              来週入荷する分、1つ取り置きしておきますね！

              [問い合わせ]
              ありがとう！じゃあお願いします〜（佐藤）

              [店主の返信]
              りょうかいです！入荷したらすぐ連絡しますね。
              いつもありがとうございます！
              """,
              "ポットって食洗機いけます？（佐藤）"),
          new Customer(
              "takahashi",
              "高橋様（とても丁寧）",
              """
              [問い合わせ]
              いつも大変お世話になっております。
              株式会社ヤマガタ総務部の高橋でございます。
              このたび、創立記念品といたしましてマグカップ50個のお見積を
              頂戴できますと幸いに存じます。
              ご多忙のところ恐れ入りますが、何卒よろしくお願い申し上げます。（高橋）

              [店主の返信]
              株式会社ヤマガタ 総務部 高橋様
              いつも格別のお引き立てを賜り、誠にありがとうございます。
              雑貨店ノルドでございます。
              お見積のご依頼を頂戴し、ありがとうございます。
              数量50個にて、本日中にお見積書をお送り申し上げます。
              納期や名入れのご希望がございましたら、
              お手数ですがあわせてお知らせくださいませ。
              今後とも変わらぬご愛顧のほど、よろしくお願い申し上げます。
              """,
              "恐れ入りますが、お見積書の宛名は社名のみとしていただくことは可能でしょうか。（高橋）"));

  private final AnthropicClient client;

  private CustomerToneReplyDemo(AnthropicClient client) {
    this.client = client;
  }

  public static void main(String[] args) throws IOException {
    String apiKey = System.getenv().getOrDefault("ANTHROPIC_API_KEY", "");
    if (apiKey.isBlank()) {
      System.out.println("ANTHROPIC_API_KEY を設定してください。");
      return;
    }

    var app = new CustomerToneReplyDemo(AnthropicOkHttpClient.fromEnv());
    boolean setup = isSetupCommand(args);

    System.out.println(setup ? "# Customer Tone Reply Demo - setup" : "# Customer Tone Reply Demo");
    if (setup) {
      app.setup();
    } else {
      app.runDemo();
    }
  }

  private void setup() {
    var environment =
        client
            .beta()
            .environments()
            .create(EnvironmentCreateParams.builder().name("customer-tone-environment").build());

    var policyStore =
        client
            .beta()
            .memoryStores()
            .create(
                MemoryStoreCreateParams.builder()
                    .name(POLICY_STORE_NAME)
                    .description("店の共通ポリシーと基本情報。全顧客の返信で必ず従う。")
                    .build());

    createPolicyMemory(policyStore.id());

    var agent =
        client
            .beta()
            .agents()
            .create(
                AgentCreateParams.builder()
                    .name("customer-tone-reply-agent")
                    .model(
                        BetaManagedAgentsModelConfigParams.builder()
                            .id(BetaManagedAgentsModel.CLAUDE_SONNET_5)
                            .build())
                    .system(agentInstructions())
                    .addTool(fileToolset())
                    .build());

    ResourceConfig config = new ResourceConfig(environment.id(), agent.id(), policyStore.id());
    saveConfig(config);

    System.out.println("セットアップが完了しました。");
    System.out.println("リソースIDを " + CONFIG_FILE + " に保存しました。");
  }

  private void runDemo() {
    ResourceConfig config = loadConfig();

    for (Customer customer : CUSTOMERS) {
      runForCustomer(config, customer);
    }
  }

  private void runForCustomer(ResourceConfig config, Customer customer) {
    String customerStoreId = findOrCreateCustomerStore(customer);

    System.out.println(customer.displayName() + " / session 1: 履歴からトーンを学習");
    String firstAnswer = runSession(config, customerStoreId, buildLearningTask(customer));
    System.out.println("A. " + firstAnswer);

    System.out.println(customer.displayName() + " / session 2: 保存済みトーンで回答");
    String secondAnswer =
        runSession(config, customerStoreId, "次の問い合わせに返信してください。\n\n[問い合わせ]\n" + COMMON_INQUIRY);
    System.out.println("A. " + secondAnswer);

    System.out.println(customer.displayName() + " / 顧客トーンストアの内容");
    printAsJson(readMemories(customerStoreId));
  }

  private String runSession(ResourceConfig config, String customerStoreId, String task) {

    var session =
        client
            .beta()
            .sessions()
            .create(
                SessionCreateParams.builder()
                    .environmentId(config.environmentId())
                    .agent(config.agentId())
                    .addResource(policyResource(config.policyStoreId()))
                    .addResource(customerResource(customerStoreId))
                    .build());

    sendUserMessage(session.id(), task);
    return collectAnswer(session.id());
  }

  private void sendUserMessage(String sessionId, String task) {
    var textBlock =
        BetaManagedAgentsTextBlock.builder()
            .type(BetaManagedAgentsTextBlock.Type.TEXT)
            .text(task)
            .build();

    var content = BetaManagedAgentsUserMessageEventParams.Content.ofText(textBlock);

    client
        .beta()
        .sessions()
        .events()
        .send(
            EventSendParams.builder()
                .sessionId(sessionId)
                .addUserMessageEvent(ImmutableList.of(content))
                .build());
  }

  private String collectAnswer(String sessionId) {
    StringBuilder answer = new StringBuilder();

    try (StreamResponse<BetaManagedAgentsStreamSessionEvents> response =
        client.beta().sessions().events().streamStreaming(sessionId)) {

      Iterator<BetaManagedAgentsStreamSessionEvents> events = response.stream().iterator();

      while (events.hasNext()) {
        BetaManagedAgentsStreamSessionEvents event = events.next();

        if (event.isAgentMessage()) {
          event.asAgentMessage().content().forEach(block -> answer.append(block.text()));
        }

        if (event.isSessionStatusIdle() && event.asSessionStatusIdle().stopReason().isEndTurn()) {
          break;
        }
      }
    }

    return answer.toString().strip();
  }

  private void createPolicyMemory(String policyStoreId) {
    client
        .beta()
        .memoryStores()
        .memories()
        .create(
            MemoryCreateParams.builder()
                .memoryStoreId(policyStoreId)
                .path("/policy.md")
                .content(
                    """
                    # 店の共通ポリシー（雑貨店ノルド）
                    - 営業時間: 11:00〜19:00（木曜定休）
                    - 不具合の連絡には、まず謝り、交換品を即送する。返送は求めない
                    - 絵文字は使わない
                    - 履歴・ポリシーにない事実（在庫・価格など）は書かない
                    """)
                .build());
  }

  private String findOrCreateCustomerStore(Customer customer) {
    String storeName = customer.memoryStoreName();

    return client
        .beta()
        .memoryStores()
        .list(MemoryStoreListParams.builder().limit(100).build())
        .data()
        .stream()
        .filter(store -> storeName.equals(store.name()))
        .findFirst()
        .map(store -> store.id())
        .orElseGet(() -> createCustomerStore(customer));
  }

  private String createCustomerStore(Customer customer) {
    return client
        .beta()
        .memoryStores()
        .create(
            MemoryStoreCreateParams.builder()
                .name(customer.memoryStoreName())
                .description("このお客様（" + customer.id() + "）専用の返信トーン手引き。")
                .build())
        .id();
  }

  private ImmutableList<MemoryView> readMemories(String storeId) {
    var result = ImmutableList.<MemoryView>builder();

    client.beta().memoryStores().memories().list(storeId).data().stream()
        .filter(item -> item.isMemory())
        .map(item -> item.asMemory())
        .forEach(
            memory -> {
              var full =
                  client
                      .beta()
                      .memoryStores()
                      .memories()
                      .retrieve(
                          memory.id(),
                          MemoryRetrieveParams.builder().memoryStoreId(storeId).build());

              result.add(new MemoryView(full.path(), full.content().orElse("")));
            });

    return result.build();
  }

  private static BetaManagedAgentsMemoryStoreResourceParam policyResource(String storeId) {
    return BetaManagedAgentsMemoryStoreResourceParam.builder()
        .type(BetaManagedAgentsMemoryStoreResourceParam.Type.MEMORY_STORE)
        .memoryStoreId(storeId)
        .access(BetaManagedAgentsMemoryStoreResourceParam.Access.READ_ONLY)
        .instructions("店の共通ポリシー。返信前に必ず読み、内容に従う。")
        .build();
  }

  private static BetaManagedAgentsMemoryStoreResourceParam customerResource(String storeId) {
    return BetaManagedAgentsMemoryStoreResourceParam.builder()
        .type(BetaManagedAgentsMemoryStoreResourceParam.Type.MEMORY_STORE)
        .memoryStoreId(storeId)
        .access(BetaManagedAgentsMemoryStoreResourceParam.Access.READ_WRITE)
        .instructions(
            "このお客様専用のトーン手引き。返信前に必ず読み、" + "履歴が与えられた場合は、このストア内の " + STYLE_GUIDE_FILE + " を更新する。")
        .build();
  }

  private static String buildLearningTask(Customer customer) {
    return """
        以下は、このお客様との過去のメール対応履歴です。
        言葉遣い・トーン・定型の特徴を、読み書き可能な顧客メモリストア内の
        %s に記録または更新してください。
        そのうえで、最後の新しい問い合わせに同じトーンで返信してください。

        [過去の履歴]
        %s

        [新しい問い合わせ]
        %s
        """
        .formatted(STYLE_GUIDE_FILE, customer.history(), customer.firstInquiry());
  }

  private static String agentInstructions() {
    return """
        あなたは小売店のメール返信エージェントです。

        返信を書く前に、セッションに添付されたメモリストアを必ず確認してください。
        - 読み取り専用ストア: 店の共通ポリシー
        - 読み書き可能ストア: 顧客ごとのトーン手引き

        共通ポリシーには厳密に従ってください。
        読み書き可能ストア内に style-guide.md がある場合は、
        その言葉遣い・トーン・定型に合わせてください。
        過去の会話履歴が与えられた場合は、特徴を style-guide.md に記録・更新してください。
        """;
  }

  private static BetaManagedAgentsAgentToolset20260401Params fileToolset() {
    return BetaManagedAgentsAgentToolset20260401Params.builder()
        .type(BetaManagedAgentsAgentToolset20260401Params.Type.AGENT_TOOLSET_20260401)
        .build();
  }

  private static void saveConfig(ResourceConfig config) {
    try (var output = new BufferedOutputStream(Files.newOutputStream(CONFIG_FILE));
        JsonGenerator generator = JSON.createGenerator(ObjectWriteContext.empty(), output)) {
      generator.writeStartObject();
      generator.writeStringProperty("environmentId", config.environmentId());
      generator.writeStringProperty("agentId", config.agentId());
      generator.writeStringProperty("policyStoreId", config.policyStoreId());
      generator.writeEndObject();
    } catch (IOException e) {
      throw new IllegalStateException(CONFIG_FILE + " に保存できません", e);
    }
  }

  private static ResourceConfig loadConfig() {
    ResourceConfig fromEnvironment = configFromEnvironment();
    if (fromEnvironment != null) {
      return fromEnvironment;
    }

    checkState(
        Files.exists(CONFIG_FILE),
        "%s がありません。先に `mvn -q compile exec:java -Dexec.args=setup` を実行してください。",
        CONFIG_FILE);

    try (var input = new BufferedInputStream(Files.newInputStream(CONFIG_FILE));
        JsonParser parser = JSON.createParser(ObjectReadContext.empty(), input)) {
      checkState(parser.nextToken() == JsonToken.START_OBJECT, "%s の形式が不正です", CONFIG_FILE);

      String environmentId = null;
      String agentId = null;
      String policyStoreId = null;
      while (parser.nextToken() != JsonToken.END_OBJECT) {
        String fieldName = parser.currentName();
        parser.nextToken();
        switch (fieldName) {
          case "environmentId" -> environmentId = parser.getValueAsString();
          case "agentId" -> agentId = parser.getValueAsString();
          case "policyStoreId" -> policyStoreId = parser.getValueAsString();
          default -> parser.skipChildren();
        }
      }
      return new ResourceConfig(environmentId, agentId, policyStoreId);
    } catch (IOException e) {
      throw new IllegalStateException(CONFIG_FILE + " を読み込めません", e);
    }
  }

  private static ResourceConfig configFromEnvironment() {
    String environmentId = System.getenv("ANTHROPIC_ENVIRONMENT_ID");
    String agentId = System.getenv("ANTHROPIC_AGENT_ID");
    String policyStoreId = System.getenv("ANTHROPIC_POLICY_STORE_ID");

    boolean noneSet =
        Strings.isNullOrEmpty(environmentId)
            && Strings.isNullOrEmpty(agentId)
            && Strings.isNullOrEmpty(policyStoreId);
    if (noneSet) {
      return null;
    }

    checkState(
        !Strings.isNullOrEmpty(environmentId)
            && !Strings.isNullOrEmpty(agentId)
            && !Strings.isNullOrEmpty(policyStoreId),
        "IDを環境変数で指定する場合は ANTHROPIC_ENVIRONMENT_ID、ANTHROPIC_AGENT_ID、"
            + "ANTHROPIC_POLICY_STORE_ID の3つをすべて設定してください");

    return new ResourceConfig(environmentId, agentId, policyStoreId);
  }

  private static void printAsJson(List<MemoryView> memories) {
    try (var json = new StringWriter();
        JsonGenerator generator = JSON.createGenerator(ObjectWriteContext.empty(), json)) {
      generator.writeStartArray();
      for (MemoryView memory : memories) {
        generator.writeStartObject();
        generator.writeStringProperty("path", memory.path());
        generator.writeStringProperty("content", memory.content());
        generator.writeEndObject();
      }
      generator.writeEndArray();
      generator.close();
      System.out.println(json.toString());
    } catch (IOException e) {
      throw new IllegalStateException("Memoryの表示用JSONを生成できません", e);
    }
  }

  private static boolean isSetupCommand(String[] args) {
    return args.length > 0 && Ascii.equalsIgnoreCase(args[0], "setup");
  }

  public record ResourceConfig(String environmentId, String agentId, String policyStoreId) {

    public ResourceConfig {
      checkState(!Strings.isNullOrEmpty(environmentId), "environmentId must not be empty");
      checkState(!Strings.isNullOrEmpty(agentId), "agentId must not be empty");
      checkState(!Strings.isNullOrEmpty(policyStoreId), "policyStoreId must not be empty");
    }
  }

  private record Customer(String id, String displayName, String history, String firstInquiry) {

    private Customer {
      checkArgument(!Strings.isNullOrEmpty(id), "customer id must not be empty");
      checkArgument(!Strings.isNullOrEmpty(displayName), "displayName must not be empty");
      checkArgument(!Strings.isNullOrEmpty(history), "history must not be empty");
      checkArgument(!Strings.isNullOrEmpty(firstInquiry), "firstInquiry must not be empty");
    }

    private String memoryStoreName() {
      return "style-" + id;
    }
  }

  public record MemoryView(String path, String content) {}
}