modaic 0.2.0__tar.gz → 0.4.0__tar.gz

{modaic-0.2.0/src/modaic.egg-info → modaic-0.4.0}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: modaic
-Version: 0.2.0
-Summary: Modular Agent Infrastructure Collective, a python framework for managing and sharing DSPy agents
+Version: 0.4.0
+Summary: Modular Agent Infrastructure Collection, a python framework for managing and sharing DSPy agents
 Author-email: Tyrin <tytodd@mit.edu>, Farouk <farouk@modaic.dev>
 License: MIT License
         

{modaic-0.2.0 → modaic-0.4.0}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "modaic"
-version = "0.2.0"
-description = "Modular Agent Infrastructure Collective, a python framework for managing and sharing DSPy agents"
+version = "0.4.0"
+description = "Modular Agent Infrastructure Collection, a python framework for managing and sharing DSPy agents"
 authors = [{ name = "Tyrin", email = "tytodd@mit.edu" }, {name = "Farouk", email = "farouk@modaic.dev"}]
 readme = "README.md"
 license = {file = "LICENSE"}
@@ -87,4 +87,5 @@ members = [
     "tests/artifacts/test_repos/nested_repo_2",
     "tests/artifacts/test_repos/nested_repo_3",
     "tests/artifacts/test_repos/multi_module_repo",
+    "tests/artifacts/test_repos/failing_repo",
 ]

{modaic-0.2.0 → modaic-0.4.0}/src/modaic/agents/rag_agent.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import Any
 
 from modaic import Indexer, PrecompiledAgent, PrecompiledConfig
 from modaic.context import Context
@@ -19,7 +19,10 @@ class RAGAgentConfig(PrecompiledConfig):
 
 @builtin_indexer(agent_name)
 class RAGIndexer(Indexer):
-    def ingest(self, config: RAGAgentConfig, contexts: List[Context]):
+    def __init__(self, config: RAGAgentConfig):
+        super().__init__(config)
+
+    def index(self, contents: Any):
         pass
 
 

{modaic-0.2.0 → modaic-0.4.0}/src/modaic/databases/sql_database.py

@@ -5,7 +5,19 @@ from typing import Any, Callable, Iterable, List, Literal, Optional, Tuple
 from urllib.parse import urlencode
 
 import pandas as pd
-from sqlalchemy import Column, CursorResult, MetaData, String, Text, create_engine, inspect, text
+from sqlalchemy import (
+    JSON,
+    Column,
+    CursorResult,
+    Index,
+    MetaData,
+    PrimaryKeyConstraint,
+    String,
+    Text,
+    create_engine,
+    inspect,
+    text,
+)
 from sqlalchemy import Table as SQLTable
 from sqlalchemy.dialects import sqlite
 from sqlalchemy.orm import sessionmaker
@@ -81,30 +93,43 @@ class SQLiteBackend(SQLDatabaseBackend):
 
 
 class SQLDatabase:
+    METADATA_TABLE_NAME = "modaic_metadata"
+
     def __init__(
         self,
         backend: SQLDatabaseBackend | str,
         engine_kwargs: dict = None,  # TODO: This may not be a smart idea, may want to enforce specific kwargs
         session_kwargs: dict = None,  # TODO: This may not be a smart idea, may want to enforce specific kwargs
+        track_metadata: bool = False,
     ):
         self.url = backend.url if isinstance(backend, SQLDatabaseBackend) else backend
         self.engine = create_engine(self.url, **(engine_kwargs or {}))
         self.metadata = MetaData()
         self.session = sessionmaker(bind=self.engine, **(session_kwargs or {}))
         self.inspector = inspect(self.engine)
-        self.preparer = IdentifierPreparer(sqlite.dialect())
+        self.preparer = self.engine.dialect.identifier_preparer
 
         # Create metadata table to store table metadata
-        self.metadata_table = SQLTable(
-            "metadata",
-            self.metadata,
-            Column("table_name", String(255), primary_key=True),
-            Column("metadata_json", Text),
+        if track_metadata:
+            self._ensure_metadata_table()
+        self.metadata.reflect(bind=self.engine)
+        self.metadata_table: Optional[Table] = (
+            self.metadata.tables[self.METADATA_TABLE_NAME] if track_metadata else None
         )
-        self.metadata.create_all(self.engine)
         self.connection = None
         self._in_transaction = False
 
+    def _ensure_metadata_table(self) -> None:
+        """Create the metadata table if missing."""
+        if not self.inspector.has_table(self.METADATA_TABLE_NAME):
+            SQLTable(
+                self.METADATA_TABLE_NAME,
+                self.metadata,
+                Column("table_name", String(255), primary_key=True),
+                Column("metadata_json", Text),
+            )
+            self.metadata.create_all(self.engine)
+
     def add_table(
         self,
         table: BaseTable,
@@ -115,17 +140,17 @@ class SQLDatabase:
         with self.connect() as connection:
             # Use the connection for to_sql to respect transaction context
             table._df.to_sql(table.name, connection, if_exists=if_exists, index=False)
-
-            # Remove existing metadata for this table if it exists
-            connection.execute(self.metadata_table.delete().where(self.metadata_table.c.table_name == table.name))
-
-            # Insert new metadata
-            connection.execute(
-                self.metadata_table.insert().values(
-                    table_name=table.name,
-                    metadata_json=json.dumps(table.metadata),
+            if self.metadata_table is not None:
+                # Remove existing metadata for this table if it exists
+                connection.execute(self.metadata_table.delete().where(self.metadata_table.c.table_name == table.name))
+
+                # Insert new metadata
+                connection.execute(
+                    self.metadata_table.insert().values(
+                        table_name=table.name,
+                        metadata_json=json.dumps(table.metadata),
+                    )
                 )
-            )
             if self._should_commit():
                 connection.commit()
 
@@ -151,7 +176,8 @@ class SQLDatabase:
             command = text(f"DROP TABLE {if_exists} {safe_name}")
             connection.execute(command)
             # Also remove metadata for this table
-            connection.execute(self.metadata_table.delete().where(self.metadata_table.c.table_name == name))
+            if self.metadata_table is not None:
+                connection.execute(self.metadata_table.delete().where(self.metadata_table.c.table_name == name))
             if self._should_commit():
                 connection.commit()
 
@@ -197,6 +223,10 @@ class SQLDatabase:
         Returns:
             Dictionary containing the table's metadata, or empty dict if not found.
         """
+        if self.metadata_table is None:
+            raise ValueError(
+                "Metadata table is not enabled. Please enable metadata tracking when initializing the SQLDatabase. with track_metadata=True."
+            )
         with self.connect() as connection:
             result = connection.execute(
                 self.metadata_table.select().where(self.metadata_table.c.table_name == name)

{modaic-0.2.0 → modaic-0.4.0}/src/modaic/exceptions.py

@@ -36,3 +36,12 @@ class BackendCompatibilityError(ModaicError):
     """Raised when a feature is not supported by a backend"""
 
     pass
+
+
+class MissingSecretError(AuthenticationError):
+    """Raised when a secret is missing"""
+
+    def __init__(self, message: str, secret_name: str):
+        self.message = message
+        self.secret_name = secret_name
+        super().__init__(message)

{modaic-0.2.0 → modaic-0.4.0}/src/modaic/module_utils.py

@@ -243,6 +243,12 @@ def init_agent_repo(repo_path: str, with_code: bool = True) -> Path:
         if src_init.exists() and not dest_init.exists():
             shutil.copy2(src_init, dest_init)
 
+    for extra_file in get_extra_files():
+        if extra_file.is_dir():
+            shutil.copytree(extra_file, repo_dir / extra_file.relative_to(project_root))
+        else:
+            shutil.copy2(extra_file, repo_dir / extra_file.relative_to(project_root))
+
     return repo_dir
 
 
@@ -272,23 +278,52 @@ def get_ignored_files() -> list[Path]:
     pyproject_path = Path("pyproject.toml")
     doc = tomlk.parse(pyproject_path.read_text(encoding="utf-8"))
 
-    # Safely get [tool.modaic.ignore]
-    ignore_table = (
+    # Safely get [tool.modaic.exclude]
+    files = (
         doc.get("tool", {})  # [tool]
         .get("modaic", {})  # [tool.modaic]
-        .get("ignore")  # [tool.modaic.ignore]
+        .get("exclude", {})  # [tool.modaic.exclude]
+        .get("files", [])  # [tool.modaic.exclude] files = ["file1", "file2"]
     )
 
-    if ignore_table is None or "files" not in ignore_table:
-        return []
+    excluded: list[Path] = []
+    for entry in files:
+        entry = Path(entry)
+        if not entry.is_absolute():
+            entry = project_root / entry
+        if entry.exists():
+            excluded.append(entry)
+    return excluded
 
-    ignored: list[Path] = []
-    for entry in ignore_table["files"]:
-        try:
-            ignored.append((project_root / entry).resolve())
-        except OSError:
-            continue
-    return ignored
+
+def get_extra_files() -> list[Path]:
+    """Return a list of extra files 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"))
+    files = (
+        doc.get("tool", {})  # [tool]
+        .get("modaic", {})  # [tool.modaic]
+        .get("include", {})  # [tool.modaic.include]
+        .get("files", [])  # [tool.modaic.include] files = ["file1", "file2"]
+    )
+    included: list[Path] = []
+    for entry in files:
+        entry = Path(entry)
+        if entry.is_absolute():
+            try:
+                entry = entry.resolve()
+                entry.relative_to(project_root.resolve())
+            except ValueError:
+                warnings.warn(
+                    f"{entry} will not be bundled because it is not inside the current working directory", stacklevel=4
+                )
+        else:
+            entry = project_root / entry
+        if entry.resolve().exists():
+            included.append(entry)
+
+    return included
 
 
 def create_pyproject_toml(repo_dir: Path, package_name: str):
@@ -304,7 +339,7 @@ def create_pyproject_toml(repo_dir: Path, package_name: str):
     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"])
+    doc_new["project"]["dependencies"] = get_final_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"])
@@ -315,29 +350,32 @@ def create_pyproject_toml(repo_dir: Path, package_name: str):
         tomlk.dump(doc_new, fp)
 
 
-def get_filtered_dependencies(dependencies: list[str]) -> list[str]:
+def get_final_dependencies(dependencies: list[str]) -> list[str]:
     """
     Get the dependencies that should be included in the bundled agent.
+    Filters out "[tool.modaic.ignore] dependencies. Adds [tool.modaic.include] dependencies.
     """
     pyproject_path = Path("pyproject.toml")
     doc = tomlk.parse(pyproject_path.read_text(encoding="utf-8"))
 
-    # Safely get [tool.modaic.ignore]
-    ignore_table = (
+    # Safely get [tool.modaic.exclude]
+    exclude_deps = (
         doc.get("tool", {})  # [tool]
         .get("modaic", {})  # [tool.modaic]
-        .get("ignore", {})  # [tool.modaic.ignore]
+        .get("exclude", {})  # [tool.modaic.exclude]
+        .get("dependencies", [])  # [tool.modaic.exclude] dependencies = ["praw", "sagemaker"]
+    )
+    include_deps = (
+        doc.get("tool", {})  # [tool]
+        .get("modaic", {})  # [tool.modaic]
+        .get("include", {})  # [tool.modaic.include]
+        .get("dependencies", [])  # [tool.modaic.include] dependencies = ["praw", "sagemaker"]
     )
 
-    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
+    if exclude_deps:
+        pattern = re.compile(r"\b(" + "|".join(map(re.escape, exclude_deps)) + r")\b")
+        dependencies = [pkg for pkg in dependencies if not pattern.search(pkg)]
+    return dependencies + include_deps
 
 
 def warn_if_local(sources: dict[str, dict]):

{modaic-0.2.0 → modaic-0.4.0}/src/modaic/precompiled.py

@@ -2,12 +2,13 @@ import inspect
 import json
 import os
 import pathlib
+import warnings
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
+    Any,
     Dict,
-    List,
     Optional,
     Type,
     TypeVar,
@@ -15,11 +16,13 @@ from typing import (
 )
 
 import dspy
+from git import config
 from pydantic import BaseModel
 
 from modaic.module_utils import create_agent_repo
 from modaic.observability import Trackable, track_modaic_obj
 
+from .exceptions import MissingSecretError
 from .hub import load_repo, push_folder_to_hub
 from .module_utils import _module_path
 
@@ -128,6 +131,7 @@ class PrecompiledConfig(BaseModel):
         return self.model_dump_json()
 
 
+# Use a metaclass to enforce super().__init__() with config
 class PrecompiledAgent(dspy.Module):
     """
     Bases: `dspy.Module`
@@ -234,15 +238,25 @@ class PrecompiledAgent(dspy.Module):
                 extra_auto_classes["AutoRetriever"] = self.retriever
         self.config.save_precompiled(path, extra_auto_classes)
         self.save(path / "agent.json")
+        _clean_secrets(path / "agent.json")
 
     @classmethod
-    def from_precompiled(cls: Type[A], path: str | Path, config_options: Optional[dict] = None, **kwargs) -> A:
+    def from_precompiled(
+        cls: Type[A],
+        path: str | Path,
+        config_options: Optional[dict] = None,
+        api_key: Optional[str | dict[str, str]] = None,
+        hf_token: Optional[str | dict[str, str]] = None,
+        **kwargs,
+    ) -> A:
         """
         Loads the agent and the config from the given path.
 
         Args:
             path: The path to load the agent and config from. Can be a local path or a path on Modaic Hub.
             config_options: A dictionary containg key-value pairs used to override the default config.
+            api_key: Your API key.
+            hf_token: Your Hugging Face token.
             **kwargs: Additional keyword arguments forwarded to the PrecompiledAgent's constructor.
 
         Returns:
@@ -260,7 +274,9 @@ class PrecompiledAgent(dspy.Module):
         agent = cls(config, **kwargs)
         agent_state_path = local_dir / "agent.json"
         if agent_state_path.exists():
-            agent.load(agent_state_path)
+            secrets = {"api_key": api_key, "hf_token": hf_token}
+            state = _get_state_with_secrets(agent_state_path, secrets)
+            agent.load_state(state)
         return agent
 
     def push_to_hub(
@@ -374,7 +390,7 @@ class Indexer(Retriever):
     config: PrecompiledConfig
 
     @abstractmethod
-    def ingest(self, contexts: List["Context"], **kwargs):
+    def index(self, contents: Any, **kwargs):
         pass
 
 
@@ -404,11 +420,88 @@ def _push_to_hub(
 
 def is_local_path(s: str | Path) -> bool:
     # absolute or relative filesystem path
+    if isinstance(s, Path):
+        return True
     s = str(s)
+
+    print("SSSS", s)
     if os.path.isabs(s) or s.startswith((".", "/", "\\")):
         return True
     parts = s.split("/")
     # hub IDs: "repo" or "user/repo"
-    if len(parts) == 1 or (len(parts) == 2 and all(parts)):
+    if len(parts) == 1:
+        raise ValueError(
+            f"Invalid repo: '{s}'. Please prefix local paths with './', '/', or '../' . And use 'user/repo' format for hub paths."
+        )
+    elif len(parts) == 2 and all(parts):
         return False
     return True
+
+
+SECRET_MASK = "********"
+COMMON_SECRETS = ["api_key", "hf_token"]
+
+
+def _clean_secrets(path: Path, extra_secrets: Optional[list[str]] = None):
+    """
+    Removes all secret keys from `lm` dict in agent.json file
+    """
+    secret_keys = COMMON_SECRETS + (extra_secrets or [])
+
+    with open(path, "r") as f:
+        d = json.load(f)
+
+    for predictor in d.values():
+        lm = predictor.get("lm", None)
+        if lm is None:
+            continue
+        for k in lm.keys():
+            if k in secret_keys:
+                lm[k] = SECRET_MASK
+
+    with open(path, "w") as f:
+        json.dump(d, f, indent=2)
+
+
+def _get_state_with_secrets(path: Path, secrets: dict[str, str | dict[str, str] | None]):
+    """`
+    Fills secret keys in `lm` dict in agent.json file
+
+    Args:
+        path: The path to the agent.json file.
+        secrets: A dictionary containing the secrets to fill in the `lm` dict.
+            - Dict[k,v] where k is the name of a secret (e.g. "api_key") and v is the value of the secret
+            - If v is a string, every lm will use v for k
+            - if v is a dict, each key of v should be the name of a named predictor
+            (e.g. "my_module.predict", "my_module.summarizer") mapping to the secret value for that predictor
+    Returns:
+        A dictionary containing the state of the agent.json file with the secrets filled in.
+    """
+    with open(path, "r") as f:
+        named_predictors = json.load(f)
+
+    def _get_secret(predictor_name: str, secret_name: str) -> Optional[str]:
+        if secret_val := secrets.get(secret_name):
+            if isinstance(secret_val, str):
+                return secret_val
+            elif isinstance(secret_val, dict):
+                return secret_val.get(predictor_name)
+        return None
+
+    for predictor_name, predictor in named_predictors.items():
+        lm = predictor.get("lm", {})
+        for kw, arg in lm.items():
+            if kw in COMMON_SECRETS and arg != "" and arg != SECRET_MASK:
+                warnings.warn(
+                    f"{str(path)} exposes the secret key {kw}. Please remove it or ensure this file is not made public."
+                )
+            secret = _get_secret(predictor_name, kw)
+            if secret is not None and arg != "" and arg != SECRET_MASK:
+                raise ValueError(
+                    f"Failed to fill insert secret value for {predictor_name}['lm']['{kw}']. It is already set to {arg}"
+                )
+            elif secret is None and kw in COMMON_SECRETS:
+                raise MissingSecretError(f"Please specify a value for {kw} in the secrets dictionary", kw)
+            elif secret is not None:
+                lm[kw] = secret
+    return named_predictors

{modaic-0.2.0 → modaic-0.4.0/src/modaic.egg-info}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: modaic
-Version: 0.2.0
-Summary: Modular Agent Infrastructure Collective, a python framework for managing and sharing DSPy agents
+Version: 0.4.0
+Summary: Modular Agent Infrastructure Collection, a python framework for managing and sharing DSPy agents
 Author-email: Tyrin <tytodd@mit.edu>, Farouk <farouk@modaic.dev>
 License: MIT License
         

modaic-0.4.0/tests/test_auto.py

@@ -0,0 +1,341 @@
+import os
+import pathlib
+import shutil
+import subprocess
+from pathlib import Path
+from typing import Union
+
+import pytest
+import tomlkit as tomlk
+
+from modaic import AutoAgent, AutoConfig, AutoRetriever
+from modaic.hub import MODAIC_CACHE, get_user_info
+from tests.testing_utils import delete_agent_repo
+
+MODAIC_TOKEN = os.getenv("MODAIC_TOKEN")
+INSTALL_TEST_REPO_DEPS = os.getenv("INSTALL_TEST_REPO_DEPS", "True").lower() == "true"
+USERNAME = get_user_info(os.environ["MODAIC_TOKEN"])["login"]
+
+
+def get_cached_agent_dir(repo_name: str) -> Path:
+    return MODAIC_CACHE / "agents" / repo_name
+
+
+def clean_modaic_cache() -> None:
+    """Remove the MODAIC cache directory if it exists.
+
+    Params:
+        None
+
+    Returns:
+        None
+    """
+    shutil.rmtree(MODAIC_CACHE, ignore_errors=True)
+
+
+def prepare_repo(repo_name: str) -> None:
+    """Clean cache and ensure remote hub repo is deleted before test run.
+
+    Params:
+        repo_name (str): The name of the test repository in artifacts/test_repos.
+
+    Returns:
+        None
+    """
+    clean_modaic_cache()
+    if not MODAIC_TOKEN:
+        pytest.skip("Skipping because MODAIC_TOKEN is not set")
+    delete_agent_repo(username=USERNAME, agent_name=repo_name)
+
+
+def run_script(repo_name: str, run_path: str = "compile.py") -> None:
+    """Run the repository's compile script inside its own uv environment.
+
+    Params:
+        repo_name (str): The name of the test repository directory to compile.
+
+    Returns:
+        None
+    """
+    env = os.environ.copy()
+    env.update(
+        {
+            "MODAIC_CACHE": "../../temp/modaic_cache",
+        }
+    )
+    repo_dir = pathlib.Path("tests/artifacts/test_repos") / repo_name
+    if INSTALL_TEST_REPO_DEPS:
+        subprocess.run(["uv", "sync"], cwd=repo_dir, check=True, env=env)
+        # Ensure the root package is available in the subproject env
+    # Run as file
+    if run_path.endswith(".py"):
+        subprocess.run(["uv", "run", run_path, USERNAME], cwd=repo_dir, check=True, env=env)
+    # Run as module
+    else:
+        subprocess.run(["uv", "run", "-m", run_path, USERNAME], cwd=repo_dir, check=True, env=env)
+    # clean cache
+    shutil.rmtree("tests/artifacts/temp/modaic_cache", ignore_errors=True)
+
+
+# recursive dict/list of dicts/lists of strs representing a folder structure
+FolderLayout = dict[str, Union[str, "FolderLayout"]] | list[Union[str, "FolderLayout"]]
+
+
+def assert_expected_files(cache_dir: Path, extra_expected_files: FolderLayout):
+    default_expected = ["agent.json", "auto_classes.json", "config.json", "pyproject.toml", "README.md", ".git"]
+    if isinstance(extra_expected_files, list):
+        expected = extra_expected_files + default_expected
+    elif isinstance(extra_expected_files, dict):
+        expected = [extra_expected_files] + default_expected
+    else:
+        raise ValueError(f"Invalid folder layout: {extra_expected_files}")
+    assert_folder_layout(cache_dir, expected)
+
+
+def assert_top_level_names(dir: Path, expected_files: FolderLayout | str, root: bool = True):
+    if isinstance(expected_files, list):
+        expected_names = []
+        for obj in expected_files:
+            if isinstance(obj, str):
+                expected_names.append(obj)
+            elif isinstance(obj, dict):
+                expected_names.extend(list(obj.keys()))
+            else:
+                raise ValueError(f"Invalid folder layout: {expected_files}")
+    elif isinstance(expected_files, dict):
+        expected_names = list(expected_files.keys())
+    elif isinstance(expected_files, str):
+        expected_names = [expected_files]
+    else:
+        raise ValueError(f"Invalid folder layout: {expected_files}")
+    expected_names = expected_names if root else expected_names + ["__init__.py"]
+    missing = set(expected_names) - set(os.listdir(dir))
+    assert missing == set(), f"Missing files, in {dir}, {missing}"
+    unexpected = set(os.listdir(dir)) - set(expected_names)
+    assert unexpected.issubset(set(["__pycache__", "__init__.py"])), (
+        f"Unexpected files in {dir}, {unexpected - set(['__pycache__', '__init__.py'])}"
+    )
+
+
+def assert_folder_layout(
+    dir: Path, expected_files: FolderLayout | str, root: bool = True, assert_top_level: bool = True
+):
+    """
+    Asserts that the files in the directory match the expected folder structure.
+    Checking that only expected files are included. Will raise assertion error if unexpected files are included.
+    Args:
+        dir: The directory to assert the files in.
+        expected_files: The expected folder structure.
+
+    Raises:
+        Assertion error if expected file not found in path or if unexpected file found in path
+    """
+    # dir is a single file folder
+    if isinstance(expected_files, str):
+        assert_top_level_names(dir, expected_files, root)
+    # dir is a folder containg multiples files or subfolders
+    elif isinstance(expected_files, list):
+        assert_top_level_names(dir, expected_files, root)
+        for file in expected_files:
+            if isinstance(file, dict):
+                assert_folder_layout(dir, file, root=False, assert_top_level=False)
+            elif not isinstance(file, str):
+                raise ValueError(f"Invalid folder layout: {expected_files}")
+    # dir contains subfolders, however don't check top level because we don't know if this is the entirety of dir or a subset
+    elif isinstance(expected_files, dict):
+        for key, value in expected_files.items():
+            assert_folder_layout(dir / key, value, root=False)
+    else:
+        raise ValueError(f"Invalid folder layout: {expected_files}")
+
+
+def assert_dependencies(cache_dir: Path, extra_expected_dependencies: list[str]):
+    expected_dependencies = extra_expected_dependencies + ["dspy", "modaic"]
+
+    pyproject_path = cache_dir / "pyproject.toml"
+    doc = tomlk.parse(pyproject_path.read_text(encoding="utf-8"))
+    actual_dependencies = doc.get("project", {}).get("dependencies", [])
+
+    missing = set(expected_dependencies) - set(actual_dependencies)
+    assert missing == set(), f"Missing dependencies, {missing}"
+    unexpected = set(actual_dependencies) - set(expected_dependencies)
+    assert unexpected == set(), f"Unexpected dependencies, {unexpected}"
+
+
+def test_simple_repo() -> None:
+    prepare_repo("simple_repo")
+    run_script("simple_repo", run_path="agent.py")
+    clean_modaic_cache()
+    config = AutoConfig.from_precompiled(f"{USERNAME}/simple_repo")
+    assert config.lm == "openai/gpt-4o"
+    assert config.output_type == "str"
+    assert config.number == 1
+    cache_dir = get_cached_agent_dir(f"{USERNAME}/simple_repo")
+    assert_expected_files(cache_dir, ["agent.py"])
+    assert_dependencies(cache_dir, ["dspy", "modaic", "praw"])
+
+    clean_modaic_cache()
+    agent = AutoAgent.from_precompiled(f"{USERNAME}/simple_repo", runtime_param="Hello")
+    assert agent.config.lm == "openai/gpt-4o"
+    assert agent.config.output_type == "str"
+    assert agent.config.number == 1
+    assert agent.runtime_param == "Hello"
+    clean_modaic_cache()
+    agent = AutoAgent.from_precompiled(
+        f"{USERNAME}/simple_repo", runtime_param="Hello", config_options={"lm": "openai/gpt-4o-mini"}
+    )
+    assert agent.config.lm == "openai/gpt-4o-mini"
+    assert agent.config.output_type == "str"
+    assert agent.config.number == 1
+    assert agent.runtime_param == "Hello"
+    # TODO: test third party deps installation
+
+
+simple_repo_with_compile_extra_files = [{"agent": ["agent.py", "mod.py"]}, "compile.py", "include_me_too.txt"]
+
+
+def test_simple_repo_with_compile():
+    prepare_repo("simple_repo_with_compile")
+    run_script("simple_repo_with_compile", run_path="compile.py")
+    clean_modaic_cache()
+    config = AutoConfig.from_precompiled(f"{USERNAME}/simple_repo_with_compile")
+    assert config.lm == "openai/gpt-4o"
+    assert config.output_type == "str"
+    assert config.number == 1
+    cache_dir = get_cached_agent_dir(f"{USERNAME}/simple_repo_with_compile")
+    assert os.path.exists(cache_dir / "config.json")
+    assert os.path.exists(cache_dir / "agent.json")
+    assert os.path.exists(cache_dir / "auto_classes.json")
+    assert os.path.exists(cache_dir / "README.md")
+    assert os.path.exists(cache_dir / "agent" / "agent.py")
+    assert os.path.exists(cache_dir / "agent" / "mod.py")
+    assert os.path.exists(cache_dir / "pyproject.toml")
+    assert os.path.exists(cache_dir / "include_me_too.txt")
+    extra_files = [{"agent": ["agent.py", "mod.py"]}, "compile.py", "include_me_too.txt"]
+    assert_expected_files(cache_dir, extra_files)
+    assert_dependencies(cache_dir, ["dspy", "modaic"])
+    clean_modaic_cache()
+    agent = AutoAgent.from_precompiled(f"{USERNAME}/simple_repo_with_compile", runtime_param="Hello")
+    assert agent.config.lm == "openai/gpt-4o"
+    assert agent.config.output_type == "str"
+    assert agent.config.number == 1
+    assert agent.runtime_param == "Hello"
+    clean_modaic_cache()
+    agent = AutoAgent.from_precompiled(
+        f"{USERNAME}/simple_repo_with_compile", runtime_param="Hello", config_options={"lm": "openai/gpt-4o-mini"}
+    )
+    assert agent.config.lm == "openai/gpt-4o-mini"
+    assert agent.config.output_type == "str"
+    assert agent.config.number == 1
+    assert agent.runtime_param == "Hello"
+    # TODO: test third party deps installation
+
+
+nested_repo_extra_files = {
+    "agent": [
+        {
+            "tools": {"google": "google_search.py", "jira": "jira_api_tools.py"},
+            "utils": ["second_degree_import.py", "used.py"],
+        },
+        "agent.py",
+        "compile.py",
+        "config.py",
+        "retriever.py",
+    ]
+}
+nested_repo_2_extra_files = [
+    {
+        "agent": [
+            {
+                "tools": {"google": "google_search.py", "jira": "jira_api_tools.py"},
+                "utils": [
+                    "second_degree_import.py",
+                    "unused_but_included.py",
+                    "used.py",
+                ],
+            },
+            "agent.py",
+            "config.py",
+            "retriever.py",
+        ]
+    },
+    {"unused_but_included_folder": [".env", "folder_content1.py", "folder_content2.txt"]},
+    "compile.py",
+]
+nested_repo_3_extra_files = {
+    "agent": [
+        {
+            "tools": [{"google": "google_search.py", "jira": "jira_api_tools.py"}, "unused_but_included2.py"],
+            "utils": ["second_degree_import.py", "unused_but_included.py", "used.py"],
+        },
+        "agent.py",
+        "config.py",
+        "retriever.py",
+    ],
+}
+
+
+@pytest.mark.parametrize(
+    "repo_name, run_path, extra_expected_files, extra_expected_dependencies",
+    [
+        (
+            "nested_repo",
+            "agent.compile",
+            nested_repo_extra_files,
+            [],
+        ),
+        (
+            "nested_repo_2",
+            "compile.py",
+            nested_repo_2_extra_files,
+            ["dspy", "modaic", "praw", "sagemaker"],
+        ),
+        (
+            "nested_repo_3",
+            "agent.agent",
+            nested_repo_3_extra_files,
+            ["dspy", "modaic"],
+        ),
+    ],
+)
+def test_nested_repo(
+    repo_name: str, run_path: str, extra_expected_files: FolderLayout, extra_expected_dependencies: list[str]
+):
+    prepare_repo(repo_name)
+    run_script(repo_name, run_path=run_path)
+    clean_modaic_cache()
+    config = AutoConfig.from_precompiled(f"{USERNAME}/{repo_name}", clients={"get_replaced": "noob"})
+    assert config.num_fetch == 1
+    assert config.lm == "openai/gpt-4o-mini"
+    assert config.embedder == "openai/text-embedding-3-small"
+    assert config.clients == {"get_replaced": "noob"}
+
+    cache_dir = get_cached_agent_dir(f"{USERNAME}/{repo_name}")
+    assert_expected_files(cache_dir, extra_expected_files)
+    assert_dependencies(cache_dir, extra_expected_dependencies)
+
+    clean_modaic_cache()
+    retriever = AutoRetriever.from_precompiled(f"{USERNAME}/{repo_name}", needed_param="hello")
+    agent = AutoAgent.from_precompiled(f"{USERNAME}/{repo_name}", retriever=retriever)
+    assert agent.config.num_fetch == 1
+    assert agent.config.lm == "openai/gpt-4o-mini"
+    assert agent.config.embedder == "openai/text-embedding-3-small"
+    assert agent.config.clients == {"mit": ["csail", "mit-media-lab"], "berkeley": ["bear"]}
+    assert retriever.needed_param == "hello"
+    assert agent.forward("my query") == "Retrieved 1 results for my query"
+    clean_modaic_cache()
+    config_options = {"lm": "openai/gpt-4o"}
+    retriever = AutoRetriever.from_precompiled(
+        f"{USERNAME}/{repo_name}", needed_param="hello", config_options=config_options
+    )
+    agent = AutoAgent.from_precompiled(f"{USERNAME}/{repo_name}", retriever=retriever, config_options=config_options)
+    assert agent.config.num_fetch == 1
+    assert agent.config.lm == "openai/gpt-4o"
+    assert agent.config.embedder == "openai/text-embedding-3-small"
+    assert agent.config.clients == {"mit": ["csail", "mit-media-lab"], "berkeley": ["bear"]}
+    assert retriever.needed_param == "hello"
+    assert agent.forward("my query") == "Retrieved 1 results for my query"
+
+
+def test_auth():
+    pass

{modaic-0.2.0 → modaic-0.4.0}/tests/test_precompiled.py

@@ -1,3 +1,4 @@
+import json
 import os
 import shutil
 from pathlib import Path
@@ -415,5 +416,44 @@ def test_precompiled_agent_with_retriever_hub(hub_repo: str):
     loaded_agent3.push_to_hub(hub_repo, with_code=False)
 
 
+class InnerSecretAgent(dspy.Module):
+    def __init__(self):
+        self.predictor = dspy.Predict(Summarize)
+        self.predictor.set_lm(lm=dspy.LM("openai/gpt-4o-mini", api_key="sk-proj-1234567890", hf_token="hf_1234567890"))
+
+    def forward(self, query: str) -> str:
+        return self.predictor(query=query)
+
+
+class SecretAgentConfig(PrecompiledConfig):
+    pass
+
+
+class SecretAgent(PrecompiledAgent):
+    config: SecretAgentConfig
+
+    def __init__(self, config: SecretAgentConfig, **kwargs):
+        super().__init__(config, **kwargs)
+        self.predictor = dspy.Predict(Summarize)
+        self.predictor.set_lm(lm=dspy.LM("openai/gpt-4o-mini", api_key="sk-proj-1234567890"))
+        self.inner = InnerSecretAgent()
+
+    def forward(self, query: str) -> str:
+        return self.inner(query=query)
+
+
+def test_precompiled_agent_with_secret(clean_folder: Path):
+    SecretAgent(SecretAgentConfig()).save_precompiled(clean_folder)
+    with open(clean_folder / "agent.json", "r") as f:
+        agent_state = json.load(f)
+    assert agent_state["inner.predictor"]["lm"]["api_key"] == "********"
+    assert agent_state["inner.predictor"]["lm"]["hf_token"] == "********"
+    assert agent_state["predictor"]["lm"]["api_key"] == "********"
+    loaded_agent = SecretAgent.from_precompiled(clean_folder, api_key="set-api-key", hf_token="set-hf-token")
+    assert loaded_agent.inner.predictor.lm.kwargs["api_key"] == "set-api-key"
+    assert loaded_agent.inner.predictor.lm.kwargs["hf_token"] == "set-hf-token"
+    assert loaded_agent.predictor.lm.kwargs["api_key"] == "set-api-key"
+
+
 def test_unauthorized_push_to_hub():
     pass

modaic-0.2.0/tests/test_auto.py

@@ -1,180 +0,0 @@
-import os
-import pathlib
-import shutil
-import subprocess
-from pathlib import Path
-
-import pytest
-
-from modaic import AutoAgent, AutoConfig, AutoRetriever
-from modaic.hub import MODAIC_CACHE, get_user_info
-from tests.testing_utils import delete_agent_repo
-
-MODAIC_TOKEN = os.getenv("MODAIC_TOKEN")
-INSTALL_TEST_REPO_DEPS = os.getenv("INSTALL_TEST_REPO_DEPS", "True").lower() == "true"
-USERNAME = get_user_info(os.environ["MODAIC_TOKEN"])["login"]
-
-
-def get_cached_agent_dir(repo_name: str) -> Path:
-    return MODAIC_CACHE / "agents" / repo_name
-
-
-def clean_modaic_cache() -> None:
-    """Remove the MODAIC cache directory if it exists.
-
-    Params:
-        None
-
-    Returns:
-        None
-    """
-    shutil.rmtree(MODAIC_CACHE, ignore_errors=True)
-
-
-def prepare_repo(repo_name: str) -> None:
-    """Clean cache and ensure remote hub repo is deleted before test run.
-
-    Params:
-        repo_name (str): The name of the test repository in artifacts/test_repos.
-
-    Returns:
-        None
-    """
-    clean_modaic_cache()
-    if not MODAIC_TOKEN:
-        pytest.skip("Skipping because MODAIC_TOKEN is not set")
-    delete_agent_repo(username=USERNAME, agent_name=repo_name)
-
-
-def run_script(repo_name: str, run_path: str = "compile.py", module_mode: bool = False) -> None:
-    """Run the repository's compile script inside its own uv environment.
-
-    Params:
-        repo_name (str): The name of the test repository directory to compile.
-
-    Returns:
-        None
-    """
-    env = os.environ.copy()
-    env.update(
-        {
-            "MODAIC_CACHE": "../../temp/modaic_cache",
-        }
-    )
-    repo_dir = pathlib.Path("tests/artifacts/test_repos") / repo_name
-    if INSTALL_TEST_REPO_DEPS:
-        subprocess.run(["uv", "sync"], cwd=repo_dir, check=True, env=env)
-        # Ensure the root package is available in the subproject env
-    if module_mode:
-        subprocess.run(["uv", "run", "-m", run_path, USERNAME], cwd=repo_dir, check=True, env=env)
-    else:
-        subprocess.run(["uv", "run", run_path, USERNAME], cwd=repo_dir, check=True, env=env)
-    # clean cache
-    shutil.rmtree("tests/artifacts/temp/modaic_cache", ignore_errors=True)
-
-
-def test_simple_repo() -> None:
-    prepare_repo("simple_repo")
-    run_script("simple_repo", run_path="agent.py")
-    clean_modaic_cache()
-    config = AutoConfig.from_precompiled(f"{USERNAME}/simple_repo")
-    assert config.lm == "openai/gpt-4o"
-    assert config.output_type == "str"
-    assert config.number == 1
-    cache_dir = get_cached_agent_dir(f"{USERNAME}/simple_repo")
-    assert os.path.exists(cache_dir / "config.json")
-    assert os.path.exists(cache_dir / "agent.json")
-    assert os.path.exists(cache_dir / "auto_classes.json")
-    assert os.path.exists(cache_dir / "README.md")
-    assert os.path.exists(cache_dir / "agent.py")
-    assert os.path.exists(cache_dir / "pyproject.toml")
-    clean_modaic_cache()
-    agent = AutoAgent.from_precompiled(f"{USERNAME}/simple_repo", runtime_param="Hello")
-    assert agent.config.lm == "openai/gpt-4o"
-    assert agent.config.output_type == "str"
-    assert agent.config.number == 1
-    assert agent.runtime_param == "Hello"
-    clean_modaic_cache()
-    agent = AutoAgent.from_precompiled(
-        f"{USERNAME}/simple_repo", runtime_param="Hello", config_options={"lm": "openai/gpt-4o-mini"}
-    )
-    assert agent.config.lm == "openai/gpt-4o-mini"
-    assert agent.config.output_type == "str"
-    assert agent.config.number == 1
-    assert agent.runtime_param == "Hello"
-    # TODO: test third party deps installation
-
-
-def test_simple_repo_with_compile():
-    prepare_repo("simple_repo_with_compile")
-    run_script("simple_repo_with_compile", run_path="compile.py")
-    clean_modaic_cache()
-    config = AutoConfig.from_precompiled(f"{USERNAME}/simple_repo_with_compile")
-    assert config.lm == "openai/gpt-4o"
-    assert config.output_type == "str"
-    assert config.number == 1
-    cache_dir = get_cached_agent_dir(f"{USERNAME}/simple_repo_with_compile")
-    assert os.path.exists(cache_dir / "config.json")
-    assert os.path.exists(cache_dir / "agent.json")
-    assert os.path.exists(cache_dir / "auto_classes.json")
-    assert os.path.exists(cache_dir / "README.md")
-    assert os.path.exists(cache_dir / "agent" / "agent.py")
-    assert os.path.exists(cache_dir / "agent" / "mod.py")
-    assert os.path.exists(cache_dir / "pyproject.toml")
-    clean_modaic_cache()
-    agent = AutoAgent.from_precompiled(f"{USERNAME}/simple_repo_with_compile", runtime_param="Hello")
-    assert agent.config.lm == "openai/gpt-4o"
-    assert agent.config.output_type == "str"
-    assert agent.config.number == 1
-    assert agent.runtime_param == "Hello"
-    clean_modaic_cache()
-    agent = AutoAgent.from_precompiled(
-        f"{USERNAME}/simple_repo_with_compile", runtime_param="Hello", config_options={"lm": "openai/gpt-4o-mini"}
-    )
-    assert agent.config.lm == "openai/gpt-4o-mini"
-    assert agent.config.output_type == "str"
-    assert agent.config.number == 1
-    assert agent.runtime_param == "Hello"
-    # TODO: test third party deps installation
-
-
-@pytest.mark.parametrize("repo_name", ["nested_repo", "nested_repo_2", "nested_repo_3"])
-def test_nested_repo(repo_name: str):
-    prepare_repo(repo_name)
-    if repo_name == "nested_repo":
-        run_script(repo_name, run_path="agent.compile", module_mode=True)
-    elif repo_name == "nested_repo_2":
-        run_script(repo_name, run_path="compile.py")
-    else:
-        run_script(repo_name, run_path="agent.agent", module_mode=True)
-    clean_modaic_cache()
-    config = AutoConfig.from_precompiled(f"{USERNAME}/{repo_name}", clients={"get_replaced": "noob"})
-    assert config.num_fetch == 1
-    assert config.lm == "openai/gpt-4o-mini"
-    assert config.embedder == "openai/text-embedding-3-small"
-    assert config.clients == {"get_replaced": "noob"}
-    clean_modaic_cache()
-    retriever = AutoRetriever.from_precompiled(f"{USERNAME}/{repo_name}", needed_param="hello")
-    agent = AutoAgent.from_precompiled(f"{USERNAME}/{repo_name}", retriever=retriever)
-    assert agent.config.num_fetch == 1
-    assert agent.config.lm == "openai/gpt-4o-mini"
-    assert agent.config.embedder == "openai/text-embedding-3-small"
-    assert agent.config.clients == {"mit": ["csail", "mit-media-lab"], "berkeley": ["bear"]}
-    assert retriever.needed_param == "hello"
-    assert agent.forward("my query") == "Retrieved 1 results for my query"
-    clean_modaic_cache()
-    config_options = {"lm": "openai/gpt-4o"}
-    retriever = AutoRetriever.from_precompiled(
-        f"{USERNAME}/{repo_name}", needed_param="hello", config_options=config_options
-    )
-    agent = AutoAgent.from_precompiled(f"{USERNAME}/{repo_name}", retriever=retriever, config_options=config_options)
-    assert agent.config.num_fetch == 1
-    assert agent.config.lm == "openai/gpt-4o"
-    assert agent.config.embedder == "openai/text-embedding-3-small"
-    assert agent.config.clients == {"mit": ["csail", "mit-media-lab"], "berkeley": ["bear"]}
-    assert retriever.needed_param == "hello"
-    assert agent.forward("my query") == "Retrieved 1 results for my query"
-
-
-def test_auth():
-    pass