File size: 6,864 Bytes
c63c390 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | from datetime import datetime, timezone
from typing import Dict, Iterable, List, Optional
from memory_policy import evaluate_memory_write
from memory_store import MemoryStore
from schemas import MemoryEvent, RetrievedMemory
class MemoryAugmentedAgent:
def __init__(self, memory_store: MemoryStore, audit_log: Optional[Iterable[dict]] = None):
self.memory_store = memory_store
self.audit_log = list(audit_log or [])
def compress_context(self, retrieved_memories: List[RetrievedMemory]) -> str:
if not retrieved_memories:
return "No relevant project memories were retrieved."
semantic = [item.memory.content for item in retrieved_memories if item.memory.memory_type == "semantic"]
episodic = [item.memory.content for item in retrieved_memories if item.memory.memory_type == "episodic"]
parts: List[str] = []
if semantic:
parts.append("Stable project context:")
parts.extend(f"- {memory}" for memory in semantic)
if episodic:
if parts:
parts.append("")
parts.append("Relevant prior decisions and events:")
parts.extend(f"- {memory}" for memory in episodic)
return "\n".join(parts)
def generate_response(self, query: str, retrieved_memories: List[RetrievedMemory]) -> str:
if not retrieved_memories:
return (
"I do not have enough remembered project context to answer this from memory. "
"The demo will not invent a project memory that was not retrieved."
)
lines = [
"Here is the most relevant remembered project context for this request:",
"",
]
for item in retrieved_memories[:4]:
lines.append(f"- {item.memory.content}")
lines.extend(
[
"",
"Memory boundary: this answer is derived only from the session memories shown in the inspection panel.",
]
)
return "\n".join(lines)
def log_run(self, result: Dict) -> dict:
retrieved_memory_ids = [item.memory.memory_id for item in result["retrieved_memories"]]
decision = result["write_decision"]
log_entry = {
"timestamp": datetime.now(timezone.utc).isoformat(),
"query_length": len(result["query"]),
"retrieved_memory_ids": retrieved_memory_ids,
"write_outcome": decision.outcome,
"write_reason": decision.reason,
"policy_flags": decision.policy_flags,
"memory_count_after_run": len(self.memory_store.memories),
}
self.audit_log.append(log_entry)
return log_entry
@staticmethod
def _frame(event: MemoryEvent, result: Dict) -> Dict:
return {"event": event, "result": dict(result)}
def run_iter(self, query: str, top_k: int = 4, save_new_memory: bool = True):
clean_query = (query or "").strip()
result: Dict = {
"query": clean_query,
"retrieved_memories": [],
"compressed_context": "",
"agent_response": "",
"write_decision": None,
"saved_memory": None,
"log_entry": None,
}
yield self._frame(
MemoryEvent(stage="query_received", message="Request received. The agent is deciding which memories are relevant."),
result,
)
retrieved = self.memory_store.retrieve(clean_query, top_k=top_k)
result["retrieved_memories"] = retrieved
yield self._frame(
MemoryEvent(
stage="memory_retrieved",
message=f"Retrieved {len(retrieved)} ranked memories from this session's memory store.",
memory_ids=[item.memory.memory_id for item in retrieved],
),
result,
)
compressed_context = self.compress_context(retrieved)
result["compressed_context"] = compressed_context
yield self._frame(
MemoryEvent(stage="context_compressed", message="Separated stable context from prior events and compressed the working context."),
result,
)
write_decision = evaluate_memory_write(clean_query)
result["write_decision"] = write_decision
yield self._frame(
MemoryEvent(
stage="write_policy_evaluated",
message=f"Memory write policy returned: {write_decision.outcome.upper()}.",
decision=write_decision.outcome,
),
result,
)
saved_memory = None
if save_new_memory and write_decision.outcome == "save":
saved_memory = self.memory_store.save_memory_from_decision(write_decision)
result["saved_memory"] = saved_memory
yield self._frame(
MemoryEvent(
stage="memory_saved",
message="The approved project memory was added to this browser session only.",
memory_ids=[saved_memory.memory_id] if saved_memory else [],
decision="save",
),
result,
)
elif write_decision.outcome == "block":
yield self._frame(
MemoryEvent(
stage="memory_blocked",
message="The application blocked this content from becoming long-term memory.",
decision="block",
),
result,
)
else:
yield self._frame(
MemoryEvent(
stage="memory_skipped",
message=(
"Memory storage is disabled for this run."
if not save_new_memory and write_decision.outcome == "save"
else "The interaction was useful for this turn but was not saved as durable memory."
),
decision="skip",
),
result,
)
result["agent_response"] = self.generate_response(clean_query, retrieved)
yield self._frame(
MemoryEvent(stage="response_ready", message="Memory-grounded response is ready."),
result,
)
result["log_entry"] = self.log_run(result)
yield self._frame(
MemoryEvent(stage="audit_recorded", message="Recorded control metadata without storing the raw query in the audit log."),
result,
)
def run(self, query: str, top_k: int = 4, save_new_memory: bool = True):
last_frame = None
for frame in self.run_iter(query=query, top_k=top_k, save_new_memory=save_new_memory):
last_frame = frame
return last_frame["result"] if last_frame else None
|