zero-agent 0.1.0__py3-none-any.whl

agentz/context/data_store.py

@@ -0,0 +1,208 @@
+"""
+Pipeline-scoped data store for sharing in-memory objects between agents.
+Enables efficient data sharing without redundant file I/O.
+"""
+
+from typing import Any, Dict, Optional, List
+from datetime import datetime
+from dataclasses import dataclass
+import threading
+from loguru import logger
+
+
+@dataclass
+class DataStoreEntry:
+    """Single entry in the pipeline data store with metadata."""
+    key: str
+    value: Any
+    timestamp: datetime
+    data_type: str
+    metadata: Dict[str, Any]
+
+    def size_mb(self) -> float:
+        """Estimate size in MB (rough approximation)."""
+        import sys
+        return sys.getsizeof(self.value) / 1024 / 1024
+
+
+class DataStore:
+    """Thread-safe data store for sharing objects within a pipeline run.
+
+    Designed for storing:
+    - DataFrames loaded from files
+    - Trained models
+    - Preprocessed data
+    - Intermediate computation results
+
+    Key naming conventions:
+    - DataFrames: f"dataframe:{file_path}"
+    - Models: f"model:{model_name}"
+    - Preprocessed data: f"preprocessed:{file_path}"
+    - Custom: any string key
+    """
+
+    def __init__(self, experiment_id: Optional[str] = None):
+        """Initialize the data store.
+
+        Args:
+            experiment_id: Optional experiment ID for tracking
+        """
+        self._store: Dict[str, DataStoreEntry] = {}
+        self._lock = threading.RLock()
+        self.experiment_id = experiment_id
+        logger.debug(f"Initialized DataStore for experiment: {experiment_id}")
+
+    def set(
+        self,
+        key: str,
+        value: Any,
+        data_type: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Store a value in the data store.
+
+        Args:
+            key: Unique key for the value
+            value: The object to store
+            data_type: Optional type descriptor (e.g., 'dataframe', 'model')
+            metadata: Optional metadata dict
+        """
+        with self._lock:
+            # Infer data type if not provided
+            if data_type is None:
+                data_type = type(value).__name__
+
+            entry = DataStoreEntry(
+                key=key,
+                value=value,
+                timestamp=datetime.now(),
+                data_type=data_type,
+                metadata=metadata or {}
+            )
+
+            self._store[key] = entry
+            logger.debug(
+                f"Stored {data_type} at key '{key}' "
+                f"(size: {entry.size_mb():.2f} MB)"
+            )
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Retrieve a value from the data store.
+
+        Args:
+            key: Key to retrieve
+            default: Default value if key not found
+
+        Returns:
+            The stored value or default
+        """
+        with self._lock:
+            entry = self._store.get(key)
+            if entry:
+                logger.debug(f"Retrieved {entry.data_type} from key '{key}'")
+                return entry.value
+            return default
+
+    def get_entry(self, key: str) -> Optional[DataStoreEntry]:
+        """Retrieve full entry with metadata.
+
+        Args:
+            key: Key to retrieve
+
+        Returns:
+            DataStoreEntry or None
+        """
+        with self._lock:
+            return self._store.get(key)
+
+    def has(self, key: str) -> bool:
+        """Check if a key exists in the store.
+
+        Args:
+            key: Key to check
+
+        Returns:
+            True if key exists
+        """
+        with self._lock:
+            return key in self._store
+
+    def delete(self, key: str) -> bool:
+        """Delete a key from the store.
+
+        Args:
+            key: Key to delete
+
+        Returns:
+            True if key was deleted, False if not found
+        """
+        with self._lock:
+            if key in self._store:
+                del self._store[key]
+                logger.debug(f"Deleted key '{key}' from data store")
+                return True
+            return False
+
+    def clear(self) -> None:
+        """Clear all data from the store."""
+        with self._lock:
+            count = len(self._store)
+            self._store.clear()
+            logger.debug(f"Cleared {count} entries from data store")
+
+    def list_keys(self, data_type: Optional[str] = None) -> List[str]:
+        """List all keys, optionally filtered by data type.
+
+        Args:
+            data_type: Optional filter by data type
+
+        Returns:
+            List of keys
+        """
+        with self._lock:
+            if data_type is None:
+                return list(self._store.keys())
+            return [
+                key for key, entry in self._store.items()
+                if entry.data_type == data_type
+            ]
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get statistics about the data store.
+
+        Returns:
+            Dictionary with stats
+        """
+        with self._lock:
+            total_size_mb = sum(entry.size_mb() for entry in self._store.values())
+            data_types = {}
+            for entry in self._store.values():
+                data_types[entry.data_type] = data_types.get(entry.data_type, 0) + 1
+
+            return {
+                "total_entries": len(self._store),
+                "total_size_mb": round(total_size_mb, 2),
+                "data_types": data_types,
+                "experiment_id": self.experiment_id,
+                "keys": list(self._store.keys())
+            }
+
+    def __len__(self) -> int:
+        """Return number of entries."""
+        with self._lock:
+            return len(self._store)
+
+    def __contains__(self, key: str) -> bool:
+        """Support 'in' operator."""
+        return self.has(key)
+
+    def __getitem__(self, key: str) -> Any:
+        """Support dictionary-style access."""
+        value = self.get(key)
+        if value is None and key not in self._store:
+            raise KeyError(f"Key '{key}' not found in data store")
+        return value
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        """Support dictionary-style assignment."""
+        self.set(key, value)

agentz/llm/llm_setup.py

@@ -0,0 +1,156 @@
+from typing import Dict, Any, Optional, Union
+
+from agents import (
+    OpenAIChatCompletionsModel,
+    OpenAIResponsesModel,
+    ModelSettings,
+)
+from agents.extensions.models.litellm_model import LitellmModel
+from openai import AsyncAzureOpenAI, AsyncOpenAI
+
+# Provider configurations - use OpenAIResponsesModel for most providers
+PROVIDER_CONFIGS = {
+    "openai": {
+        "base_url": "https://api.openai.com/v1",
+        "model_class": OpenAIResponsesModel,
+    },
+    "deepseek": {
+        "base_url": "https://api.deepseek.com/v1",
+        "model_class": OpenAIResponsesModel,
+    },
+    "openrouter": {
+        "base_url": "https://openrouter.ai/api/v1",
+        "model_class": OpenAIResponsesModel,
+    },
+    "gemini": {
+        "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
+        "model_class": OpenAIChatCompletionsModel,
+        "use_litellm": False,
+    },
+    "anthropic": {
+        "base_url": "https://api.anthropic.com/v1/",
+        "model_class": OpenAIResponsesModel,
+    },
+    "perplexity": {
+        "base_url": "https://api.perplexity.ai/chat/completions",
+        "model_class": OpenAIResponsesModel,
+    },
+    "huggingface": {
+        "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
+        "model_class": OpenAIResponsesModel,
+    },
+    "local": {
+        "base_url": None,  # Will be provided in config
+        "model_class": OpenAIChatCompletionsModel,
+        "default_api_key": "ollama",
+    },
+    "azureopenai": {
+        "model_class": OpenAIChatCompletionsModel,
+        "requires_azure": True,
+    },
+    "bedrock": {
+        "model_class": LitellmModel,
+        "use_litellm": True,
+    },
+}
+
+
+class LLMConfig:
+    """Direct configuration system - no environment variables."""
+
+    def __init__(self, config: Dict[str, Any], full_config: Optional[Dict[str, Any]] = None):
+        """
+        Initialize LLM configuration from direct config.
+
+        Args:
+            config: Dictionary containing:
+                - provider: str (e.g., "openai", "gemini", "deepseek")
+                - api_key: str
+                - model: str (optional, will use defaults)
+                - base_url: str (optional for custom endpoints)
+                - azure_config: dict (for Azure OpenAI)
+                - aws_config: dict (for Bedrock)
+                - model_settings: dict (optional, for temperature etc.)
+            full_config: Optional full configuration including agent prompts, pipeline settings
+        """
+        self.provider = config["provider"]
+        self.api_key = config["api_key"]
+        self.model_name = config.get("model", self._get_default_model())
+        self.config = config
+        self.full_config = full_config
+
+        # Validate provider
+        if self.provider not in PROVIDER_CONFIGS:
+            valid = list(PROVIDER_CONFIGS.keys())
+            raise ValueError(f"Invalid provider: {self.provider}. Available: {valid}")
+
+        # Create main model (used for all purposes - reasoning, main, fast)
+        self.main_model = self._create_model()
+        self.reasoning_model = self.main_model
+        self.fast_model = self.main_model
+
+        # Model settings from config or defaults
+        model_settings_config = self.config.get("model_settings", {})
+        self.default_model_settings = ModelSettings(
+            temperature=model_settings_config.get("temperature", 0.1)
+        )
+
+        # Set tracing if OpenAI key provided
+        if self.provider == "openai" and self.api_key:
+            from agents import set_tracing_export_api_key
+            set_tracing_export_api_key(self.api_key)
+
+    def _get_default_model(self) -> str:
+        """Get default model for provider."""
+        defaults = {
+            "openai": "gpt-4.1",
+            "gemini": "gemini-2.5-flash",
+            "deepseek": "deepseek-chat",
+            "anthropic": "claude-3-5-sonnet-20241022",
+            "bedrock": "anthropic.claude-3-5-sonnet-20241022-v2:0",
+            "perplexity": "llama-3.1-sonar-large-128k-online",
+            "openrouter": "meta-llama/llama-3.2-3b-instruct:free",
+        }
+        return defaults.get(self.provider, "gpt-4.1")
+
+    def _create_model(self):
+        """Create model instance using direct configuration."""
+        provider_config = PROVIDER_CONFIGS[self.provider]
+        model_class = provider_config["model_class"]
+
+        if provider_config.get("use_litellm"):
+            return model_class(model=self.model_name, api_key=self.api_key, base_url=provider_config.get("base_url"))
+
+        elif self.provider == "azureopenai":
+            azure_config = self.config.get("azure_config", {})
+            client = AsyncAzureOpenAI(
+                api_key=self.api_key,
+                azure_endpoint=azure_config.get("endpoint"),
+                azure_deployment=azure_config.get("deployment"),
+                api_version=azure_config.get("api_version", "2023-12-01-preview"),
+            )
+            return model_class(model=self.model_name, openai_client=client)
+
+        else:
+            # Standard OpenAI-compatible providers
+            base_url = self.config.get("base_url", provider_config["base_url"])
+            api_key = self.api_key or provider_config.get("default_api_key", "key")
+
+            client = AsyncOpenAI(
+                api_key=api_key,
+                base_url=base_url,
+            )
+            return model_class(model=self.model_name, openai_client=client)
+
+def get_base_url(model: Union[OpenAIChatCompletionsModel, OpenAIResponsesModel]) -> str:
+    """Utility function to get the base URL for a given model"""
+    return str(model._client._base_url)
+
+def model_supports_json_and_tool_calls(
+    model: Union[OpenAIChatCompletionsModel, OpenAIResponsesModel],
+) -> bool:
+    """Utility function to check if a model supports structured output"""
+    structured_output_providers = ["openai.com", "anthropic.com"]
+    return any(
+        provider in get_base_url(model) for provider in structured_output_providers
+    )

agentz/mcp/manager.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from contextlib import AsyncExitStack
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+from agents.mcp import MCPServer, MCPServerSse, MCPServerStdio
+
+
+class MCPConfigurationError(ValueError):
+    """Raised when an MCP server configuration is invalid."""
+
+
+@dataclass(frozen=True)
+class MCPServerSpec:
+    """Lightweight specification describing how to build an MCP server."""
+
+    type: str
+    options: Dict[str, Any]
+
+    def __post_init__(self) -> None:
+        if not self.type:
+            raise MCPConfigurationError("MCP server configuration requires a 'type'.")
+
+
+class MCPRegistry:
+    """Registry responsible for storing MCP server specifications."""
+
+    def __init__(self, specs: Optional[Dict[str, MCPServerSpec]] = None) -> None:
+        self._specs = specs or {}
+
+    @classmethod
+    def from_config(cls, config: Optional[Mapping[str, Any]]) -> "MCPRegistry":
+        """Create a registry from configuration mapping."""
+        if config is None:
+            return cls()
+
+        servers_config = config.get("servers", {})
+        if servers_config and not isinstance(servers_config, Mapping):
+            raise MCPConfigurationError("'servers' must be a mapping of server definitions.")
+
+        specs: Dict[str, MCPServerSpec] = {}
+        for name, server_cfg in (servers_config or {}).items():
+            if not isinstance(server_cfg, Mapping):
+                raise MCPConfigurationError(f"MCP server '{name}' configuration must be a mapping.")
+
+            server_type = str(server_cfg.get("type") or server_cfg.get("transport") or "").strip().lower()
+            if not server_type:
+                raise MCPConfigurationError(f"MCP server '{name}' must define a 'type' or 'transport'.")
+
+            options = {k: v for k, v in server_cfg.items() if k not in {"type", "transport"}}
+            specs[name] = MCPServerSpec(type=server_type, options=options)
+
+        return cls(specs)
+
+    def register(self, name: str, spec: MCPServerSpec) -> None:
+        """Register (or overwrite) a spec by name."""
+        self._specs[name] = spec
+
+    def get(self, name: str) -> MCPServerSpec:
+        try:
+            return self._specs[name]
+        except KeyError as exc:
+            raise MCPConfigurationError(
+                f"MCP server '{name}' is not defined; add it to the configuration."
+            ) from exc
+
+    def as_dict(self) -> Dict[str, MCPServerSpec]:
+        return dict(self._specs)
+
+    def contains(self, name: str) -> bool:
+        return name in self._specs
+
+
+SERVER_TYPE_MAP = {
+    "stdio": MCPServerStdio,
+    "sse": MCPServerSse,
+}
+
+
+class MCPManagerSession:
+    """Async context manager that keeps MCP server connections alive."""
+
+    def __init__(self, registry: MCPRegistry):
+        self._registry = registry
+        self._stack = AsyncExitStack()
+        self._servers: Dict[str, MCPServer] = {}
+
+    async def __aenter__(self) -> "MCPManagerSession":
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        await self._stack.aclose()
+        self._servers.clear()
+
+    async def get_server(self, name: str, overrides: Optional[Mapping[str, Any]] = None) -> MCPServer:
+        """Return a connected MCP server instance by name, creating it on demand."""
+        if name in self._servers:
+            return self._servers[name]
+
+        spec = self._registry.get(name)
+        options = dict(spec.options)
+        if overrides:
+            options.update(overrides)
+
+        try:
+            server_cls = SERVER_TYPE_MAP[spec.type]
+        except KeyError as exc:
+            raise MCPConfigurationError(
+                f"Unsupported MCP server type '{spec.type}' for '{name}'. "
+                f"Supported types: {', '.join(SERVER_TYPE_MAP)}."
+            ) from exc
+
+        server_ctx = server_cls(**options)
+        server = await self._stack.enter_async_context(server_ctx)
+        self._servers[name] = server
+        return server
+
+
+class MCPManager:
+    """Entry point that provides MCP manager sessions."""
+
+    def __init__(self, registry: MCPRegistry):
+        self._registry = registry
+
+    @classmethod
+    def from_config(cls, config: Optional[Mapping[str, Any]]) -> "MCPManager":
+        registry = MCPRegistry.from_config(config)
+        return cls(registry)
+
+    def ensure_server(self, name: str, spec: MCPServerSpec) -> None:
+        """Add a default server if one isn't already configured."""
+        if not self._registry.contains(name):
+            self._registry.register(name, spec)
+
+    def session(self) -> MCPManagerSession:
+        """Create a new MCPManagerSession for a pipeline run."""
+        return MCPManagerSession(self._registry)
+
+    def list_servers(self) -> Dict[str, MCPServerSpec]:
+        return self._registry.as_dict()

agentz/mcp/patches.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import logging
+import shutil
+import subprocess
+from pathlib import Path
+from typing import Iterable
+
+
+logger = logging.getLogger(__name__)
+
+
+PATCH_MARKER = "__browsermcp_original_close"
+BUGGY_SNIPPET = "  server.close = async () => {\n    await server.close();"
+PATCHED_SNIPPET = """  const __browsermcp_original_close = server.close.bind(server);
+  server.close = async () => {
+    if (server.__browsermcp_closing) {
+      return;
+    }
+    server.__browsermcp_closing = true;
+    await __browsermcp_original_close();
+"""
+
+
+def _candidate_paths(base_dir: Path) -> Iterable[Path]:
+    if not base_dir.exists():
+        return []
+    return base_dir.glob("**/node_modules/@browsermcp/mcp/dist/index.js")
+
+
+def _prime_browsermcp_cache() -> bool:
+    npx_path = shutil.which("npx")
+    if not npx_path:
+        logger.warning("Unable to locate `npx`; skipping Browser MCP patch.")
+        return False
+
+    try:
+        subprocess.run(
+            [
+                npx_path,
+                "--yes",
+                "--package",
+                "@browsermcp/mcp@latest",
+                "node",
+                "-e",
+                "process.exit(0)",
+            ],
+            check=True,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        return True
+    except subprocess.CalledProcessError as exc:
+        logger.warning("Failed to prime Browser MCP cache via npx: %s", exc)
+        return False
+
+
+def apply_browsermcp_close_patch() -> None:
+    """Patch the Browser MCP CLI to avoid recursive `server.close` calls."""
+    cache_root = Path.home() / ".npm" / "_npx"
+    candidate_paths = list(_candidate_paths(cache_root))
+    if not candidate_paths:
+        if not _prime_browsermcp_cache():
+            return
+        candidate_paths = list(_candidate_paths(cache_root))
+
+    for path in candidate_paths:
+        try:
+            original = path.read_text(encoding="utf-8")
+        except OSError as exc:
+            logger.warning("Failed to read %s while applying Browser MCP patch: %s", path, exc)
+            continue
+
+        if PATCH_MARKER in original:
+            continue
+
+        if BUGGY_SNIPPET not in original:
+            continue
+
+        patched = original.replace(BUGGY_SNIPPET, PATCHED_SNIPPET, 1)
+        if patched == original:
+            continue
+
+        try:
+            path.write_text(patched, encoding="utf-8")
+            logger.info("Patched Browser MCP close handler in %s", path)
+        except OSError as exc:
+            logger.warning("Failed to patch %s: %s", path, exc)

agentz/mcp/servers/chrome_devtools/server.py

@@ -0,0 +1,14 @@
+import asyncio
+import shutil
+
+from agents import Agent, Runner, trace
+from agents.mcp import MCPServer, MCPServerStdio
+
+
+def ChromeDevToolsMCP():
+    server = MCPServerStdio(
+        cache_tools_list=True,  # Cache the tools list, for demonstration
+        params={"command": "npx", "args": ["-y", "chrome-devtools-mcp@latest"]},
+    ) 
+    return server
+