multi-tenant Phase 1: user registry + per-user memory
- CDK: add agent-claw-users DynamoDB table (actor_id PK, RETAIN policy) - CDK: grant agent-runner read/write on users table; add USERS_TABLE_NAME env - CDK: fix cdk.json app field (was object, must be command string) - CDK: add UsersTableName output - agent-runner: get_or_create_user() auto-registers users on first contact (stores display_name, telegram_username, created_at, allowed) - agent-runner: pass user_profile in AgentCore payload - prompt_builder: split base prompt (cached) from per-user context (injected per-call) removes USER.md/MEMORY.md from shared load; user name/username injected dynamically - main.py: extract user_profile from payload, build user_context string for prompt
This commit is contained in:
daniel
@@ -0,0 +1,151 @@
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
import boto3
|
||||
from typing import Any
|
||||
|
||||
# AWS clients
|
||||
_ddb = None
|
||||
_agentcore = None
|
||||
|
||||
|
||||
def get_ddb():
|
||||
global _ddb
|
||||
if _ddb is None:
|
||||
_ddb = boto3.resource('dynamodb')
|
||||
return _ddb
|
||||
|
||||
|
||||
def get_agentcore():
|
||||
global _agentcore
|
||||
if _agentcore is None:
|
||||
_agentcore = boto3.client('bedrock-agentcore', region_name='us-east-1')
|
||||
return _agentcore
|
||||
|
||||
|
||||
def get_or_create_user(actor_id: str, from_info: dict) -> dict:
|
||||
"""Look up user in registry, auto-registering on first contact."""
|
||||
table_name = os.environ.get('USERS_TABLE_NAME', '')
|
||||
if not table_name:
|
||||
return {'actor_id': actor_id, 'display_name': from_info.get('from_name', actor_id)}
|
||||
table = get_ddb().Table(table_name)
|
||||
response = table.get_item(Key={'actor_id': actor_id})
|
||||
item = response.get('Item')
|
||||
if item:
|
||||
return item
|
||||
now = int(time.time())
|
||||
item = {
|
||||
'actor_id': actor_id,
|
||||
'display_name': from_info.get('from_name') or actor_id,
|
||||
'telegram_username': from_info.get('from_username', ''),
|
||||
'created_at': str(now),
|
||||
'allowed': True,
|
||||
}
|
||||
table.put_item(Item=item)
|
||||
print(f'[agent-runner] Registered new user: {actor_id}')
|
||||
return item
|
||||
|
||||
|
||||
def get_or_create_session(actor_id: str) -> str:
|
||||
"""Look up active session for actor, or create a new one."""
|
||||
table = get_ddb().Table(os.environ['SESSION_TABLE_NAME'])
|
||||
|
||||
response = table.get_item(Key={'actor_id': actor_id})
|
||||
item = response.get('Item')
|
||||
|
||||
now = int(time.time())
|
||||
ttl_8hr = now + (8 * 3600)
|
||||
|
||||
if item and item.get('ttl', 0) > now:
|
||||
# Active session exists — extend TTL
|
||||
table.update_item(
|
||||
Key={'actor_id': actor_id},
|
||||
UpdateExpression='SET #ttl = :ttl',
|
||||
ExpressionAttributeNames={'#ttl': 'ttl'},
|
||||
ExpressionAttributeValues={':ttl': ttl_8hr},
|
||||
)
|
||||
return item['session_id']
|
||||
|
||||
# Create new session
|
||||
session_id = str(uuid.uuid4())
|
||||
table.put_item(Item={
|
||||
'actor_id': actor_id,
|
||||
'session_id': session_id,
|
||||
'created_at': str(now),
|
||||
'ttl': ttl_8hr,
|
||||
})
|
||||
return session_id
|
||||
|
||||
|
||||
def handler(event, context):
|
||||
# ── Parse SQS records (FIFO — all from same actor) ───────────────────
|
||||
records = []
|
||||
for record in event.get('Records', []):
|
||||
try:
|
||||
records.append(json.loads(record['body']))
|
||||
except (json.JSONDecodeError, KeyError):
|
||||
continue
|
||||
|
||||
if not records:
|
||||
return
|
||||
|
||||
first = records[0]
|
||||
channel = first.get('channel', 'telegram')
|
||||
chat_id = first.get('chat_id', '')
|
||||
actor_id = f"{channel}:{chat_id}"
|
||||
|
||||
# ── User registry ─────────────────────────────────────────────────────
|
||||
from_info = first.get('messages', [{}])[0]
|
||||
user_profile = get_or_create_user(actor_id, from_info)
|
||||
|
||||
# ── Get or create AgentCore session ──────────────────────────────────
|
||||
session_id = get_or_create_session(actor_id)
|
||||
print(f"[agent-runner] actor={actor_id} session={session_id} user={user_profile.get('display_name', '')}")
|
||||
|
||||
# ── Bundle messages ───────────────────────────────────────────────────
|
||||
if len(records) == 1:
|
||||
prompt = records[0]['messages'][0]['text']
|
||||
else:
|
||||
lines = [
|
||||
f"[{i+1}] {r['messages'][0]['text']}"
|
||||
for i, r in enumerate(records)
|
||||
]
|
||||
prompt = f"You have {len(records)} queued messages:\n" + "\n".join(lines)
|
||||
|
||||
# ── Build payload for AgentCore Runtime 1 ────────────────────────────
|
||||
payload: dict[str, Any] = {
|
||||
'prompt': prompt,
|
||||
'actor_id': actor_id,
|
||||
'session_id': session_id,
|
||||
'user_profile': {
|
||||
'display_name': user_profile.get('display_name', actor_id),
|
||||
'telegram_username': user_profile.get('telegram_username', ''),
|
||||
'allowed': user_profile.get('allowed', True),
|
||||
},
|
||||
'channel_adapter': {
|
||||
'type': channel,
|
||||
'target_id': str(chat_id),
|
||||
'bot_token_secret_arn': os.environ.get('TELEGRAM_BOT_TOKEN_SECRET_ARN', ''),
|
||||
},
|
||||
}
|
||||
|
||||
# ── Invoke AgentCore Runtime 1 ────────────────────────────────────────
|
||||
runtime_arn = os.environ.get('RUNTIME_1_ARN', '')
|
||||
if not runtime_arn or runtime_arn == 'PLACEHOLDER_SET_AFTER_RUNTIME_DEPLOY':
|
||||
print(f"[agent-runner] RUNTIME_1_ARN not set — skipping AgentCore invoke")
|
||||
print(f"[agent-runner] Would have sent: {json.dumps(payload)[:200]}")
|
||||
return
|
||||
|
||||
client = get_agentcore()
|
||||
response = client.invoke_agent_runtime(
|
||||
agentRuntimeArn=runtime_arn,
|
||||
runtimeSessionId=session_id,
|
||||
payload=json.dumps(payload).encode(),
|
||||
)
|
||||
|
||||
# Consume streaming response (agent delivers to Telegram via send_message tool)
|
||||
for chunk in response.get('response', []):
|
||||
pass # intentional no-op — agent handles delivery internally
|
||||
|
||||
print(f"[agent-runner] Completed session={session_id} actor={actor_id}")
|
||||
Reference in New Issue
Block a user