autobyteus 1.2.1__py3-none-any.whl → 1.3.0__py3-none-any.whl

autobyteus/memory/models/raw_trace_item.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+from .memory_types import MemoryType
+
+
+@dataclass
+class RawTraceItem:
+    id: str
+    ts: float
+    turn_id: str
+    seq: int
+    trace_type: str
+    content: str
+    source_event: str
+    media: Optional[Dict[str, List[str]]] = None
+    tool_name: Optional[str] = None
+    tool_call_id: Optional[str] = None
+    tool_args: Optional[Dict[str, Any]] = None
+    tool_result: Optional[Any] = None
+    tool_error: Optional[str] = None
+    correlation_id: Optional[str] = None
+    tags: List[str] = field(default_factory=list)
+    tool_result_ref: Optional[str] = None
+
+    @property
+    def memory_type(self) -> MemoryType:
+        return MemoryType.RAW_TRACE
+
+    def to_dict(self) -> Dict[str, Any]:
+        data: Dict[str, Any] = {
+            "id": self.id,
+            "ts": self.ts,
+            "turn_id": self.turn_id,
+            "seq": self.seq,
+            "trace_type": self.trace_type,
+            "content": self.content,
+            "source_event": self.source_event,
+        }
+        if self.media is not None:
+            data["media"] = self.media
+        if self.tool_name is not None:
+            data["tool_name"] = self.tool_name
+        if self.tool_call_id is not None:
+            data["tool_call_id"] = self.tool_call_id
+        if self.tool_args is not None:
+            data["tool_args"] = self.tool_args
+        if self.tool_result is not None:
+            data["tool_result"] = self.tool_result
+        if self.tool_error is not None:
+            data["tool_error"] = self.tool_error
+        if self.correlation_id is not None:
+            data["correlation_id"] = self.correlation_id
+        if self.tags:
+            data["tags"] = self.tags
+        if self.tool_result_ref is not None:
+            data["tool_result_ref"] = self.tool_result_ref
+        return data
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "RawTraceItem":
+        return cls(
+            id=data["id"],
+            ts=data["ts"],
+            turn_id=data["turn_id"],
+            seq=data["seq"],
+            trace_type=data["trace_type"],
+            content=data.get("content", ""),
+            source_event=data["source_event"],
+            media=data.get("media"),
+            tool_name=data.get("tool_name"),
+            tool_call_id=data.get("tool_call_id"),
+            tool_args=data.get("tool_args"),
+            tool_result=data.get("tool_result"),
+            tool_error=data.get("tool_error"),
+            correlation_id=data.get("correlation_id"),
+            tags=data.get("tags", []),
+            tool_result_ref=data.get("tool_result_ref"),
+        )

autobyteus/memory/models/semantic_item.py

@@ -0,0 +1,41 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, List
+
+from .memory_types import MemoryType
+
+
+@dataclass
+class SemanticItem:
+    id: str
+    ts: float
+    fact: str
+    tags: List[str] = field(default_factory=list)
+    confidence: float = 0.0
+    salience: float = 0.0
+
+    @property
+    def memory_type(self) -> MemoryType:
+        return MemoryType.SEMANTIC
+
+    def to_dict(self) -> Dict[str, Any]:
+        data: Dict[str, Any] = {
+            "id": self.id,
+            "ts": self.ts,
+            "fact": self.fact,
+            "confidence": self.confidence,
+            "salience": self.salience,
+        }
+        if self.tags:
+            data["tags"] = self.tags
+        return data
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SemanticItem":
+        return cls(
+            id=data["id"],
+            ts=data["ts"],
+            fact=data.get("fact", ""),
+            tags=data.get("tags", []),
+            confidence=data.get("confidence", 0.0),
+            salience=data.get("salience", 0.0),
+        )

autobyteus/memory/models/tool_interaction.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Optional, Dict
+
+
+class ToolInteractionStatus(Enum):
+    PENDING = "pending"
+    SUCCESS = "success"
+    ERROR = "error"
+
+
+@dataclass
+class ToolInteraction:
+    tool_call_id: str
+    turn_id: Optional[str]
+    tool_name: Optional[str]
+    arguments: Optional[Dict[str, Any]]
+    result: Optional[Any]
+    error: Optional[str]
+    status: ToolInteractionStatus

autobyteus/memory/path_resolver.py

@@ -0,0 +1,27 @@
+import os
+from pathlib import Path
+from typing import Mapping, Optional, Union
+
+
+def resolve_memory_base_dir(
+    override_dir: Optional[str] = None,
+    env: Optional[Mapping[str, str]] = None,
+    fallback_dir: Optional[Union[str, Path]] = None,
+) -> str:
+    override_value = override_dir.strip() if override_dir else ""
+    if override_value:
+        return override_value
+
+    env_values = env if env is not None else os.environ
+    env_value = env_values.get("AUTOBYTEUS_MEMORY_DIR", "").strip()
+    if env_value:
+        return env_value
+
+    if fallback_dir is not None:
+        return str(fallback_dir)
+
+    return str(Path.cwd() / "memory")
+
+
+def resolve_agent_memory_dir(base_dir: Union[str, Path], agent_id: str) -> str:
+    return str(Path(base_dir) / "agents" / agent_id)

autobyteus/memory/policies/__init__.py

@@ -0,0 +1,5 @@
+from autobyteus.memory.policies.compaction_policy import CompactionPolicy
+
+__all__ = [
+    "CompactionPolicy",
+]

autobyteus/memory/policies/compaction_policy.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class CompactionPolicy:
+    trigger_ratio: float = 0.8
+    raw_tail_turns: int = 4
+    max_item_chars: int = 2000
+    safety_margin_tokens: int = 256
+
+    def should_compact(self, prompt_tokens: int, input_budget: int) -> bool:
+        if input_budget <= 0:
+            return True
+        if prompt_tokens >= input_budget:
+            return True
+        return prompt_tokens >= int(self.trigger_ratio * input_budget)

autobyteus/memory/restore/__init__.py

@@ -0,0 +1 @@
+# restore package

autobyteus/memory/restore/working_context_snapshot_bootstrapper.py

@@ -0,0 +1,61 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from autobyteus.memory.working_context_snapshot_serializer import WorkingContextSnapshotSerializer
+from autobyteus.memory.compaction_snapshot_builder import CompactionSnapshotBuilder
+from autobyteus.memory.store.working_context_snapshot_store import WorkingContextSnapshotStore
+
+
+@dataclass
+class WorkingContextSnapshotBootstrapOptions:
+    max_episodic: int = 3
+    max_semantic: int = 20
+    raw_tail_turns: Optional[int] = None
+
+
+class WorkingContextSnapshotBootstrapper:
+    def __init__(
+        self,
+        working_context_snapshot_store: Optional[WorkingContextSnapshotStore] = None,
+        snapshot_builder: Optional[CompactionSnapshotBuilder] = None,
+    ) -> None:
+        self.working_context_snapshot_store = working_context_snapshot_store
+        self.snapshot_builder = snapshot_builder or CompactionSnapshotBuilder()
+
+    def bootstrap(self, memory_manager, system_prompt: str, options: WorkingContextSnapshotBootstrapOptions) -> None:
+        store = self._resolve_store(memory_manager)
+        agent_id = self._resolve_agent_id(memory_manager, store)
+
+        if store and agent_id and store.exists(agent_id):
+            payload = store.read(agent_id)
+            if payload and WorkingContextSnapshotSerializer.validate(payload):
+                snapshot, _meta = WorkingContextSnapshotSerializer.deserialize(payload)
+                memory_manager.reset_working_context_snapshot(snapshot.build_messages())
+                return
+
+        bundle = memory_manager.retriever.retrieve(
+            max_episodic=options.max_episodic,
+            max_semantic=options.max_semantic,
+        )
+        tail_turns = options.raw_tail_turns
+        if tail_turns is None:
+            policy = getattr(memory_manager, "compaction_policy", None)
+            tail_turns = getattr(policy, "raw_tail_turns", 0) if policy else 0
+        raw_tail = memory_manager.get_raw_tail(tail_turns or 0, exclude_turn_id=None)
+        snapshot_messages = self.snapshot_builder.build(
+            system_prompt=system_prompt,
+            bundle=bundle,
+            raw_tail=raw_tail,
+        )
+        memory_manager.reset_working_context_snapshot(snapshot_messages)
+
+    def _resolve_store(self, memory_manager) -> Optional[WorkingContextSnapshotStore]:
+        if self.working_context_snapshot_store is not None:
+            return self.working_context_snapshot_store
+        return getattr(memory_manager, "working_context_snapshot_store", None)
+
+    def _resolve_agent_id(self, memory_manager, store: Optional[WorkingContextSnapshotStore]) -> Optional[str]:
+        if store and getattr(store, "agent_id", None):
+            return store.agent_id
+        store_obj = getattr(memory_manager, "store", None)
+        return getattr(store_obj, "agent_id", None)

autobyteus/memory/retrieval/__init__.py

@@ -0,0 +1,7 @@
+from autobyteus.memory.retrieval.memory_bundle import MemoryBundle
+from autobyteus.memory.retrieval.retriever import Retriever
+
+__all__ = [
+    "MemoryBundle",
+    "Retriever",
+]

autobyteus/memory/retrieval/memory_bundle.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass, field
+from typing import List
+
+from autobyteus.memory.models.episodic_item import EpisodicItem
+from autobyteus.memory.models.semantic_item import SemanticItem
+
+
+@dataclass
+class MemoryBundle:
+    episodic: List[EpisodicItem] = field(default_factory=list)
+    semantic: List[SemanticItem] = field(default_factory=list)

autobyteus/memory/retrieval/retriever.py

@@ -0,0 +1,13 @@
+from autobyteus.memory.models.memory_types import MemoryType
+from autobyteus.memory.retrieval.memory_bundle import MemoryBundle
+from autobyteus.memory.store.base_store import MemoryStore
+
+
+class Retriever:
+    def __init__(self, store: MemoryStore):
+        self.store = store
+
+    def retrieve(self, max_episodic: int, max_semantic: int) -> MemoryBundle:
+        episodic = self.store.list(MemoryType.EPISODIC, limit=max_episodic)
+        semantic = self.store.list(MemoryType.SEMANTIC, limit=max_semantic)
+        return MemoryBundle(episodic=episodic, semantic=semantic)

autobyteus/memory/store/__init__.py

@@ -0,0 +1,9 @@
+from autobyteus.memory.store.base_store import MemoryStore
+from autobyteus.memory.store.file_store import FileMemoryStore
+from autobyteus.memory.store.working_context_snapshot_store import WorkingContextSnapshotStore
+
+__all__ = [
+    "MemoryStore",
+    "FileMemoryStore",
+    "WorkingContextSnapshotStore",
+]

autobyteus/memory/store/base_store.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod
+from typing import Iterable, List, Optional
+
+from autobyteus.memory.models.memory_types import MemoryType
+
+
+class MemoryStore(ABC):
+    @abstractmethod
+    def add(self, items: Iterable[object]) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def list(self, memory_type: MemoryType, limit: Optional[int] = None) -> List[object]:
+        raise NotImplementedError

autobyteus/memory/store/file_store.py

@@ -0,0 +1,98 @@
+import json
+from pathlib import Path
+from typing import Iterable, List, Optional, Union
+
+from autobyteus.memory.models.memory_types import MemoryType
+from autobyteus.memory.models.raw_trace_item import RawTraceItem
+from autobyteus.memory.models.episodic_item import EpisodicItem
+from autobyteus.memory.models.semantic_item import SemanticItem
+from autobyteus.memory.store.base_store import MemoryStore
+
+
+class FileMemoryStore(MemoryStore):
+    def __init__(self, base_dir: Union[str, Path], agent_id: str):
+        self.base_dir = Path(base_dir)
+        self.agent_id = agent_id
+        self.agent_dir = self.base_dir / "agents" / agent_id
+        self.agent_dir.mkdir(parents=True, exist_ok=True)
+
+    def add(self, items: Iterable[object]) -> None:
+        for item in items:
+            memory_type = getattr(item, "memory_type", None)
+            if memory_type is None:
+                raise ValueError("Memory item missing memory_type")
+            path = self._get_file_path(memory_type)
+            record = item.to_dict() if hasattr(item, "to_dict") else item
+            with path.open("a", encoding="utf-8") as handle:
+                handle.write(json.dumps(record) + "\n")
+
+    def list(self, memory_type: MemoryType, limit: Optional[int] = None) -> List[object]:
+        path = self._get_file_path(memory_type)
+        if not path.exists():
+            return []
+        with path.open("r", encoding="utf-8") as handle:
+            lines = [line.strip() for line in handle.readlines() if line.strip()]
+        if limit is not None:
+            lines = lines[-limit:]
+        return [self._deserialize(memory_type, json.loads(line)) for line in lines]
+
+    def list_raw_trace_dicts(self) -> List[dict]:
+        path = self._get_file_path(MemoryType.RAW_TRACE)
+        if not path.exists():
+            return []
+        with path.open("r", encoding="utf-8") as handle:
+            return [json.loads(line) for line in handle if line.strip()]
+
+    def read_archive_raw_traces(self) -> List[dict]:
+        path = self._get_archive_path()
+        if not path.exists():
+            return []
+        with path.open("r", encoding="utf-8") as handle:
+            return [json.loads(line) for line in handle if line.strip()]
+
+    def prune_raw_traces(self, keep_turn_ids: set[str], archive: bool = True) -> None:
+        raw_items = self.list_raw_trace_dicts()
+        if not raw_items:
+            return
+
+        keep = []
+        removed = []
+        for item in raw_items:
+            if item.get("turn_id") in keep_turn_ids:
+                keep.append(item)
+            else:
+                removed.append(item)
+
+        raw_path = self._get_file_path(MemoryType.RAW_TRACE)
+        tmp_path = raw_path.with_suffix(".jsonl.tmp")
+        with tmp_path.open("w", encoding="utf-8") as handle:
+            for item in keep:
+                handle.write(json.dumps(item) + "\n")
+        tmp_path.replace(raw_path)
+
+        if archive and removed:
+            archive_path = self._get_archive_path()
+            with archive_path.open("a", encoding="utf-8") as handle:
+                for item in removed:
+                    handle.write(json.dumps(item) + "\n")
+
+    def _get_file_path(self, memory_type: MemoryType) -> Path:
+        if memory_type == MemoryType.RAW_TRACE:
+            return self.agent_dir / "raw_traces.jsonl"
+        if memory_type == MemoryType.EPISODIC:
+            return self.agent_dir / "episodic.jsonl"
+        if memory_type == MemoryType.SEMANTIC:
+            return self.agent_dir / "semantic.jsonl"
+        raise ValueError(f"Unknown memory type: {memory_type}")
+
+    def _deserialize(self, memory_type: MemoryType, data: dict) -> object:
+        if memory_type == MemoryType.RAW_TRACE:
+            return RawTraceItem.from_dict(data)
+        if memory_type == MemoryType.EPISODIC:
+            return EpisodicItem.from_dict(data)
+        if memory_type == MemoryType.SEMANTIC:
+            return SemanticItem.from_dict(data)
+        raise ValueError(f"Unknown memory type: {memory_type}")
+
+    def _get_archive_path(self) -> Path:
+        return self.agent_dir / "raw_traces_archive.jsonl"

autobyteus/memory/store/working_context_snapshot_store.py

@@ -0,0 +1,28 @@
+import json
+from pathlib import Path
+from typing import Optional, Union, Dict, Any
+
+
+class WorkingContextSnapshotStore:
+    def __init__(self, base_dir: Union[str, Path], agent_id: str) -> None:
+        self.base_dir = Path(base_dir)
+        self.agent_id = agent_id
+
+    def exists(self, agent_id: str) -> bool:
+        return self._get_path(agent_id).exists()
+
+    def read(self, agent_id: str) -> Optional[Dict[str, Any]]:
+        path = self._get_path(agent_id)
+        if not path.exists():
+            return None
+        with path.open("r", encoding="utf-8") as handle:
+            return json.load(handle)
+
+    def write(self, agent_id: str, payload: Dict[str, Any]) -> None:
+        path = self._get_path(agent_id)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with path.open("w", encoding="utf-8") as handle:
+            json.dump(payload, handle)
+
+    def _get_path(self, agent_id: str) -> Path:
+        return self.base_dir / "agents" / agent_id / "working_context_snapshot.json"

autobyteus/memory/tool_interaction_builder.py

@@ -0,0 +1,46 @@
+from typing import Dict, List
+
+from autobyteus.memory.models.raw_trace_item import RawTraceItem
+from autobyteus.memory.models.tool_interaction import ToolInteraction, ToolInteractionStatus
+
+
+def build_tool_interactions(raw_traces: List[RawTraceItem]) -> List[ToolInteraction]:
+    interactions: Dict[str, ToolInteraction] = {}
+
+    for trace in raw_traces:
+        if trace.trace_type not in {"tool_call", "tool_result"}:
+            continue
+
+        tool_call_id = trace.tool_call_id
+        if not tool_call_id:
+            continue
+
+        interaction = interactions.get(tool_call_id)
+        if interaction is None:
+            interaction = ToolInteraction(
+                tool_call_id=tool_call_id,
+                turn_id=trace.turn_id,
+                tool_name=trace.tool_name,
+                arguments=None,
+                result=None,
+                error=None,
+                status=ToolInteractionStatus.PENDING,
+            )
+            interactions[tool_call_id] = interaction
+
+        if trace.trace_type == "tool_call":
+            interaction.tool_name = trace.tool_name
+            interaction.arguments = trace.tool_args
+            if interaction.status == ToolInteractionStatus.PENDING and interaction.error:
+                interaction.status = ToolInteractionStatus.ERROR
+            continue
+
+        if trace.trace_type == "tool_result":
+            interaction.tool_name = trace.tool_name or interaction.tool_name
+            interaction.result = trace.tool_result
+            interaction.error = trace.tool_error
+            interaction.status = (
+                ToolInteractionStatus.ERROR if trace.tool_error else ToolInteractionStatus.SUCCESS
+            )
+
+    return list(interactions.values())

autobyteus/memory/turn_tracker.py

@@ -0,0 +1,9 @@
+class TurnTracker:
+    def __init__(self, start: int = 1):
+        if start < 1:
+            raise ValueError("start must be >= 1")
+        self._counter = start - 1
+
+    def next_turn_id(self) -> str:
+        self._counter += 1
+        return f"turn_{self._counter:04d}"

autobyteus/memory/working_context_snapshot.py

@@ -0,0 +1,69 @@
+import time
+from typing import List, Optional, Iterable, Any
+
+from autobyteus.llm.utils.messages import (
+    Message,
+    MessageRole,
+    ToolCallPayload,
+    ToolCallSpec,
+    ToolResultPayload,
+)
+
+
+class WorkingContextSnapshot:
+    def __init__(self, initial_messages: Optional[Iterable[Message]] = None):
+        self._messages: List[Message] = list(initial_messages) if initial_messages else []
+        self.epoch_id: int = 1
+        self.last_compaction_ts: Optional[float] = None
+
+    def append_message(self, message: Message) -> None:
+        self._messages.append(message)
+
+    def append_user(self, content: str) -> None:
+        self._messages.append(Message(role=MessageRole.USER, content=content))
+
+    def append_assistant(self, content: Optional[str], reasoning: Optional[str] = None) -> None:
+        self._messages.append(
+            Message(
+                role=MessageRole.ASSISTANT,
+                content=content,
+                reasoning_content=reasoning,
+            )
+        )
+
+    def append_tool_calls(self, tool_calls: List[ToolCallSpec]) -> None:
+        self._messages.append(
+            Message(
+                role=MessageRole.ASSISTANT,
+                content=None,
+                tool_payload=ToolCallPayload(tool_calls=tool_calls),
+            )
+        )
+
+    def append_tool_result(
+        self,
+        tool_call_id: str,
+        tool_name: str,
+        tool_result: Any,
+        tool_error: Optional[str] = None,
+    ) -> None:
+        self._messages.append(
+            Message(
+                role=MessageRole.TOOL,
+                content=None,
+                tool_payload=ToolResultPayload(
+                    tool_call_id=tool_call_id,
+                    tool_name=tool_name,
+                    tool_result=tool_result,
+                    tool_error=tool_error,
+                ),
+            )
+        )
+
+    def build_messages(self) -> List[Message]:
+        return list(self._messages)
+
+    def reset(self, snapshot_messages: Iterable[Message], last_compaction_ts: Optional[float] = None) -> None:
+        self._messages = list(snapshot_messages)
+        self.epoch_id += 1
+        self.last_compaction_ts = last_compaction_ts if last_compaction_ts is not None else time.time()