طراحی معماری حافظه برای عوامل هوش مصنوعی دولتی

# context_budget.py
# Instruments an agent's session log so you can see context pressure build
# up turn by turn, instead of discovering it when a request silently truncates
# or the provider returns a context-length error mid-conversation.

from dataclasses import dataclass, field

# In production, count with the model's real tokenizer (for example tiktoken's
# cl100k_base or the o200k_base encoding). We wrap it so a missing encoding or
# offline build degrades to a conservative heuristic instead of crashing the
# request path, which is what you actually want in a hot loop.
def count_tokens(text: str) -> int:
    try:
        import tiktoken
        return len(tiktoken.get_encoding("o200k_base").encode(text))
    except Exception:
        # ~4 characters per token is a safe upper-bound heuristic for English.
        return (len(text) + 3) // 4

@dataclass
class SessionLog:
    budget: int                       # hard ceiling for prompt tokens
    events: list = field(default_factory=list)

    def add(self, role: str, text: str) -> None:
        self.events.append((role, text, count_tokens(text)))

    def used(self) -> int:
        return sum(t for _, _, t in self.events)

    def pressure(self) -> float:
        # Fraction of the budget consumed. Alert teams above ~0.75 so there is
        # headroom for the model's own output tokens, not just the prompt.
        return self.used() / self.budget

log = SessionLog(budget=2000)
for turn in range(1, 13):
    log.add("user", "Walk me through the failure we saw in the nightly run. " * turn)
    log.add("assistant", "Here is the breakdown of the root cause and the fix. " * turn)
    if log.pressure() > 0.75:
        print(f"turn {turn:>2}: used={log.used():>5} tokens  "
              f"pressure={log.pressure():.0%}  ACTION: evict or summarize")
    else:
        print(f"turn {turn:>2}: used={log.used():>5} tokens  pressure={log.pressure():.0%}")

# namespaces.py
# Memory records are addressed by a hierarchical namespace, not a flat key.
# The hierarchy is what enforces multi-tenant isolation: a retrieval scoped to
# one tenant can never resolve a key that belongs to another, because the
# tenant id is a structural prefix of every key rather than a filter applied
# after the fetch.

import re

# Components are validated so a caller cannot smuggle a separator into an id and
# escape its own subtree. This is the single most common isolation bug in
# home-grown memory stores: an unescaped "https://dzone.com/" in a user id reads as a new level.
_SAFE = re.compile(r"^[A-Za-z0-9_-]+$")

def make_namespace(tenant: str, actor: str, scope: str) -> str:
    for part in (tenant, actor, scope):
        if not _SAFE.match(part):
            raise ValueError(f"unsafe namespace component: {part!r}")
    return f"/{tenant}/{actor}/{scope}"

class MemoryStore:
    def __init__(self):
        self._records: dict[str, list[str]] = {}

    def put(self, namespace: str, value: str) -> None:
        self._records.setdefault(namespace, []).append(value)

    def query(self, namespace_prefix: str) -> list[str]:
        # Only keys under the caller's prefix are visible. The prefix check is
        # the isolation boundary, so it runs before any ranking or scoring.
        hits = []
        for key, values in self._records.items():
            if key == namespace_prefix or key.startswith(namespace_prefix + "https://dzone.com/"):
                hits.extend(values)
        return hits

store = MemoryStore()
store.put(make_namespace("tenant-a", "user-42", "preferences"), "prefers metric units")
store.put(make_namespace("tenant-a", "user-42", "facts"), "deploys to two regions")
store.put(make_namespace("tenant-b", "user-42", "preferences"), "prefers imperial units")

print("tenant-a/user-42 sees:", store.query("/tenant-a/user-42"))
print("tenant-b/user-42 sees:", store.query("/tenant-b/user-42"))

# A caller trying to break out of its subtree by embedding a separator is
# rejected at construction time, before any record is written or read.
try:
    make_namespace("tenant-a", "user-42/../tenant-b", "preferences")
except ValueError as exc:
    print("isolation guard:", exc)

# extraction_worker.py
# Raw session events are cheap to append and useless to search. The extraction
# worker runs off the request path, reads batches of raw events and distills
# them into durable long-term records: facts, preferences and rolling
# summaries. Keeping this asynchronous is the point. Distillation calls a model
# and takes hundreds of milliseconds, so doing it inline would tax every turn.

import queue
import threading
import time

raw_events: "queue.Queue[dict]" = queue.Queue()
long_term: list[dict] = []

def distill(batch: list[dict]) -> list[dict]:
    # PLACEHOLDER for the model call. In production this is a structured-output
    # request to an extraction model that returns typed records. The rule-based
    # stub below keeps the worker deterministic and runnable for tests; swap it
    # for the real call and keep the same return contract.
    records = []
    for ev in batch:
        text = ev["text"].lower()
        if "prefer" in text:
            records.append({"type": "preference", "value": ev["text"], "src": ev["id"]})
        elif any(k in text for k in ("region", "deploy", "cluster", "schema")):
            records.append({"type": "fact", "value": ev["text"], "src": ev["id"]})
    return records

def worker(batch_size: int = 4, idle_flush: float = 0.2) -> None:
    batch: list[dict] = []
    last = time.monotonic()
    while True:
        try:
            ev = raw_events.get(timeout=idle_flush)
            if ev is None:                      # poison pill: drain then stop
                if batch:
                    long_term.extend(distill(batch))
                break
            batch.append(ev)
        except queue.Empty:
            pass
        # Flush on a full batch or after an idle gap, so a quiet session still
        # gets its memory persisted instead of waiting for more traffic.
        if len(batch) >= batch_size or (batch and time.monotonic() - last > idle_flush):
            long_term.extend(distill(batch))
            batch, last = [], time.monotonic()

t = threading.Thread(target=worker, daemon=True)
t.start()
for i, msg in enumerate([
    "I prefer metric units in reports",
    "we deploy to two regions",
    "the order schema changed last week",
    "thanks, that helps",
    "I prefer dark themed dashboards",
]):
    raw_events.put({"id": f"e{i}", "text": msg})
raw_events.put(None)
t.join(timeout=2)
for rec in long_term:
    print(rec)

# agent_turn.py
# One agent turn that reads from both memory layers and respects a token
# budget. Long-term records come first because they are dense and stable;
# recent raw turns fill whatever budget remains. This ordering matters: models
# attend less reliably to the middle of a long prompt, so the highest-value
# memory belongs near the edges, not buried in the transcript.

def count_tokens(text: str) -> int:
    return (len(text) + 3) // 4   # heuristic; production swaps in the real tokenizer

def assemble_context(long_term: list[str], recent: list[str], budget: int) -> list[str]:
    context, used = [], 0
    # Long-term memory is admitted first, up to half the budget, so a long
    # transcript can never crowd out durable facts the agent depends on.
    for rec in long_term:
        cost = count_tokens(rec)
        if used + cost > budget // 2:
            break
        context.append(rec)
        used += cost
    # Recent turns fill the remainder, newest first, dropping the oldest.
    for turn in reversed(recent):
        cost = count_tokens(turn)
        if used + cost > budget:
            break
        context.append(turn)
        used += cost
    return context

def run_turn(user_msg, long_term, recent, budget=400):
    # Reserve room for the user message before admitting any memory, otherwise
    # the assembled prompt overshoots the budget once the message is appended.
    reserve = count_tokens(f"user: {user_msg}") + 1
    context = assemble_context(long_term, recent, budget - reserve)
    prompt = "\n".join(context + [f"user: {user_msg}"])
    # PLACEHOLDER for the model call. The returned context window is what the
    # agent actually reasons over; assert it before spending tokens on inference.
    assert count_tokens(prompt) 

# ranking.py
# Retrieval that scores only on semantic similarity will happily surface a
# year-old preference the user has since reversed. Blending a recency decay
# into the score fixes that, but the decay rate is a real tuning knob: too
# aggressive and the agent forgets stable facts, too gentle and stale memory
# wins. This compares pure similarity against a recency-weighted score on the
# same candidate set so the trade-off is visible.

import numpy as np

# query vector and candidate vectors stand in for embeddings; in production
# these come from an embedding model and a vector index, but the ranking math
# is identical regardless of where the vectors originate.
query = np.array([0.9, 0.1, 0.0])
candidates = [
    {"text": "prefers metric units",   "vec": np.array([0.88, 0.12, 0.0]), "age_days": 210},
    {"text": "prefers imperial units",  "vec": np.array([0.86, 0.14, 0.0]), "age_days": 3},
    {"text": "works on payments",        "vec": np.array([0.20, 0.80, 0.1]), "age_days": 12},
]

def cosine(a, b):
    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))

def rank(decay_per_day: float):
    scored = []
    for c in candidates:
        sim = cosine(query, c["vec"])
        # Exponential recency decay. A fresher record with slightly lower
        # similarity can now outrank a near-duplicate that is months old.
        recency = np.exp(-decay_per_day * c["age_days"])
        blended = sim if decay_per_day == 0 else 0.6 * sim + 0.4 * recency
        scored.append((round(float(blended), 3), c["text"]))
    return sorted(scored, reverse=True)

print("pure semantic:        ", rank(decay_per_day=0.0)[0])
print("recency-weighted 0.01:", rank(decay_per_day=0.01)[0])
print("recency-weighted 0.05:", rank(decay_per_day=0.05)[0])