agent-memory-engine 0.1.0__py3-none-any.whl

agent_memory/interfaces/mcp_server.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+from dataclasses import asdict
+
+from agent_memory.client import MemoryClient
+from agent_memory.config import AgentMemoryConfig
+
+
+class _FallbackMCPServer:
+    def __init__(self, name: str) -> None:
+        self.name = name
+        self.tools: dict[str, object] = {}
+
+    def tool(self, name: str | None = None):
+        def decorator(func):
+            self.tools[name or func.__name__] = func
+            return func
+
+        return decorator
+
+
+def _serialize_trace(report) -> dict[str, object]:
+    return {
+        "focus": {"id": report.focus.id, "content": report.focus.content},
+        "ancestors": [{"id": item.id, "content": item.content} for item in report.ancestors],
+        "descendants": [{"id": item.id, "content": item.content} for item in report.descendants],
+        "relations": [
+            {
+                "source_id": edge.source_id,
+                "target_id": edge.target_id,
+                "relation_type": edge.relation_type.value,
+            }
+            for edge in report.relations
+        ],
+        "evolution_events": report.evolution_events,
+    }
+
+
+def _make_server() -> object:
+    try:
+        from fastmcp import FastMCP
+
+        return FastMCP("agent-memory")
+    except ImportError:
+        return _FallbackMCPServer("agent-memory")
+
+
+def create_mcp_server(client: MemoryClient) -> object:
+    server = _make_server()
+
+    @server.tool("memory_store")
+    def memory_store(content: str, source_id: str, memory_type: str = "semantic") -> dict[str, object]:
+        item = client.add(content, source_id=source_id, memory_type=memory_type)  # type: ignore[arg-type]
+        return {"id": item.id, "content": item.content, "trust_score": item.trust_score}
+
+    @server.tool("memory_search")
+    def memory_search(query: str, limit: int = 5) -> list[dict[str, object]]:
+        return [
+            {
+                "id": result.item.id,
+                "content": result.item.content,
+                "score": result.score,
+                "matched_by": result.matched_by,
+            }
+            for result in client.search(query, limit=limit)
+        ]
+
+    @server.tool("memory_ingest_conversation")
+    def memory_ingest_conversation(turns: list[dict[str, str]], source_id: str) -> list[dict[str, object]]:
+        items = client.ingest_conversation(
+            [client.turn_model(**turn) for turn in turns],
+            source_id=source_id,
+        )
+        return [{"id": item.id, "content": item.content} for item in items]
+
+    @server.tool("memory_trace")
+    def memory_trace(memory_id: str, max_depth: int = 10) -> dict[str, object]:
+        return _serialize_trace(client.trace_graph(memory_id, max_depth=max_depth))
+
+    @server.tool("memory_health")
+    def memory_health() -> dict[str, object]:
+        return asdict(client.health())
+
+    @server.tool("memory_audit")
+    def memory_audit(limit: int = 50) -> list[dict[str, object]]:
+        return client.audit_events(limit=limit)
+
+    @server.tool("memory_evolution")
+    def memory_evolution(memory_id: str, limit: int = 50) -> list[dict[str, object]]:
+        return client.evolution_events(memory_id=memory_id, limit=limit)
+
+    @server.tool("memory_update")
+    def memory_update(memory_id: str, content: str) -> dict[str, object]:
+        item = client.get(memory_id)
+        if item is None:
+            raise ValueError(f"Memory {memory_id} not found")
+        updated = client.update(item, content=content)
+        return {"id": updated.id, "content": updated.content}
+
+    @server.tool("memory_delete")
+    def memory_delete(memory_id: str) -> dict[str, object]:
+        return {"deleted": client.delete(memory_id)}
+
+    @server.tool("memory_maintain")
+    def memory_maintain() -> dict[str, object]:
+        return asdict(client.maintain())
+
+    @server.tool("memory_export")
+    def memory_export(path: str) -> dict[str, object]:
+        exported = client.export_jsonl(path)
+        return {"path": path, "exported": exported}
+
+    return server
+
+
+def main() -> int:
+    client = MemoryClient(config=AgentMemoryConfig.from_env())
+    server = create_mcp_server(client)
+    run = getattr(server, "run", None)
+    if callable(run):
+        run()
+        return 0
+    print("FastMCP is not installed. Install with `pip install -e .[mcp]`.")
+    return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

agent_memory/interfaces/rest_api.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from dataclasses import asdict
+
+from agent_memory.client import MemoryClient
+from agent_memory.interfaces.mcp_server import _serialize_trace
+
+
+class _FallbackRestApp:
+    def __init__(self) -> None:
+        self.routes: dict[str, object] = {}
+
+    def add_api_route(self, path: str, endpoint, methods: list[str]) -> None:
+        for method in methods:
+            self.routes[f"{method.upper()} {path}"] = endpoint
+
+
+def _make_app() -> object:
+    try:
+        from fastapi import FastAPI
+
+        return FastAPI(title="agent-memory")
+    except ImportError:
+        return _FallbackRestApp()
+
+
+def create_rest_app(client: MemoryClient) -> object:
+    app = _make_app()
+
+    def create_memory(payload: dict[str, object]) -> dict[str, object]:
+        item = client.add(
+            str(payload["content"]),
+            source_id=str(payload.get("source_id", "rest")),
+            memory_type=str(payload.get("memory_type", "semantic")),  # type: ignore[arg-type]
+        )
+        return {"id": item.id, "content": item.content}
+
+    def search_memory(query: str, limit: int = 5) -> list[dict[str, object]]:
+        return [
+            {"id": result.item.id, "content": result.item.content, "score": result.score}
+            for result in client.search(query, limit=limit)
+        ]
+
+    def read_health() -> dict[str, object]:
+        return asdict(client.health())
+
+    def read_audit(limit: int = 50) -> list[dict[str, object]]:
+        return client.audit_events(limit=limit)
+
+    def read_evolution(memory_id: str, limit: int = 50) -> list[dict[str, object]]:
+        return client.evolution_events(memory_id=memory_id, limit=limit)
+
+    def trace_memory(memory_id: str, max_depth: int = 10) -> dict[str, object]:
+        return _serialize_trace(client.trace_graph(memory_id, max_depth=max_depth))
+
+    def run_maintenance() -> dict[str, object]:
+        return asdict(client.maintain())
+
+    def export_memory(path: str) -> dict[str, object]:
+        exported = client.export_jsonl(path)
+        return {"path": path, "exported": exported}
+
+    app.add_api_route("/memories", create_memory, methods=["POST"])
+    app.add_api_route("/search", search_memory, methods=["GET"])
+    app.add_api_route("/health", read_health, methods=["GET"])
+    app.add_api_route("/audit", read_audit, methods=["GET"])
+    app.add_api_route("/evolution", read_evolution, methods=["GET"])
+    app.add_api_route("/trace", trace_memory, methods=["GET"])
+    app.add_api_route("/maintain", run_maintenance, methods=["POST"])
+    app.add_api_route("/export", export_memory, methods=["POST"])
+    return app

agent_memory/llm/__init__.py

@@ -0,0 +1,5 @@
+from agent_memory.llm.ollama_client import OllamaClient
+from agent_memory.llm.openai_client import OpenAIClient
+
+__all__ = ["OllamaClient", "OpenAIClient"]
+

agent_memory/llm/base.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from typing import Any
+from typing import Protocol
+
+
+class LLMClient(Protocol):
+    def complete(self, prompt: str, system_prompt: str | None = None) -> str:
+        ...
+
+    def generate_json(
+        self,
+        *,
+        prompt: str,
+        schema: dict[str, Any],
+        schema_name: str,
+        system_prompt: str | None = None,
+    ) -> dict[str, Any]:
+        ...
+
+
+class LLMClientError(RuntimeError):
+    pass

agent_memory/llm/ollama_client.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+from urllib import request
+
+from agent_memory.llm.base import LLMClientError
+
+
+class OllamaClient:
+    def __init__(self, model: str = "llama3.1", base_url: str | None = None) -> None:
+        self.model = model
+        self.base_url = (base_url or os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")).rstrip("/")
+
+    def complete(self, prompt: str, system_prompt: str | None = None) -> str:
+        payload: dict[str, Any] = {
+            "model": self.model,
+            "prompt": prompt,
+            "stream": False,
+        }
+        if system_prompt:
+            payload["system"] = system_prompt
+        response = self._request_json("/api/generate", payload)
+        if "response" not in response:
+            raise LLMClientError(f"Ollama response missing text: {response}")
+        return str(response["response"])
+
+    def generate_json(
+        self,
+        *,
+        prompt: str,
+        schema: dict[str, Any],
+        schema_name: str,
+        system_prompt: str | None = None,
+    ) -> dict[str, Any]:
+        payload: dict[str, Any] = {
+            "model": self.model,
+            "prompt": prompt,
+            "stream": False,
+            "format": schema,
+        }
+        if system_prompt:
+            payload["system"] = system_prompt
+        response = self._request_json("/api/generate", payload)
+        text = str(response.get("response", "")).strip()
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError as exc:
+            raise LLMClientError(f"Ollama response was not valid JSON: {text}") from exc
+
+    def _request_json(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
+        body = json.dumps(payload).encode("utf-8")
+        req = request.Request(
+            f"{self.base_url}{path}",
+            data=body,
+            method="POST",
+            headers={"Content-Type": "application/json"},
+        )
+        try:
+            with request.urlopen(req) as response:
+                return json.loads(response.read().decode("utf-8"))
+        except Exception as exc:
+            raise LLMClientError(f"Ollama request failed: {exc}") from exc

agent_memory/llm/openai_client.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+from urllib import request
+
+from agent_memory.llm.base import LLMClientError
+
+
+class OpenAIClient:
+    def __init__(
+        self,
+        model: str = "gpt-4o-mini",
+        api_key: str | None = None,
+        base_url: str = "https://api.openai.com/v1",
+    ) -> None:
+        self.model = model
+        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
+        self.base_url = base_url.rstrip("/")
+
+    def complete(self, prompt: str, system_prompt: str | None = None) -> str:
+        payload = {
+            "model": self.model,
+            "input": self._build_input(prompt=prompt, system_prompt=system_prompt),
+        }
+        response = self._request_json("/responses", payload)
+        return self._extract_output_text(response)
+
+    def generate_json(
+        self,
+        *,
+        prompt: str,
+        schema: dict[str, Any],
+        schema_name: str,
+        system_prompt: str | None = None,
+    ) -> dict[str, Any]:
+        payload = {
+            "model": self.model,
+            "input": self._build_input(prompt=prompt, system_prompt=system_prompt),
+            "text": {
+                "format": {
+                    "type": "json_schema",
+                    "name": schema_name,
+                    "strict": True,
+                    "schema": schema,
+                }
+            },
+        }
+        response = self._request_json("/responses", payload)
+        text = self._extract_output_text(response)
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError as exc:
+            raise LLMClientError(f"OpenAI response was not valid JSON: {text}") from exc
+
+    def _build_input(self, *, prompt: str, system_prompt: str | None) -> list[dict[str, str]]:
+        messages: list[dict[str, str]] = []
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": prompt})
+        return messages
+
+    def _request_json(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
+        if not self.api_key:
+            raise LLMClientError("OPENAI_API_KEY is required to use OpenAIClient.")
+        body = json.dumps(payload).encode("utf-8")
+        req = request.Request(
+            f"{self.base_url}{path}",
+            data=body,
+            method="POST",
+            headers={
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json",
+            },
+        )
+        try:
+            with request.urlopen(req) as response:
+                return json.loads(response.read().decode("utf-8"))
+        except Exception as exc:
+            raise LLMClientError(f"OpenAI request failed: {exc}") from exc
+
+    def _extract_output_text(self, response: dict[str, Any]) -> str:
+        for output in response.get("output", []):
+            if output.get("type") != "message":
+                continue
+            for item in output.get("content", []):
+                if item.get("type") == "refusal":
+                    raise LLMClientError(item.get("refusal", "Model refused request."))
+                if item.get("type") in {"output_text", "text"} and item.get("text"):
+                    return str(item["text"])
+        if response.get("output_text"):
+            return str(response["output_text"])
+        raise LLMClientError(f"Could not extract text from OpenAI response: {response}")

agent_memory/models.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any
+
+
+def utc_now() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+class MemoryType(str, Enum):
+    SEMANTIC = "semantic"
+    EPISODIC = "episodic"
+    PROCEDURAL = "procedural"
+
+
+class MemoryLayer(str, Enum):
+    SHORT_TERM = "short_term"
+    LONG_TERM = "long_term"
+
+
+class QueryIntent(str, Enum):
+    FACTUAL = "factual"
+    TEMPORAL = "temporal"
+    CAUSAL = "causal"
+    EXPLORATORY = "exploratory"
+    PROCEDURAL = "procedural"
+    GENERAL = "general"
+
+
+class RelationType(str, Enum):
+    DERIVED_FROM = "derived_from"
+    SUPERSEDES = "supersedes"
+    SUPPORTS = "supports"
+    CONTRADICTS = "contradicts"
+    RELATED_TO = "related_to"
+
+
+class ConflictResolution(str, Enum):
+    SUPERSEDE = "supersede"
+    KEEP_BOTH = "keep_both"
+    MANUAL = "manual"
+
+
+@dataclass(slots=True)
+class MemoryItem:
+    id: str
+    content: str
+    memory_type: MemoryType
+    embedding: list[float]
+    created_at: datetime = field(default_factory=utc_now)
+    last_accessed: datetime = field(default_factory=utc_now)
+    access_count: int = 0
+    valid_from: datetime | None = None
+    valid_until: datetime | None = None
+    trust_score: float = 0.75
+    importance: float = 0.5
+    layer: MemoryLayer = MemoryLayer.SHORT_TERM
+    decay_rate: float = 0.1
+    source_id: str = "manual"
+    causal_parent_id: str | None = None
+    supersedes_id: str | None = None
+    entity_refs: list[str] = field(default_factory=list)
+    tags: list[str] = field(default_factory=list)
+    deleted_at: datetime | None = None
+
+
+@dataclass(slots=True)
+class MemoryDraft:
+    content: str
+    memory_type: MemoryType = MemoryType.SEMANTIC
+    importance: float = 0.5
+    trust_score: float = 0.7
+    source_id: str = "conversation"
+    causal_parent_id: str | None = None
+    supersedes_id: str | None = None
+    entity_refs: list[str] = field(default_factory=list)
+    tags: list[str] = field(default_factory=list)
+
+
+@dataclass(slots=True)
+class ConversationTurn:
+    role: str
+    content: str
+    timestamp: datetime | None = None
+
+
+@dataclass(slots=True)
+class SearchResult:
+    item: MemoryItem
+    score: float
+    matched_by: list[str]
+
+
+@dataclass(slots=True)
+class RetrievalPlan:
+    intent: QueryIntent
+    strategies: list[str]
+    filters: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(slots=True)
+class RelationEdge:
+    source_id: str
+    target_id: str
+    relation_type: RelationType
+    created_at: datetime = field(default_factory=utc_now)
+
+
+@dataclass(slots=True)
+class ConflictRecord:
+    existing_id: str
+    candidate_id: str | None
+    confidence: float
+    resolution: ConflictResolution
+    reason: str
+
+
+@dataclass(slots=True)
+class MaintenanceReport:
+    promoted: int = 0
+    demoted: int = 0
+    decayed: int = 0
+    conflicts_found: int = 0
+    conflicts_resolved: int = 0
+    consolidated: int = 0
+
+
+@dataclass(slots=True)
+class HealthReport:
+    total_memories: int
+    stale_ratio: float
+    orphan_ratio: float
+    unresolved_conflicts: int
+    average_trust_score: float
+    database_size_bytes: int
+    audit_events: int
+    suggestions: list[str] = field(default_factory=list)
+
+
+@dataclass(slots=True)
+class TraceReport:
+    focus: MemoryItem
+    ancestors: list[MemoryItem] = field(default_factory=list)
+    descendants: list[MemoryItem] = field(default_factory=list)
+    relations: list[RelationEdge] = field(default_factory=list)
+    evolution_events: list[dict[str, Any]] = field(default_factory=list)

agent_memory/storage/__init__.py

@@ -0,0 +1,4 @@
+from agent_memory.storage.sqlite_backend import SQLiteBackend
+
+__all__ = ["SQLiteBackend"]
+

agent_memory/storage/base.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import Protocol
+
+from agent_memory.models import MemoryItem, RelationEdge
+
+
+class StorageBackend(Protocol):
+    def add_memory(self, item: MemoryItem) -> MemoryItem:
+        ...
+
+    def get_memory(self, memory_id: str) -> MemoryItem | None:
+        ...
+
+    def update_memory(self, item: MemoryItem) -> MemoryItem:
+        ...
+
+    def soft_delete_memory(self, memory_id: str) -> bool:
+        ...
+
+    def search_full_text(self, query: str, limit: int = 10, memory_type: str | None = None) -> list[tuple[MemoryItem, float]]:
+        ...
+
+    def search_by_entities(self, entities: list[str], limit: int = 10, memory_type: str | None = None) -> list[tuple[MemoryItem, float]]:
+        ...
+
+    def search_by_vector(
+        self,
+        embedding: list[float],
+        limit: int = 10,
+        memory_type: str | None = None,
+    ) -> list[tuple[MemoryItem, float]]:
+        ...
+
+    def touch_memory(self, memory_id: str) -> None:
+        ...
+
+    def trace_ancestors(self, memory_id: str, max_depth: int = 10) -> list[MemoryItem]:
+        ...
+
+    def list_memories(self, include_deleted: bool = False) -> list[MemoryItem]:
+        ...
+
+    def add_relation(self, edge: RelationEdge) -> None:
+        ...
+
+    def list_relations(self, memory_id: str | None = None) -> list[RelationEdge]:
+        ...
+
+    def trace_descendants(self, memory_id: str, max_depth: int = 10) -> list[MemoryItem]:
+        ...
+
+    def relation_exists_between(
+        self,
+        left_id: str,
+        right_id: str,
+        relation_types: list[str] | None = None,
+    ) -> bool:
+        ...