agent_os_kernel 3.1.0__py3-none-any.whl

emk/hf_utils.py

@@ -0,0 +1,421 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Hugging Face Hub utilities for EMK.
+
+This module provides functions to push and pull episode data and experiment
+results to/from the Hugging Face Hub for sharing and reproducibility.
+
+Requirements:
+    pip install agent-os-kernel[full]  # includes emk with huggingface support
+
+Example:
+    >>> from emk.hf_utils import upload_episodes_to_hub
+    >>> upload_episodes_to_hub(
+    ...     episodes=my_episodes,
+    ...     repo_id="microsoft/emk-experiments",
+    ...     filename="episodes.jsonl"
+    ... )
+
+Note:
+    You must be logged in to Hugging Face Hub to push data:
+    >>> huggingface_hub.login()
+    or set the HF_TOKEN environment variable.
+"""
+
+from __future__ import annotations
+
+import json
+import tempfile
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+if TYPE_CHECKING:
+    from emk.schema import Episode
+
+# Lazy import to avoid requiring huggingface_hub at import time
+_HF_HUB_AVAILABLE = None
+
+
+def _check_hf_hub() -> None:
+    """Check if huggingface_hub is available and raise helpful error if not."""
+    global _HF_HUB_AVAILABLE
+    
+    if _HF_HUB_AVAILABLE is None:
+        try:
+            import huggingface_hub  # noqa: F401
+            _HF_HUB_AVAILABLE = True
+        except ImportError:
+            _HF_HUB_AVAILABLE = False
+    
+    if not _HF_HUB_AVAILABLE:
+        raise ImportError(
+            "huggingface_hub is required for this functionality. "
+            "Install it with: pip install agent-os-kernel[full]"
+        )
+
+
+def upload_episodes_to_hub(
+    episodes: List["Episode"],
+    repo_id: str,
+    filename: str = "episodes.jsonl",
+    *,
+    commit_message: Optional[str] = None,
+    private: bool = False,
+    token: Optional[str] = None,
+    branch: Optional[str] = None,
+) -> str:
+    """
+    Upload episodes to a Hugging Face Hub dataset repository.
+    
+    This function serializes episodes to JSONL format and uploads them
+    to the specified Hugging Face Hub repository.
+    
+    Args:
+        episodes: List of Episode objects to upload.
+        repo_id: The Hugging Face Hub repository ID (e.g., "username/repo-name").
+        filename: Name of the file in the repository (default: "episodes.jsonl").
+        commit_message: Custom commit message (auto-generated if not provided).
+        private: Whether the repository should be private (default: False).
+        token: Hugging Face API token (uses cached token if not provided).
+        branch: Branch to upload to (default: main).
+    
+    Returns:
+        str: URL of the uploaded file.
+    
+    Raises:
+        ImportError: If huggingface_hub is not installed.
+        ValueError: If episodes list is empty.
+    
+    Example:
+        >>> from emk import Episode
+        >>> from emk.hf_utils import upload_episodes_to_hub
+        >>> episodes = [Episode(goal="Test", action="Run", result="Pass", reflection="Good")]
+        >>> url = upload_episodes_to_hub(
+        ...     episodes=episodes,
+        ...     repo_id="microsoft/emk-test-data"
+        ... )
+        >>> print(f"Uploaded to: {url}")
+    """
+    _check_hf_hub()
+    from huggingface_hub import HfApi
+    
+    if not episodes:
+        raise ValueError("Episodes list cannot be empty")
+    
+    api = HfApi(token=token)
+    
+    # Create repository if it doesn't exist
+    api.create_repo(
+        repo_id=repo_id,
+        repo_type="dataset",
+        private=private,
+        exist_ok=True,
+    )
+    
+    # Serialize episodes to JSONL
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f:
+        for episode in episodes:
+            f.write(episode.to_json() + "\n")
+        temp_path = f.name
+    
+    try:
+        # Generate commit message if not provided
+        if commit_message is None:
+            commit_message = f"Upload {len(episodes)} episodes via emk"
+        
+        # Upload file
+        result = api.upload_file(
+            path_or_fileobj=temp_path,
+            path_in_repo=filename,
+            repo_id=repo_id,
+            repo_type="dataset",
+            commit_message=commit_message,
+            revision=branch,
+        )
+        
+        return result
+    finally:
+        # Cleanup temp file
+        Path(temp_path).unlink(missing_ok=True)
+
+
+def download_episodes_from_hub(
+    repo_id: str,
+    filename: str = "episodes.jsonl",
+    *,
+    token: Optional[str] = None,
+    revision: Optional[str] = None,
+) -> List["Episode"]:
+    """
+    Download episodes from a Hugging Face Hub dataset repository.
+    
+    Args:
+        repo_id: The Hugging Face Hub repository ID (e.g., "username/repo-name").
+        filename: Name of the file in the repository (default: "episodes.jsonl").
+        token: Hugging Face API token (uses cached token if not provided).
+        revision: Git revision (branch, tag, or commit) to download from.
+    
+    Returns:
+        List[Episode]: List of Episode objects loaded from the repository.
+    
+    Raises:
+        ImportError: If huggingface_hub is not installed.
+        FileNotFoundError: If the file doesn't exist in the repository.
+    
+    Example:
+        >>> from emk.hf_utils import download_episodes_from_hub
+        >>> episodes = download_episodes_from_hub(
+        ...     repo_id="microsoft/emk-test-data"
+        ... )
+        >>> print(f"Downloaded {len(episodes)} episodes")
+    """
+    _check_hf_hub()
+    from huggingface_hub import hf_hub_download
+    
+    # Import Episode here to avoid circular imports
+    from emk.schema import Episode
+    
+    # Download file
+    local_path = hf_hub_download(
+        repo_id=repo_id,
+        filename=filename,
+        repo_type="dataset",
+        token=token,
+        revision=revision,
+    )
+    
+    # Parse episodes
+    episodes = []
+    with open(local_path, "r") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                episodes.append(Episode.from_json(line))
+    
+    return episodes
+
+
+def push_experiment_results(
+    results: Dict[str, Any],
+    repo_id: str,
+    filename: str = "results.json",
+    *,
+    commit_message: Optional[str] = None,
+    private: bool = False,
+    token: Optional[str] = None,
+    append_timestamp: bool = True,
+) -> str:
+    """
+    Push experiment results to Hugging Face Hub.
+    
+    This is useful for tracking experiment runs and sharing reproducible
+    results with the research community.
+    
+    Args:
+        results: Dictionary of experiment results to upload.
+        repo_id: The Hugging Face Hub repository ID.
+        filename: Name of the results file (default: "results.json").
+        commit_message: Custom commit message.
+        private: Whether the repository should be private.
+        token: Hugging Face API token.
+        append_timestamp: Whether to append timestamp to filename (default: True).
+    
+    Returns:
+        str: URL of the uploaded file.
+    
+    Example:
+        >>> from emk.hf_utils import push_experiment_results
+        >>> results = {
+        ...     "accuracy": 0.95,
+        ...     "latency_ms": 12.5,
+        ...     "episodes_processed": 1000
+        ... }
+        >>> url = push_experiment_results(
+        ...     results=results,
+        ...     repo_id="microsoft/emk-experiments"
+        ... )
+    """
+    _check_hf_hub()
+    from huggingface_hub import HfApi
+    
+    api = HfApi(token=token)
+    
+    # Create repository if it doesn't exist
+    api.create_repo(
+        repo_id=repo_id,
+        repo_type="dataset",
+        private=private,
+        exist_ok=True,
+    )
+    
+    # Add metadata to results
+    results_with_meta = {
+        "_uploaded_at": datetime.now(timezone.utc).isoformat(),
+        "_emk_version": _get_emk_version(),
+        **results,
+    }
+    
+    # Modify filename with timestamp if requested
+    if append_timestamp:
+        timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
+        stem = Path(filename).stem
+        suffix = Path(filename).suffix or ".json"
+        filename = f"{stem}_{timestamp}{suffix}"
+    
+    # Serialize to JSON
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(results_with_meta, f, indent=2, default=str)
+        temp_path = f.name
+    
+    try:
+        if commit_message is None:
+            commit_message = f"Upload experiment results via emk"
+        
+        result = api.upload_file(
+            path_or_fileobj=temp_path,
+            path_in_repo=filename,
+            repo_id=repo_id,
+            repo_type="dataset",
+            commit_message=commit_message,
+        )
+        
+        return result
+    finally:
+        Path(temp_path).unlink(missing_ok=True)
+
+
+def create_dataset_card(
+    repo_id: str,
+    description: str,
+    *,
+    num_episodes: Optional[int] = None,
+    tags: Optional[List[str]] = None,
+    license: str = "mit",
+    token: Optional[str] = None,
+) -> str:
+    """
+    Create or update a dataset card (README.md) for an EMK dataset.
+    
+    Args:
+        repo_id: The Hugging Face Hub repository ID.
+        description: Description of the dataset.
+        num_episodes: Number of episodes in the dataset (optional).
+        tags: List of tags for the dataset (optional).
+        license: License identifier (default: "mit").
+        token: Hugging Face API token.
+    
+    Returns:
+        str: URL of the dataset card.
+    
+    Example:
+        >>> from emk.hf_utils import create_dataset_card
+        >>> url = create_dataset_card(
+        ...     repo_id="microsoft/emk-agent-logs",
+        ...     description="Agent experience logs from production system",
+        ...     num_episodes=10000,
+        ...     tags=["agents", "episodic-memory", "nlp"]
+        ... )
+    """
+    _check_hf_hub()
+    from huggingface_hub import HfApi
+    
+    api = HfApi(token=token)
+    
+    # Build tags list
+    all_tags = ["emk", "episodic-memory", "agent-experiences"]
+    if tags:
+        all_tags.extend(tags)
+    tags_yaml = "\n".join(f"- {tag}" for tag in all_tags)
+    
+    # Build dataset card content
+    card_content = f"""---
+license: {license}
+tags:
+{tags_yaml}
+library_name: emk
+---
+
+# {repo_id.split('/')[-1]}
+
+{description}
+
+## Dataset Information
+
+- **Format**: JSONL (newline-delimited JSON)
+- **Schema**: EMK Episode (Goal → Action → Result → Reflection)
+- **Library**: [emk](https://github.com/microsoft/agent-governance-toolkit)
+"""
+
+    if num_episodes:
+        card_content += f"- **Episodes**: {num_episodes:,}\n"
+    
+    card_content += """
+## Usage
+
+```python
+from emk.hf_utils import download_episodes_from_hub
+
+episodes = download_episodes_from_hub(
+    repo_id="{repo_id}"
+)
+
+for episode in episodes[:5]:
+    print(f"Goal: {episode.goal}")
+    print(f"Result: {episode.result}")
+    print("---")
+```
+
+## Episode Schema
+
+Each episode contains:
+
+| Field | Type | Description |
+|-------|------|-------------|
+| `goal` | string | The agent's intended objective |
+| `action` | string | The action taken |
+| `result` | string | The outcome |
+| `reflection` | string | Agent's analysis or learning |
+| `timestamp` | datetime | When the episode was created |
+| `metadata` | object | Additional context |
+| `episode_id` | string | Unique SHA-256 identifier |
+
+## License
+
+This dataset is released under the {license.upper()} license.
+""".format(repo_id=repo_id, license=license)
+    
+    # Upload README
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".md", delete=False) as f:
+        f.write(card_content)
+        temp_path = f.name
+    
+    try:
+        result = api.upload_file(
+            path_or_fileobj=temp_path,
+            path_in_repo="README.md",
+            repo_id=repo_id,
+            repo_type="dataset",
+            commit_message="Create/update dataset card via emk",
+        )
+        return result
+    finally:
+        Path(temp_path).unlink(missing_ok=True)
+
+
+def _get_emk_version() -> str:
+    """Get the current emk version."""
+    try:
+        from emk import __version__
+        return __version__
+    except ImportError:
+        return "unknown"
+
+
+__all__ = [
+    "upload_episodes_to_hub",
+    "download_episodes_from_hub",
+    "push_experiment_results",
+    "create_dataset_card",
+]

emk/indexer.py

@@ -0,0 +1,83 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# Public Preview — basic context/memory management
+"""
+Indexer — simple tag-based filtering for episodes.
+"""
+
+from typing import List, Set, Dict, Any
+import hashlib
+import re
+
+from emk.schema import Episode
+
+
+class Indexer:
+    """Tag extraction and simple metadata-based filtering for episodes."""
+
+    @staticmethod
+    def extract_tags(text: str, min_length: int = 3) -> Set[str]:
+        """Extract potential search tags from *text*."""
+        words = re.findall(r'\b\w+\b', text.lower())
+        stop_words = {
+            'the', 'is', 'at', 'which', 'on', 'and', 'a', 'an',
+            'as', 'are', 'was', 'were', 'been', 'be', 'have', 'has',
+            'had', 'do', 'does', 'did', 'will', 'would', 'should',
+            'could', 'may', 'might', 'must', 'can', 'to', 'from',
+            'in', 'out', 'up', 'down', 'for', 'with', 'by', 'of',
+        }
+        return {w for w in words if len(w) >= min_length and w not in stop_words}
+
+    @staticmethod
+    def generate_episode_tags(episode: Episode) -> List[str]:
+        """Generate searchable tags from an episode."""
+        combined = f"{episode.goal} {episode.action} {episode.result} {episode.reflection}"
+        tags = Indexer.extract_tags(combined)
+        for key in episode.metadata.keys():
+            tags.add(key.lower())
+        return sorted(tags)
+
+    @staticmethod
+    def compute_content_hash(episode: Episode) -> str:
+        """Return the content hash (episode_id) of the episode."""
+        return episode.episode_id
+
+    @staticmethod
+    def enrich_metadata(episode: Episode, auto_tags: bool = True) -> Dict[str, Any]:
+        """Enrich episode metadata with tags and length metrics."""
+        enriched = episode.metadata.copy()
+        if auto_tags and 'tags' not in enriched:
+            enriched['tags'] = Indexer.generate_episode_tags(episode)
+        enriched['goal_length'] = len(episode.goal)
+        enriched['action_length'] = len(episode.action)
+        enriched['result_length'] = len(episode.result)
+        enriched['reflection_length'] = len(episode.reflection)
+        return enriched
+
+    @staticmethod
+    def create_search_text(episode: Episode) -> str:
+        """Create a concatenated search text from an episode."""
+        parts = [
+            f"Goal: {episode.goal}",
+            f"Action: {episode.action}",
+            f"Result: {episode.result}",
+            f"Reflection: {episode.reflection}",
+        ]
+        if episode.metadata:
+            metadata_str = ", ".join(f"{k}: {v}" for k, v in episode.metadata.items())
+            parts.append(f"Context: {metadata_str}")
+        return " | ".join(parts)
+
+    @staticmethod
+    def filter_by_tags(
+        episodes: List[Episode],
+        required_tags: Set[str],
+    ) -> List[Episode]:
+        """Return episodes whose auto-generated tags include all *required_tags*."""
+        required_lower = {t.lower() for t in required_tags}
+        results: List[Episode] = []
+        for ep in episodes:
+            ep_tags = set(Indexer.generate_episode_tags(ep))
+            if required_lower.issubset(ep_tags):
+                results.append(ep)
+        return results