codex-agent-framework 0.1.1__py3-none-any.whl

codex_agent/memory.py

@@ -0,0 +1,318 @@
+from datetime import datetime
+import json
+import os
+from typing import List
+from xml.sax.saxutils import escape, quoteattr
+
+import numpy as np
+from modict import modict
+
+from .message import DeveloperMessage, Message
+from .utils import short_id, token_count, timestamp
+
+
+class MemoryEntry(DeveloperMessage):
+    title = ""
+    source = "manual"
+    metadata = modict.factory(dict)
+
+    def to_response_format(self):
+        return None
+
+
+def format_memory_results_xml(results, root=None):
+    entries = "\n".join(
+        (
+            f"<memory id={quoteattr(str(entry.id))} "
+            f"title={quoteattr(str(entry.get('title') or ''))}>"
+            f"\n{escape(str(entry.content))}\n"
+            f"</memory>"
+        )
+        for entry in results
+    )
+    if root is None:
+        return entries
+    return f"<{root}>\n{entries}\n</{root}>"
+
+
+class RAGMemoryConfig(modict):
+    _config = modict.config(enforce_json=True)
+
+    precision: int = 5
+    dimensions: int = 128
+    half_life_days: float = 2
+    half_life_steps: float = 2
+    moment_exponent: float = 3
+    max_tokens: int | None = None
+
+
+class RAGMemory(modict):
+    """Timestamped embedding retrieval over an explicit message collection."""
+
+    settings: RAGMemoryConfig = modict.factory(RAGMemoryConfig)
+    messages: List[Message] = modict.factory(list)
+    file: str | None = None
+
+    def __init__(self, *args, agent=None, file=None, messages=None, **kwargs):
+        config_data = {}
+        for key in list(kwargs):
+            if key in RAGMemoryConfig.__fields__:
+                config_data[key] = kwargs.pop(key)
+        super().__init__(*args, **kwargs)
+        self.set_attr("_agent", agent)
+        if config_data:
+            self.settings.update(config_data)
+        if file is not None:
+            self.file = file
+        if messages is not None:
+            self.messages = list(messages)
+        elif self.file:
+            self.load()
+
+    @property
+    def agent(self):
+        return self.get_attr("_agent", None)
+
+    def add(self, message):
+        self.messages.append(message)
+        return message
+
+    def load(self, file=None):
+        self.file = file or self.file
+        if not self.file or not os.path.isfile(self.file):
+            return self
+        with open(self.file, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        self.settings.update(data.get("settings") or {})
+        self.messages = [
+            MemoryEntry(entry)
+            for entry in data.get("messages", [])
+        ]
+        return self
+
+    def dump(self, file=None):
+        self.file = file or self.file
+        if not self.file:
+            return self
+        parent = os.path.dirname(self.file)
+        if parent:
+            os.makedirs(parent, exist_ok=True)
+        data = modict(
+            type="memory",
+            settings=self.settings,
+            messages=self.messages,
+        )
+        with open(self.file, "w", encoding="utf-8") as f:
+            json.dump(data, f, ensure_ascii=False, indent=2)
+        return self
+
+    def remember(self, content, title="", source="manual", metadata=None, memory_id=None):
+        entry = MemoryEntry(
+            id=memory_id or short_id(),
+            content=content,
+            title=title,
+            source=source,
+            metadata=metadata or {},
+            timestamp=timestamp(),
+        )
+        self.ensure_embedding(entry)
+        self.add(entry)
+        self.dump()
+        return entry
+
+    def get(self, memory_id):
+        for entry in self.messages:
+            if entry.get("id") == memory_id:
+                return entry
+        return None
+
+    def find(self, ref):
+        entry = self.get(ref)
+        if entry is not None:
+            return entry
+        for entry in self.messages:
+            if entry.get("title") == ref:
+                return entry
+        return None
+
+    def edit(self, memory_id, content=None, title=None, source=None, metadata=None):
+        entry = self.find(memory_id)
+        if entry is None:
+            raise KeyError(memory_id)
+        if content is not None and content != entry.get("content"):
+            entry.content = content
+            entry.embedding = None
+            self.ensure_embedding(entry)
+        if title is not None:
+            entry.title = title
+        if source is not None:
+            entry.source = source
+        if metadata is not None:
+            entry.metadata = metadata
+        entry.timestamp = timestamp()
+        self.dump()
+        return entry
+
+    def delete(self, memory_id):
+        entry = self.find(memory_id)
+        if entry is None:
+            return None
+        self.messages = [
+            message for message in self.messages
+            if message.get("id") != entry.get("id")
+        ]
+        self.dump()
+        return entry
+
+    def ensure_embedding(self, message):
+        if message.get("embedding") is not None:
+            return message.embedding
+        if self.agent is None:
+            raise ValueError("Message has no embedding and no agent was provided to embed it.")
+        return self.agent.ai.embed_message(
+            message,
+            precision=self.settings.precision,
+            dimensions=self.settings.dimensions,
+        )
+
+    def ensure_embeddings(self, messages):
+        for message in messages:
+            self.ensure_embedding(message)
+
+    def similarity(self, msg1, msg2):
+        self.ensure_embedding(msg1)
+        self.ensure_embedding(msg2)
+        return float((np.dot(msg1.embedding, msg2.embedding) + 1) / 2)
+
+    @staticmethod
+    def recency_weight(k: int, n: int, half_life: float) -> float:
+        if n == 1:
+            return 1
+        if not 0 <= k < n:
+            raise ValueError("k must be between in range(n)")
+        if half_life <= 0:
+            raise ValueError("half_life must be positive")
+
+        decay = np.log(2) / half_life
+        ratio = np.exp(-decay)
+        numerator = (1 - ratio) * ratio**k
+        denominator = 1 - ratio**n
+        return float(numerator / denominator)
+
+    def relevance_score(self, message, context, exponent=None):
+        if not context:
+            return 0
+        exponent = exponent or self.settings.moment_exponent
+        n = len(context)
+        similarities = [self.similarity(message, ctx_msg) for ctx_msg in reversed(context)]
+        weights = [self.recency_weight(k, n, self.settings.half_life_steps) for k in range(n)]
+        score = sum(
+            weight * similarity**exponent
+            for weight, similarity in zip(weights, similarities)
+        )
+        return float(score ** (1.0 / exponent))
+
+    def temporal_score(self, message, now=None):
+        now = now or datetime.now()
+        timestamp = message.get("timestamp")
+        msg_time = self.parse_timestamp(timestamp)
+        delta_days = (now - msg_time).total_seconds() / (3600 * 24)
+        decay = np.log(2) / self.settings.half_life_days
+        return float(np.exp(-decay * delta_days))
+
+    @staticmethod
+    def parse_timestamp(value):
+        if isinstance(value, datetime):
+            return value
+        return datetime.strptime(value, "%Y%m%dT%H%M%S.%f")
+
+    def filter_by_timestamp(self, messages, since=None, until=None):
+        if since is None and until is None:
+            return list(messages)
+        since_time = self.parse_timestamp(since) if since is not None else None
+        until_time = self.parse_timestamp(until) if until is not None else None
+        filtered = []
+        for message in messages:
+            msg_time = self.parse_timestamp(message.get("timestamp"))
+            if since_time is not None and msg_time < since_time:
+                continue
+            if until_time is not None and msg_time > until_time:
+                continue
+            filtered.append(message)
+        return filtered
+
+    def score(self, message, context, recent_is_better=True, now=None):
+        score = self.relevance_score(message, context=context)
+        if recent_is_better:
+            score += self.temporal_score(message, now=now)
+        return score
+
+    def rank(self, context, messages=None, recent_is_better=True, now=None):
+        messages = list(self.messages if messages is None else messages)
+        self.ensure_embeddings(context)
+        self.ensure_embeddings(messages)
+        return sorted(
+            messages,
+            key=lambda message: self.score(
+                message,
+                context=context,
+                recent_is_better=recent_is_better,
+                now=now,
+            ),
+            reverse=True,
+        )
+
+    def retrieve(
+        self,
+        context,
+        messages=None,
+        limit=None,
+        max_tokens=None,
+        recent_is_better=True,
+        exclude_session_id=None,
+        since=None,
+        until=None,
+        now=None,
+    ):
+        if messages is None:
+            messages = self.filter_by_timestamp(self.messages, since=since, until=until)
+        else:
+            messages = self.filter_by_timestamp(messages, since=since, until=until)
+        results = self.rank(
+            context,
+            messages=messages,
+            recent_is_better=recent_is_better,
+            now=now,
+        )
+        if exclude_session_id is not None:
+            results = [
+                entry for entry in results
+                if entry.get("session_id") != exclude_session_id
+            ]
+
+        selected = []
+        total = 0
+        token_budget = self.settings.max_tokens if max_tokens is None else max_tokens
+        for entry in results:
+            if limit is not None and len(selected) >= limit:
+                break
+            if token_budget is not None:
+                tokens = token_count(str(entry.get("content") or ""))
+                if total + tokens > token_budget:
+                    break
+                total += tokens
+            selected.append(entry)
+        return selected
+
+    def search(self, query, limit=5, recent_is_better=False, since=None, until=None, **kwargs):
+        query_message = Message(content=query)
+        self.ensure_embedding(query_message)
+        return self.retrieve(
+            [query_message],
+            limit=limit,
+            recent_is_better=recent_is_better,
+            since=since,
+            until=until,
+            **kwargs,
+        )