[アップデート] AgentCore Memoryの名前空間に独自変数を使えるようになったので試してみた

"""Compare legacy and flexible namespaces on a disposable AgentCore Memory."""
import argparse
from datetime import datetime, timezone
import json
import os
from pathlib import Path
import time
import uuid

import boto3
from botocore.exceptions import ClientError

ROOT = Path(__file__).resolve().parents[1]
VARIANT = os.environ.get('MEMORY_LAB_VARIANT', 'flexible')
if VARIANT not in ('legacy', 'flexible'):
    raise ValueError('MEMORY_LAB_VARIANT must be legacy or flexible')
STATE_PATH = ROOT / f'.local/memory-namespace-{VARIANT}-state.json'
LOG_PATH = ROOT / '.local/memory-namespace-events.jsonl'
REGION = 'us-east-1'
ACTOR = 'staff001'
TEMPLATE = '/tenant/{tenant}/env/{env}/actor/{actorId}/'
CASES = [
    ('ramen', 'prod', '22時', '食券機'),
    ('cafe', 'prod', '21時', 'エスプレッソマシン'),
    ('ramen', 'training', '20時', 'レジ練習機'),
]
session = boto3.Session(region_name=REGION)
control = session.client('bedrock-agentcore-control')
data = session.client('bedrock-agentcore')
state = json.loads(STATE_PATH.read_text()) if STATE_PATH.exists() else {}

def save():
    STATE_PATH.parent.mkdir(exist_ok=True, mode=0o700)
    fd = os.open(STATE_PATH, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, 'w') as f:
        json.dump(state, f, indent=2, default=str)

def record(action, **values):
    item = {'time': datetime.now(timezone.utc).isoformat(), 'variant': VARIANT, 'action': action, **values}
    LOG_PATH.parent.mkdir(exist_ok=True, mode=0o700)
    with LOG_PATH.open('a') as f:
        f.write(json.dumps(item, ensure_ascii=False, default=str) + '\n')
    print(json.dumps(item, ensure_ascii=False, default=str), flush=True)

def setup():
    if state.get('deleted'):
        raise RuntimeError('Previous lab was deleted; archive its state before a new run.')
    if 'memory_id' not in state:
        options = ({
            'namespaceKeys': [
                {'key': 'tenant', 'validation': {'allowedValues': ['ramen', 'cafe']}},
                {'key': 'env', 'validation': {'allowedValues': ['prod', 'training']}},
            ],
            'memoryStrategies': [{'semanticMemoryStrategy': {
                'name': 'FlexibleByTenantEnv', 'namespaceTemplates': [TEMPLATE],
            }}],
        } if VARIANT == 'flexible' else {
            'memoryStrategies': [{'semanticMemoryStrategy': {
                'name': 'LegacyByActor', 'namespaces': ['/legacy/actor/{actorId}/'],
            }}],
        })
        memory = control.create_memory(
            name='NamespaceBlog_' + uuid.uuid4().hex[:12],
            description='Disposable flexible namespace comparison',
            eventExpiryDuration=3,
            tags={'Project': 'memory-namespace-blog', 'Purpose': 'disposable-test'},
            **options,
        )['memory']
        state['memory_id'] = memory['id']
        save()
        record('created', status=memory['status'])
    status()

def status():
    memory = control.get_memory(memoryId=state['memory_id'])['memory']
    state['strategies'] = {s['name']: s['strategyId'] for s in memory['strategies']}
    save()
    record('status', status=memory['status'], namespaceKeys=memory.get('namespaceKeys'),
           strategies=[{k: s.get(k) for k in ('name', 'status', 'namespaces', 'namespaceTemplates')}
                       for s in memory['strategies']])
    return memory

def create_event(session_id, text, variables):
    kwargs = dict(memoryId=state['memory_id'], actorId=ACTOR, sessionId=session_id,
                  eventTimestamp=datetime.now(timezone.utc), clientToken=str(uuid.uuid4()),
                  payload=[{'conversational': {'role': 'USER', 'content': {'text': text}}},
                           {'conversational': {'role': 'ASSISTANT', 'content': {
                               'text': '承知しました。閉店作業の手順を覚えておきます。'}}}])
    if variables is not None and VARIANT == 'flexible':
        kwargs['extractionConfig'] = {'namespaceVariables': variables}
    return data.create_event(**kwargs)

def seed():
    memory = status()
    if memory['status'] != 'ACTIVE' or any(s['status'] != 'ACTIVE' for s in memory['strategies']):
        raise RuntimeError('Wait until memory and strategies are ACTIVE.')
    for tenant, env, closing_time, item in CASES:
        key = f'{tenant}-{env}'
        if key in state.get('events', {}):
            continue
        text = (f'閉店作業は{closing_time}から始めます。'
                f'閉店作業では{item}の片付けが必要です。'
                '今後、閉店作業について質問したときはこの手順を思い出してください。')
        result = create_event(key, text, {'tenant': tenant, 'env': env})
        state.setdefault('events', {})[key] = result['event']['eventId']
        save()
        record('event_created', case=key, closing_time=closing_time, item=item,
               http_status=result['ResponseMetadata']['HTTPStatusCode'])

def wait_ready():
    deadline = time.monotonic() + 600
    while time.monotonic() < deadline:
        memory = status()
        if memory['status'] == 'FAILED' or any(s['status'] == 'FAILED' for s in memory['strategies']):
            raise RuntimeError(memory.get('failureReason', 'Memory strategy failed'))
        if memory['status'] == 'ACTIVE' and all(s['status'] == 'ACTIVE' for s in memory['strategies']):
            return
        time.sleep(15)
    raise TimeoutError('Memory did not become ACTIVE within 10 minutes.')

def list_records(namespace, strategy_id=None):
    args = dict(memoryId=state['memory_id'], namespace=namespace, maxResults=100)
    if strategy_id:
        args['memoryStrategyId'] = strategy_id
    records = []
    while True:
        response = data.list_memory_records(**args)
        records.extend(response.get('memoryRecordSummaries', []))
        if not response.get('nextToken'):
            break
        args['nextToken'] = response['nextToken']
    return records

def check():
    paths = (['/legacy/actor/staff001/'] if VARIANT == 'legacy' else [
        f'/tenant/{t}/env/{e}/actor/staff001/' for t, e, _, _ in CASES])
    results = {}
    for path in paths:
        records = list_records(path)
        results[path] = [{k: r.get(k) for k in ('content', 'namespaces', 'memoryStrategyId')}
                         for r in records]
        record('records', namespace=path, count=len(records), records=results[path])
    state['last_records'] = results
    save()
    return results

def retrieve():
    query = '閉店作業は何時から始めますか？'
    if VARIANT == 'legacy':
        namespace = '/legacy/actor/staff001/'
        response = data.retrieve_memory_records(memoryId=state['memory_id'], namespace=namespace,
            searchCriteria={'searchQuery': query, 'topK': 10})
        records = response.get('memoryRecordSummaries', [])
        text = json.dumps([r['content'] for r in records], ensure_ascii=False)
        assert all(t in text and item in text for _, _, t, item in CASES), records
        record('retrieved', namespace=namespace, records=[
            {k: r.get(k) for k in ('content', 'namespaces', 'score')} for r in records],
            all_cases_present=True, assertions_passed=True)
        return
    for tenant, env, closing_time, item in CASES:
        namespace = f'/tenant/{tenant}/env/{env}/actor/staff001/'
        response = data.retrieve_memory_records(memoryId=state['memory_id'], namespace=namespace,
            searchCriteria={'searchQuery': query, 'topK': 10})
        records = response.get('memoryRecordSummaries', [])
        text = json.dumps([r['content'] for r in records], ensure_ascii=False)
        assert closing_time in text and item in text, (namespace, records)
        assert all(other not in text for _, _, _, other in CASES if other != item), (namespace, records)
        assert all(t not in text for _, _, t, _ in CASES if t != closing_time), (namespace, records)
        assert all(r['namespaces'] == [namespace] for r in records), records
        record('retrieved', namespace=namespace, records=[
            {k: r.get(k) for k in ('content', 'namespaces', 'score')} for r in records],
            expected_case=f'{closing_time}・{item}', assertions_passed=True)

def probes():
    if 'missing_env' not in state:
        result = create_event('missing-env',
            '閉店作業のあとに日報を提出します。日報の提出期限は23時です。'
            '今後、日報について質問したときはこの期限を思い出してください。', {'tenant': 'ramen'})
        state['missing_env'] = result['event']['eventId']
        save()
        stored = data.get_event(memoryId=state['memory_id'], actorId=ACTOR,
                                sessionId='missing-env', eventId=state['missing_env'])
        record('missing_env', http_status=result['ResponseMetadata']['HTTPStatusCode'],
               persisted_in_short_term=bool(stored['event']['payload']))
    for label, variables in [
        ('invalid_allowed_value', {'tenant': 'ramen', 'env': 'staging'}),
        ('uppercase_value', {'tenant': 'Ramen', 'env': 'prod'}),
    ]:
        try:
            result = create_event(label.replace('_', '-'), '新メニューの試食は毎週金曜です。', variables)
        except ClientError as e:
            record(label, code=e.response['Error']['Code'], message=e.response['Error']['Message'])
        else:
            record(label, unexpected_success=True, http_status=result['ResponseMetadata']['HTTPStatusCode'])

def jobs():
    args = dict(memoryId=state['memory_id'], maxResults=50)
    result = []
    while True:
        response = data.list_memory_extraction_jobs(**args)
        result.extend(response.get('jobs', []))
        if not response.get('nextToken'):
            break
        args['nextToken'] = response['nextToken']
    record('extraction_jobs', response=result)

def cleanup():
    try:
        response = control.delete_memory(memoryId=state['memory_id'])
        record('delete_requested', status=response.get('status'))
    except control.exceptions.ResourceNotFoundException:
        pass
    for _ in range(60):
        try:
            control.get_memory(memoryId=state['memory_id'])
        except control.exceptions.ResourceNotFoundException:
            state['deleted'] = True
            save()
            record('deleted', confirmed=True)
            return
        time.sleep(5)
    raise TimeoutError('Deletion still pending; rerun cleanup.')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('command', choices=['setup', 'status', 'wait_ready', 'seed', 'check', 'retrieve', 'probes', 'jobs', 'cleanup'])
    globals()[parser.parse_args().command]()

// namespace=/legacy/actor/staff001/（6件）
[
  {
    "content": {
      "text": "閉店作業は21時から始まる。"
    },
    "namespaces": [
      "/legacy/actor/staff001/"
    ]
  },
  {
    "content": {
      "text": "閉店作業では食券機の片付けが必要である。"
    },
    "namespaces": [
      "/legacy/actor/staff001/"
    ]
  },
  {
    "content": {
      "text": "閉店作業は20時から始まる。"
    },
    "namespaces": [
      "/legacy/actor/staff001/"
    ]
  },
  {
    "content": {
      "text": "閉店作業は22時から始まる。"
    },
    "namespaces": [
      "/legacy/actor/staff001/"
    ]
  },
  {
    "content": {
      "text": "閉店作業ではエスプレッソマシンの片付けが必要である。"
    },
    "namespaces": [
      "/legacy/actor/staff001/"
    ]
  },
  {
    "content": {
      "text": "閉店作業ではレジ練習機の片付けが必要である。"
    },
    "namespaces": [
      "/legacy/actor/staff001/"
    ]
  }
]

// namespace=/tenant/ramen/env/prod/actor/staff001/（2件）
[
  {
    "content": {
      "text": "閉店作業は22時から始まる。"
    },
    "namespaces": [
      "/tenant/ramen/env/prod/actor/staff001/"
    ]
  },
  {
    "content": {
      "text": "閉店作業では食券機の片付けが必要である。"
    },
    "namespaces": [
      "/tenant/ramen/env/prod/actor/staff001/"
    ]
  }
]

// namespace=/tenant/cafe/env/prod/actor/staff001/（2件）
[
  {
    "content": {
      "text": "閉店作業は21時から始まる。"
    },
    "namespaces": [
      "/tenant/cafe/env/prod/actor/staff001/"
    ]
  },
  {
    "content": {
      "text": "閉店作業ではエスプレッソマシンの片付けが必要である。"
    },
    "namespaces": [
      "/tenant/cafe/env/prod/actor/staff001/"
    ]
  }
]

// namespace=/tenant/ramen/env/training/actor/staff001/（2件）
[
  {
    "content": {
      "text": "閉店作業ではレジ練習機の片付けが必要である。"
    },
    "namespaces": [
      "/tenant/ramen/env/training/actor/staff001/"
    ]
  },
  {
    "content": {
      "text": "閉店作業は20時から始まる。"
    },
    "namespaces": [
      "/tenant/ramen/env/training/actor/staff001/"
    ]
  }
]

[
  {
    "content": {
      "text": "閉店作業は22時から始まる。"
    },
    "namespaces": [
      "/legacy/actor/staff001/"
    ],
    "score": 0.7048244
  },
  {
    "content": {
      "text": "閉店作業は20時から始まる。"
    },
    "namespaces": [
      "/legacy/actor/staff001/"
    ],
    "score": 0.67026436
  },
  {
    "content": {
      "text": "閉店作業は21時から始まる。"
    },
    "namespaces": [
      "/legacy/actor/staff001/"
    ],
    "score": 0.6684311
  },
  {
    "content": {
      "text": "閉店作業ではエスプレッソマシンの片付けが必要である。"
    },
    "namespaces": [
      "/legacy/actor/staff001/"
    ],
    "score": 0.47130024
  },
  {
    "content": {
      "text": "閉店作業ではレジ練習機の片付けが必要である。"
    },
    "namespaces": [
      "/legacy/actor/staff001/"
    ],
    "score": 0.46928594
  },
  {
    "content": {
      "text": "閉店作業では食券機の片付けが必要である。"
    },
    "namespaces": [
      "/legacy/actor/staff001/"
    ],
    "score": 0.4545679
  }
]

// namespace=/tenant/ramen/env/prod/actor/staff001/ を指定した結果
[
  {
    "content": {
      "text": "閉店作業は22時から始まる。"
    },
    "namespaces": [
      "/tenant/ramen/env/prod/actor/staff001/"
    ],
    "score": 0.7048244
  },
  {
    "content": {
      "text": "閉店作業では食券機の片付けが必要である。"
    },
    "namespaces": [
      "/tenant/ramen/env/prod/actor/staff001/"
    ],
    "score": 0.45456797
  }
]

// namespace=/tenant/cafe/env/prod/actor/staff001/ を指定した結果
[
  {
    "content": {
      "text": "閉店作業は21時から始まる。"
    },
    "namespaces": [
      "/tenant/cafe/env/prod/actor/staff001/"
    ],
    "score": 0.66843104
  },
  {
    "content": {
      "text": "閉店作業ではエスプレッソマシンの片付けが必要である。"
    },
    "namespaces": [
      "/tenant/cafe/env/prod/actor/staff001/"
    ],
    "score": 0.4713003
  }
]

// namespace=/tenant/ramen/env/training/actor/staff001/ を指定した結果
[
  {
    "content": {
      "text": "閉店作業は20時から始まる。"
    },
    "namespaces": [
      "/tenant/ramen/env/training/actor/staff001/"
    ],
    "score": 0.67026436
  },
  {
    "content": {
      "text": "閉店作業ではレジ練習機の片付けが必要である。"
    },
    "namespaces": [
      "/tenant/ramen/env/training/actor/staff001/"
    ],
    "score": 0.46928594
  }
]

{"variant": "flexible", "action": "missing_env", "http_status": 201, "persisted_in_short_term": true}
{"variant": "flexible", "action": "invalid_allowed_value", "code": "ValidationException", "message": "Value 'staging' for namespace variable 'env' failed validation"}
{"variant": "flexible", "action": "uppercase_value", "code": "ValidationException", "message": "1 validation error detected: Value at 'extractionConfig.namespaceVariables' failed to satisfy constraint: Map value must satisfy constraint: [Member must have length less than or equal to 64, Member must have length greater than or equal to 1, Member must satisfy regular expression pattern: [a-z0-9][a-z0-9-_]*]"}

"""Search tenant-scoped AgentCore Memory records from a Strands agent tool."""
import json
import os
from pathlib import Path

import boto3
from strands import Agent, tool
from strands.models import BedrockModel

REGION = 'us-east-1'
data = boto3.client('bedrock-agentcore', region_name=REGION)

state_path = Path(__file__).resolve().parents[1] / '.local/memory-namespace-flexible-state.json'
MEMORY_ID = os.environ.get('MEMORY_ID') or json.loads(state_path.read_text())['memory_id']

def make_memory_tools(memory_id: str, tenant: str, env: str, actor_id: str):
    # JWTなどのユーザー属性から判別したテナント・環境をクロージャで固定する
    namespace = f'/tenant/{tenant}/env/{env}/actor/{actor_id}/'

    @tool
    def search_store_memory(query: str) -> str:
        """いま担当している現場の記憶（閉店作業のルールなど）を検索する"""
        response = data.retrieve_memory_records(
            memoryId=memory_id,
            namespace=namespace,
            searchCriteria={'searchQuery': query, 'topK': 5},
        )
        records = [r['content']['text'] for r in response['memoryRecordSummaries']]
        return '\n'.join(records) if records else '該当する記憶はありません'

    return [search_store_memory]

if __name__ == '__main__':
    agent = Agent(
        model=BedrockModel(model_id='us.anthropic.claude-haiku-4-5-20251001-v1:0', region_name=REGION),
        tools=make_memory_tools(MEMORY_ID, tenant='ramen', env='prod', actor_id='staff001'),
    )
    agent('閉店作業は何時から始めますか？')

閉店作業の開始時間についての情報を検索させていただきます。
Tool #1: search_store_memory
閉店作業は**22時から始まります**。

閉店作業では食券機の片付けなどの業務も含まれています。

"""Let the LLM choose the tenant and env as tool arguments (compare with memory_tool_agent.py)."""
import json
import os
from pathlib import Path

import boto3
from strands import Agent, tool
from strands.models import BedrockModel

REGION = 'us-east-1'
MODEL_ID = 'us.anthropic.claude-haiku-4-5-20251001-v1:0'
data = boto3.client('bedrock-agentcore', region_name=REGION)

state_path = Path(__file__).resolve().parents[1] / '.local/memory-namespace-flexible-state.json'
MEMORY_ID = os.environ.get('MEMORY_ID') or json.loads(state_path.read_text())['memory_id']

@tool
def search_store_memory(query: str, tenant: str, env: str) -> str:
    """現場の記憶を検索する。tenantはramenまたはcafe、envはprodまたはtraining"""
    namespace = f'/tenant/{tenant}/env/{env}/actor/staff001/'
    print(f'[tool called with tenant={tenant} env={env}]')
    response = data.retrieve_memory_records(
        memoryId=MEMORY_ID,
        namespace=namespace,
        searchCriteria={'searchQuery': query, 'topK': 5},
    )
    records = [r['content']['text'] for r in response['memoryRecordSummaries']]
    return '\n'.join(records) if records else '該当する記憶はありません'

def build_agent() -> Agent:
    return Agent(
        model=BedrockModel(model_id=MODEL_ID, region_name=REGION),
        tools=[search_store_memory],
    )

if __name__ == '__main__':
    build_agent()('カフェの閉店作業は何時から始めますか？')
    print('=== 次: あいまいな聞き方 ===')
    build_agent()('閉店作業は何時から始めますか？')

カフェの閉店作業の時間を検索します。
Tool #1: search_store_memory
[tool called with tenant=cafe env=prod]
カフェの閉店作業は**21時から始まります**。

閉店作業ではエスプレッソマシンの片付けなどが必要とのことです。
=== 次: あいまいな聞き方 ===
閉店作業の開始時刻を確認するために、現場の記憶を検索させていただきます。

ただし、どの店舗の情報を探すのか、また本番環境（prod）と訓練環境（training）のどちらを確認するのか教えていただけますか？

- **テナント**: ramen（ラーメン店）または cafe（カフェ）のどちらですか？
- **環境**: prod（本番）または training（訓練）のどちらですか？