| """ |
| Memory Store — ChromaDB-backed trajectory memory for self-improving tool agents. |
| |
| Stores past episode experiences (query, tool sequence, reward, lesson) |
| and retrieves similar past experiences to guide future decisions. |
| """ |
|
|
| import json |
| import os |
| from datetime import datetime |
| from typing import Any |
|
|
| import chromadb |
| from chromadb.config import Settings |
|
|
|
|
| class MemoryStore: |
| """Persistent memory for agent trajectories using ChromaDB.""" |
|
|
| def __init__(self, persist_dir: str = "./data/chroma_data", collection_name: str = "tool_experiences"): |
| self.persist_dir = persist_dir |
| os.makedirs(persist_dir, exist_ok=True) |
|
|
| self.client = chromadb.PersistentClient( |
| path=persist_dir, |
| settings=Settings(anonymized_telemetry=False), |
| ) |
| self.collection = self.client.get_or_create_collection( |
| name=collection_name, |
| metadata={"hnsw:space": "cosine"}, |
| ) |
|
|
| def count(self) -> int: |
| return self.collection.count() |
|
|
| def store_experience( |
| self, |
| query: str, |
| scenario_id: int, |
| tool_sequence: list[str], |
| reward: float, |
| lesson: str, |
| should_refuse: bool = False, |
| difficulty: str = "medium", |
| episode: int = 0, |
| extra_metadata: dict[str, Any] | None = None, |
| ) -> str: |
| """Store one episode's experience in memory.""" |
| entry_id = f"ep{episode}_s{scenario_id}_{datetime.now().strftime('%H%M%S')}" |
|
|
| outcome = "correct" if reward > 0.7 else "partial" if reward > 0.3 else "wrong" |
|
|
| metadata = { |
| "scenario_id": str(scenario_id), |
| "tool_sequence": json.dumps(tool_sequence), |
| "reward": float(reward), |
| "outcome": outcome, |
| "lesson": lesson, |
| "should_refuse": str(should_refuse), |
| "difficulty": difficulty, |
| "episode": str(episode), |
| "timestamp": datetime.now().isoformat(), |
| } |
| if extra_metadata: |
| for k, v in extra_metadata.items(): |
| metadata[k] = str(v) |
|
|
| self.collection.upsert( |
| documents=[query], |
| metadatas=[metadata], |
| ids=[entry_id], |
| ) |
| return entry_id |
|
|
| def retrieve_lessons( |
| self, |
| query: str, |
| n_results: int = 3, |
| min_reward: float | None = None, |
| ) -> list[dict]: |
| """Retrieve similar past experiences for a given query.""" |
| if self.count() == 0: |
| return [] |
|
|
| n = min(n_results, self.count()) |
| results = self.collection.query( |
| query_texts=[query], |
| n_results=n, |
| ) |
|
|
| experiences = [] |
| if not results["metadatas"] or not results["metadatas"][0]: |
| return [] |
|
|
| for i, meta in enumerate(results["metadatas"][0]): |
| reward = float(meta.get("reward", 0)) |
| if min_reward is not None and reward < min_reward: |
| continue |
|
|
| experiences.append({ |
| "query": results["documents"][0][i] if results["documents"] else "", |
| "tool_sequence": json.loads(meta.get("tool_sequence", "[]")), |
| "reward": reward, |
| "outcome": meta.get("outcome", "unknown"), |
| "lesson": meta.get("lesson", ""), |
| "should_refuse": meta.get("should_refuse", "False") == "True", |
| "similarity": results["distances"][0][i] if results["distances"] else 0.0, |
| "difficulty": meta.get("difficulty", "medium"), |
| }) |
|
|
| return experiences |
|
|
| def get_tool_preference_scores( |
| self, |
| query: str, |
| tool_names: list[str], |
| n_results: int = 5, |
| ) -> dict[str, float]: |
| """Convert retrieved memories into per-tool preference scores.""" |
| experiences = self.retrieve_lessons(query, n_results=n_results) |
| if not experiences: |
| return {t: 0.0 for t in tool_names} |
|
|
| scores: dict[str, float] = {t: 0.0 for t in tool_names} |
| total_weight = 0.0 |
|
|
| for exp in experiences: |
| sim = 1.0 - exp["similarity"] |
| reward = exp["reward"] |
|
|
| if reward > 0.5: |
| weight = sim * reward |
| else: |
| weight = sim * (reward - 1.0) |
|
|
| for tool in exp["tool_sequence"]: |
| if tool in scores: |
| scores[tool] += weight |
|
|
| total_weight += abs(weight) |
|
|
| if total_weight > 0: |
| scores = {t: s / total_weight for t, s in scores.items()} |
|
|
| return scores |
|
|
| def format_lessons_for_prompt( |
| self, |
| query: str, |
| n_results: int = 3, |
| ) -> str: |
| """Format retrieved lessons as a string for prompt injection.""" |
| experiences = self.retrieve_lessons(query, n_results=n_results) |
| if not experiences: |
| return "" |
|
|
| positive = [e for e in experiences if e["reward"] > 0.5] |
| negative = [e for e in experiences if e["reward"] <= 0.5] |
|
|
| lines = [] |
| for exp in positive: |
| tools = " → ".join(exp["tool_sequence"]) if exp["tool_sequence"] else "REFUSE" |
| lines.append( |
| f" [reward={exp['reward']:.2f}] {exp['lesson']} (tools: {tools})" |
| ) |
|
|
| for exp in negative: |
| tools = " → ".join(exp["tool_sequence"]) if exp["tool_sequence"] else "REFUSE" |
| lines.append( |
| f" [AVOID, reward={exp['reward']:.2f}] {exp['lesson']} (tools: {tools})" |
| ) |
|
|
| if not lines: |
| return "" |
|
|
| return "LESSONS FROM PAST EXPERIENCE:\n" + "\n".join(lines) |
|
|
| def get_all_experiences(self, limit: int = 100) -> list[dict]: |
| """Get all stored experiences for analysis/export.""" |
| if self.count() == 0: |
| return [] |
|
|
| results = self.collection.get(limit=limit, include=["documents", "metadatas"]) |
| experiences = [] |
| for i, meta in enumerate(results["metadatas"]): |
| experiences.append({ |
| "id": results["ids"][i], |
| "query": results["documents"][i], |
| "tool_sequence": json.loads(meta.get("tool_sequence", "[]")), |
| "reward": float(meta.get("reward", 0)), |
| "outcome": meta.get("outcome", ""), |
| "lesson": meta.get("lesson", ""), |
| "episode": meta.get("episode", "0"), |
| "difficulty": meta.get("difficulty", ""), |
| "timestamp": meta.get("timestamp", ""), |
| }) |
| return experiences |
|
|
| def get_stats(self) -> dict: |
| """Get summary statistics of the memory store.""" |
| if self.count() == 0: |
| return {"total": 0, "avg_reward": 0.0, "correct": 0, "wrong": 0} |
|
|
| all_exp = self.get_all_experiences(limit=1000) |
| rewards = [e["reward"] for e in all_exp] |
| return { |
| "total": len(all_exp), |
| "avg_reward": sum(rewards) / len(rewards) if rewards else 0.0, |
| "correct": sum(1 for e in all_exp if e["outcome"] == "correct"), |
| "partial": sum(1 for e in all_exp if e["outcome"] == "partial"), |
| "wrong": sum(1 for e in all_exp if e["outcome"] == "wrong"), |
| "episodes": len(set(e["episode"] for e in all_exp)), |
| } |
|
|
| def clear(self): |
| """Clear all stored experiences.""" |
| self.client.delete_collection(self.collection.name) |
| self.collection = self.client.get_or_create_collection( |
| name=self.collection.name, |
| metadata={"hnsw:space": "cosine"}, |
| ) |
|
|
