agent_os_kernel 3.1.0__py3-none-any.whl

emk/store.py

@@ -0,0 +1,281 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# Public Preview — basic context/memory management
+"""
+Store — Abstract interfaces and mutable file-based implementation for episodic memory.
+"""
+
+from abc import ABC, abstractmethod
+from typing import List, Optional, Dict, Any
+from pathlib import Path
+import json
+import numpy as np
+
+from emk.schema import Episode
+
+
+class VectorStoreAdapter(ABC):
+    """Abstract interface for vector store implementations."""
+
+    @abstractmethod
+    def store(self, episode: Episode, embedding: Optional[np.ndarray] = None) -> str:
+        """Store an episode. Returns the episode_id."""
+        pass
+
+    @abstractmethod
+    def retrieve(
+        self,
+        query_embedding: Optional[np.ndarray] = None,
+        filters: Optional[Dict[str, Any]] = None,
+        limit: int = 10,
+    ) -> List[Episode]:
+        """Retrieve episodes, optionally filtered."""
+        pass
+
+    @abstractmethod
+    def get_by_id(self, episode_id: str) -> Optional[Episode]:
+        """Retrieve a specific episode by ID."""
+        pass
+
+    def update(self, episode_id: str, episode: Episode) -> bool:
+        """Update an existing episode. Returns True if found and updated."""
+        raise NotImplementedError
+
+    def delete(self, episode_id: str) -> bool:
+        """Delete an episode by ID. Returns True if found and deleted."""
+        raise NotImplementedError
+
+    def retrieve_failures(
+        self,
+        query_embedding: Optional[np.ndarray] = None,
+        filters: Optional[Dict[str, Any]] = None,
+        limit: int = 10,
+    ) -> List[Episode]:
+        """Retrieve episodes marked as failures."""
+        failure_filters = filters.copy() if filters else {}
+        failure_filters["is_failure"] = True
+        return self.retrieve(query_embedding=query_embedding, filters=failure_filters, limit=limit)
+
+    def retrieve_successes(
+        self,
+        query_embedding: Optional[np.ndarray] = None,
+        filters: Optional[Dict[str, Any]] = None,
+        limit: int = 10,
+    ) -> List[Episode]:
+        """Retrieve episodes that are NOT failures."""
+        batch_size = min(limit * 3, 1000)
+        all_episodes = self.retrieve(query_embedding=query_embedding, filters=filters, limit=batch_size)
+        return [ep for ep in all_episodes if not ep.is_failure()][:limit]
+
+    def retrieve_with_anti_patterns(
+        self,
+        query_embedding: Optional[np.ndarray] = None,
+        filters: Optional[Dict[str, Any]] = None,
+        limit: int = 10,
+        include_failures: bool = True,
+    ) -> Dict[str, List[Episode]]:
+        """Retrieve both successes and failures."""
+        result: Dict[str, List[Episode]] = {"successes": [], "failures": []}
+        result["successes"] = self.retrieve_successes(query_embedding=query_embedding, filters=filters, limit=limit)
+        if include_failures:
+            result["failures"] = self.retrieve_failures(query_embedding=query_embedding, filters=filters, limit=limit)
+        return result
+
+
+class FileAdapter(VectorStoreAdapter):
+    """
+    Mutable JSONL-based file storage adapter.
+
+    Supports store, retrieve, get_by_id, update and delete.
+    """
+
+    def __init__(self, filepath: str = "episodes.jsonl"):
+        self.filepath = Path(filepath)
+        self.filepath.parent.mkdir(parents=True, exist_ok=True)
+        if not self.filepath.exists():
+            self.filepath.touch()
+
+    # -- helpers -----------------------------------------------------------
+
+    def _read_all(self) -> List[Episode]:
+        episodes: List[Episode] = []
+        if not self.filepath.exists() or self.filepath.stat().st_size == 0:
+            return episodes
+        with open(self.filepath, "r") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    episodes.append(Episode.from_json(line))
+                except Exception:
+                    continue
+        return episodes
+
+    def _write_all(self, episodes: List[Episode]) -> None:
+        with open(self.filepath, "w") as f:
+            for ep in episodes:
+                f.write(ep.to_json() + "\n")
+
+    # -- core API ----------------------------------------------------------
+
+    def store(self, episode: Episode, embedding: Optional[np.ndarray] = None) -> str:
+        with open(self.filepath, "a") as f:
+            f.write(episode.to_json() + "\n")
+        return episode.episode_id
+
+    def retrieve(
+        self,
+        query_embedding: Optional[np.ndarray] = None,
+        filters: Optional[Dict[str, Any]] = None,
+        limit: int = 10,
+    ) -> List[Episode]:
+        episodes = self._read_all()
+        if filters:
+            episodes = [
+                ep for ep in episodes
+                if all(ep.metadata.get(k) == v for k, v in filters.items())
+            ]
+        episodes.reverse()
+        return episodes[:limit]
+
+    def get_by_id(self, episode_id: str) -> Optional[Episode]:
+        for ep in self._read_all():
+            if ep.episode_id == episode_id:
+                return ep
+        return None
+
+    def update(self, episode_id: str, episode: Episode) -> bool:
+        """Replace the episode with the given ID."""
+        episodes = self._read_all()
+        for i, ep in enumerate(episodes):
+            if ep.episode_id == episode_id:
+                episodes[i] = episode
+                self._write_all(episodes)
+                return True
+        return False
+
+    def delete(self, episode_id: str) -> bool:
+        """Remove the episode with the given ID."""
+        episodes = self._read_all()
+        new_episodes = [ep for ep in episodes if ep.episode_id != episode_id]
+        if len(new_episodes) == len(episodes):
+            return False
+        self._write_all(new_episodes)
+        return True
+
+
+try:
+    import chromadb
+    from chromadb.config import Settings
+
+    class ChromaDBAdapter(VectorStoreAdapter):
+        """
+        ChromaDB-based vector storage adapter.
+
+        Uses ChromaDB for vector similarity search and efficient retrieval
+        of episodes based on embeddings.
+        """
+
+        def __init__(
+            self,
+            collection_name: str = "episodes",
+            persist_directory: str = "./chroma_data",
+        ):
+            self.client = chromadb.Client(Settings(
+                persist_directory=persist_directory,
+                anonymized_telemetry=False,
+            ))
+            self.collection = self.client.get_or_create_collection(
+                name=collection_name,
+                metadata={"description": "Episodic memory storage"},
+            )
+
+        def store(self, episode: Episode, embedding: Optional[np.ndarray] = None) -> str:
+            if embedding is None:
+                text = f"{episode.goal} {episode.action} {episode.result} {episode.reflection}"
+            else:
+                text = None
+
+            self.collection.add(
+                ids=[episode.episode_id],
+                documents=[text] if text else None,
+                embeddings=[embedding.tolist()] if embedding is not None else None,
+                metadatas=[{
+                    "goal": episode.goal,
+                    "action": episode.action,
+                    "result": episode.result,
+                    "reflection": episode.reflection,
+                    "timestamp": episode.timestamp.isoformat(),
+                    **episode.metadata,
+                }],
+            )
+            return episode.episode_id
+
+        def retrieve(
+            self,
+            query_embedding: Optional[np.ndarray] = None,
+            filters: Optional[Dict[str, Any]] = None,
+            limit: int = 10,
+        ) -> List[Episode]:
+            if query_embedding is not None:
+                results = self.collection.query(
+                    query_embeddings=[query_embedding.tolist()],
+                    n_results=limit,
+                    where=filters,
+                )
+            else:
+                results = self.collection.get(limit=limit, where=filters)
+
+            episodes = []
+            if 'metadatas' in results and results['metadatas']:
+                if isinstance(results['metadatas'][0], list) if results['metadatas'] else False:
+                    metadatas = results['metadatas'][0]
+                else:
+                    metadatas = results['metadatas']
+            else:
+                metadatas = []
+
+            for metadata in metadatas:
+                episode_metadata = {
+                    k: v for k, v in metadata.items()
+                    if k not in ('goal', 'action', 'result', 'reflection', 'timestamp')
+                }
+                try:
+                    episode = Episode(
+                        goal=metadata.get("goal", ""),
+                        action=metadata.get("action", ""),
+                        result=metadata.get("result", ""),
+                        reflection=metadata.get("reflection", ""),
+                        timestamp=metadata.get("timestamp", None),
+                        metadata=episode_metadata,
+                    )
+                    episodes.append(episode)
+                except Exception:
+                    continue
+            return episodes
+
+        def get_by_id(self, episode_id: str) -> Optional[Episode]:
+            try:
+                results = self.collection.get(ids=[episode_id])
+                if not results['ids']:
+                    return None
+                metadata = results['metadatas'][0]
+                episode_metadata = {
+                    k: v for k, v in metadata.items()
+                    if k not in ('goal', 'action', 'result', 'reflection', 'timestamp')
+                }
+                return Episode(
+                    goal=metadata.get("goal", ""),
+                    action=metadata.get("action", ""),
+                    result=metadata.get("result", ""),
+                    reflection=metadata.get("reflection", ""),
+                    timestamp=metadata.get("timestamp", None),
+                    metadata=episode_metadata,
+                )
+            except Exception:
+                return None
+
+except ImportError:
+    # ChromaDB not installed — adapter will not be available
+    pass

iatp/__init__.py

@@ -0,0 +1,166 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Inter-Agent Trust Protocol (IATP).
+
+A sidecar-based service mesh for preventing cascading hallucinations
+in autonomous agent networks. IATP provides the missing "governance layer"
+for multi-agent LLM systems.
+
+Core Features:
+    - **Discovery**: Capability manifest exchange via `/.well-known/agent-manifest`
+    - **Trust**: Dynamic trust scoring and security validation
+    - **Reversibility**: Enforced transaction rollback requirements
+    - **Privacy**: Data retention and handling policy enforcement
+    - **Telemetry**: Distributed tracing via Flight Recorder
+
+Quick Start:
+    >>> from iatp import CapabilityManifest, TrustLevel, SidecarProxy
+    >>> manifest = CapabilityManifest(
+    ...     agent_id="my-agent",
+    ...     trust_level=TrustLevel.TRUSTED,
+    ...     capabilities=AgentCapabilities(reversibility=ReversibilityLevel.FULL),
+    ...     privacy_contract=PrivacyContract(retention=RetentionPolicy.EPHEMERAL),
+    ... )
+    >>> print(f"Trust Score: {manifest.calculate_trust_score()}/10")
+
+CLI Usage:
+    .. code-block:: bash
+
+        # Verify a manifest
+        iatp verify manifest.json
+
+        # Run the sidecar
+        uvicorn iatp.main:app --host 0.0.0.0 --port 8081
+
+Docker:
+    .. code-block:: bash
+
+        docker run -p 8081:8081 -e IATP_AGENT_URL=http://my-agent:8000 iatp-sidecar
+
+For more information, see:
+    - Documentation: https://github.com/microsoft/agent-governance-toolkit
+    - Paper: paper/whitepaper.md
+
+Example:
+    Basic trust score calculation::
+
+        from iatp import (
+            CapabilityManifest,
+            AgentCapabilities,
+            PrivacyContract,
+            TrustLevel,
+            ReversibilityLevel,
+            RetentionPolicy,
+        )
+
+        manifest = CapabilityManifest(
+            agent_id="secure-bank-agent",
+            trust_level=TrustLevel.VERIFIED_PARTNER,
+            capabilities=AgentCapabilities(
+                reversibility=ReversibilityLevel.FULL,
+                idempotency=True,
+            ),
+            privacy_contract=PrivacyContract(
+                retention=RetentionPolicy.EPHEMERAL,
+            ),
+        )
+
+        score = manifest.calculate_trust_score()
+        print(f"Trust Score: {score}/10")  # Output: Trust Score: 10/10
+"""
+
+from __future__ import annotations
+
+__version__ = "3.1.0"
+__author__ = "Microsoft Corporation"
+__license__ = "MIT"
+
+# Core Models
+# Attestation & Reputation
+from iatp.attestation import AttestationValidator, ReputationManager
+from iatp.models import (
+    AgentCapabilities,
+    AttestationRecord,
+    CapabilityManifest,
+    PrivacyContract,
+    QuarantineSession,
+    ReputationEvent,
+    ReputationScore,
+    RetentionPolicy,
+    ReversibilityLevel,
+    TracingContext,
+    TrustLevel,
+)
+
+# Engines
+from iatp.policy_engine import IATPPolicyEngine
+from iatp.recovery import IATPRecoveryEngine
+
+# Security & Privacy
+from iatp.security import PrivacyScrubber, SecurityValidator
+
+# Sidecar Components
+from iatp.sidecar import SidecarProxy, create_sidecar
+
+# Telemetry & Tracing
+from iatp.telemetry import FlightRecorder, TraceIDGenerator
+
+# IPC Pipes - Typed inter-agent communication (v0.4.0)
+from iatp.ipc_pipes import (
+    TypedPipe,
+    PipeMessage,
+    PipeConfig,
+    PipeState,
+    PolicyCheckPipe,
+    Pipeline,
+    AgentPipelineStage,
+    create_pipeline,
+    pipe_agents,
+)
+
+# Public API exports - controls what is visible via `from iatp import *`
+__all__ = [
+    # Package metadata
+    "__version__",
+    "__author__",
+    "__license__",
+    # Models - Core data structures for capability manifests
+    "CapabilityManifest",
+    "AgentCapabilities",
+    "PrivacyContract",
+    "TrustLevel",
+    "ReversibilityLevel",
+    "RetentionPolicy",
+    "QuarantineSession",
+    "TracingContext",
+    "AttestationRecord",
+    "ReputationScore",
+    "ReputationEvent",
+    # Sidecar - The proxy that wraps agents
+    "SidecarProxy",
+    "create_sidecar",
+    # Security - Validation and privacy enforcement
+    "SecurityValidator",
+    "PrivacyScrubber",
+    # Attestation & Reputation
+    "AttestationValidator",
+    "ReputationManager",
+    # Telemetry - Distributed tracing and audit logging
+    "FlightRecorder",
+    "TraceIDGenerator",
+    # Policy Engine - Rule-based policy evaluation
+    "IATPPolicyEngine",
+    # Recovery Engine - Compensating transaction support (scak integration)
+    "IATPRecoveryEngine",
+    # IPC Pipes - Typed inter-agent communication (v0.4.0)
+    "TypedPipe",
+    "PipeMessage",
+    "PipeConfig",
+    "PipeState",
+    "PolicyCheckPipe",
+    "Pipeline",
+    "AgentPipelineStage",
+    "create_pipeline",
+    "pipe_agents",
+]