modaic 0.1.0__py3-none-any.whl

modaic/module_utils.py

@@ -0,0 +1,341 @@
+import importlib.util
+import os
+import re
+import shutil
+import sys
+import sysconfig
+import warnings
+from pathlib import Path
+from types import ModuleType
+from typing import Dict
+
+import tomlkit as tomlk
+
+from .utils import compute_cache_dir
+
+MODAIC_CACHE = compute_cache_dir()
+AGENTS_CACHE = Path(MODAIC_CACHE) / "agents"
+EDITABLE_MODE = os.getenv("EDITABLE_MODE", "false").lower() == "true"
+
+
+def is_builtin(module_name: str) -> bool:
+    """Check whether a module name refers to a built-in module.
+
+    Args:
+      module_name: The fully qualified module name.
+
+    Returns:
+      bool: True if the module is a Python built-in.
+    """
+
+    return module_name in sys.builtin_module_names
+
+
+def is_stdlib(module_name: str) -> bool:
+    """Check whether a module belongs to the Python standard library.
+
+    Args:
+      module_name: The fully qualified module name.
+
+    Returns:
+      bool: True if the module is part of the stdlib (including built-ins).
+    """
+
+    try:
+        spec = importlib.util.find_spec(module_name)
+    except ValueError:
+        return False
+    except Exception:
+        return False
+    if not spec:
+        return False
+    if spec.origin == "built-in":
+        return True
+    origin = spec.origin or ""
+    stdlib_dir = Path(sysconfig.get_paths()["stdlib"]).resolve()
+    try:
+        origin_path = Path(origin).resolve()
+    except OSError:
+        return False
+    return stdlib_dir in origin_path.parents or origin_path == stdlib_dir
+
+
+def is_builtin_or_frozen(mod: ModuleType) -> bool:
+    """Check whether a module object is built-in or frozen.
+
+    Args:
+      mod: The module object.
+
+    Returns:
+      bool: True if the module is built-in or frozen.
+    """
+
+    spec = getattr(mod, "__spec__", None)
+    origin = getattr(spec, "origin", None)
+    name = getattr(mod, "__name__", None)
+    return (name in sys.builtin_module_names) or (origin in ("built-in", "frozen"))
+
+
+def get_internal_imports() -> Dict[str, ModuleType]:
+    """Return only internal modules currently loaded in sys.modules.
+
+    Internal modules are defined as those not installed in site/dist packages
+    (covers virtualenv `.venv` cases as well).
+
+    If the environment variable `EDITABLE_MODE` is set to "true" (case-insensitive),
+    modules located under `src/modaic/` are also excluded.
+
+    Args:
+      None
+
+    Returns:
+      Dict[str, ModuleType]: Mapping of module names to module objects that are
+      not located under any "site-packages" or "dist-packages" directory.
+    """
+
+    internal: Dict[str, ModuleType] = {}
+
+    seen: set[int] = set()
+    for name, module in list(sys.modules.items()):
+        if module is None:
+            continue
+        module_id = id(module)
+        if module_id in seen:
+            continue
+        seen.add(module_id)
+
+        if is_builtin_or_frozen(module):
+            continue
+
+        module_file = getattr(module, "__file__", None)
+        if not module_file:
+            continue
+        try:
+            module_path = Path(module_file).resolve()
+        except OSError:
+            continue
+
+        if is_builtin(name) or is_stdlib(name):
+            continue
+        if is_external_package(module_path):
+            continue
+        if EDITABLE_MODE:
+            posix_path = module_path.as_posix().lower()
+            if "src/modaic" in posix_path:
+                continue
+        normalized_name = name
+
+        internal[normalized_name] = module
+
+    return internal
+
+
+def resolve_project_root() -> Path:
+    """
+    Return the project root directory, must be a directory containing a pyproject.toml file.
+
+    Raises:
+        FileNotFoundError: If pyproject.toml is not found in the current directory.
+    """
+    pyproject_path = Path("pyproject.toml")
+    if not pyproject_path.exists():
+        raise FileNotFoundError("pyproject.toml not found in current directory")
+    return pyproject_path.resolve().parent
+
+
+def is_path_ignored(target_path: Path, ignored_paths: list[Path]) -> bool:
+    """Return True if target_path matches or is contained within any ignored path."""
+    try:
+        absolute_target = target_path.resolve()
+    except OSError:
+        return False
+    for ignored in ignored_paths:
+        if absolute_target == ignored:
+            return True
+        try:
+            absolute_target.relative_to(ignored)
+            return True
+        except Exception:
+            pass
+    return False
+
+
+def copy_module_layout(base_dir: Path, name_parts: list[str]) -> None:
+    """
+    Create ancestor package directories and ensure each contains an __init__.py file.
+    Example:
+        Given a base_dir of "/tmp/modaic" and name_parts of ["agent","indexer"],
+        creates the following layout:
+        | /tmp/modaic/
+        |   | agent/
+        |   |   | __init__.py
+        |   | indexer/
+        |   |   | __init__.py
+    """
+    current = base_dir
+    for part in name_parts:
+        current = current / part
+        current.mkdir(parents=True, exist_ok=True)
+        init_file = current / "__init__.py"
+        if not init_file.exists():
+            init_file.touch()
+
+
+def is_external_package(path: Path) -> bool:
+    """Return True if the path is under site-packages or dist-packages."""
+    parts = {p.lower() for p in path.parts}
+    return "site-packages" in parts or "dist-packages" in parts
+
+
+def init_agent_repo(repo_path: str, with_code: bool = True) -> Path:
+    """Create a local repository staging directory for agent modules and files, excluding ignored files and folders."""
+    repo_dir = Path(AGENTS_CACHE) / repo_path
+    repo_dir.mkdir(parents=True, exist_ok=True)
+
+    internal_imports = get_internal_imports()
+    ignored_paths = get_ignored_files()
+
+    seen_files: set[Path] = set()
+
+    readme_src = Path("README.md")
+    if readme_src.exists() and not is_path_ignored(readme_src, ignored_paths):
+        readme_dest = repo_dir / "README.md"
+        shutil.copy2(readme_src, readme_dest)
+    else:
+        warnings.warn("README.md not found in current directory. Please add one when pushing to the hub.", stacklevel=4)
+
+    if not with_code:
+        return repo_dir
+
+    for module_name, module in internal_imports.items():
+        module_file = getattr(module, "__file__", None)
+        if not module_file:
+            continue
+        try:
+            src_path = Path(module_file).resolve()
+        except OSError:
+            continue
+        if src_path.suffix != ".py":
+            continue
+        if is_path_ignored(src_path, ignored_paths):
+            continue
+        if src_path in seen_files:
+            continue
+        seen_files.add(src_path)
+
+        # Split modul_name to get the relative path
+        name_parts = module_name.split(".")
+        if src_path.name == "__init__.py":
+            copy_module_layout(repo_dir, name_parts)
+            dest_path = repo_dir.joinpath(*name_parts) / "__init__.py"
+        else:
+            if len(name_parts) > 1:
+                copy_module_layout(repo_dir, name_parts[:-1])
+            else:
+                repo_dir.mkdir(parents=True, exist_ok=True)
+            # use the file name to name the file
+            dest_path = repo_dir.joinpath(*name_parts[:-1]) / src_path.name
+        dest_path.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(src_path, dest_path)
+    return repo_dir
+
+
+def create_agent_repo(repo_path: str, with_code: bool = True) -> Path:
+    """
+    Create a temporary directory inside the Modaic cache. Containing everything that will be pushed to the hub. This function adds the following files:
+    - All internal modules used to run the agent
+    - The pyproject.toml
+    - The README.md
+    """
+    package_name = repo_path.split("/")[-1]
+    repo_dir = init_agent_repo(repo_path, with_code=with_code)
+    if with_code:
+        create_pyproject_toml(repo_dir, package_name)
+
+    return repo_dir
+
+
+def get_ignored_files() -> list[Path]:
+    """Return a list of absolute Paths that should be excluded from staging."""
+    project_root = resolve_project_root()
+    pyproject_path = Path("pyproject.toml")
+    doc = tomlk.parse(pyproject_path.read_text(encoding="utf-8"))
+
+    # Safely get [tool.modaic.ignore]
+    ignore_table = (
+        doc.get("tool", {})  # [tool]
+        .get("modaic", {})  # [tool.modaic]
+        .get("ignore")  # [tool.modaic.ignore]
+    )
+
+    if ignore_table is None or "files" not in ignore_table:
+        return []
+
+    ignored: list[Path] = []
+    for entry in ignore_table["files"]:
+        try:
+            ignored.append((project_root / entry).resolve())
+        except OSError:
+            continue
+    return ignored
+
+
+def create_pyproject_toml(repo_dir: Path, package_name: str):
+    """
+    Create a new pyproject.toml for the bundled agent in the temp directory.
+    """
+    old = Path("pyproject.toml").read_text(encoding="utf-8")
+    new = repo_dir / "pyproject.toml"
+
+    doc_old = tomlk.parse(old)
+    doc_new = tomlk.document()
+
+    if "project" not in doc_old:
+        raise KeyError("No [project] table in old TOML")
+    doc_new["project"] = doc_old["project"]
+    doc_new["project"]["dependencies"] = get_filtered_dependencies(doc_old["project"]["dependencies"])
+    if "tool" in doc_old and "uv" in doc_old["tool"] and "sources" in doc_old["tool"]["uv"]:
+        doc_new["tool"] = {"uv": {"sources": doc_old["tool"]["uv"]["sources"]}}
+        warn_if_local(doc_new["tool"]["uv"]["sources"])
+
+    doc_new["project"]["name"] = package_name
+
+    with open(new, "w") as fp:
+        tomlk.dump(doc_new, fp)
+
+
+def get_filtered_dependencies(dependencies: list[str]) -> list[str]:
+    """
+    Get the dependencies that should be included in the bundled agent.
+    """
+    pyproject_path = Path("pyproject.toml")
+    doc = tomlk.parse(pyproject_path.read_text(encoding="utf-8"))
+
+    # Safely get [tool.modaic.ignore]
+    ignore_table = (
+        doc.get("tool", {})  # [tool]
+        .get("modaic", {})  # [tool.modaic]
+        .get("ignore", {})  # [tool.modaic.ignore]
+    )
+
+    if "dependencies" not in ignore_table:
+        return dependencies
+
+    ignored_dependencies = ignore_table["dependencies"]
+    if not ignored_dependencies:
+        return dependencies
+    pattern = re.compile(r"\b(" + "|".join(map(re.escape, ignored_dependencies)) + r")\b")
+    filtered_dependencies = [pkg for pkg in dependencies if not pattern.search(pkg)]
+    return filtered_dependencies
+
+
+def warn_if_local(sources: dict[str, dict]):
+    """
+    Warn if the agent is bundled with a local package.
+    """
+    for source, config in sources.items():
+        if "path" in config:
+            warnings.warn(
+                f"Bundling agent with local package {source} installed from {config['path']}. This is not recommended.",
+                stacklevel=5,
+            )

modaic/observability.py

@@ -0,0 +1,275 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from functools import wraps
+from typing import Any, Callable, Dict, Optional, TypeVar, cast
+
+import opik
+from opik import Opik, config
+from typing_extensions import Concatenate, ParamSpec
+
+from .utils import validate_project_name
+
+P = ParamSpec("P")  # params of the function
+R = TypeVar("R")  # return type of the function
+T = TypeVar("T", bound="Trackable")  # an instance of a class that inherits from Trackable
+
+
+@dataclass
+class ModaicSettings:
+    """Global settings for Modaic observability."""
+
+    tracing: bool = False
+    repo: Optional[str] = None
+    project: Optional[str] = None
+    base_url: str = "https://api.modaic.dev"
+    modaic_token: Optional[str] = None
+    default_tags: Dict[str, str] = field(default_factory=dict)
+    log_inputs: bool = True
+    log_outputs: bool = True
+    max_input_size: int = 10000
+    max_output_size: int = 10000
+
+
+# global settings instance
+_settings = ModaicSettings()
+_opik_client: Optional[Opik] = None
+_configured = False
+
+
+def configure(
+    tracing: bool = True,
+    repo: Optional[str] = None,
+    project: Optional[str] = None,
+    base_url: str = "https://api.modaic.dev",
+    modaic_token: Optional[str] = None,
+    default_tags: Optional[Dict[str, str]] = None,
+    log_inputs: bool = True,
+    log_outputs: bool = True,
+    max_input_size: int = 10000,
+    max_output_size: int = 10000,
+    **opik_kwargs,
+) -> None:
+    """Configure Modaic observability settings globally.
+
+    Args:
+        tracing: Whether observability is enabled
+        repo: Default repository name (e.g., 'user/repo')
+        project: Default project name
+        base_url: Opik server URL
+        modaic_token: Authentication token for Opik
+        default_tags: Default tags to apply to all traces
+        log_inputs: Whether to log function inputs
+        log_outputs: Whether to log function outputs
+        max_input_size: Maximum size of logged inputs
+        max_output_size: Maximum size of logged outputs
+        **opik_kwargs: Additional arguments passed to opik.configure()
+    """
+    global _settings, _opik_client, _configured
+    if project and not repo:
+        raise ValueError("You cannot specify a project without a repo")
+
+    # update global settings
+    _settings.tracing = tracing
+    _settings.repo = repo
+    _settings.project = project
+    _settings.base_url = base_url
+    _settings.modaic_token = modaic_token
+    _settings.default_tags = default_tags or {}
+    _settings.log_inputs = log_inputs
+    _settings.log_outputs = log_outputs
+    _settings.max_input_size = max_input_size
+    _settings.max_output_size = max_output_size
+
+    if tracing:
+        # configure Opik
+        opik_config = {"use_local": True, "url": base_url, "force": True, "automatic_approvals": True, **opik_kwargs}
+
+        opik.configure(**opik_config)
+
+        # create client with project if specified
+        project_name = f"{repo}-{project}" if project else repo
+        if project_name:
+            _opik_client = Opik(host=base_url, project_name=project_name)
+
+        config.update_session_config("track_disable", not tracing)
+
+    _configured = True
+
+
+def _get_effective_settings(
+    repo: Optional[str] = None, project: Optional[str] = None, tags: Optional[Dict[str, str]] = None
+) -> Dict[str, Any]:
+    """Get effective settings by merging global and local parameters."""
+    effective_repo = repo or _settings.repo
+    effective_project = (f"{repo}-{project}" if project else effective_repo) or _settings.project
+
+    # validate project name if provided
+    if effective_project:
+        validate_project_name(effective_project)
+
+    # merge tags
+    effective_tags = {**_settings.default_tags}
+    if tags:
+        effective_tags.update(tags)
+
+    return {"repo": effective_repo, "project": effective_project, "tags": effective_tags}
+
+
+def _truncate_data(data: Any, max_size: int) -> Any:
+    """Truncate data if it exceeds max_size when serialized."""
+    try:
+        import json
+
+        serialized = json.dumps(data, default=str)
+        if len(serialized) > max_size:
+            return f"<Data truncated: {len(serialized)} chars>"
+        return data
+    except Exception:
+        # if serialization fails, convert to string and truncate
+        str_data = str(data)
+        if len(str_data) > max_size:
+            return str_data[:max_size] + "..."
+        return str_data
+
+
+def track(  # noqa: ANN201
+    name: Optional[str] = None,
+    repo: Optional[str] = None,
+    project: Optional[str] = None,
+    tags: Optional[Dict[str, str]] = None,
+    span_type: str = "general",
+    capture_input: Optional[bool] = None,
+    capture_output: Optional[bool] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **opik_kwargs,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """Decorator to track function calls with Opik.
+
+    Args:
+        name: Custom name for the tracked operation
+        repo: Repository name (overrides global setting)
+        project: Project name (overrides global setting)
+        tags: Additional tags for this operation
+        span_type: Type of span ('general', 'tool', 'llm', 'guardrail')
+        capture_input: Whether to capture input (overrides global setting)
+        capture_output: Whether to capture output (overrides global setting)
+        metadata: Additional metadata
+        **opik_kwargs: Additional arguments passed to opik.track
+    """
+
+    def decorator(func: Callable) -> Callable:
+        if not _settings.tracing:
+            return func
+
+        # get effective settings
+        settings = _get_effective_settings(repo, project, tags)
+
+        # determine capture settings
+        should_capture_input = capture_input if capture_input is not None else _settings.log_inputs
+        should_capture_output = capture_output if capture_output is not None else _settings.log_outputs
+
+        # build opik.track arguments
+        track_args: Dict[str, Any] = {
+            "type": span_type,
+            "capture_input": should_capture_input,
+            "capture_output": should_capture_output,
+            **opik_kwargs,
+        }
+
+        # add project if available
+        if settings["project"]:
+            track_args["project_name"] = settings["project"]
+
+        if name:
+            track_args["name"] = name
+
+        # add tags and metadata
+        if settings["tags"] or metadata:
+            combined_metadata = {**(metadata or {})}
+            if settings["tags"]:
+                combined_metadata["tags"] = settings["tags"]
+            if settings["repo"]:
+                combined_metadata["repo"] = settings["repo"]
+            track_args["metadata"] = combined_metadata
+
+        # apply opik.track decorator
+        # Return function with type annotations persisted for static type checking
+        return cast(Callable[P, R], opik.track(**track_args)(func))
+
+    return decorator
+
+
+class Trackable:
+    """Base class for objects that support automatic tracking.
+
+    Manages the attributes repo, project, and commit for classes that subclass it.
+    All Modaic classes except PrecompiledAgent should inherit from this class.
+    """
+
+    def __init__(self, repo: Optional[str] = None, project: Optional[str] = None, commit: Optional[str] = None):
+        self.repo = repo
+        self.project = project
+        self.commit = commit
+
+    def set_repo_project(self, repo: Optional[str] = None, project: Optional[str] = None, trace: bool = True):
+        """Update the repo and project for this trackable object."""
+        if repo is not None:
+            self.repo = repo
+
+        self.project = f"{self.repo}-{project}" if project else self.repo
+
+        # configure global tracing
+        if trace and (repo or project):
+            configure(tracing=trace, repo=repo or self.repo, project=project or self.project)
+
+
+MethodDecorator = Callable[
+    [Callable[Concatenate[T, P], R]],
+    Callable[Concatenate[T, P], R],
+]
+
+
+def track_modaic_obj(func: Callable[Concatenate[T, P], R]) -> Callable[Concatenate[T, P], R]:
+    """Method decorator for Trackable objects to automatically track method calls.
+
+    Uses self.repo and self.project to automatically set repository and project
+    for modaic.track, then wraps the function with modaic.track.
+
+    Usage:
+        class Retriever(Trackable):
+            @track_modaic_obj
+            def retrieve(self, query: str):
+                ...
+    """
+
+    @wraps(func)
+    def wrapper(self: T, *args: P.args, **kwargs: P.kwargs) -> R:
+        # self should be a Trackable instance
+        # TODO: may want to get rid of this type check for hot paths
+        if not isinstance(self, Trackable):
+            raise ValueError("@track_modaic_obj can only be used on methods of Trackable subclasses")
+
+        # get repo and project from self
+        repo = getattr(self, "repo", None)
+        project = getattr(self, "project", None)
+
+        # check if tracking is enabled globally
+        if not _settings.tracing:
+            # binds the method to self so it can be called with args and kwars, also type cast's it to callable with type vars for static type checking
+            bound = cast(Callable[P, R], func.__get__(self, type(self)))
+            return bound(*args, **kwargs)
+
+        # create tracking decorator with automatic name generation
+        tracker = track(
+            name=f"{self.__class__.__name__}.{func.__name__}", repo=repo, project=project, span_type="general"
+        )
+
+        # apply tracking and call method
+        # type casts the 'track' decorator static type checking
+        tracked_func = cast(MethodDecorator, tracker)(func)
+        # binds the method to self so it can be called with args and kwars, also type cast's it to callable with type vars for static type checking
+        bound_tracked = cast(Callable[P, R], tracked_func.__get__(self, type(self)))
+        return bound_tracked(*args, **kwargs)
+
+    return cast(Callable[Concatenate[T, P], R], wrapper)