diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md new file mode 100644 index 000000000000..39e3ccde1c36 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md @@ -0,0 +1,24 @@ +# Release History + +## 1.0.0b1 (2026-05-24) + +### Features Added + +- Initial beta release. +- `load_config(*, config_dir, required)` — single-call config loader with 4-priority resolution and graceful fallback. +- `load_skills_from_dir(path)` — load skills from a directory on demand (not loaded inline by `load_config`). +- `OptimizationConfig` dataclass with instructions, model, temperature, skills, skills_dir, tool_definitions, source, and candidate_id. +- `OptimizationConfig.apply_tool_descriptions(tools)` — patch `__doc__`, `.description`, and `input_model` parameter descriptions on @tool-decorated functions from optimized tool definitions. +- `OptimizationConfig.compose_instructions()` — append skill catalog to instructions. +- `CandidateConfig` — typed representation of the resolver API payload. +- `Skill` — learned skill model (name, description, body). +- 4-priority resolution order: + 1. Inline JSON via `OPTIMIZATION_CONFIG` env var. + 2. Resolver API via `OPTIMIZATION_CANDIDATE_ID` + `OPTIMIZATION_RESOLVE_ENDPOINT` (endpoint is the full job-scoped URL). + 3. Local directory layout (`OPTIMIZATION_LOCAL_DIR` or `config_dir` param, defaults to `.agent_configs/`). + 4. `required=True` raises `ValueError`; `required=False` returns `None`. +- Local directory layout: `metadata.yaml` + `instructions.md` + `tools.json` + `skills/` per candidate, with `baseline/` fallback. +- Tool definitions use the OpenAI function-calling list format exclusively. +- Skill loading from `SKILL.md` files with YAML frontmatter. +- Resolver API persists fetched configs and skill files to local directory for offline use. +- Path traversal (zip-slip) protection on skill file downloads from the API. diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/LICENSE b/sdk/agentserver/azure-ai-agentserver-optimization/LICENSE new file mode 100644 index 000000000000..4c3581d3b052 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/LICENSE @@ -0,0 +1,21 @@ +Copyright (c) Microsoft Corporation. + +MIT License + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/MANIFEST.in b/sdk/agentserver/azure-ai-agentserver-optimization/MANIFEST.in new file mode 100644 index 000000000000..0614561fbf5c --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/MANIFEST.in @@ -0,0 +1,8 @@ +include *.md +include LICENSE +recursive-include tests *.py +recursive-include samples *.py *.md +include azure/__init__.py +include azure/ai/__init__.py +include azure/ai/agentserver/__init__.py +include azure/ai/agentserver/optimization/py.typed diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/README.md b/sdk/agentserver/azure-ai-agentserver-optimization/README.md new file mode 100644 index 000000000000..9d0204d30072 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/README.md @@ -0,0 +1,220 @@ +# Azure AI Agent Server Optimization client library for Python + +The `azure-ai-agentserver-optimization` package provides a drop-in config loader for optimization-ready Azure AI Hosted Agents. A single `load_config()` call resolves optimization parameters (instructions, model, temperature, skills, tool definitions) from multiple sources with graceful fallback — your agent works unchanged when not running under optimization. + +## Getting started + +### Install the package + +```bash +pip install azure-ai-agentserver-optimization +``` + +### Prerequisites + +- Python 3.10 or later + +## Key concepts + +### Resolution Order + +`load_config()` resolves from four sources in order — first match wins: + +| Priority | Source | Env vars required | Description | +|----------|--------|-------------------|-------------| +| 1 | **Inline JSON** | `OPTIMIZATION_CONFIG` | Full config as a JSON string. Used by temporary agent versions during evaluation. | +| 2 | **Resolver API** | `OPTIMIZATION_CANDIDATE_ID`, `OPTIMIZATION_RESOLVE_ENDPOINT` | Fetches the candidate config from the remote optimization service and persists it to the local directory. The endpoint should be the full job-scoped URL. | +| 3 | **Local directory** | `OPTIMIZATION_LOCAL_DIR` (optional, defaults to `.agent_configs/`) | Reads from `//` or `baseline/` as fallback. | +| 4 | **No config** | *(none)* | `required=True` (default) raises `ValueError`; `required=False` returns `None`. | + +Any unexpected error is caught and logged — `load_config()` never crashes (only `ValueError` from `required=True` propagates). + +### Environment Variables + +| Variable | Description | +|----------|-------------| +| `OPTIMIZATION_CONFIG` | Inline JSON config (Priority 1). | +| `OPTIMIZATION_CANDIDATE_ID` | Candidate ID for resolver API or local folder lookup. | +| `OPTIMIZATION_RESOLVE_ENDPOINT` | Full job-scoped URL of the optimization service (e.g. `https://host/api/projects/proj/agents_optimization_job/job-1`). | +| `OPTIMIZATION_LOCAL_DIR` | Path to the local config directory (default: `.agent_configs/`). | + +### Local Directory Layout + +The local config directory **must** exist with at least a `baseline/` folder before the agent can start to do optimization. When running under optimization, the resolver API (Priority 2) persists candidate configs into the same layout. If a `/` folder exists it takes precedence; otherwise `baseline/` is used as the fallback. + +``` +.agent_configs/ +├── baseline/ # required — default candidate used at startup +│ ├── metadata.yaml # model, temperature, file pointers +│ ├── instructions.md # system prompt +│ ├── tools.json # tool definitions (OpenAI function-calling list format) +│ └── skills/ # learned skills +│ └── / +│ └── SKILL.md +└── / # optional — created by resolver API, same layout as baseline/ + ├── metadata.yaml + ├── instructions.md + ├── tools.json + └── skills/ + └── / + └── SKILL.md +``` + +#### `metadata.yaml` + +Points to the other files and sets model parameters. All fields are optional: + +```yaml +model: gpt-4o +temperature: 0.7 +instruction_file: instructions.md +skill_dir: skills +tool_file: tools.json +``` + +#### `instructions.md` + +The system prompt for the agent — plain text or Markdown: + +```markdown +You are a travel assistant. Help users search flights, book hotels, +and answer questions about company travel policy. +``` + +#### `tools.json` + +Tool definitions in the OpenAI function-calling list format: + +```json +[ + { + "type": "function", + "function": { + "name": "lookup_policy", + "description": "Look up the company travel policy.", + "parameters": { + "type": "object", + "properties": { + "dept": {"type": "string", "description": "Department name"} + } + } + } + } +] +``` + +#### `skills/*/SKILL.md` + +Each skill lives in its own subfolder with a `SKILL.md` file. An optional YAML frontmatter block provides the name and description; the rest is the skill body: + +```markdown +--- +name: budget-checker +description: Check whether a trip is within budget. +--- + +Compare the trip cost against the department's remaining travel budget +and return APPROVED or DENIED with a reason. +``` + +### OptimizationConfig Properties + +| Property | Type | Description | +|----------|------|-------------| +| `instructions` | `str \| None` | System prompt (optimized or default). | +| `model` | `str \| None` | Model deployment name. | +| `temperature` | `float \| None` | Sampling temperature. | +| `skills` | `list[Skill]` | Learned skills (from inline config). | +| `skills_dir` | `str \| None` | Path to skills directory (for on-demand loading). | +| `tool_definitions` | `list[dict]` | Optimized tool definitions (OpenAI function-calling format). | +| `source` | `str` | Where the config was loaded from. | +| `candidate_id` | `str \| None` | Candidate ID (when resolved via API or local folder). | +| `has_skills` | `bool` | Whether skills are available (inline or via skills_dir). | + +### Public API + +| Export | Type | Description | +|--------|------|-------------| +| `load_config(*, config_dir, required)` | function | Load optimization config with 4-priority resolution. | +| `load_skills_from_dir(path)` | function | Load skills from a directory of `SKILL.md` files. | +| `OptimizationConfig` | dataclass | Resolved config with instructions, model, temperature, skills_dir, tool_definitions. | +| `OptimizationConfig.apply_tool_descriptions(tools)` | method | Patch `__doc__`, `.description`, and parameter descriptions on tool functions. | +| `OptimizationConfig.compose_instructions()` | method | Return instructions with skill catalog appended. | +| `CandidateConfig` | dataclass | Typed representation of the resolver API response. | +| `Skill` | dataclass | A learned skill (name, description, body). | + +## Examples + +### Basic usage + +```python +from azure.ai.agentserver.optimization import load_config + +config = load_config() # uses .agent_configs/baseline/ +config = load_config(config_dir="my_configs") # custom directory +config = load_config(required=False) # returns None if no config found + +print(config.instructions) # optimized system prompt +print(config.model) # optimized model name +print(config.temperature) # optimized temperature +print(config.tool_definitions) # optimized tool definitions (list) +print(config.source) # "env:OPTIMIZATION_CONFIG", "api:candidate:abc", "local:...", etc. +``` + +### Apply optimized tool descriptions + +```python +from azure.ai.agentserver.optimization import load_config + +config = load_config() + +# Your @tool-decorated functions +def search_flights(origin: str, destination: str): + """Search for flights.""" + ... + +def book_hotel(city: str): + """Book a hotel room.""" + ... + +# Patches __doc__ and .description on matching tools with optimized descriptions +config.apply_tool_descriptions([search_flights, book_hotel]) +``` + +### Load skills on demand + +```python +from pathlib import Path +from azure.ai.agentserver.optimization import load_config, load_skills_from_dir + +config = load_config() + +# Skills are not loaded inline — load them when needed +if config.skills_dir: + skills = load_skills_from_dir(Path(config.skills_dir)) + for skill in skills: + print(f"{skill.name}: {skill.description}") +``` + +## Troubleshooting + +Enable debug logging to see resolution details: + +```python +import logging +logging.getLogger("azure.ai.agentserver.optimization").setLevel(logging.DEBUG) +``` + +Common issues: +- **Config not loading from resolver API** — ensure both env vars are set: `OPTIMIZATION_CANDIDATE_ID` and `OPTIMIZATION_RESOLVE_ENDPOINT`. +- **Local directory not found** — check that `OPTIMIZATION_LOCAL_DIR` points to an existing directory, or ensure `.agent_configs/` exists relative to your main script. +- **`load_config()` raises ValueError** — no config source was found; either set up a baseline folder or pass `required=False`. + +## Next steps + +- [Azure SDK for Python documentation](https://learn.microsoft.com/azure/developer/python/) +- [Contributing guide](https://github.com/Azure/azure-sdk-for-python/blob/main/CONTRIBUTING.md) + +## Contributing + +This project welcomes contributions and suggestions. See [CONTRIBUTING.md](https://github.com/Azure/azure-sdk-for-python/blob/main/CONTRIBUTING.md) for details. diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/azure/__init__.py b/sdk/agentserver/azure-ai-agentserver-optimization/azure/__init__.py new file mode 100644 index 000000000000..8db66d3d0f0f --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/azure/__init__.py @@ -0,0 +1 @@ +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/__init__.py b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/__init__.py new file mode 100644 index 000000000000..8db66d3d0f0f --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/__init__.py @@ -0,0 +1 @@ +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/__init__.py b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/__init__.py new file mode 100644 index 000000000000..8db66d3d0f0f --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/__init__.py @@ -0,0 +1 @@ +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/__init__.py b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/__init__.py new file mode 100644 index 000000000000..4763a8c69568 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/__init__.py @@ -0,0 +1,37 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- + +"""Agent Optimization — Config loader for optimization-ready hosted agents. + +One import, one call:: + + from azure.ai.agentserver.optimization import load_config + + config = load_config() # uses .agent_configs/baseline/ + config = load_config(config_dir="my_configs") # custom directory + config = load_config(required=False) # returns None fields instead of raising + +Resolution order (first match wins): + 1. OPTIMIZATION_CONFIG env var → inline JSON (used by temp agent versions) + 2. OPTIMIZATION_CANDIDATE_ID + ENDPOINT → resolver API → full config + skills + 3. Local directory (config_dir or .agent_configs/) → metadata.yaml + instructions.md + tools.json + skills/ + 4. No config found → raises ValueError (or returns empty config if required=False) +""" + +from azure.ai.agentserver.optimization._config import load_config, load_skills_from_dir +from azure.ai.agentserver.optimization._models import ( + CandidateConfig, + OptimizationConfig, + Skill, +) +from azure.ai.agentserver.optimization._version import VERSION + +__all__ = [ + "CandidateConfig", + "OptimizationConfig", + "Skill", + "load_config", + "load_skills_from_dir", +] +__version__ = VERSION diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_config.py b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_config.py new file mode 100644 index 000000000000..dabc868a9b88 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_config.py @@ -0,0 +1,479 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- + +"""Config loader — resolves optimization config from multiple sources. + +The local directory uses a reserved folder structure:: + + / (default: .agent_configs/) + ├── baseline/ (fallback candidate) + │ ├── metadata.yaml (model, temperature, file pointers) + │ ├── instructions.md (system prompt) + │ ├── tools.json (tool definitions — list format) + │ └── skills/ (learned skills) + │ └── / + │ └── SKILL.md + └── / (same layout as baseline/) + ├── metadata.yaml + ├── instructions.md + ├── tools.json + └── skills/ + └── / + └── SKILL.md + +All folder and file names are defined as constants on +:class:`~OptimizationConfig` (e.g. ``METADATA_FILE``, ``SKILLS_DIR``, +``BASELINE_DIR``). +""" + +from __future__ import annotations + +import json +import logging +import os +from pathlib import Path +from typing import Any + +from azure.ai.agentserver.optimization._models import ( + CandidateConfig, + MetadataConfig, + OptimizationConfig, + Skill, +) +from azure.ai.agentserver.optimization._resolver import resolve_candidate + +logger = logging.getLogger("azure.ai.agentserver.optimization") + + +def load_config( + *, + config_dir: str | Path | None = None, + required: bool = True, +) -> OptimizationConfig | None: + """Load optimization config with graceful fallback. + + Resolution order (first match wins): + + 1. **Inline JSON** — ``OPTIMIZATION_CONFIG`` env var contains the + full config as a JSON string. Used by temporary agent versions + during evaluation; this path is being deprecated. + 2. **Resolver API** — ``OPTIMIZATION_CANDIDATE_ID`` and + ``OPTIMIZATION_RESOLVE_ENDPOINT`` are both set. The endpoint + should be the full job-scoped URL. Fetches the candidate + config from the remote optimization service and persists it + to the local directory. + 3. **Local directory** — reads from + ``//`` (or ``/baseline/`` + as fallback). Defaults to ``.agent_configs/`` relative to the + main script, overridable via ``OPTIMIZATION_LOCAL_DIR`` env var. + 4. When none of the above match: + + - ``required=True`` (default) → raises ``ValueError``. + - ``required=False`` → returns ``None``. + + :keyword config_dir: Path to the agent config directory. When ``None``, + falls back to the ``OPTIMIZATION_LOCAL_DIR`` env var, then + to ``.agent_configs/`` next to the main script. + :paramtype config_dir: str | Path | None + :keyword required: If ``True`` (default), raise ``ValueError`` when no + config source is found. Set to ``False`` during initial + setup or testing. + :paramtype required: bool + :return: The resolved optimization config, or ``None`` when not found + and *required* is ``False``. + :rtype: OptimizationConfig | None + :raises ValueError: When *required* is ``True`` and no config source + (env var, resolver API, or local directory) provides a + valid config. + """ + try: + return _load_config_inner(config_dir=config_dir, required=required) + except ValueError: + raise + except Exception as exc: # noqa: BLE001 # pylint: disable=broad-exception-caught + logger.error("Unexpected error loading optimization config: %s", exc) + return None + + +def _load_config_inner( + *, + config_dir: str | Path | None, + required: bool, +) -> OptimizationConfig | None: + """Internal config loader — may raise on unexpected errors. + + :keyword config_dir: Path to the agent config directory. + :paramtype config_dir: str | Path | None + :keyword required: Whether to raise on missing config. + :paramtype required: bool + :return: Resolved config or ``None``. + :rtype: OptimizationConfig | None + """ + # ── Priority 1: Inline JSON env var (used by temp agent versions, deprecating) ─ + env_var = OptimizationConfig.ENV_CONFIG + raw_config = os.environ.get(env_var, "").strip() + if raw_config: + try: + cfg = json.loads(raw_config) + candidate = CandidateConfig.from_dict(cfg) + logger.warning( + "Loaded optimization config from %s env var (%d chars instructions)", + env_var, + len(candidate.instructions or ""), + ) + return OptimizationConfig( + instructions=candidate.instructions, + model=candidate.model, + temperature=candidate.temperature, + skills=candidate.skills, + tool_definitions=candidate.tool_definitions, + source=f"env:{env_var}", + ) + except (json.JSONDecodeError, TypeError) as exc: + logger.warning("Bad %s env var: %s", env_var, exc) + + # ── Priority 2: Candidate ID → resolver API ────────────────────── + candidate_id = os.environ.get(OptimizationConfig.ENV_CANDIDATE_ID, "").strip() + endpoint = ( + os.environ.get(OptimizationConfig.ENV_RESOLVE_ENDPOINT, "").strip().rstrip("/") + ) + if candidate_id and endpoint: + local_dir = _resolve_local_dir(config_dir) + resolved = resolve_candidate( + candidate_id, endpoint=endpoint, local_dir=local_dir + ) + if resolved is not None: + candidate = CandidateConfig.from_dict(resolved) + logger.warning( + "Loaded optimization config from resolver API for candidate %s", + candidate_id, + ) + return OptimizationConfig( + instructions=candidate.instructions, + model=candidate.model, + temperature=candidate.temperature, + skills=candidate.skills, + skills_dir=resolved.get("skills_dir"), + tool_definitions=candidate.tool_definitions, + source=f"api:candidate:{candidate_id}", + candidate_id=candidate_id, + ) + logger.warning( + "Failed to resolve candidate %s — falling through to local/defaults", + candidate_id, + ) + + # ── Priority 3: Local directory (.agent_configs/) ────────── + local_config = _load_local_dir(candidate_id or None, config_dir) + if local_config is not None: + logger.warning( + "Loaded optimization config from local directory: %s (candidate_id=%s)", + local_config.source, + local_config.candidate_id, + ) + return local_config + + # ── Priority 4: No config found ─────────────────────────────────── + if required: + local_dir = _resolve_local_dir(config_dir) + raise ValueError( + "No optimization config found. Prepare a baseline folder at " + f"'{local_dir / OptimizationConfig.BASELINE_DIR}' with a " + "metadata.yaml file, or pass required=False." + ) + logger.warning("No optimization config found — returning None") + return None + + +def _resolve_local_dir(config_dir: str | Path | None = None) -> Path: + """Resolve the local optimization directory path. + + Priority: *config_dir* argument → ``OPTIMIZATION_LOCAL_DIR`` env + var → ``OptimizationConfig.DEFAULT_LOCAL_DIR`` (``.agent_configs``). + + :param config_dir: Explicit config directory path. + :type config_dir: str | Path | None + :return: Resolved directory path. + :rtype: Path + """ + if config_dir is not None: + local_dir = Path(config_dir) + explicitly_set = True + else: + local_dir_env = os.environ.get(OptimizationConfig.ENV_LOCAL_DIR, "").strip() + explicitly_set = bool(local_dir_env) + local_dir = ( + Path(local_dir_env) + if explicitly_set + else Path(OptimizationConfig.DEFAULT_LOCAL_DIR) + ) + + if not local_dir.is_absolute(): + import sys + + main_mod = sys.modules.get("__main__") + main_file = getattr(main_mod, "__file__", None) if main_mod else None + if main_file is not None: + local_dir = Path(main_file).resolve().parent / local_dir + if explicitly_set and not local_dir.is_dir(): + logger.warning( + "OPTIMIZATION_LOCAL_DIR is set to %r but the directory does not exist", + str(local_dir), + ) + return local_dir + + +def _load_local_dir( + candidate_id: str | None, + config_dir: str | Path | None, +) -> OptimizationConfig | None: + """Load optimization config from a local directory. + + :param candidate_id: Candidate identifier, or ``None`` for baseline. + :type candidate_id: str | None + :param config_dir: Explicit config directory path. + :type config_dir: str | Path | None + :return: Loaded config or ``None`` if directory does not exist. + :rtype: OptimizationConfig | None + """ + local_dir = _resolve_local_dir(config_dir) + if not local_dir.is_dir(): + return None + + candidate_path = _resolve_candidate_folder(local_dir, candidate_id) + if candidate_path is None: + return None + + metadata_file = candidate_path / OptimizationConfig.METADATA_FILE + + return _load_candidate_from_metadata(candidate_path, metadata_file, candidate_id) + + +def _load_candidate_from_metadata( + candidate_path: Path, + metadata_file: Path, + candidate_id: str | None, +) -> OptimizationConfig | None: + """Load candidate config from metadata.yaml + instructions.md layout. + + If ``metadata_file`` does not exist, all default paths + (instructions.md, skills/, tools.json) are used. + + :param candidate_path: Path to the candidate folder. + :type candidate_path: Path + :param metadata_file: Path to the metadata.yaml file. + :type metadata_file: Path + :param candidate_id: Candidate identifier. + :type candidate_id: str | None + :return: Loaded config or ``None``. + :rtype: OptimizationConfig | None + """ + if metadata_file.is_file(): + try: + import yaml # type: ignore[import-untyped] + except ImportError: + raw = _parse_simple_yaml(metadata_file) + else: + try: + raw = yaml.safe_load(metadata_file.read_text(encoding="utf-8")) or {} + except (yaml.YAMLError, OSError) as exc: + logger.warning("Failed to read %s: %s", metadata_file, exc) + raw = {} + else: + raw = {} + + meta = MetadataConfig.from_dict(raw) + + # Read instructions from the referenced file + instructions_path = candidate_path / meta.instruction_file + if instructions_path.is_file(): + instructions: str | None = instructions_path.read_text(encoding="utf-8").strip() + else: + instructions = None + + # Resolve skills directory + skills_dir: str | None + skills_path = candidate_path / meta.skill_dir + if skills_path.resolve().is_dir(): + skills_dir = str(skills_path.resolve()) + else: + skills_dir = None + + # Load tool definitions + tool_file_path = candidate_path / meta.tool_file + tool_definitions = _load_tool_definitions(tool_file_path) + + return OptimizationConfig( + instructions=instructions, + model=meta.model, + temperature=meta.temperature, + skills_dir=skills_dir, + tool_definitions=tool_definitions, + source=f"local:{candidate_path}", + candidate_id=candidate_id, + ) + + +def _load_tool_definitions(tool_file: Path) -> list[dict]: + """Load tool definitions from a tools.json file. + + Expects the OpenAI function-calling list format:: + + [{"type": "function", "function": {"name": "...", "description": "...", "parameters": {...}}}] + + :param tool_file: Path to the tools.json file. + :type tool_file: Path + :return: List of tool definition dicts. + :rtype: list[dict] + """ + if not tool_file.is_file(): + return [] + try: + raw = tool_file.read_text(encoding="utf-8") + data = json.loads(raw) + if isinstance(data, list): + return data + return [] + except (json.JSONDecodeError, OSError) as exc: + logger.warning("Failed to read tools file %s: %s", tool_file, exc) + return [] + + +def _parse_simple_yaml(path: Path) -> dict: + """Minimal key: value parser for metadata.yaml when PyYAML is not installed. + + Coerces numeric-looking values to float/int and recognizes + null/true/false literals. + + :param path: Path to the YAML file. + :type path: Path + :return: Parsed key-value mapping. + :rtype: dict + """ + result: dict = {} + try: + for line in path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line or line.startswith("#"): + continue + if ":" in line: + key, _, value = line.partition(":") + result[key.strip()] = _coerce_yaml_value(value.strip()) + except OSError as exc: + logger.warning("Failed to read %s: %s", path, exc) + return result + + +def _coerce_yaml_value(value: str) -> Any: + """Coerce a YAML scalar string to the appropriate Python type. + + :param value: Raw YAML scalar string. + :type value: str + :return: Coerced Python value. + :rtype: Any + """ + if not value or value in ("null", "~"): + return None + if value.lower() == "true": + return True + if value.lower() == "false": + return False + try: + return int(value) + except ValueError: + pass + try: + return float(value) + except ValueError: + pass + return value + + +def _resolve_candidate_folder(local_dir: Path, candidate_id: str | None) -> Path | None: + """Pick the candidate folder from the local optimization dir. + + Returns ``local_dir/`` if it exists, otherwise falls + back to ``local_dir/baseline/``. Returns ``None`` if neither exists. + + :param local_dir: Root optimization directory. + :type local_dir: Path + :param candidate_id: Candidate identifier. + :type candidate_id: str | None + :return: Resolved candidate folder path, or ``None``. + :rtype: Path | None + """ + if candidate_id: + exact = local_dir / candidate_id + if exact.is_dir(): + return exact + baseline = local_dir / OptimizationConfig.BASELINE_DIR + return baseline if baseline.is_dir() else None + + +def load_skills_from_dir(skills_dir: Path) -> list[Skill]: + """Load skills from a directory of skill folders. + + Expected layout:: + + skills/ + └── / + └── SKILL.md + + :param skills_dir: Path to the skills directory. + :type skills_dir: Path + :return: List of loaded skills. + :rtype: list[Skill] + """ + if not skills_dir.is_dir(): + return [] + + skills: list[Skill] = [] + for skill_folder in sorted(skills_dir.iterdir()): + if not skill_folder.is_dir(): + continue + skill_file = skill_folder / OptimizationConfig.SKILL_FILE + if not skill_file.is_file(): + continue + try: + content = skill_file.read_text(encoding="utf-8").strip() + frontmatter, body = _parse_skill_frontmatter(content) + name = frontmatter.get("name", skill_folder.name) + description = frontmatter.get("description", "") + if not frontmatter and body: + lines = body.split("\n", 1) + description = lines[0].lstrip("#").strip() + body = lines[1].strip() if len(lines) > 1 else "" + skills.append(Skill(name=name, description=description, body=body)) + except OSError as exc: + logger.warning("Failed to read skill %s: %s", skill_file, exc) + + return skills + + +def _parse_skill_frontmatter(content: str) -> tuple[dict, str]: + """Extract YAML frontmatter and body from a SKILL.md file. + + :param content: Raw SKILL.md content. + :type content: str + :return: Tuple of (frontmatter dict, body text). + :rtype: tuple[dict, str] + """ + if not content.startswith("---"): + return {}, content + + end = content.find("---", 3) + if end == -1: + return {}, content + + fm_text = content[3:end].strip() + body = content[end + 3 :].strip() + + frontmatter: dict = {} + for line in fm_text.splitlines(): + line = line.strip() + if ":" in line: + key, _, value = line.partition(":") + frontmatter[key.strip()] = value.strip() + + return frontmatter, body diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_models.py b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_models.py new file mode 100644 index 000000000000..4d2c204d8200 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_models.py @@ -0,0 +1,254 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- + +"""Data models for the optimization config system.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field, fields +from typing import Any, ClassVar + +logger = logging.getLogger(__name__) + + +@dataclass +class Skill: + """A learned skill discovered during optimization. + + Matches the API contract:: + + {"name": "budget-checker", "description": "...", "body": "..."} + """ + + name: str + description: str + body: str = "" + + +@dataclass +class CandidateConfig: + """Typed representation of the candidate config payload from the API. + + This mirrors the wire format produced by the optimization service's + ``to_hosted_agent_config_payload()``:: + + { + "name": "travel", + "instructions": "You are a travel assistant...", + "model": "gpt-4o", + "temperature": 0.7, + "skills": [{"name": "...", "description": "...", "body": "..."}], + "tools": [{"type": "function", "function": {"name": "...", ...}}] + } + """ + + name: str | None = None + instructions: str | None = None + model: str | None = None + temperature: float | None = None + skills: list[Skill] = field(default_factory=list) + tool_definitions: list[dict] = field(default_factory=list) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> CandidateConfig: + """Parse from a raw API response / JSON dict. + + :param data: Raw API response dict. + :type data: dict[str, Any] + :return: Parsed candidate config. + :rtype: CandidateConfig + """ + tools = data.get("tools", []) + return cls( + name=data.get("name"), + instructions=data.get("instructions"), + model=data.get("model"), + temperature=data.get("temperature"), + skills=_parse_skills(data.get("skills", [])), + tool_definitions=tools if isinstance(tools, list) else [], + ) + + +@dataclass +class MetadataConfig: + """Schema for metadata.yaml in the local directory layout. + + Example metadata.yaml:: + + model: gpt-4o + temperature: 0.7 + instruction_file: instructions.md + skill_dir: skills + tool_file: tools.json + """ + + model: str | None = None + temperature: float | None = None + instruction_file: str = "instructions.md" + skill_dir: str = "skills" + tool_file: str = "tools.json" + + @classmethod + def from_dict(cls, data: dict) -> MetadataConfig: + """Create from a parsed YAML dict, ignoring unknown keys. + + :param data: Parsed YAML dict. + :type data: dict + :return: Metadata config. + :rtype: MetadataConfig + """ + known = {f.name for f in fields(cls)} + filtered = {k: v for k, v in data.items() if k in known} + return cls(**filtered) + + +@dataclass +class OptimizationConfig: # pylint: disable=too-many-instance-attributes + """Resolved optimization config. + + When not running under optimization, fields are ``None`` unless a + local config directory (baseline) supplies values. + """ + + ENV_CANDIDATE_ID: ClassVar[str] = "OPTIMIZATION_CANDIDATE_ID" + ENV_CONFIG: ClassVar[str] = "OPTIMIZATION_CONFIG" + ENV_LOCAL_DIR: ClassVar[str] = "OPTIMIZATION_LOCAL_DIR" + ENV_RESOLVE_ENDPOINT: ClassVar[str] = "OPTIMIZATION_RESOLVE_ENDPOINT" + DEFAULT_LOCAL_DIR: ClassVar[str] = ".agent_configs" + + METADATA_FILE: ClassVar[str] = "metadata.yaml" + INSTRUCTIONS_FILE: ClassVar[str] = "instructions.md" + TOOLS_FILE: ClassVar[str] = "tools.json" + SKILLS_DIR: ClassVar[str] = "skills" + SKILL_FILE: ClassVar[str] = "SKILL.md" + BASELINE_DIR: ClassVar[str] = "baseline" + + instructions: str | None = None + model: str | None = None + temperature: float | None = None + skills: list[Skill] = field(default_factory=list) + skills_dir: str | None = None + tool_definitions: list[dict] = field(default_factory=list) + source: str = "defaults" + candidate_id: str | None = None + + @property + def has_skills(self) -> bool: + return len(self.skills) > 0 or self.skills_dir is not None + + def apply_tool_descriptions(self, tools: list) -> list: + """Apply optimized tool definitions to a list of tool functions. + + For each tool function whose name matches a definition in + :attr:`tool_definitions`, patches ``__doc__`` and ``.description`` + with the optimized description, and patches parameter descriptions + on the ``input_model`` if present. + + :param tools: List of @tool-decorated functions. + :type tools: list + :return: The same list of tools (mutated in place). + :rtype: list + """ + if not self.tool_definitions: + return tools + # Build name → function-definition lookup + lookup: dict[str, dict] = {} + for item in self.tool_definitions: + if not isinstance(item, dict): + continue + func = item.get("function", {}) + if isinstance(func, dict) and func.get("name"): + lookup[func["name"]] = func + if not lookup: + return tools + for tool_fn in tools: + tool_name = getattr(tool_fn, "__name__", None) or getattr( + tool_fn, "name", None + ) + if not tool_name or tool_name not in lookup: + continue + func_def = lookup[tool_name] + description = func_def.get("description", "") + if description: + try: + tool_fn.description = description + except AttributeError: + pass + tool_fn.__doc__ = description + logger.debug("Applied optimized description for tool '%s'", tool_name) + _patch_input_model_params(tool_fn, func_def, tool_name) + return tools + + def compose_instructions(self) -> str: + """Return instructions with skill catalog appended (if any). + + :return: Instructions text with skills appended. + :rtype: str + """ + base = self.instructions or "" + if not self.skills: + return base + + lines = [base, "", "## Available Skills"] if base else ["## Available Skills"] + for s in self.skills: + lines.append(f"- **{s.name}**: {s.description}") + return "\n".join(lines) + + +# ── Parsing / patching helpers ──────────────────────────────────────── + + +def _patch_input_model_params(tool_fn: Any, func_def: dict, tool_name: str) -> None: + """Patch parameter descriptions on a tool's input_model. + + :param tool_fn: The tool function to patch. + :type tool_fn: Any + :param func_def: Function definition dict with parameters schema. + :type func_def: dict + :param tool_name: Tool name for logging. + :type tool_name: str + """ + params_schema = func_def.get("parameters", {}) + if not isinstance(params_schema, dict): + return + props = params_schema.get("properties", {}) + if not isinstance(props, dict): + return + input_model = getattr(tool_fn, "input_model", None) + if not input_model or not hasattr(input_model, "model_fields"): + return + patched = False + for param_name, param_val in props.items(): + if not isinstance(param_val, dict) or "description" not in param_val: + continue + if param_name in input_model.model_fields: + input_model.model_fields[param_name].description = param_val["description"] + patched = True + if patched: + input_model.model_rebuild(force=True) + logger.debug( + "Applied optimized parameter descriptions for tool '%s'", tool_name + ) + + +def _parse_skills(raw: list) -> list[Skill]: + """Parse skills from API/env config JSON. + + :param raw: Raw skills list from API response. + :type raw: list + :return: Parsed list of Skill objects. + :rtype: list[Skill] + """ + skills: list[Skill] = [] + for item in raw: + if isinstance(item, dict) and item.get("name"): + skills.append( + Skill( + name=item["name"], + description=item.get("description", ""), + body=item.get("body", ""), + ) + ) + return skills diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_resolver.py b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_resolver.py new file mode 100644 index 000000000000..09cf08dfc4ba --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_resolver.py @@ -0,0 +1,357 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- + +"""Candidate config resolution via the optimization service API. + +Fetches candidate config and skill files from the remote optimization +service and persists them into the standard local directory layout:: + + / + └── / + ├── metadata.yaml + ├── instructions.md + ├── tools.json + └── skills/ + └── / + └── SKILL.md +""" + +from __future__ import annotations + +import json +import logging +import pathlib +import shutil +from typing import Any + +from azure.core import PipelineClient +from azure.core.pipeline.policies import BearerTokenCredentialPolicy, RetryPolicy +from azure.core.rest import HttpRequest + +from azure.ai.agentserver.optimization._models import OptimizationConfig + +logger = logging.getLogger("azure.ai.agentserver.optimization") + +_downloaded: set[str] = set() + +# API path and version constants +_API_VERSION = "2025-11-15-preview" +_AUTH_SCOPE = "https://ai.azure.com/.default" + + +def resolve_candidate( + candidate_id: str, + endpoint: str, + local_dir: pathlib.Path | None = None, +) -> dict[str, Any] | None: + """Resolve a candidate's full config from the optimization service. + + ``endpoint`` should be the full job-scoped URL, + The resolver appends ``/candidates/{candidate_id}/config``. + + Downloads config and skills into ``local_dir//`` + following the standard local directory layout. + Returns ``None`` if the call fails. + + :param candidate_id: Candidate identifier. + :type candidate_id: str + :param endpoint: Full job-scoped endpoint URL. + :type endpoint: str + :param local_dir: Local directory for persisting config. + :type local_dir: pathlib.Path | None + :return: Candidate config dict, or ``None`` on failure. + :rtype: dict[str, Any] | None + """ + if candidate_id in _downloaded: + if local_dir is not None and (local_dir / candidate_id).is_dir(): + logger.warning("Candidate %s already downloaded — skipping", candidate_id) + return None + logger.warning( + "Candidate %s was downloaded but folder is missing — re-downloading", + candidate_id, + ) + _downloaded.discard(candidate_id) + + client = _build_client(endpoint) + + # ── Step 1: Fetch config ───────────────────────────────────────── + config = _api_get_json( + client, + f"/candidates/{candidate_id}/config", + params={"api-version": _API_VERSION}, + ) + if config is None: + client.close() + return None + + logger.info( + "Resolved candidate %s: model=%s, instructions=%d chars, skills=%d, tool_definitions=%d", + candidate_id, + config.get("model", "?"), + len(config.get("instructions", "")), + len(config.get("skills", [])), + len(config.get("tools", [])), + ) + + # ── Step 2: Persist to local directory layout ──────────────────── + if local_dir is not None: + candidate_path = local_dir / candidate_id + try: + _persist_to_local_layout(candidate_path, config) + _download_skill_files(client, candidate_id, candidate_path) + except OSError as exc: + logger.warning( + "Failed to persist candidate %s to disk: %s", candidate_id, exc + ) + # Point skills_dir to the downloaded skills folder + skills_path = candidate_path / OptimizationConfig.SKILLS_DIR + if skills_path.is_dir(): + config["skills_dir"] = str(skills_path) + + client.close() + _downloaded.add(candidate_id) + return config + + +def _persist_to_local_layout( + candidate_path: pathlib.Path, config: dict[str, Any] +) -> None: + """Write config into the standard local directory layout. + + Produces the same structure that ``_load_local_dir`` reads:: + + / + ├── metadata.yaml + ├── instructions.md + ├── tools.json + └── skills/ + └── / + └── SKILL.md + + If the folder already exists it is removed and re-created. + + :param candidate_path: Target directory for the candidate layout. + :type candidate_path: pathlib.Path + :param config: Candidate config dict from the API. + :type config: dict[str, Any] + """ + if candidate_path.is_dir(): + logger.info("Overwriting existing candidate folder: %s", candidate_path) + shutil.rmtree(candidate_path) + + candidate_path.mkdir(parents=True, exist_ok=True) + + # metadata.yaml + meta_lines: list[str] = [] + if config.get("model"): + meta_lines.append(f"model: {config['model']}") + if config.get("temperature") is not None: + meta_lines.append(f"temperature: {config['temperature']}") + meta_lines.append(f"instruction_file: {OptimizationConfig.INSTRUCTIONS_FILE}") + meta_lines.append(f"skill_dir: {OptimizationConfig.SKILLS_DIR}") + meta_lines.append(f"tool_file: {OptimizationConfig.TOOLS_FILE}") + meta_file = candidate_path / OptimizationConfig.METADATA_FILE + meta_file.write_text("\n".join(meta_lines) + "\n", encoding="utf-8") + + # instructions.md + instructions = config.get("instructions", "") + if instructions: + instr_file = candidate_path / OptimizationConfig.INSTRUCTIONS_FILE + instr_file.write_text(instructions, encoding="utf-8") + + # tools.json — write tool definitions in list format + tools_list = config.get("tools") + if tools_list and isinstance(tools_list, list): + tools_file = candidate_path / OptimizationConfig.TOOLS_FILE + tools_file.write_text( + json.dumps(tools_list, indent=2, ensure_ascii=False), encoding="utf-8" + ) + + # skills/ — write inline skills as //SKILL.md + inline_skills = config.get("skills", []) + if inline_skills and isinstance(inline_skills, list): + skills_dir = candidate_path / OptimizationConfig.SKILLS_DIR + for skill in inline_skills: + if not isinstance(skill, dict) or not skill.get("name"): + continue + skill_name = skill["name"] + skill_folder = skills_dir / skill_name + skill_folder.mkdir(parents=True, exist_ok=True) + # Build SKILL.md with YAML frontmatter + lines: list[str] = ["---"] + lines.append(f"name: {skill_name}") + if skill.get("description"): + lines.append(f"description: {skill['description']}") + lines.append("---") + if skill.get("body"): + lines.append("") + lines.append(skill["body"]) + skill_file = skill_folder / OptimizationConfig.SKILL_FILE + skill_file.write_text("\n".join(lines) + "\n", encoding="utf-8") + logger.info( + "Persisted %d inline skill(s) to %s", len(inline_skills), skills_dir + ) + + logger.info("Persisted config to local layout: %s", candidate_path) + + +def _download_skill_files( + client: PipelineClient, + candidate_id: str, + candidate_path: pathlib.Path, +) -> None: + """Fetch manifest and download skill files into candidate_path/skills//SKILL.md. + + :param client: Azure PipelineClient for API calls. + :type client: PipelineClient + :param candidate_id: Candidate identifier. + :type candidate_id: str + :param candidate_path: Local directory for the candidate. + :type candidate_path: pathlib.Path + """ + manifest = _api_get_json( + client, + f"/candidates/{candidate_id}", + params={"api-version": _API_VERSION}, + ) + if manifest is None: + logger.debug("Could not fetch manifest for candidate %s", candidate_id) + return + + files = manifest.get("files", []) + skill_files = [f for f in files if _is_skill_file(f)] + if not skill_files: + logger.debug("No skill files in manifest for candidate %s", candidate_id) + return + + logger.info( + "Downloading %d skill file(s) for candidate %s", + len(skill_files), + candidate_id, + ) + + skills_dir = candidate_path / OptimizationConfig.SKILLS_DIR + for file_entry in skill_files: + file_path = file_entry.get("path", "") + if not file_path: + continue + + content = _api_get_text( + client, + f"/candidates/{candidate_id}/files", + params={"path": file_path, "api-version": _API_VERSION}, + ) + if content is None: + logger.warning("Failed to download skill file: %s", file_path) + continue + + # file_path is like "skills/math/SKILL.md" → write to skills_dir/math/SKILL.md + rel_path = file_path + prefix = OptimizationConfig.SKILLS_DIR + "/" + if rel_path.startswith(prefix): + rel_path = rel_path[len(prefix) :] + + out_path = (skills_dir / rel_path).resolve() + if not str(out_path).startswith(str(skills_dir.resolve())): + logger.warning( + "Path traversal detected in skill file path: %r — skipping", file_path + ) + continue + + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(content, encoding="utf-8") + logger.info(" → %s (%d bytes)", out_path, len(content)) + + +def _is_skill_file(file_entry: dict) -> bool: + """Check if a manifest entry is a skill file. + + :param file_entry: Manifest entry dict. + :type file_entry: dict + :return: ``True`` if the entry represents a skill file. + :rtype: bool + """ + path = file_entry.get("path", "") + file_type = file_entry.get("type", "") + return file_type == "skill" or path.startswith("skills/") + + +# ── HTTP helpers (azure.core transport) ────────────────────────────── + + +def _build_client(endpoint: str) -> PipelineClient: + """Create a PipelineClient with credential-based auth and retry. + + :param endpoint: Base URL for the API. + :type endpoint: str + :return: Configured pipeline client. + :rtype: PipelineClient + """ + policies: list = [RetryPolicy()] + try: + from azure.identity import DefaultAzureCredential + + credential = DefaultAzureCredential() + policies.insert(0, BearerTokenCredentialPolicy(credential, _AUTH_SCOPE)) + except Exception: # noqa: BLE001 # pylint: disable=broad-exception-caught + logger.debug( + "azure-identity not available or credentials failed — proceeding without auth" + ) + return PipelineClient(base_url=endpoint, policies=policies) + + +def _api_get_json( + client: PipelineClient, path: str, params: dict[str, str] | None = None +) -> dict[str, Any] | None: + """GET a JSON endpoint, return parsed dict or None on failure. + + :param client: Azure PipelineClient. + :type client: PipelineClient + :param path: API path to append to the base URL. + :type path: str + :param params: Query parameters. + :type params: dict[str, str] | None + :return: Parsed response dict or ``None``. + :rtype: dict[str, Any] | None + """ + url = f"{client._base_url.rstrip('/')}{path}" # pylint: disable=protected-access + request = HttpRequest("GET", url, params=params) + logger.debug("GET %s", url) + try: + response = client.send_request(request) + if response.status_code != 200: + logger.error("GET %s returned %d", url, response.status_code) + return None + return response.json() + except Exception as exc: # noqa: BLE001 # pylint: disable=broad-exception-caught + logger.error("GET %s failed: %s", url, exc) + return None + + +def _api_get_text( + client: PipelineClient, path: str, params: dict[str, str] | None = None +) -> str | None: + """GET an endpoint, return response body as text or None on failure. + + :param client: Azure PipelineClient. + :type client: PipelineClient + :param path: API path to append to the base URL. + :type path: str + :param params: Query parameters. + :type params: dict[str, str] | None + :return: Response body text or ``None``. + :rtype: str | None + """ + url = f"{client._base_url.rstrip('/')}{path}" # pylint: disable=protected-access + request = HttpRequest("GET", url, params=params) + logger.debug("GET %s", url) + try: + response = client.send_request(request) + if response.status_code != 200: + logger.error("GET %s returned %d", url, response.status_code) + return None + return response.text() + except Exception as exc: # noqa: BLE001 # pylint: disable=broad-exception-caught + logger.error("GET %s failed: %s", url, exc) + return None diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_version.py b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_version.py new file mode 100644 index 000000000000..67d209a8cafd --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_version.py @@ -0,0 +1,5 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- + +VERSION = "1.0.0b1" diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/py.typed b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/py.typed new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/cspell.json b/sdk/agentserver/azure-ai-agentserver-optimization/cspell.json new file mode 100644 index 000000000000..b3814147ea7e --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/cspell.json @@ -0,0 +1,16 @@ +{ + "ignoreWords": [ + "agentserver", + "autouse", + "cand", + "dataclass", + "frontmatter", + "paramtype", + "pathlib", + "rtype" + ], + "ignorePaths": [ + "*.json", + "*.rst" + ] +} diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/dev_requirements.txt b/sdk/agentserver/azure-ai-agentserver-optimization/dev_requirements.txt new file mode 100644 index 000000000000..c29511d76ecb --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/dev_requirements.txt @@ -0,0 +1,6 @@ +-e ../../../eng/tools/azure-sdk-tools +pytest +pytest-mock +pyyaml>=6.0 +azure-core>=1.31.0 +azure-identity>=1.15.0 diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/pyproject.toml b/sdk/agentserver/azure-ai-agentserver-optimization/pyproject.toml new file mode 100644 index 000000000000..13e7aaec3ca5 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/pyproject.toml @@ -0,0 +1,62 @@ +[project] +name = "azure-ai-agentserver-optimization" +dynamic = ["version", "readme"] +description = "Optimization config loader for Azure AI Hosted Agents" +requires-python = ">=3.10" +authors = [ + { name = "Microsoft Corporation", email = "azpysdkhelp@microsoft.com" }, +] +license = "MIT" +classifiers = [ + "Development Status :: 4 - Beta", + "Programming Language :: Python", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", +] +keywords = ["azure", "azure sdk", "agent", "agentserver", "optimization"] + +dependencies = [ + "pyyaml>=6.0", + "azure-core>=1.31.0", + "azure-identity>=1.15.0", +] + +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project.urls] +repository = "https://github.com/Azure/azure-sdk-for-python" + +[tool.setuptools.packages.find] +exclude = [ + "tests*", + "samples*", + "doc*", + "azure", + "azure.ai", + "azure.ai.agentserver", +] + +[tool.setuptools.dynamic] +version = { attr = "azure.ai.agentserver.optimization._version.VERSION" } +readme = { file = ["README.md"], content-type = "text/markdown" } + +[tool.setuptools.package-data] +"azure.ai.agentserver.optimization" = ["py.typed"] + +[tool.azure-sdk-build] +breaking = false +mypy = true +pyright = true +verifytypes = false +latestdependency = false +pylint = true +type_check_samples = false + +[tool.uv.sources] diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/tests/conftest.py b/sdk/agentserver/azure-ai-agentserver-optimization/tests/conftest.py new file mode 100644 index 000000000000..e806f9eaf652 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/conftest.py @@ -0,0 +1,23 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Shared fixtures for azure-ai-agentserver-optimization tests.""" + +import pytest + + +ENV_VARS = [ + "OPTIMIZATION_CONFIG", + "OPTIMIZATION_CANDIDATE_ID", + "OPTIMIZATION_LOCAL_DIR", + "OPTIMIZATION_RESOLVE_ENDPOINT", + "MODEL_DEPLOYMENT_NAME", +] + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch): + """Ensure optimization env vars are cleared before each test.""" + for var in ENV_VARS: + monkeypatch.delenv(var, raising=False) + yield diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py new file mode 100644 index 000000000000..d212eafb82b3 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py @@ -0,0 +1,1362 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for load_config — priority resolution, fallback, and edge cases.""" + +import json + +import pytest + +from azure.ai.agentserver.optimization import ( + CandidateConfig, + OptimizationConfig, + Skill, + load_config, +) +from azure.ai.agentserver.optimization._models import ( + MetadataConfig, + _parse_skills, +) +from azure.ai.agentserver.optimization._config import ( + _load_tool_definitions, + _parse_simple_yaml, + _parse_skill_frontmatter, + _resolve_candidate_folder, + _resolve_local_dir, +) +from azure.ai.agentserver.optimization._resolver import _downloaded + + +@pytest.fixture(autouse=True) +def clear_downloaded(): + _downloaded.clear() + yield + _downloaded.clear() + + +# ── Defaults (Priority 4) ─────────────────────────────────────────── + + +class TestDefaults: + """When no env vars or config dir are set.""" + + def test_required_true_raises_by_default(self): + with pytest.raises(ValueError, match="No optimization config found"): + load_config() + + def test_required_false_returns_none(self): + config = load_config(required=False) + assert config is None + + def test_falls_back_to_model_deployment_name_env(self, monkeypatch): + """MODEL_DEPLOYMENT_NAME is only used by priority 1/2, not by required=False.""" + monkeypatch.setenv("MODEL_DEPLOYMENT_NAME", "gpt-4o-mini") + config = load_config(required=False) + assert config is None + + def test_config_dir_loads_baseline(self, monkeypatch, tmp_path): + """config_dir parameter points to a custom agent config folder.""" + baseline = tmp_path / "baseline" + baseline.mkdir() + (baseline / "instructions.md").write_text("Custom dir prompt.") + (baseline / "metadata.yaml").write_text("model: gpt-4o\n") + config = load_config(config_dir=tmp_path) + assert config.instructions == "Custom dir prompt." + assert config.model == "gpt-4o" + assert config.source.startswith("local:") + + +# ── Inline JSON env var (Priority 1) ──────────────────────────────── + + +class TestEnvConfig: + """OPTIMIZATION_CONFIG env var overrides everything.""" + + def test_loads_from_env_config(self, monkeypatch): + payload = { + "instructions": "Optimized prompt.", + "model": "gpt-4o", + "temperature": 0.3, + } + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) + config = load_config() + assert config.instructions == "Optimized prompt." + assert config.model == "gpt-4o" + assert config.temperature == 0.3 + assert config.source == "env:OPTIMIZATION_CONFIG" + + def test_env_config_with_skills(self, monkeypatch): + payload = { + "instructions": "With skills.", + "skills": [ + {"name": "math", "description": "Math skill", "body": "do math"}, + {"name": "code", "description": "Code skill"}, + ], + } + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) + config = load_config() + assert len(config.skills) == 2 + assert config.skills[0].name == "math" + assert config.skills[0].body == "do math" + assert config.skills[1].name == "code" + assert config.skills[1].body == "" + assert config.has_skills + + def test_env_config_with_tool_descriptions(self, monkeypatch): + payload = { + "instructions": "With tools.", + "tools": [ + { + "type": "function", + "function": { + "name": "lookup_travel_policy", + "description": "Look up the company travel policy.", + "parameters": {"type": "object", "properties": {}}, + }, + }, + { + "type": "function", + "function": { + "name": "check_department_budget", + "description": "Check remaining budget.", + "parameters": { + "type": "object", + "properties": { + "dept": {"type": "string", "description": "Department name"} + }, + }, + }, + }, + ], + } + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) + config = load_config() + assert len(config.tool_definitions) == 2 + assert config.tool_definitions[0]["function"]["name"] == "lookup_travel_policy" + + + def test_env_config_with_tools_list(self, monkeypatch): + """OpenAI function-calling list format is supported.""" + payload = { + "instructions": "With tools list.", + "tools": [ + { + "type": "function", + "function": { + "name": "lookup_policy", + "description": "Look up policy", + "parameters": { + "type": "object", + "properties": { + "dept": {"type": "string", "description": "Department name"} + }, + }, + }, + } + ], + } + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) + config = load_config() + assert len(config.tool_definitions) == 1 + assert config.tool_definitions[0]["function"]["name"] == "lookup_policy" + + + def test_bad_json_falls_through(self, monkeypatch): + monkeypatch.setenv("OPTIMIZATION_CONFIG", "not-json{{{") + config = load_config(required=False) + assert config is None + + def test_empty_env_var_ignored(self, monkeypatch): + monkeypatch.setenv("OPTIMIZATION_CONFIG", " ") + config = load_config(required=False) + assert config is None + + def test_partial_config_fills_none(self, monkeypatch): + payload = {"model": "gpt-4o"} + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) + config = load_config() + assert config.instructions is None + assert config.model == "gpt-4o" + assert config.temperature is None + + def test_env_config_takes_priority_over_candidate_id(self, monkeypatch): + payload = {"instructions": "From env."} + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "some-candidate") + monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") + config = load_config() + assert config.source == "env:OPTIMIZATION_CONFIG" + + +# ── Candidate ID / Resolver (Priority 2) ──────────────────────────── + + +class TestCandidateResolver: + """OPTIMIZATION_CANDIDATE_ID + ENDPOINT triggers resolver API.""" + + def test_candidate_id_calls_resolver(self, monkeypatch): + resolved = { + "instructions": "Resolved prompt.", + "model": "gpt-4o", + "temperature": 0.2, + "skills": [{"name": "s1", "description": "d1"}], + } + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-123") + monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") + monkeypatch.setattr( + "azure.ai.agentserver.optimization._config.resolve_candidate", + lambda cid, endpoint, local_dir=None: resolved, + ) + config = load_config() + assert config.source == "api:candidate:cand-123" + assert config.instructions == "Resolved prompt." + assert config.candidate_id == "cand-123" + assert len(config.skills) == 1 + + def test_resolver_failure_falls_to_defaults(self, monkeypatch): + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "bad-id") + monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") + monkeypatch.setattr( + "azure.ai.agentserver.optimization._config.resolve_candidate", + lambda cid, endpoint, local_dir=None: None, + ) + config = load_config(required=False) + assert config is None + + def test_missing_endpoint_skips_resolver(self, monkeypatch): + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-1") + # No ENDPOINT set + config = load_config(required=False) + assert config is None + + def test_resolver_falls_to_local_dir(self, monkeypatch, tmp_path): + """When resolver returns None, falls to local dir (priority 3).""" + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-local") + monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + monkeypatch.setattr( + "azure.ai.agentserver.optimization._config.resolve_candidate", + lambda cid, endpoint, local_dir=None: None, + ) + # Set up local dir with this candidate + candidate_dir = tmp_path / "cand-local" + candidate_dir.mkdir() + (candidate_dir / "metadata.yaml").write_text("model: local-model\n") + (candidate_dir / "instructions.md").write_text("Local instructions.") + + config = load_config() + assert config.source.startswith("local:") + assert config.instructions == "Local instructions." + assert config.model == "local-model" + + +# ── Local directory (Priority 3) ──────────────────────────────────── + + +class TestLocalDir: + """OPTIMIZATION_LOCAL_DIR triggers local directory loading.""" + + def test_loads_from_baseline(self, monkeypatch, tmp_path): + candidate_dir = tmp_path / "baseline" + candidate_dir.mkdir() + (candidate_dir / "metadata.yaml").write_text( + "model: gpt-4o\ntemperature: 0.4\n" + ) + (candidate_dir / "instructions.md").write_text("Baseline instructions.") + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert config.instructions == "Baseline instructions." + assert config.model == "gpt-4o" + assert config.temperature == 0.4 + assert config.source.startswith("local:") + + def test_candidate_id_folder_takes_priority_over_baseline(self, monkeypatch, tmp_path): + baseline = tmp_path / "baseline" + baseline.mkdir() + (baseline / "metadata.yaml").write_text("model: baseline\n") + (baseline / "instructions.md").write_text("Baseline.") + + candidate = tmp_path / "cand-123" + candidate.mkdir() + (candidate / "metadata.yaml").write_text("model: candidate\n") + (candidate / "instructions.md").write_text("Candidate.") + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-123") + config = load_config() + assert config.model == "candidate" + assert config.instructions == "Candidate." + + def test_falls_to_baseline_when_candidate_folder_missing(self, monkeypatch, tmp_path): + baseline = tmp_path / "baseline" + baseline.mkdir() + (baseline / "metadata.yaml").write_text("model: baseline\n") + (baseline / "instructions.md").write_text("Baseline.") + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "nonexistent-id") + config = load_config() + assert config.model == "baseline" + assert config.instructions == "Baseline." + + def test_loads_without_metadata_yaml(self, monkeypatch, tmp_path): + """Without metadata.yaml, uses default file paths.""" + candidate_dir = tmp_path / "baseline" + candidate_dir.mkdir() + # No metadata.yaml — should use defaults (instructions.md, skills/, tools.json) + (candidate_dir / "instructions.md").write_text("No metadata instructions.") + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert config.instructions == "No metadata instructions." + + def test_loads_skills_from_local_dir(self, monkeypatch, tmp_path): + candidate_dir = tmp_path / "baseline" + candidate_dir.mkdir() + (candidate_dir / "metadata.yaml").write_text("skill_dir: skills\n") + (candidate_dir / "instructions.md").write_text("With skills.") + + skills_dir = candidate_dir / "skills" / "math" + skills_dir.mkdir(parents=True) + (skills_dir / "SKILL.md").write_text( + "---\nname: math\ndescription: Do math\n---\nBody here." + ) + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert config.skills_dir is not None + assert config.skills == [] + + from pathlib import Path + from azure.ai.agentserver.optimization import load_skills_from_dir + skills = load_skills_from_dir(Path(config.skills_dir)) + assert len(skills) == 1 + assert skills[0].name == "math" + assert skills[0].description == "Do math" + assert skills[0].body == "Body here." + + def test_loads_tools_dict_from_local_dir(self, monkeypatch, tmp_path): + candidate_dir = tmp_path / "baseline" + candidate_dir.mkdir() + (candidate_dir / "metadata.yaml").write_text("tool_file: tools.json\n") + (candidate_dir / "instructions.md").write_text("With tools.") + tools = [ + {"type": "function", "function": {"name": "search", "description": "Search stuff", "parameters": {"type": "object", "properties": {"q": {"type": "string", "description": "query"}}}}} + ] + (candidate_dir / "tools.json").write_text(json.dumps(tools)) + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert len(config.tool_definitions) == 1 + assert config.tool_definitions[0]["function"]["name"] == "search" + + def test_loads_tools_list_from_local_dir(self, monkeypatch, tmp_path): + """OpenAI function-calling list format in tools.json.""" + candidate_dir = tmp_path / "baseline" + candidate_dir.mkdir() + (candidate_dir / "instructions.md").write_text("With tools list.") + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + }, + }, + }, + } + ] + (candidate_dir / "tools.json").write_text(json.dumps(tools)) + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert len(config.tool_definitions) == 1 + assert config.tool_definitions[0]["function"]["name"] == "get_weather" + + def test_missing_instructions_returns_none(self, monkeypatch, tmp_path): + candidate_dir = tmp_path / "baseline" + candidate_dir.mkdir() + (candidate_dir / "metadata.yaml").write_text("model: gpt-4o\n") + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert config.instructions is None + assert config.model == "gpt-4o" + + def test_nonexistent_local_dir_falls_to_defaults(self, monkeypatch): + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", "/nonexistent/path") + config = load_config(required=False) + assert config is None + + def test_no_candidate_no_baseline_falls_to_defaults(self, monkeypatch, tmp_path): + """Empty local dir with no baseline falls through.""" + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config(required=False) + assert config is None + + def test_metadata_instruction_file_outside_candidate(self, monkeypatch, tmp_path): + """instruction_file can point outside the candidate folder.""" + candidate_dir = tmp_path / "baseline" + candidate_dir.mkdir() + shared = tmp_path / "shared_instructions.md" + shared.write_text("Shared prompt.") + (candidate_dir / "metadata.yaml").write_text( + "instruction_file: ../shared_instructions.md\n" + ) + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert config.instructions == "Shared prompt." + + def test_metadata_skill_dir_outside_candidate(self, monkeypatch, tmp_path): + """skill_dir can point outside the candidate folder.""" + candidate_dir = tmp_path / "baseline" + candidate_dir.mkdir() + (candidate_dir / "instructions.md").write_text("ok") + shared_skills = tmp_path / "shared_skills" / "math" + shared_skills.mkdir(parents=True) + (shared_skills / "SKILL.md").write_text( + "---\nname: math\ndescription: Do math\n---\nBody." + ) + (candidate_dir / "metadata.yaml").write_text( + "skill_dir: ../shared_skills\n" + ) + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert config.skills_dir is not None + + def test_metadata_tool_file_outside_candidate(self, monkeypatch, tmp_path): + """tool_file can point outside the candidate folder.""" + candidate_dir = tmp_path / "baseline" + candidate_dir.mkdir() + (candidate_dir / "instructions.md").write_text("ok") + shared_tools = tmp_path / "shared_tools.json" + shared_tools.write_text('[{"type": "function", "function": {"name": "search", "description": "Search"}}]') + (candidate_dir / "metadata.yaml").write_text( + "tool_file: ../shared_tools.json\n" + ) + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert len(config.tool_definitions) == 1 + assert config.tool_definitions[0]["function"]["name"] == "search" + + +# ── _resolve_local_dir ────────────────────────────────────────────── + + +class TestResolveLocalDir: + """Tests for _resolve_local_dir.""" + + def test_defaults_to_agent_configs(self, monkeypatch): + monkeypatch.delenv("OPTIMIZATION_LOCAL_DIR", raising=False) + local_dir = _resolve_local_dir() + assert local_dir.name == ".agent_configs" + + def test_uses_env_var(self, monkeypatch, tmp_path): + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + local_dir = _resolve_local_dir() + assert local_dir == tmp_path + + def test_config_dir_param_takes_priority_over_env(self, monkeypatch, tmp_path): + """config_dir argument wins over OPTIMIZATION_LOCAL_DIR env var.""" + env_dir = tmp_path / "env_dir" + env_dir.mkdir() + param_dir = tmp_path / "param_dir" + param_dir.mkdir() + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(env_dir)) + local_dir = _resolve_local_dir(str(param_dir)) + assert local_dir == param_dir + + def test_config_dir_param_as_path_object(self, tmp_path): + """config_dir can be a Path object.""" + from pathlib import Path + local_dir = _resolve_local_dir(Path(tmp_path)) + assert local_dir == tmp_path + + def test_env_var_whitespace_stripped(self, monkeypatch, tmp_path): + """Leading/trailing whitespace in OPTIMIZATION_LOCAL_DIR is stripped.""" + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", f" {tmp_path} ") + local_dir = _resolve_local_dir() + assert local_dir == tmp_path + + def test_empty_env_var_uses_default(self, monkeypatch): + """Empty OPTIMIZATION_LOCAL_DIR falls back to default.""" + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", " ") + local_dir = _resolve_local_dir() + assert local_dir.name == ".agent_configs" + + +# ── _resolve_candidate_folder ─────────────────────────────────────── + + +class TestResolveCandidateFolder: + """Tests for _resolve_candidate_folder.""" + + def test_exact_candidate_found(self, tmp_path): + (tmp_path / "cand-1").mkdir() + result = _resolve_candidate_folder(tmp_path, "cand-1") + assert result == tmp_path / "cand-1" + + def test_falls_to_baseline(self, tmp_path): + (tmp_path / "baseline").mkdir() + result = _resolve_candidate_folder(tmp_path, "nonexistent") + assert result == tmp_path / "baseline" + + def test_no_candidate_id_uses_baseline(self, tmp_path): + (tmp_path / "baseline").mkdir() + result = _resolve_candidate_folder(tmp_path, None) + assert result == tmp_path / "baseline" + + def test_returns_none_when_nothing_exists(self, tmp_path): + result = _resolve_candidate_folder(tmp_path, "nonexistent") + assert result is None + + def test_returns_none_no_id_no_baseline(self, tmp_path): + result = _resolve_candidate_folder(tmp_path, None) + assert result is None + + def test_empty_string_candidate_id_uses_baseline(self, tmp_path): + """Empty string candidate_id is falsy, falls to baseline.""" + (tmp_path / "baseline").mkdir() + result = _resolve_candidate_folder(tmp_path, "") + assert result == tmp_path / "baseline" + + +# ── Graceful error handling ───────────────────────────────────────── + + +class TestGracefulErrorHandling: + """load_config never crashes — always returns a valid config (unless required ValueError).""" + + def test_unexpected_exception_returns_none(self, monkeypatch): + """Any unexpected error in _load_config_inner returns None.""" + monkeypatch.setattr( + "azure.ai.agentserver.optimization._config._load_config_inner", + lambda **kwargs: (_ for _ in ()).throw(RuntimeError("boom")), + ) + config = load_config(required=False) + assert config is None + + def test_load_config_never_raises_on_corrupt_env(self, monkeypatch): + """Even with corrupted env vars, load_config(required=False) returns None.""" + monkeypatch.setenv("OPTIMIZATION_CONFIG", "{invalid") + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "x") + monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://nope") + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", "/nonexistent") + config = load_config(required=False) + assert config is None + + def test_required_true_raises_value_error(self): + """required=True (default) raises ValueError when no config found.""" + with pytest.raises(ValueError, match="No optimization config found"): + load_config() + + def test_required_true_not_caught_by_broad_except(self, monkeypatch): + """ValueError from required=True is NOT swallowed by the outer try/except.""" + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", "/nonexistent") + with pytest.raises(ValueError, match="baseline"): + load_config(required=True) + + +# ── OptimizationConfig dataclass ──────────────────────────────────── + + +class TestOptimizationConfig: + """Unit tests for OptimizationConfig properties and methods.""" + + def test_compose_instructions_no_skills(self): + config = OptimizationConfig( + instructions="Base prompt.", model=None, temperature=None + ) + assert config.compose_instructions() == "Base prompt." + + def test_compose_instructions_with_skills(self): + config = OptimizationConfig( + instructions="Base prompt.", + model=None, + temperature=None, + skills=[ + Skill(name="math", description="Math operations"), + Skill(name="code", description="Code generation"), + ], + ) + result = config.compose_instructions() + assert "Base prompt." in result + assert "## Available Skills" in result + assert "- **math**: Math operations" in result + assert "- **code**: Code generation" in result + + def test_has_skills_with_list(self): + config = OptimizationConfig( + instructions="", model=None, temperature=None, + skills=[Skill(name="s", description="d")], + ) + assert config.has_skills + + def test_has_skills_with_dir(self): + config = OptimizationConfig( + instructions="", model=None, temperature=None, + skills_dir="/some/dir", + ) + assert config.has_skills + + def test_no_skills(self): + config = OptimizationConfig( + instructions="", model=None, temperature=None, + ) + assert not config.has_skills + + def test_constants(self): + assert OptimizationConfig.DEFAULT_LOCAL_DIR == ".agent_configs" + assert OptimizationConfig.METADATA_FILE == "metadata.yaml" + assert OptimizationConfig.INSTRUCTIONS_FILE == "instructions.md" + assert OptimizationConfig.TOOLS_FILE == "tools.json" + assert OptimizationConfig.SKILLS_DIR == "skills" + assert OptimizationConfig.SKILL_FILE == "SKILL.md" + assert OptimizationConfig.BASELINE_DIR == "baseline" + + +# ── CandidateConfig ───────────────────────────────────────────────── + + +class TestCandidateConfig: + """Tests for CandidateConfig.from_dict parsing.""" + + def test_full_payload(self): + payload = { + "name": "travel-agent-v2", + "instructions": "You are a travel assistant.", + "model": "gpt-4o", + "temperature": 0.7, + "skills": [ + {"name": "budget-checker", "description": "Check budget", "body": "# Budget"}, + {"name": "policy-reviewer", "description": "Review policy"}, + ], + "tools": [ + {"type": "function", "function": {"name": "lookup_travel_policy", "description": "Look up travel policy."}}, + {"type": "function", "function": {"name": "get_flight_alternatives", "description": "Find cheaper flights.", "parameters": {"type": "object", "properties": {"destination": {"type": "string", "description": "The destination city"}}}}}, + ], + } + candidate = CandidateConfig.from_dict(payload) + assert candidate.name == "travel-agent-v2" + assert candidate.instructions == "You are a travel assistant." + assert candidate.model == "gpt-4o" + assert candidate.temperature == 0.7 + assert len(candidate.skills) == 2 + assert candidate.skills[0].name == "budget-checker" + assert candidate.skills[0].body == "# Budget" + assert len(candidate.tool_definitions) == 2 + assert candidate.tool_definitions[1]["function"]["name"] == "get_flight_alternatives" + + def test_minimal_payload(self): + candidate = CandidateConfig.from_dict({}) + assert candidate.name is None + assert candidate.instructions is None + assert candidate.model is None + assert candidate.temperature is None + assert candidate.skills == [] + assert candidate.tool_definitions == [] + + def test_tools_list_format(self): + payload = { + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City"}, + }, + }, + }, + }, + ], + } + candidate = CandidateConfig.from_dict(payload) + assert len(candidate.tool_definitions) == 1 + assert candidate.tool_definitions[0]["function"]["name"] == "get_weather" + + def test_tools_non_list_ignored(self): + """Non-list tools value is coerced to empty list.""" + candidate = CandidateConfig.from_dict({"tools": "not a list"}) + assert candidate.tool_definitions == [] + + def test_tools_dict_ignored(self): + """Dict tools value is coerced to empty list.""" + candidate = CandidateConfig.from_dict({"tools": {"a": "b"}}) + assert candidate.tool_definitions == [] + + def test_tools_none_ignored(self): + """None tools value results in empty list.""" + candidate = CandidateConfig.from_dict({"tools": None}) + assert candidate.tool_definitions == [] + + def test_no_job_id_field(self): + """OptimizationConfig no longer carries a job_id field.""" + candidate = CandidateConfig.from_dict({"instructions": "test"}) + assert not hasattr(candidate, "job_id") + + +# ── MetadataConfig ────────────────────────────────────────────────── + + +class TestMetadataConfig: + """Unit tests for MetadataConfig.from_dict.""" + + def test_from_dict_basic(self): + meta = MetadataConfig.from_dict({"model": "gpt-4o", "temperature": 0.5}) + assert meta.model == "gpt-4o" + assert meta.temperature == 0.5 + + def test_from_dict_ignores_unknown(self): + meta = MetadataConfig.from_dict({"model": "gpt-4o", "unknown_key": "value"}) + assert meta.model == "gpt-4o" + assert not hasattr(meta, "unknown_key") + + def test_from_dict_defaults(self): + meta = MetadataConfig.from_dict({}) + assert meta.model is None + assert meta.temperature is None + assert meta.instruction_file == "instructions.md" + assert meta.skill_dir == "skills" + assert meta.tool_file == "tools.json" + + def test_custom_file_paths(self): + meta = MetadataConfig.from_dict({ + "instruction_file": "prompt.txt", + "skill_dir": "my_skills", + "tool_file": "my_tools.json", + }) + assert meta.instruction_file == "prompt.txt" + assert meta.skill_dir == "my_skills" + assert meta.tool_file == "my_tools.json" + + +# ── _parse_skills ─────────────────────────────────────────────────── + + +class TestParseSkills: + """Tests for _parse_skills edge cases.""" + + def test_skips_non_dict_items(self): + result = _parse_skills(["not-a-dict", 42, None]) + assert result == [] + + def test_skips_items_without_name(self): + result = _parse_skills([{"description": "no name"}]) + assert result == [] + + def test_parses_valid_skills(self): + result = _parse_skills([ + {"name": "a", "description": "desc-a", "body": "body-a"}, + {"name": "b"}, + ]) + assert len(result) == 2 + assert result[0].name == "a" + assert result[0].body == "body-a" + assert result[1].description == "" + + def test_empty_list(self): + assert _parse_skills([]) == [] + + def test_mixed_valid_invalid(self): + result = _parse_skills([ + {"name": "valid", "description": "ok"}, + "garbage", + {"no_name": True}, + {"name": "also-valid"}, + ]) + assert len(result) == 2 + + +# ── _load_tool_definitions (file loading) ─────────────────────────── + + +class TestLoadToolDefinitions: + """Tests for _load_tool_definitions from tools.json.""" + + def test_load_list_format(self, tmp_path): + tool_file = tmp_path / "tools.json" + tools = [ + {"type": "function", "function": {"name": "f1", "description": "Func 1"}}, + ] + tool_file.write_text(json.dumps(tools)) + result = _load_tool_definitions(tool_file) + assert len(result) == 1 + assert result[0]["function"]["name"] == "f1" + + def test_missing_file_returns_empty(self, tmp_path): + result = _load_tool_definitions(tmp_path / "nonexistent.json") + assert result == [] + + def test_bad_json_returns_empty(self, tmp_path): + tool_file = tmp_path / "tools.json" + tool_file.write_text("not json") + result = _load_tool_definitions(tool_file) + assert result == [] + + def test_non_list_returns_empty(self, tmp_path): + tool_file = tmp_path / "tools.json" + tool_file.write_text('"just a string"') + result = _load_tool_definitions(tool_file) + assert result == [] + + def test_dict_format_returns_empty(self, tmp_path): + tool_file = tmp_path / "tools.json" + tool_file.write_text(json.dumps({"t": "simple description"})) + result = _load_tool_definitions(tool_file) + assert result == [] + + def test_empty_list_returns_empty(self, tmp_path): + tool_file = tmp_path / "tools.json" + tool_file.write_text("[]") + result = _load_tool_definitions(tool_file) + assert result == [] + + def test_multiple_tools(self, tmp_path): + tool_file = tmp_path / "tools.json" + tools = [ + {"type": "function", "function": {"name": "a", "description": "A"}}, + {"type": "function", "function": {"name": "b", "description": "B"}}, + ] + tool_file.write_text(json.dumps(tools)) + result = _load_tool_definitions(tool_file) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + +# ── Skill frontmatter parsing ─────────────────────────────────────── + + +class TestSkillFrontmatter: + """Tests for _parse_skill_frontmatter.""" + + def test_no_frontmatter(self): + fm, body = _parse_skill_frontmatter("Just a body.") + assert fm == {} + assert body == "Just a body." + + def test_with_frontmatter(self): + content = "---\nname: test\ndescription: A test\n---\nBody text." + fm, body = _parse_skill_frontmatter(content) + assert fm["name"] == "test" + assert fm["description"] == "A test" + assert body == "Body text." + + def test_unclosed_frontmatter(self): + content = "---\nname: broken" + fm, body = _parse_skill_frontmatter(content) + assert fm == {} + + def test_empty_frontmatter(self): + content = "---\n---\nBody." + fm, body = _parse_skill_frontmatter(content) + assert fm == {} + assert body == "Body." + + +# ── Simple YAML parser ────────────────────────────────────────────── + + +class TestSimpleYaml: + """Tests for the fallback YAML parser (no PyYAML).""" + + def test_basic_parsing(self, tmp_path): + f = tmp_path / "test.yaml" + f.write_text("model: gpt-4o\ntemperature: 0.5\n") + result = _parse_simple_yaml(f) + assert result["model"] == "gpt-4o" + assert result["temperature"] == 0.5 + assert isinstance(result["temperature"], float) + + def test_skips_comments_and_blanks(self, tmp_path): + f = tmp_path / "test.yaml" + f.write_text("# comment\n\nmodel: gpt-4o\n") + result = _parse_simple_yaml(f) + assert result == {"model": "gpt-4o"} + + def test_missing_file(self, tmp_path): + result = _parse_simple_yaml(tmp_path / "nope.yaml") + assert result == {} + + def test_colon_in_value(self, tmp_path): + f = tmp_path / "test.yaml" + f.write_text("url: http://example.com\n") + result = _parse_simple_yaml(f) + assert result["url"] == "http://example.com" + + def test_coerces_int(self, tmp_path): + f = tmp_path / "test.yaml" + f.write_text("count: 42\n") + result = _parse_simple_yaml(f) + assert result["count"] == 42 + assert isinstance(result["count"], int) + + def test_coerces_null(self, tmp_path): + f = tmp_path / "test.yaml" + f.write_text("value: null\nempty: ~\nblank:\n") + result = _parse_simple_yaml(f) + assert result["value"] is None + assert result["empty"] is None + assert result["blank"] is None + + def test_coerces_bool(self, tmp_path): + f = tmp_path / "test.yaml" + f.write_text("enabled: true\ndisabled: false\n") + result = _parse_simple_yaml(f) + assert result["enabled"] is True + assert result["disabled"] is False + + +# ── apply_tool_descriptions ────────────────────────────────────────── + + +class TestApplyToolDescriptions: + """Tests for OptimizationConfig.apply_tool_descriptions.""" + + @staticmethod + def _tool_def(name, description="", parameters=None): + """Helper to build a tool definition dict.""" + func = {"name": name, "description": description} + if parameters: + func["parameters"] = parameters + return {"type": "function", "function": func} + + def _make_config(self, tool_definitions=None): + return OptimizationConfig( + instructions="test", + model=None, + temperature=None, + tool_definitions=tool_definitions or [], + ) + + def test_patches_docstring(self): + def lookup_policy(): + """Original description.""" + + config = self._make_config([ + self._tool_def("lookup_policy", "Optimized description."), + ]) + result = config.apply_tool_descriptions([lookup_policy]) + assert result is not None + assert lookup_policy.__doc__ == "Optimized description." + + def test_returns_same_list(self): + def my_tool(): + """Original.""" + + config = self._make_config([ + self._tool_def("my_tool", "New."), + ]) + tools = [my_tool] + result = config.apply_tool_descriptions(tools) + assert result is tools + + def test_skips_tools_not_in_definitions(self): + def unknown_tool(): + """Should not change.""" + + config = self._make_config([ + self._tool_def("other_tool", "Something."), + ]) + config.apply_tool_descriptions([unknown_tool]) + assert unknown_tool.__doc__ == "Should not change." + + def test_skips_empty_description(self): + def my_tool(): + """Original doc.""" + + config = self._make_config([ + self._tool_def("my_tool", ""), + ]) + config.apply_tool_descriptions([my_tool]) + assert my_tool.__doc__ == "Original doc." + + def test_noop_when_no_tool_definitions(self): + def my_tool(): + """Keep me.""" + + config = self._make_config() + config.apply_tool_descriptions([my_tool]) + assert my_tool.__doc__ == "Keep me." + + def test_handles_objects_without_name(self): + """Non-function items without __name__ are silently skipped.""" + config = self._make_config([ + self._tool_def("something", "X."), + ]) + obj = object() + # Should not raise + config.apply_tool_descriptions([obj]) + + def test_multiple_tools_selective_patch(self): + def tool_a(): + """A original.""" + + def tool_b(): + """B original.""" + + def tool_c(): + """C original.""" + + config = self._make_config([ + self._tool_def("tool_a", "A optimized."), + self._tool_def("tool_c", "C optimized."), + ]) + config.apply_tool_descriptions([tool_a, tool_b, tool_c]) + assert tool_a.__doc__ == "A optimized." + assert tool_b.__doc__ == "B original." + assert tool_c.__doc__ == "C optimized." + + def test_uses_name_attribute_fallback(self): + """Falls back to .name when __name__ is not available.""" + class FakeTool: + name = "my_tool" + __doc__ = "Original." + + tool = FakeTool() + config = self._make_config([ + self._tool_def("my_tool", "Patched via .name attr."), + ]) + config.apply_tool_descriptions([tool]) + assert tool.__doc__ == "Patched via .name attr." + + def test_patches_description_attribute(self): + """Patches .description attribute for AIFunction/ToolProtocol objects.""" + def search_flights(): + """Original.""" + + search_flights.description = "Original." # type: ignore[attr-defined] + + config = self._make_config([ + self._tool_def("search_flights", "Optimized flights search."), + ]) + config.apply_tool_descriptions([search_flights]) + assert search_flights.description == "Optimized flights search." # type: ignore[attr-defined] + assert search_flights.__doc__ == "Optimized flights search." + + def test_description_attribute_not_set_when_readonly(self): + """If .description is read-only, only __doc__ is patched.""" + class ReadOnlyTool: + __name__ = "my_tool" + __doc__ = "Original doc." + + @property + def description(self): + return "read-only" + + tool = ReadOnlyTool() + config = self._make_config([ + self._tool_def("my_tool", "Patched."), + ]) + config.apply_tool_descriptions([tool]) + assert tool.__doc__ == "Patched." + assert tool.description == "read-only" # unchanged + + def test_patches_input_model_param_descriptions(self): + """Patches parameter descriptions on input_model.model_fields.""" + class FakeField: + def __init__(self, desc): + self.description = desc + + class FakeInputModel: + model_fields = { + "destination": FakeField("Old dest description"), + "date": FakeField("Old date description"), + } + _rebuild_called = False + + @classmethod + def model_rebuild(cls, force=False): + cls._rebuild_called = True + + def search_flights(destination: str, date: str): + """Search.""" + + search_flights.input_model = FakeInputModel # type: ignore[attr-defined] + + config = self._make_config([ + self._tool_def("search_flights", "Find flights.", parameters={ + "type": "object", + "properties": { + "destination": {"type": "string", "description": "The travel destination city"}, + }, + }), + ]) + config.apply_tool_descriptions([search_flights]) + assert search_flights.__doc__ == "Find flights." + assert FakeInputModel.model_fields["destination"].description == "The travel destination city" + assert FakeInputModel.model_fields["date"].description == "Old date description" + assert FakeInputModel._rebuild_called is True + + def test_skips_param_patch_when_no_input_model(self): + """No crash when tool has no input_model attribute.""" + def my_tool(): + """Original.""" + + config = self._make_config([ + self._tool_def("my_tool", "New.", parameters={ + "type": "object", + "properties": {"x": {"type": "string", "description": "Some param"}}, + }), + ]) + config.apply_tool_descriptions([my_tool]) + assert my_tool.__doc__ == "New." + + def test_skips_unknown_params_in_input_model(self): + """Parameters not in model_fields are silently ignored.""" + class FakeField: + def __init__(self, desc): + self.description = desc + + class FakeInputModel: + model_fields = {"known": FakeField("Known param")} + _rebuild_called = False + + @classmethod + def model_rebuild(cls, force=False): + cls._rebuild_called = True + + def my_tool(): + """Doc.""" + + my_tool.input_model = FakeInputModel # type: ignore[attr-defined] + + config = self._make_config([ + self._tool_def("my_tool", "New doc.", parameters={ + "type": "object", + "properties": {"unknown_param": {"type": "string", "description": "Should be ignored"}}, + }), + ]) + config.apply_tool_descriptions([my_tool]) + assert FakeInputModel.model_fields["known"].description == "Known param" + assert FakeInputModel._rebuild_called is False + + def test_no_rebuild_when_no_params_patched(self): + """model_rebuild is NOT called if no parameters were actually patched.""" + class FakeField: + def __init__(self, desc): + self.description = desc + + class FakeInputModel: + model_fields = {"x": FakeField("X")} + _rebuild_called = False + + @classmethod + def model_rebuild(cls, force=False): + cls._rebuild_called = True + + def my_tool(): + """Doc.""" + + my_tool.input_model = FakeInputModel # type: ignore[attr-defined] + + config = self._make_config([ + self._tool_def("my_tool", "New."), + ]) + config.apply_tool_descriptions([my_tool]) + assert FakeInputModel._rebuild_called is False + + def test_empty_tools_list(self): + """Passing an empty list is fine.""" + config = self._make_config([ + self._tool_def("tool", "X."), + ]) + result = config.apply_tool_descriptions([]) + assert result == [] + + def test_non_dict_items_in_tool_definitions_ignored(self): + """Non-dict entries in tool_definitions are silently skipped.""" + def my_tool(): + """Original.""" + + config = self._make_config([ + "not a dict", + 42, + self._tool_def("my_tool", "Patched."), + ]) + config.apply_tool_descriptions([my_tool]) + assert my_tool.__doc__ == "Patched." + + def test_tool_definition_without_function_key_ignored(self): + """Tool definition dict without 'function' key is skipped.""" + def my_tool(): + """Original.""" + + config = self._make_config([ + {"type": "function"}, # missing 'function' key + self._tool_def("my_tool", "Patched."), + ]) + config.apply_tool_descriptions([my_tool]) + assert my_tool.__doc__ == "Patched." + + +# ── OptimizationConfig field checks ───────────────────────────────── + + +class TestOptimizationConfigFields: + """Verify OptimizationConfig fields after recent refactors.""" + + def test_no_job_id_field(self): + """job_id was removed — OptimizationConfig should not have it.""" + config = OptimizationConfig(instructions="test") + assert not hasattr(config, "job_id") + + def test_no_env_job_id_classvar(self): + """ENV_JOB_ID ClassVar was removed.""" + assert not hasattr(OptimizationConfig, "ENV_JOB_ID") + + def test_has_candidate_id(self): + config = OptimizationConfig(instructions="test", candidate_id="cand-1") + assert config.candidate_id == "cand-1" + + def test_tool_definitions_is_list(self): + config = OptimizationConfig( + instructions="test", + tool_definitions=[{"type": "function", "function": {"name": "f"}}], + ) + assert isinstance(config.tool_definitions, list) + assert len(config.tool_definitions) == 1 + + def test_default_tool_definitions_empty(self): + config = OptimizationConfig(instructions="test") + assert config.tool_definitions == [] + + +# ── Priority ordering ─────────────────────────────────────────────── + + +class TestPriorityOrdering: + """Verify the priority resolution: env > resolver > local > defaults.""" + + def test_env_beats_resolver(self, monkeypatch): + """Priority 1 (env) wins over Priority 2 (resolver).""" + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps({"instructions": "from env"})) + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-1") + monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") + config = load_config() + assert config.source == "env:OPTIMIZATION_CONFIG" + assert config.instructions == "from env" + + def test_resolver_beats_local(self, monkeypatch, tmp_path): + """Priority 2 (resolver) wins over Priority 3 (local).""" + # Set up local dir with baseline + baseline = tmp_path / "baseline" + baseline.mkdir() + (baseline / "instructions.md").write_text("From local.") + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + + # Set up resolver to succeed + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-1") + monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") + monkeypatch.setattr( + "azure.ai.agentserver.optimization._config.resolve_candidate", + lambda cid, endpoint, local_dir=None: {"instructions": "From resolver."}, + ) + config = load_config() + assert config.source == "api:candidate:cand-1" + assert config.instructions == "From resolver." + + def test_local_is_last_resort(self, monkeypatch, tmp_path): + """Priority 3 (local) used when no env var and no resolver.""" + baseline = tmp_path / "baseline" + baseline.mkdir() + (baseline / "instructions.md").write_text("From local.") + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert config.source.startswith("local:") + assert config.instructions == "From local." + + def test_candidate_id_without_endpoint_falls_to_local(self, monkeypatch, tmp_path): + """candidate_id without endpoint skips resolver, uses local.""" + baseline = tmp_path / "baseline" + baseline.mkdir() + (baseline / "instructions.md").write_text("Local fallback.") + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-1") + # No OPTIMIZATION_RESOLVE_ENDPOINT set + config = load_config() + assert config.source.startswith("local:") + + def test_resolver_provides_skills_dir(self, monkeypatch): + """Resolver response with skills_dir is passed through to OptimizationConfig.""" + resolved = { + "instructions": "With skills.", + "skills_dir": "/some/path/skills", + } + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-1") + monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") + monkeypatch.setattr( + "azure.ai.agentserver.optimization._config.resolve_candidate", + lambda cid, endpoint, local_dir=None: resolved, + ) + config = load_config() + assert config.skills_dir == "/some/path/skills" + + def test_resolver_provides_tool_definitions(self, monkeypatch): + """Resolver response with tools list is parsed into tool_definitions.""" + resolved = { + "instructions": "With tools.", + "tools": [ + {"type": "function", "function": {"name": "search", "description": "Find stuff"}}, + ], + } + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-1") + monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") + monkeypatch.setattr( + "azure.ai.agentserver.optimization._config.resolve_candidate", + lambda cid, endpoint, local_dir=None: resolved, + ) + config = load_config() + assert len(config.tool_definitions) == 1 + assert config.tool_definitions[0]["function"]["name"] == "search" + + +# ── config_dir parameter ──────────────────────────────────────────── + + +class TestConfigDirParam: + """Tests for the config_dir parameter of load_config.""" + + def test_config_dir_string(self, tmp_path): + baseline = tmp_path / "baseline" + baseline.mkdir() + (baseline / "instructions.md").write_text("String dir.") + config = load_config(config_dir=str(tmp_path)) + assert config.instructions == "String dir." + + def test_config_dir_path_object(self, tmp_path): + from pathlib import Path + baseline = tmp_path / "baseline" + baseline.mkdir() + (baseline / "instructions.md").write_text("Path dir.") + config = load_config(config_dir=Path(tmp_path)) + assert config.instructions == "Path dir." + + def test_config_dir_nonexistent_required_false(self, tmp_path): + config = load_config(config_dir=tmp_path / "nope", required=False) + assert config is None + + def test_config_dir_nonexistent_required_true(self, tmp_path): + with pytest.raises(ValueError, match="baseline"): + load_config(config_dir=tmp_path / "nope") + + def test_config_dir_with_candidate(self, tmp_path, monkeypatch): + """config_dir + OPTIMIZATION_CANDIDATE_ID selects the right folder.""" + cand_dir = tmp_path / "my-cand" + cand_dir.mkdir() + (cand_dir / "instructions.md").write_text("Candidate instructions.") + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "my-cand") + config = load_config(config_dir=tmp_path) + assert config.instructions == "Candidate instructions." diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_integration.py b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_integration.py new file mode 100644 index 000000000000..276f8ea7d0ef --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_integration.py @@ -0,0 +1,476 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Integration tests — exercise the full public API end-to-end.""" + +import json +from pathlib import Path + +import pytest + +from azure.ai.agentserver.optimization import ( + OptimizationConfig, + Skill, + load_config, + load_skills_from_dir, +) +from azure.ai.agentserver.optimization._resolver import _downloaded + + +@pytest.fixture(autouse=True) +def clear_downloaded(): + _downloaded.clear() + yield + _downloaded.clear() + + +class TestLoadConfigAndApplyTools: + """End-to-end: load_config → apply_tool_descriptions.""" + + def test_env_config_apply_tool_descriptions(self, monkeypatch): + """Load from OPTIMIZATION_CONFIG env and apply to tool functions.""" + cfg = { + "instructions": "Optimized prompt.", + "model": "gpt-4o", + "temperature": 0.5, + "tools": [ + { + "type": "function", + "function": { + "name": "search_flights", + "description": "Find the cheapest flight options.", + "parameters": { + "type": "object", + "properties": { + "destination": {"type": "string", "description": "City name"}, + }, + }, + }, + }, + { + "type": "function", + "function": { + "name": "book_hotel", + "description": "Reserve a hotel room.", + "parameters": {"type": "object", "properties": {}}, + }, + }, + ], + } + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(cfg)) + + def search_flights(destination: str): + """Original search flights doc.""" + + def book_hotel(city: str): + """Original book hotel doc.""" + + def unrelated_tool(): + """Should stay unchanged.""" + + config = load_config() + assert config.source == "env:OPTIMIZATION_CONFIG" + assert config.instructions == "Optimized prompt." + assert len(config.tool_definitions) == 2 + + tools = config.apply_tool_descriptions([search_flights, book_hotel, unrelated_tool]) + assert search_flights.__doc__ == "Find the cheapest flight options." + assert book_hotel.__doc__ == "Reserve a hotel room." + assert unrelated_tool.__doc__ == "Should stay unchanged." + assert tools == [search_flights, book_hotel, unrelated_tool] + + def test_env_config_openai_tools_list_apply(self, monkeypatch): + """OpenAI function-calling format loads and applies correctly.""" + cfg = { + "instructions": "Agent prompt.", + "tools": [ + { + "type": "function", + "function": { + "name": "lookup_policy", + "description": "Look up travel policy.", + "parameters": { + "type": "object", + "properties": { + "dept": {"type": "string", "description": "Department name"}, + }, + }, + }, + } + ], + } + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(cfg)) + + def lookup_policy(dept: str): + """Old doc.""" + + config = load_config() + config.apply_tool_descriptions([lookup_policy]) + assert lookup_policy.__doc__ == "Look up travel policy." + + +class TestLoadConfigAndLoadSkills: + """End-to-end: load_config → load_skills_from_dir.""" + + def test_local_dir_skills_workflow(self, monkeypatch, tmp_path): + """Full local directory workflow: config sets skills_dir, user loads skills.""" + candidate_dir = tmp_path / "baseline" + candidate_dir.mkdir() + (candidate_dir / "metadata.yaml").write_text( + "model: gpt-4o\ntemperature: 0.7\nskill_dir: skills\n" + ) + (candidate_dir / "instructions.md").write_text("You are a travel agent.") + + # Create two skills + (candidate_dir / "skills" / "budget" ).mkdir(parents=True) + (candidate_dir / "skills" / "budget" / "SKILL.md").write_text( + "---\nname: budget-checker\ndescription: Check trip budget\n---\nCalculate costs." + ) + (candidate_dir / "skills" / "routing").mkdir(parents=True) + (candidate_dir / "skills" / "routing" / "SKILL.md").write_text( + "---\nname: route-planner\ndescription: Plan optimal route\n---\nFind shortest path." + ) + + # Create tools.json + tools_data = [ + {"type": "function", "function": {"name": "search", "description": "Search destinations.", "parameters": {"type": "object", "properties": {"q": {"type": "string", "description": "Query"}}}}}, + ] + (candidate_dir / "tools.json").write_text(json.dumps(tools_data)) + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + + # Verify config loaded + assert config.instructions == "You are a travel agent." + assert config.model == "gpt-4o" + assert config.temperature == 0.7 + assert "local:" in config.source + assert config.skills_dir is not None + assert config.skills == [] # skills not loaded inline + + # User calls load_skills_from_dir + skills = load_skills_from_dir(Path(config.skills_dir)) + assert len(skills) == 2 + names = {s.name for s in skills} + assert "budget-checker" in names + assert "route-planner" in names + + # Verify tool definitions also loaded + assert len(config.tool_definitions) == 1 + assert config.tool_definitions[0]["function"]["name"] == "search" + + def test_no_skills_dir_returns_empty(self): + """load_skills_from_dir on non-existent dir returns empty list.""" + skills = load_skills_from_dir(Path("/nonexistent/path")) + assert skills == [] + + def test_skills_dir_with_no_skill_files(self, tmp_path): + """Directory exists but has no valid skill folders.""" + skills = load_skills_from_dir(tmp_path) + assert skills == [] + + def test_skills_without_frontmatter(self, tmp_path): + """Skills with plain markdown (no frontmatter) use folder name and first line.""" + skill_dir = tmp_path / "my-skill" + skill_dir.mkdir() + (skill_dir / "SKILL.md").write_text("# Summarize Emails\nCondense inbox messages.") + + skills = load_skills_from_dir(tmp_path) + assert len(skills) == 1 + assert skills[0].name == "my-skill" + assert skills[0].description == "Summarize Emails" + assert skills[0].body == "Condense inbox messages." + + +class TestFullWorkflow: + """Complete end-to-end: load → skills → tools → compose instructions.""" + + def test_complete_agent_setup(self, monkeypatch, tmp_path): + """Simulate a full agent startup with optimization.""" + candidate_dir = tmp_path / "candidate-v2" + candidate_dir.mkdir() + (candidate_dir / "metadata.yaml").write_text( + "model: gpt-4o-mini\ntemperature: 0.3\n" + ) + (candidate_dir / "instructions.md").write_text( + "You are a concise travel booking assistant." + ) + tools_data = [ + {"type": "function", "function": {"name": "search_flights", "description": "Find flights between cities."}}, + {"type": "function", "function": {"name": "book_flight", "description": "Book the selected flight."}}, + ] + (candidate_dir / "tools.json").write_text(json.dumps(tools_data)) + skills_dir = candidate_dir / "skills" / "rebooking" + skills_dir.mkdir(parents=True) + (skills_dir / "SKILL.md").write_text( + "---\nname: rebooking\ndescription: Handle rebooking requests\n---\n" + "Steps:\n1. Cancel old flight\n2. Search alternatives\n3. Book new one" + ) + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "candidate-v2") + + # Step 1: Load config + config = load_config() + assert config.instructions == "You are a concise travel booking assistant." + assert config.model == "gpt-4o-mini" + assert config.temperature == 0.3 + + # Step 2: Apply tool descriptions + def search_flights(origin, dest): + """Old doc.""" + + def book_flight(flight_id): + """Old doc.""" + + config.apply_tool_descriptions([search_flights, book_flight]) + assert search_flights.__doc__ == "Find flights between cities." + assert book_flight.__doc__ == "Book the selected flight." + + # Step 3: Load skills + assert config.skills_dir is not None + skills = load_skills_from_dir(Path(config.skills_dir)) + assert len(skills) == 1 + assert skills[0].name == "rebooking" + + # Step 4: Compose instructions (manually with skills since they're loaded separately) + # Simulate what compose_instructions would do + config_with_skills = OptimizationConfig( + instructions=config.instructions, + model=config.model, + temperature=config.temperature, + skills=skills, + skills_dir=config.skills_dir, + tool_definitions=config.tool_definitions, + source=config.source, + candidate_id=config.candidate_id, + ) + composed = config_with_skills.compose_instructions() + assert "You are a concise travel booking assistant." in composed + assert "rebooking" in composed + assert "Handle rebooking requests" in composed + + def test_defaults_workflow_no_optimization(self): + """When no optimization is configured, returns None.""" + config = load_config(required=False) + assert config is None + + +class TestResolverPersistLoadRoundTrip: + """End-to-end: resolver persists → load_config reads back.""" + + def test_resolver_persist_and_reload(self, monkeypatch, tmp_path): + """Simulate resolver writing to disk, then load_config reading it via local dir.""" + from azure.ai.agentserver.optimization._resolver import _persist_to_local_layout + + api_response = { + "instructions": "Optimized agent prompt.", + "model": "gpt-4o-mini", + "temperature": 0.2, + "tools": [ + { + "type": "function", + "function": { + "name": "search", + "description": "Search the web.", + "parameters": { + "type": "object", + "properties": {"q": {"type": "string", "description": "Query"}}, + }, + }, + } + ], + "skills": [ + {"name": "summarize", "description": "Summarize text", "body": "Condense."}, + ], + } + candidate_path = tmp_path / "cand-resolved" + _persist_to_local_layout(candidate_path, api_response) + + # Now load via local dir + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-resolved") + config = load_config() + + assert config.instructions == "Optimized agent prompt." + assert config.model == "gpt-4o-mini" + assert config.temperature == 0.2 + assert config.source.startswith("local:") + assert len(config.tool_definitions) == 1 + assert config.tool_definitions[0]["function"]["name"] == "search" + # Skills are loaded via skills_dir, not inline + assert config.skills_dir is not None + + skills = load_skills_from_dir(Path(config.skills_dir)) + assert len(skills) == 1 + assert skills[0].name == "summarize" + assert skills[0].body == "Condense." + + +class TestComposeInstructions: + """Tests for OptimizationConfig.compose_instructions edge cases.""" + + def test_compose_with_no_base_instructions(self): + """compose_instructions with None base instructions.""" + config = OptimizationConfig( + instructions=None, + skills=[Skill(name="s1", description="Skill one")], + ) + result = config.compose_instructions() + assert "## Available Skills" in result + assert "**s1**: Skill one" in result + + def test_compose_empty_instructions_with_skills(self): + config = OptimizationConfig( + instructions="", + skills=[Skill(name="s1", description="d1")], + ) + result = config.compose_instructions() + assert result.startswith("## Available Skills") + + def test_compose_no_instructions_no_skills(self): + config = OptimizationConfig(instructions=None) + assert config.compose_instructions() == "" + + def test_compose_empty_instructions_no_skills(self): + config = OptimizationConfig(instructions="") + assert config.compose_instructions() == "" + + def test_compose_preserves_multiline_base(self): + config = OptimizationConfig( + instructions="Line 1.\nLine 2.", + skills=[Skill(name="s1", description="d1")], + ) + result = config.compose_instructions() + assert "Line 1.\nLine 2." in result + assert "## Available Skills" in result + + +class TestApplyToolDescriptionsEndToEnd: + """Integration: load_config → apply_tool_descriptions with parameter patching.""" + + def test_parameter_patching_e2e(self, monkeypatch): + """Full flow: env config → apply → verify parameter descriptions patched.""" + + class FakeField: + def __init__(self, desc): + self.description = desc + + class FakeInputModel: + model_fields = { + "destination": FakeField("Original dest"), + "date": FakeField("Original date"), + } + _rebuild_called = False + + @classmethod + def model_rebuild(cls, force=False): + cls._rebuild_called = True + + cfg = { + "instructions": "Agent.", + "tools": [ + { + "type": "function", + "function": { + "name": "search_flights", + "description": "Find flights.", + "parameters": { + "type": "object", + "properties": { + "destination": {"type": "string", "description": "The city to fly to"}, + }, + }, + }, + } + ], + } + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(cfg)) + + def search_flights(destination: str, date: str): + """Original.""" + + search_flights.input_model = FakeInputModel # type: ignore[attr-defined] + + config = load_config() + config.apply_tool_descriptions([search_flights]) + + assert search_flights.__doc__ == "Find flights." + assert FakeInputModel.model_fields["destination"].description == "The city to fly to" + assert FakeInputModel.model_fields["date"].description == "Original date" + assert FakeInputModel._rebuild_called is True + + +class TestConfigDirIntegration: + """Integration tests for config_dir parameter.""" + + def test_config_dir_with_full_setup(self, tmp_path): + """config_dir param with complete directory layout.""" + baseline = tmp_path / "baseline" + baseline.mkdir() + (baseline / "metadata.yaml").write_text( + "model: gpt-4o\ntemperature: 0.8\n" + ) + (baseline / "instructions.md").write_text("Custom agent.") + tools = [ + {"type": "function", "function": {"name": "t1", "description": "Tool one"}}, + ] + (baseline / "tools.json").write_text(json.dumps(tools)) + skill_dir = baseline / "skills" / "helper" + skill_dir.mkdir(parents=True) + (skill_dir / "SKILL.md").write_text( + "---\nname: helper\ndescription: A helper skill\n---\nHelp." + ) + + config = load_config(config_dir=tmp_path) + assert config.instructions == "Custom agent." + assert config.model == "gpt-4o" + assert config.temperature == 0.8 + assert len(config.tool_definitions) == 1 + assert config.skills_dir is not None + + skills = load_skills_from_dir(Path(config.skills_dir)) + assert len(skills) == 1 + assert skills[0].name == "helper" + + def test_config_dir_overrides_env_local_dir(self, monkeypatch, tmp_path): + """config_dir param takes priority over OPTIMIZATION_LOCAL_DIR env var.""" + env_dir = tmp_path / "env_configs" + env_dir.mkdir() + (env_dir / "baseline").mkdir() + (env_dir / "baseline" / "instructions.md").write_text("From env dir.") + + param_dir = tmp_path / "param_configs" + param_dir.mkdir() + (param_dir / "baseline").mkdir() + (param_dir / "baseline" / "instructions.md").write_text("From param dir.") + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(env_dir)) + config = load_config(config_dir=param_dir) + assert config.instructions == "From param dir." + + +class TestNoJobIdAnywhere: + """Verify job_id is completely removed from the system.""" + + def test_no_job_id_in_env_vars_list(self): + """OPTIMIZATION_JOB_ID should not be recognized.""" + assert not hasattr(OptimizationConfig, "ENV_JOB_ID") + + def test_no_job_id_on_config(self): + config = OptimizationConfig(instructions="test") + assert not hasattr(config, "job_id") + + def test_config_from_env_no_job_id(self, monkeypatch): + cfg = {"instructions": "test"} + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(cfg)) + config = load_config() + assert not hasattr(config, "job_id") + + def test_config_from_local_no_job_id(self, monkeypatch, tmp_path): + baseline = tmp_path / "baseline" + baseline.mkdir() + (baseline / "instructions.md").write_text("ok") + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert not hasattr(config, "job_id") diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py new file mode 100644 index 000000000000..3e841958bf2a --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py @@ -0,0 +1,687 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for the candidate resolver module.""" + +import json +from unittest.mock import patch, MagicMock + +import pytest + +from azure.ai.agentserver.optimization._resolver import ( + resolve_candidate, + _downloaded, + _persist_to_local_layout, + _download_skill_files, + _is_skill_file, + _build_client, +) +from azure.ai.agentserver.optimization._models import OptimizationConfig + + +@pytest.fixture(autouse=True) +def clear_downloaded(): + """Clear the downloaded set before each test.""" + _downloaded.clear() + yield + _downloaded.clear() + + +@pytest.fixture() +def mock_client(): + """Return a MagicMock that stands in for PipelineClient.""" + client = MagicMock() + client._base_url = "http://fake-endpoint" + return client + + +ENDPOINT = "http://fake-endpoint" + + +# ── resolve_candidate ─────────────────────────────────────────────── + + +class TestResolveCandidate: + """Tests for resolve_candidate function.""" + + def test_returns_none_on_api_failure(self): + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client"), + patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=None, + ), + ): + result = resolve_candidate("cand-1", endpoint=ENDPOINT) + assert result is None + + def test_returns_config_on_success(self): + config = { + "instructions": "Optimized.", + "model": "gpt-4o", + "temperature": 0.2, + "skills": [], + } + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client"), + patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=config, + ), + ): + result = resolve_candidate("cand-1", endpoint=ENDPOINT) + assert result is not None + assert result["instructions"] == "Optimized." + assert result["model"] == "gpt-4o" + + def test_uses_correct_path(self): + """Verify the API route follows /candidates/{candidateId}/config.""" + called_args: list = [] + + def capture_call(client, path, params=None): + called_args.append((path, params)) + return {"instructions": "ok"} + + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client"), + patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + side_effect=capture_call, + ), + ): + resolve_candidate("cand-abc", endpoint="http://api.test") + assert called_args[0][0] == "/candidates/cand-abc/config" + + def test_marks_downloaded_after_success(self): + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client"), + patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value={"instructions": "ok"}, + ), + ): + resolve_candidate("cand-mark", endpoint=ENDPOINT) + assert "cand-mark" in _downloaded + + def test_skips_if_already_downloaded_and_folder_exists(self, tmp_path): + """Already-downloaded candidate with existing folder is skipped.""" + (tmp_path / "cand-skip").mkdir() + _downloaded.add("cand-skip") + + result = resolve_candidate( + "cand-skip", endpoint=ENDPOINT, local_dir=tmp_path, + ) + assert result is None + + def test_redownloads_if_folder_missing(self): + """If downloaded but folder is gone, re-download.""" + _downloaded.add("cand-gone") + config = {"instructions": "re-downloaded"} + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client"), + patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=config, + ), + ): + result = resolve_candidate( + "cand-gone", endpoint=ENDPOINT, local_dir=None, + ) + # local_dir is None → can't check folder → should re-download + assert result is not None + assert result["instructions"] == "re-downloaded" + + def test_does_not_mark_downloaded_on_api_failure(self): + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client"), + patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=None, + ), + ): + resolve_candidate("cand-fail", endpoint=ENDPOINT) + assert "cand-fail" not in _downloaded + + +# ── _persist_to_local_layout ──────────────────────────────────────── + + +class TestPersistToLocalLayout: + """Tests for _persist_to_local_layout.""" + + def test_writes_metadata_yaml(self, tmp_path): + candidate_path = tmp_path / "cand-1" + config = {"model": "gpt-4o", "temperature": 0.5} + _persist_to_local_layout(candidate_path, config) + + meta = (candidate_path / "metadata.yaml").read_text() + assert "model: gpt-4o" in meta + assert "temperature: 0.5" in meta + assert f"instruction_file: {OptimizationConfig.INSTRUCTIONS_FILE}" in meta + assert f"skill_dir: {OptimizationConfig.SKILLS_DIR}" in meta + assert f"tool_file: {OptimizationConfig.TOOLS_FILE}" in meta + + def test_writes_instructions_md(self, tmp_path): + candidate_path = tmp_path / "cand-2" + config = {"instructions": "Hello world."} + _persist_to_local_layout(candidate_path, config) + + instr = (candidate_path / "instructions.md").read_text() + assert instr == "Hello world." + + def test_no_instructions_file_when_empty(self, tmp_path): + candidate_path = tmp_path / "cand-3" + config = {"model": "gpt-4o"} + _persist_to_local_layout(candidate_path, config) + + assert not (candidate_path / "instructions.md").exists() + + def test_writes_tools_json_list_format_from_tools_key(self, tmp_path): + candidate_path = tmp_path / "cand-4" + config = { + "tools": [ + {"type": "function", "function": {"name": "search", "description": "Search it", "parameters": {"type": "object", "properties": {"q": {"type": "string", "description": "query"}}}}}, + ] + } + _persist_to_local_layout(candidate_path, config) + + tools = json.loads((candidate_path / "tools.json").read_text()) + assert isinstance(tools, list) + assert tools[0]["function"]["name"] == "search" + + def test_writes_tools_json_multiple_tools(self, tmp_path): + candidate_path = tmp_path / "cand-5" + config = { + "tools": [ + {"type": "function", "function": {"name": "lookup", "description": "Look up policy"}}, + {"type": "function", "function": {"name": "search", "description": "Search"}}, + ] + } + _persist_to_local_layout(candidate_path, config) + + tools = json.loads((candidate_path / "tools.json").read_text()) + assert len(tools) == 2 + assert tools[0]["function"]["name"] == "lookup" + + def test_writes_tools_json_list_format(self, tmp_path): + candidate_path = tmp_path / "cand-6" + config = { + "tools": [ + {"type": "function", "function": {"name": "f1", "description": "Func 1"}}, + ] + } + _persist_to_local_layout(candidate_path, config) + + tools = json.loads((candidate_path / "tools.json").read_text()) + assert isinstance(tools, list) + assert tools[0]["function"]["name"] == "f1" + + def test_no_tools_file_when_tools_is_not_list(self, tmp_path): + """Non-list tools value does not produce tools.json.""" + candidate_path = tmp_path / "cand-7" + config = { + "tools": "not a list", + } + _persist_to_local_layout(candidate_path, config) + + assert not (candidate_path / "tools.json").exists() + + def test_no_tools_file_when_no_tools(self, tmp_path): + candidate_path = tmp_path / "cand-8" + config = {"instructions": "No tools here."} + _persist_to_local_layout(candidate_path, config) + + assert not (candidate_path / "tools.json").exists() + + def test_overwrites_existing_folder(self, tmp_path): + candidate_path = tmp_path / "cand-overwrite" + candidate_path.mkdir() + (candidate_path / "old_file.txt").write_text("stale") + + config = {"instructions": "Fresh.", "model": "gpt-4o"} + _persist_to_local_layout(candidate_path, config) + + assert not (candidate_path / "old_file.txt").exists() + assert (candidate_path / "metadata.yaml").exists() + assert (candidate_path / "instructions.md").read_text() == "Fresh." + + def test_metadata_without_model_and_temperature(self, tmp_path): + candidate_path = tmp_path / "cand-minimal" + config = {"instructions": "Minimal."} + _persist_to_local_layout(candidate_path, config) + + meta = (candidate_path / "metadata.yaml").read_text() + assert "model:" not in meta + assert "temperature:" not in meta + + def test_writes_inline_skills(self, tmp_path): + """Inline skills are persisted as skills//SKILL.md with frontmatter.""" + candidate_path = tmp_path / "cand-skills" + config = { + "instructions": "With skills.", + "skills": [ + {"name": "math", "description": "Do math", "body": "Calculate things."}, + {"name": "code", "description": "Write code"}, + ], + } + _persist_to_local_layout(candidate_path, config) + + math_skill = candidate_path / "skills" / "math" / "SKILL.md" + assert math_skill.exists() + content = math_skill.read_text() + assert "name: math" in content + assert "description: Do math" in content + assert "Calculate things." in content + + code_skill = candidate_path / "skills" / "code" / "SKILL.md" + assert code_skill.exists() + code_content = code_skill.read_text() + assert "name: code" in code_content + + def test_skips_skills_without_name(self, tmp_path): + """Skills without a name are skipped during persist.""" + candidate_path = tmp_path / "cand-no-name" + config = { + "skills": [ + {"description": "no name"}, + {"name": "valid", "description": "has name"}, + ], + } + _persist_to_local_layout(candidate_path, config) + skills_dir = candidate_path / "skills" + # Only the valid skill should be written + assert not (skills_dir / "no name").exists() + assert (skills_dir / "valid" / "SKILL.md").exists() + + def test_skips_non_dict_skills(self, tmp_path): + """Non-dict skill entries are skipped.""" + candidate_path = tmp_path / "cand-bad-skills" + config = {"skills": ["not a dict", 42]} + _persist_to_local_layout(candidate_path, config) + assert not (candidate_path / "skills").exists() + + +# ── _persist + resolve round-trip ──────────────────────────────────── + + +class TestPersistRoundTrip: + """Ensure persisted layout can be read back by _load_local_dir.""" + + def test_round_trip(self, monkeypatch, tmp_path): + from azure.ai.agentserver.optimization import load_config + + config = { + "instructions": "Round-trip test.", + "model": "gpt-4o", + "temperature": 0.3, + "tools": [ + {"type": "function", "function": {"name": "search", "description": "Find things", "parameters": {"type": "object", "properties": {"q": {"type": "string", "description": "query"}}}}}, + ], + } + candidate_path = tmp_path / "cand-rt" + _persist_to_local_layout(candidate_path, config) + + # Now load via local dir + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-rt") + loaded = load_config() + assert loaded.instructions == "Round-trip test." + assert loaded.model == "gpt-4o" + assert loaded.temperature == 0.3 + assert len(loaded.tool_definitions) == 1 + assert loaded.tool_definitions[0]["function"]["name"] == "search" + assert loaded.source.startswith("local:") + + +# ── _download_skill_files ─────────────────────────────────────────── + + +class TestDownloadSkillFiles: + """Tests for _download_skill_files.""" + + def test_downloads_skill_files(self, tmp_path, mock_client): + candidate_path = tmp_path / "cand-sk" + candidate_path.mkdir() + manifest = { + "files": [ + {"path": "skills/math/SKILL.md", "type": "skill"}, + ] + } + + def mock_json(client, path, params=None): + return manifest + + def mock_text(client, path, params=None): + return "# Math Skill\nDo math." + + with ( + patch("azure.ai.agentserver.optimization._resolver._api_get_json", side_effect=mock_json), + patch("azure.ai.agentserver.optimization._resolver._api_get_text", side_effect=mock_text), + ): + _download_skill_files(mock_client, "cand-sk", candidate_path) + + skill_file = candidate_path / "skills" / "math" / "SKILL.md" + assert skill_file.exists() + assert "Math Skill" in skill_file.read_text() + + def test_skips_when_no_manifest(self, tmp_path, mock_client): + candidate_path = tmp_path / "cand-no-manifest" + candidate_path.mkdir() + with patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=None, + ): + _download_skill_files(mock_client, "cand-no-manifest", candidate_path) + assert not (candidate_path / "skills").exists() + + def test_skips_when_no_skill_files_in_manifest(self, tmp_path, mock_client): + candidate_path = tmp_path / "cand-no-skills" + candidate_path.mkdir() + manifest = {"files": [{"path": "other.txt", "type": "config"}]} + with patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=manifest, + ): + _download_skill_files(mock_client, "cand-no-skills", candidate_path) + assert not (candidate_path / "skills").exists() + + def test_skips_empty_path_entries(self, tmp_path, mock_client): + candidate_path = tmp_path / "cand-empty-path" + candidate_path.mkdir() + manifest = {"files": [{"path": "", "type": "skill"}]} + with patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=manifest, + ): + _download_skill_files(mock_client, "cand-empty-path", candidate_path) + assert not (candidate_path / "skills").exists() + + def test_handles_download_failure(self, tmp_path, mock_client): + candidate_path = tmp_path / "cand-dl-fail" + candidate_path.mkdir() + manifest = {"files": [{"path": "skills/bad/SKILL.md", "type": "skill"}]} + with ( + patch("azure.ai.agentserver.optimization._resolver._api_get_json", return_value=manifest), + patch("azure.ai.agentserver.optimization._resolver._api_get_text", return_value=None), + ): + _download_skill_files(mock_client, "cand-dl-fail", candidate_path) + # No crash, skill file simply not written + assert not (candidate_path / "skills" / "bad" / "SKILL.md").exists() + + def test_rejects_traversal_in_file_path(self, tmp_path, mock_client): + """File paths with '../' are rejected (zip-slip prevention).""" + candidate_path = tmp_path / "cand-traversal" + candidate_path.mkdir() + manifest = {"files": [{"path": "skills/../../etc/passwd", "type": "skill"}]} + with ( + patch("azure.ai.agentserver.optimization._resolver._api_get_json", return_value=manifest), + patch("azure.ai.agentserver.optimization._resolver._api_get_text", return_value="malicious"), + ): + _download_skill_files(mock_client, "cand-traversal", candidate_path) + # Malicious file must NOT be written outside skills_dir + assert not (tmp_path / "etc" / "passwd").exists() + assert not (candidate_path / "skills" / ".." / ".." / "etc" / "passwd").exists() + + def test_downloads_multiple_skill_files(self, tmp_path, mock_client): + """Multiple skill files are downloaded correctly.""" + candidate_path = tmp_path / "cand-multi" + candidate_path.mkdir() + manifest = { + "files": [ + {"path": "skills/math/SKILL.md", "type": "skill"}, + {"path": "skills/code/SKILL.md", "type": "skill"}, + ] + } + call_count = {"json": 0, "text": 0} + + def mock_json(client, path, params=None): + call_count["json"] += 1 + return manifest + + def mock_text(client, path, params=None): + call_count["text"] += 1 + if "math" in str(params): + return "# Math\nDo math." + return "# Code\nWrite code." + + with ( + patch("azure.ai.agentserver.optimization._resolver._api_get_json", side_effect=mock_json), + patch("azure.ai.agentserver.optimization._resolver._api_get_text", side_effect=mock_text), + ): + _download_skill_files(mock_client, "cand-multi", candidate_path) + + assert (candidate_path / "skills" / "math" / "SKILL.md").exists() + assert (candidate_path / "skills" / "code" / "SKILL.md").exists() + + def test_uses_correct_api_paths(self, tmp_path, mock_client): + """Verify correct API paths: /candidates/{id} for manifest, /candidates/{id}/files for downloads.""" + candidate_path = tmp_path / "cand-paths" + candidate_path.mkdir() + called_paths: list = [] + + def mock_json(client, path, params=None): + called_paths.append(("json", path)) + return {"files": [{"path": "skills/s1/SKILL.md", "type": "skill"}]} + + def mock_text(client, path, params=None): + called_paths.append(("text", path)) + return "content" + + with ( + patch("azure.ai.agentserver.optimization._resolver._api_get_json", side_effect=mock_json), + patch("azure.ai.agentserver.optimization._resolver._api_get_text", side_effect=mock_text), + ): + _download_skill_files(mock_client, "cand-x", candidate_path) + + assert called_paths[0] == ("json", "/candidates/cand-x") + assert called_paths[1] == ("text", "/candidates/cand-x/files") + + def test_empty_files_list_in_manifest(self, tmp_path, mock_client): + """Empty files list in manifest → no downloads.""" + candidate_path = tmp_path / "cand-empty-files" + candidate_path.mkdir() + with patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value={"files": []}, + ): + _download_skill_files(mock_client, "cand-empty-files", candidate_path) + assert not (candidate_path / "skills").exists() + + +class TestNormalCandidateId: + def test_normal_candidate_id_allowed(self, tmp_path): + """Normal candidate IDs pass the guard.""" + config = {"instructions": "ok", "model": "gpt-4o"} + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client"), + patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=config, + ), + ): + result = resolve_candidate( + "valid-candidate-123", endpoint=ENDPOINT, local_dir=tmp_path + ) + assert result is not None + + def test_candidate_id_with_dots_allowed(self): + """candidate_id containing dots is fine.""" + config = {"instructions": "ok"} + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client"), + patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=config, + ), + ): + result = resolve_candidate( + "candidate.v2.1", endpoint=ENDPOINT + ) + assert result is not None + + +# ── resolve_candidate with local_dir ───────────────────────────────── + + +class TestResolveCandidateWithLocalDir: + """Tests for resolve_candidate persisting to disk.""" + + def test_persists_config_to_disk(self, tmp_path): + config = { + "instructions": "Persisted.", + "model": "gpt-4o", + "temperature": 0.5, + } + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client"), + patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=config, + ), + ): + result = resolve_candidate("cand-persist", endpoint=ENDPOINT, local_dir=tmp_path) + assert result is not None + assert (tmp_path / "cand-persist" / "metadata.yaml").exists() + assert (tmp_path / "cand-persist" / "instructions.md").read_text() == "Persisted." + + def test_sets_skills_dir_when_skills_exist(self, tmp_path): + """skills_dir is set in returned config when skills folder exists after persist.""" + config = { + "instructions": "ok", + "skills": [{"name": "math", "description": "Do math", "body": "# Math"}], + } + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client"), + patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=config, + ), + ): + result = resolve_candidate("cand-skills", endpoint=ENDPOINT, local_dir=tmp_path) + assert result is not None + assert "skills_dir" in result + assert "cand-skills" in result["skills_dir"] + + def test_no_skills_dir_when_no_skills(self, tmp_path): + """skills_dir is not set when no skills are persisted.""" + config = {"instructions": "ok", "model": "gpt-4o"} + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client"), + patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=config, + ), + ): + result = resolve_candidate("cand-no-skills", endpoint=ENDPOINT, local_dir=tmp_path) + assert result is not None + assert "skills_dir" not in result + + def test_local_dir_none_skips_persist(self): + """When local_dir is None, no files are written.""" + config = {"instructions": "ok"} + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client"), + patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=config, + ), + ): + result = resolve_candidate("cand-no-dir", endpoint=ENDPOINT, local_dir=None) + assert result is not None + assert result["instructions"] == "ok" + # No skills_dir because no local_dir + assert "skills_dir" not in result + + def test_endpoint_trailing_slash_stripped(self): + """Endpoint URL trailing slash doesn't break API paths.""" + called_urls: list = [] + + def capture_build(endpoint): + m = MagicMock() + m._base_url = endpoint + return m + + def capture_json(client, path, params=None): + called_urls.append(f"{client._base_url}{path}") + return {"instructions": "ok"} + + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client", side_effect=capture_build), + patch("azure.ai.agentserver.optimization._resolver._api_get_json", side_effect=capture_json), + ): + resolve_candidate("cand-1", endpoint="http://host/job/123") + # Verify the URL construction + assert "/candidates/cand-1/config" in called_urls[0] + + +# ── _is_skill_file ────────────────────────────────────────────────── + + +class TestIsSkillFile: + """Tests for _is_skill_file.""" + + def test_type_skill(self): + assert _is_skill_file({"path": "anything", "type": "skill"}) + + def test_path_starts_with_skills(self): + assert _is_skill_file({"path": "skills/math/SKILL.md", "type": ""}) + + def test_not_a_skill(self): + assert not _is_skill_file({"path": "config.json", "type": "config"}) + + def test_empty_entry(self): + assert not _is_skill_file({}) + + +# ── Persist IO error handling ──────────────────────────────────────── + + +class TestPersistErrorHandling: + """Ensure IO errors during persist don't crash resolve_candidate.""" + + def test_persist_oserror_does_not_crash(self, tmp_path): + config = {"instructions": "ok", "model": "gpt-4o"} + with ( + patch("azure.ai.agentserver.optimization._resolver._build_client"), + patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=config, + ), + patch( + "azure.ai.agentserver.optimization._resolver._persist_to_local_layout", + side_effect=OSError("disk full"), + ), + ): + result = resolve_candidate( + "cand-io", endpoint=ENDPOINT, local_dir=tmp_path, + ) + # Config is still returned from API even if persist fails + assert result is not None + assert result["instructions"] == "ok" + assert "cand-io" in _downloaded + + +# ── HTTP helpers ───────────────────────────────────────────────────── + + +class TestBuildClient: + """Tests for _build_client.""" + + def test_returns_pipeline_client(self): + from azure.core import PipelineClient + + client = _build_client("http://example.com") + assert isinstance(client, PipelineClient) + client.close() + + def test_works_without_credentials(self): + """If DefaultAzureCredential fails, client is still created (no auth).""" + from azure.core import PipelineClient + + with patch( + "azure.identity.DefaultAzureCredential", + side_effect=Exception("No cred"), + ): + client = _build_client("http://example.com") + assert isinstance(client, PipelineClient) + client.close() diff --git a/sdk/agentserver/ci.yml b/sdk/agentserver/ci.yml index 11be178cf2a6..e6f654f329ad 100644 --- a/sdk/agentserver/ci.yml +++ b/sdk/agentserver/ci.yml @@ -48,3 +48,5 @@ extends: safeName: azureaiagentserverresponses - name: azure-ai-agentserver-ghcopilot safeName: azureaiagentserverghcopilot + - name: azure-ai-agentserver-optimization + safeName: azureaiagentserveroptimization