用 AlloyDB AI 与 Valkey 构建两级 Agent 记忆：让长期任务不再靠塞满上下文

python -m venv .venv
source .venv/bin/activate
pip install redis
export VALKEY_HOST=127.0.0.1
export VALKEY_PORT=6379

import json
import os
import time
from redis import Redis

client = Redis(
    host=os.environ.get("VALKEY_HOST", "127.0.0.1"),
    port=int(os.environ.get("VALKEY_PORT", "6379")),
    decode_responses=True,
    ssl=os.environ.get("VALKEY_TLS", "false").lower() == "true",
)

MAX_MESSAGES = 20
SESSION_TTL_SECONDS = 24 * 60 * 60


def append_message(user_id: str, session_id: str, role: str, content: str) -> None:
    key = f"agent:session:{user_id}:{session_id}:messages"
    message = json.dumps(
        {"role": role, "content": content, "timestamp": int(time.time())},
        ensure_ascii=False,
    )

    with client.pipeline(transaction=True) as pipe:
        pipe.rpush(key, message)
        pipe.ltrim(key, -MAX_MESSAGES, -1)
        pipe.expire(key, SESSION_TTL_SECONDS)
        pipe.execute()


def load_messages(user_id: str, session_id: str) -> list[dict]:
    key = f"agent:session:{user_id}:{session_id}:messages"
    return [json.loads(item) for item in client.lrange(key, 0, -1)]


if __name__ == "__main__":
    append_message("user-42", "trip-2025", "user", "我只接受直飞航班。")
    append_message("user-42", "trip-2025", "assistant", "已记录，本次行程仅筛选直飞航班。")
    print(json.dumps(load_messages("user-42", "trip-2025"), ensure_ascii=False, indent=2))

CREATE EXTENSION IF NOT EXISTS google_ml_integration CASCADE;
CREATE EXTENSION IF NOT EXISTS vector CASCADE;
CREATE EXTENSION IF NOT EXISTS rum CASCADE;

CREATE TABLE IF NOT EXISTS agent_entities (
    entity_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    user_id TEXT NOT NULL,
    entity_name TEXT NOT NULL,
    project_id TEXT NOT NULL DEFAULT 'global',
    session_id TEXT,
    scope TEXT NOT NULL DEFAULT 'session'
        CHECK (scope IN ('session', 'project', 'global')),
    memory_type TEXT NOT NULL DEFAULT 'rule'
        CHECK (memory_type IN ('rule', 'preference', 'episode', 'summary')),
    summary TEXT NOT NULL,
    source_message_id TEXT,
    confidence NUMERIC(4, 3),
    summary_embedding VECTOR(768),
    summary_tsv TSVECTOR GENERATED ALWAYS AS (
        to_tsvector('english', entity_name || ' ' || summary)
    ) STORED,
    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    UNIQUE (user_id, project_id, entity_name)
);

CREATE INDEX IF NOT EXISTS agent_entities_scope_idx
    ON agent_entities (user_id, project_id, scope, updated_at DESC);

CREATE INDEX IF NOT EXISTS agent_entities_embedding_idx
    ON agent_entities USING hnsw (summary_embedding vector_cosine_ops);

CREATE INDEX IF NOT EXISTS agent_entities_tsv_idx
    ON agent_entities USING rum (summary_tsv rum_tsvector_ops);

CALL ai.initialize_embeddings(
    model_id => 'text-embedding-005',
    table_name => 'agent_entities',
    content_column => 'summary',
    embedding_column => 'summary_embedding',
    incremental_refresh_mode => 'transactional',
    batch_size => 10
);

INSERT INTO agent_entities (
    user_id,
    project_id,
    entity_name,
    scope,
    memory_type,
    summary,
    source_message_id,
    confidence
)
VALUES (
    'user-42',
    'trip-2025',
    'flight_preference',
    'global',
    'preference',
    'User accepts non-stop flights only.',
    'msg-0182',
    0.990
)
ON CONFLICT (user_id, project_id, entity_name)
DO UPDATE SET
    summary = EXCLUDED.summary,
    source_message_id = EXCLUDED.source_message_id,
    confidence = EXCLUDED.confidence,
    updated_at = NOW();

System policy: 1,500 tokens
Retrieved hard rules: 1,000 tokens
Relevant episodic memories: 2,000 tokens
Recent Valkey session window: 4,000 tokens
Current user request and tool schema: remaining budget