From 3ceb9144df846d76a47199a3cea3237eecdeb6e9 Mon Sep 17 00:00:00 2001 From: zyysurely Date: Fri, 22 May 2026 20:28:49 -0700 Subject: [PATCH 1/5] [feat] new azure-ai-agentserver-optimization package for load_config --- .../CHANGELOG.md | 14 + .../azure-ai-agentserver-optimization/LICENSE | 21 + .../MANIFEST.in | 8 + .../README.md | 128 ++ .../azure/__init__.py | 1 + .../azure/ai/__init__.py | 1 + .../azure/ai/agentserver/__init__.py | 1 + .../ai/agentserver/optimization/__init__.py | 42 + .../ai/agentserver/optimization/_config.py | 387 +++++++ .../ai/agentserver/optimization/_models.py | 250 ++++ .../ai/agentserver/optimization/_resolver.py | 249 ++++ .../ai/agentserver/optimization/_version.py | 5 + .../ai/agentserver/optimization/py.typed | 0 .../dev_requirements.txt | 2 + .../pyproject.toml | 61 + .../tests/conftest.py | 23 + .../tests/test_config.py | 1028 +++++++++++++++++ .../tests/test_resolver.py | 431 +++++++ sdk/agentserver/ci.yml | 2 + 19 files changed, 2654 insertions(+) create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/LICENSE create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/MANIFEST.in create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/README.md create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/azure/__init__.py create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/__init__.py create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/__init__.py create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/__init__.py create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_config.py create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_models.py create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_resolver.py create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_version.py create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/py.typed create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/dev_requirements.txt create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/pyproject.toml create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/tests/conftest.py create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py 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..e93f27a39c91 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md @@ -0,0 +1,14 @@ +# Release History + +## 0.1.0b1 (Unreleased) + +### Features Added + +- Initial beta release. +- `load_config()` — single-call config loader with graceful fallback. +- `OptimizationConfig` dataclass with instructions, model, temperature, skills, and tool definitions. +- Candidate resolution via `OPTIMIZATION_CANDIDATE_ID` env var and remote resolver API. +- Inline JSON config via `OPTIMIZATION_CONFIG` env var. +- Local directory layout support (`metadata.yaml` + `instructions.md` + `skills/`). +- Skill loading from `SKILL.md` files with YAML frontmatter. +- Tool definition loading from `tools.json`. 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..e3e1b173ec53 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/README.md @@ -0,0 +1,128 @@ +# azure-ai-agentserver-optimization + +Optimization config loader for Azure AI Hosted Agents. + +Provides a single `load_config()` call that resolves optimization parameters (instructions, model, temperature, skills, tool definitions) from multiple sources with graceful fallback — your agent works unchanged when not running under optimization. + +## Installation + +```bash +pip install azure-ai-agentserver-optimization +``` + +## Quick Start + +```python +from azure.ai.agentserver.optimization import load_config + +config = load_config(default_instructions="You are a helpful assistant.") + +# Use config in your agent +print(config.instructions) # optimized or default +print(config.model) # optimized or default +print(config.temperature) # optimized or default +print(config.skills) # learned skills (empty list if none) +print(config.tool_descriptions) # optimized tool descriptions (empty dict if none) +print(config.source) # "api:candidate:abc", "env:OPTIMIZATION_CONFIG", "local:...", or "defaults" +``` + +## 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_JOB_ID`, `OPTIMIZATION_RESOLVE_ENDPOINT` | Fetches the candidate config from the remote optimization service and persists it to the local directory. | +| 3 | **Local directory** | `OPTIMIZATION_LOCAL_DIR` (optional, defaults to `.agent_configs/`) | Reads from `//` or `baseline/` as fallback. | +| 4 | **Defaults** | *(none)* | Returns the caller-supplied defaults unchanged — the agent works normally. | + +Any unexpected error is caught and logged — `load_config()` always returns a valid `OptimizationConfig`. + +## Environment Variables + +| Variable | Description | +|----------|-------------| +| `OPTIMIZATION_CONFIG` | Inline JSON config (Priority 1). | +| `OPTIMIZATION_CANDIDATE_ID` | Candidate ID for resolver API or local folder lookup. | +| `OPTIMIZATION_JOB_ID` | Job ID for the resolver API. | +| `OPTIMIZATION_RESOLVE_ENDPOINT` | Base URL of the optimization service. | +| `OPTIMIZATION_LOCAL_DIR` | Path to the local config directory (default: `.agent_configs/`). | +| `MODEL_DEPLOYMENT_NAME` | Fallback model name when no model is resolved or specified. | + +## Local Directory Layout + +When using the local directory (Priority 3) or after the resolver API persists a candidate (Priority 2), the directory uses the following structure: + +``` +.agent_configs/ +├── baseline/ # fallback candidate +│ ├── metadata.yaml # model, temperature, file pointers +│ ├── instructions.md # system prompt +│ ├── tools.json # tool descriptions (dict or OpenAI list format) +│ └── skills/ # learned skills +│ └── / +│ └── SKILL.md +└── / # same layout as baseline/ + ├── metadata.yaml + ├── instructions.md + ├── tools.json + └── skills/ + └── / + └── SKILL.md +``` + +## Tool Description Formats + +`tools.json` and the inline JSON config support three formats: + +**Dict format** (`tool_descriptions`): +```json +{ + "lookup_policy": { + "description": "Look up the company travel policy.", + "parameters": {"dept": "Department name"} + } +} +``` + +**Legacy camelCase** (`toolDescriptions`) — same structure, different key. `tool_descriptions` takes priority when both are present. + +**OpenAI function-calling list** (`tools`): +```json +[ + { + "type": "function", + "function": { + "name": "lookup_policy", + "description": "Look up the company travel policy.", + "parameters": { + "type": "object", + "properties": { + "dept": {"type": "string", "description": "Department name"} + } + } + } + } +] +``` + +## OptimizationConfig Properties + +| Property | Type | Description | +|----------|------|-------------| +| `instructions` | `str` | System prompt (optimized or default). | +| `model` | `str \| None` | Model deployment name. | +| `temperature` | `float \| None` | Sampling temperature. | +| `skills` | `list[Skill]` | Learned skills. | +| `skills_dir` | `str \| None` | Path to skills directory. | +| `tool_descriptions` | `dict[str, ToolDescription]` | Optimized tool descriptions. | +| `source` | `str` | Where the config was loaded from. | +| `candidate_id` | `str \| None` | Candidate ID (when resolved via API). | +| `job_id` | `str \| None` | Job ID (when resolved via API). | +| `has_skills` | `bool` | Whether skills are available. | +| `has_tool_descriptions` | `bool` | Whether tool descriptions are available. | + +## 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..0288fa3806f9 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/__init__.py @@ -0,0 +1,42 @@ +# --------------------------------------------------------- +# 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(default_instructions="You are a helpful assistant.") + # config.instructions — optimized or default + # config.model — optimized or default + # config.temperature — optimized or default + # config.skills — learned skills (empty if none) + # config.tool_definitions — optimized tool definitions (empty if none) + # config.source — "api:candidate:abc", "env:config", or "defaults" + +Resolution order (first match wins): + 1. OPTIMIZATION_CONFIG env var → inline JSON (used by temp agent versions) + 2. OPTIMIZATION_CANDIDATE_ID + JOB_ID + ENDPOINT → resolver API → full config + skills + 3. Local directory (.agent_configs/) → metadata.yaml + instructions.md + tools.json + skills/ + 4. Defaults → your hardcoded values (agent works normally) +""" + +from azure.ai.agentserver.optimization._config import load_config +from azure.ai.agentserver.optimization._models import ( + CandidateConfig, + OptimizationConfig, + Skill, + ToolDescription, +) +from azure.ai.agentserver.optimization._version import VERSION + +__all__ = [ + "CandidateConfig", + "OptimizationConfig", + "Skill", + "ToolDescription", + "load_config", +] +__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..332f4db1aaeb --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_config.py @@ -0,0 +1,387 @@ +# --------------------------------------------------------- +# 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 descriptions — dict or 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, + ToolDescription, + _parse_tools_list, +) +from azure.ai.agentserver.optimization._resolver import resolve_candidate + +logger = logging.getLogger("azure.ai.agentserver.optimization") + + +def load_config( + *, + default_instructions: str = "You are a helpful assistant.", + default_model: str | None = None, + default_temperature: float | None = None, + default_skills_dir: str | None = None, +) -> OptimizationConfig: + """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``, + ``OPTIMIZATION_JOB_ID``, and ``OPTIMIZATION_RESOLVE_ENDPOINT`` + are all set. 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). + The local directory defaults to ``.agent_configs/`` relative to + the main script, overridable via ``OPTIMIZATION_LOCAL_DIR``. + 4. **Defaults** — returns the caller-supplied defaults unchanged. + The agent works exactly as if optimization were not installed. + + Safe to call at module load time. Any unexpected error is caught + and logged — the caller always gets a valid config back. + """ + try: + return _load_config_inner( + default_instructions=default_instructions, + default_model=default_model, + default_temperature=default_temperature, + default_skills_dir=default_skills_dir, + ) + except Exception as exc: # noqa: BLE001 + logger.error("Unexpected error loading optimization config — returning defaults: %s", exc) + model = default_model or os.environ.get("MODEL_DEPLOYMENT_NAME") + return OptimizationConfig( + instructions=default_instructions, + model=model, + temperature=default_temperature, + skills_dir=default_skills_dir, + source="defaults", + ) + + +def _load_config_inner( + *, + default_instructions: str, + default_model: str | None, + default_temperature: float | None, + default_skills_dir: str | None, +) -> OptimizationConfig: + """Internal config loader — may raise on unexpected errors.""" + # ── 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 or default_instructions, + model=candidate.model or default_model, + temperature=candidate.temperature if candidate.temperature is not None else default_temperature, + skills=candidate.skills, + skills_dir=cfg.get("skills_dir", default_skills_dir), + tool_descriptions=candidate.tool_descriptions, + 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() + job_id = os.environ.get(OptimizationConfig.ENV_JOB_ID, "").strip() + endpoint = os.environ.get(OptimizationConfig.ENV_RESOLVE_ENDPOINT, "").strip().rstrip("/") + if candidate_id and job_id and endpoint: + local_dir = _resolve_local_dir() + resolved = resolve_candidate(candidate_id, job_id=job_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 or default_instructions, + model=candidate.model or default_model, + temperature=candidate.temperature if candidate.temperature is not None else default_temperature, + skills=candidate.skills, + skills_dir=resolved.get("skills_dir", default_skills_dir), + tool_descriptions=candidate.tool_descriptions, + source=f"api:candidate:{candidate_id}", + candidate_id=candidate_id, + job_id=job_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, default_instructions, + default_model, default_temperature, default_skills_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: Defaults ───────────────────────────────────────── + model = default_model or os.environ.get("MODEL_DEPLOYMENT_NAME") + return OptimizationConfig( + instructions=default_instructions, + model=model, + temperature=default_temperature, + skills_dir=default_skills_dir, + source="defaults", + ) + + +def _resolve_local_dir() -> Path: + """Resolve the local optimization directory path. + + Falls back to :pyattr:`OptimizationConfig.DEFAULT_LOCAL_DIR` + (``".agent_configs"``) when the env var is not set. + """ + local_dir_env = os.environ.get(OptimizationConfig.ENV_LOCAL_DIR, "").strip() + local_dir = Path(local_dir_env) if local_dir_env 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 + return local_dir + + +def _load_local_dir( + candidate_id: str | None, + default_instructions: str, + default_model: str | None, + default_temperature: float | None, + default_skills_dir: str | None, +) -> OptimizationConfig | None: + """Load optimization config from a local directory.""" + local_dir = _resolve_local_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, + default_instructions, default_model, default_temperature, default_skills_dir, + ) + + +def _load_candidate_from_metadata( + candidate_path: Path, + metadata_file: Path, + candidate_id: str | None, + default_instructions: str, + default_model: str | None, + default_temperature: float | None, + default_skills_dir: 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. + """ + 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 = instructions_path.read_text(encoding="utf-8").strip() + else: + instructions = default_instructions + + # Resolve skills directory + skills_path = (candidate_path / meta.skill_dir).resolve() + skills = _load_skills_from_dir(skills_path) if skills_path.is_dir() else [] + skills_dir = str(skills_path) if skills_path.is_dir() else default_skills_dir + + # Load tool descriptions + tool_descriptions = _load_tool_descriptions(candidate_path / meta.tool_file) + + return OptimizationConfig( + instructions=instructions, + model=meta.model or default_model, + temperature=meta.temperature if meta.temperature is not None else default_temperature, + skills=skills, + skills_dir=skills_dir, + tool_descriptions=tool_descriptions, + source=f"local:{candidate_path}", + candidate_id=candidate_id, + ) + + +def _load_tool_descriptions(tool_file: Path) -> dict[str, ToolDescription]: + """Load tool descriptions from a tools.json file. + + Supports both dict format ``{name: {description, parameters}}`` + and OpenAI function-calling list format ``[{type, function: {name, ...}}]``. + """ + if not tool_file.is_file(): + return {} + try: + raw = tool_file.read_text(encoding="utf-8") + data = json.loads(raw) + if isinstance(data, dict): + return { + name: ToolDescription.from_dict(v) if isinstance(v, dict) else ToolDescription(description=str(v)) + for name, v in data.items() + } + if isinstance(data, list): + return _parse_tools_list(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.""" + 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()] = value.strip() + except OSError as exc: + logger.warning("Failed to read %s: %s", path, exc) + return result + + +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. + """ + 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 + """ + 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.""" + 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..f88b4eb82e1b --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_models.py @@ -0,0 +1,250 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- + +"""Data models for the optimization config system.""" + +from __future__ import annotations + +from dataclasses import dataclass, field, fields +from typing import Any, ClassVar + + +@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 ToolDescription: + """Description-only projection of a tool, optimized by the service. + + The optimizer patches *descriptions* (human-readable text) — it does + **not** change the tool's JSON-Schema (type, required, etc.) because + the hosted agent owns the static definition. + + Matches the API contract:: + + { + "description": "Find cheaper flight alternatives.", + "parameters": {"destination": "The travel destination city"} + } + """ + + description: str + parameters: dict[str, str] = field(default_factory=dict) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> ToolDescription: + return cls( + description=data.get("description", ""), + parameters=data.get("parameters", {}), + ) + + +@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": "..."}], + "tool_descriptions": {"lookup_policy": {"description": "...", "parameters": {}}} + } + """ + + name: str | None = None + instructions: str | None = None + model: str | None = None + temperature: float | None = None + skills: list[Skill] = field(default_factory=list) + tool_descriptions: dict[str, ToolDescription] = field(default_factory=dict) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> CandidateConfig: + """Parse from a raw API response / JSON dict.""" + 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_descriptions=_parse_tool_descriptions(data), + ) + + +@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.""" + 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: + """Resolved optimization config. + + When not running under optimization, all fields contain the defaults + you passed to :func:`load_config` — your agent works unchanged. + """ + + 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" + ENV_JOB_ID: ClassVar[str] = "OPTIMIZATION_JOB_ID" + 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 + model: str | None + temperature: float | None + skills: list[Skill] = field(default_factory=list) + skills_dir: str | None = None + tool_descriptions: dict[str, ToolDescription] = field(default_factory=dict) + source: str = "defaults" + candidate_id: str | None = None + job_id: str | None = None + + @property + def has_skills(self) -> bool: + return len(self.skills) > 0 or self.skills_dir is not None + + @property + def has_tool_descriptions(self) -> bool: + return len(self.tool_descriptions) > 0 + + def get_tool_description(self, tool_name: str) -> ToolDescription | None: + """Look up the optimized description for a specific tool.""" + return self.tool_descriptions.get(tool_name) + + def get_tool_param_description(self, tool_name: str, param_name: str) -> str | None: + """Look up the optimized description for a specific tool parameter.""" + td = self.tool_descriptions.get(tool_name) + if td is None: + return None + return td.parameters.get(param_name) + + def compose_instructions(self) -> str: + """Return instructions with skill catalog appended (if any).""" + if not self.skills: + return self.instructions + + lines = [self.instructions, "", "## Available Skills"] + for s in self.skills: + lines.append(f"- **{s.name}**: {s.description}") + return "\n".join(lines) + + +# ── Parsing helpers (used by CandidateConfig.from_dict) ────────────── + + +def _parse_skills(raw: list) -> list[Skill]: + """Parse skills from API/env config JSON.""" + 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 + + +def _parse_tool_descriptions(data: dict[str, Any]) -> dict[str, ToolDescription]: + """Parse tool descriptions from an API response dict. + + Supports three formats: + - ``tool_descriptions`` / ``toolDescriptions``: ``{name: {description, parameters}}`` + - ``tools``: OpenAI function-calling list ``[{type: function, function: {name, description, parameters}}]`` + + ``tool_descriptions`` wins over ``toolDescriptions`` wins over ``tools``. + """ + raw = data.get("tool_descriptions") or data.get("toolDescriptions") + if isinstance(raw, dict): + return { + name: ToolDescription.from_dict(v) if isinstance(v, dict) else ToolDescription(description=str(v)) + for name, v in raw.items() + } + + tools_list = data.get("tools") + if isinstance(tools_list, list): + return _parse_tools_list(tools_list) + + return {} + + +def _parse_tools_list(tools: list) -> dict[str, ToolDescription]: + """Parse tool descriptions from OpenAI function-calling list format. + + Expected shape:: + + [{"type": "function", "function": {"name": "...", "description": "...", "parameters": {...}}}] + + Extracts per-parameter descriptions from ``parameters.properties..description``. + """ + result: dict[str, ToolDescription] = {} + for item in tools: + if not isinstance(item, dict): + continue + func = item.get("function", {}) + if not isinstance(func, dict): + continue + name = func.get("name") + if not name: + continue + description = func.get("description", "") + params_schema = func.get("parameters", {}) + param_descs: dict[str, str] = {} + if isinstance(params_schema, dict): + props = params_schema.get("properties", {}) + if isinstance(props, dict): + for pname, pval in props.items(): + if isinstance(pval, dict) and "description" in pval: + param_descs[pname] = pval["description"] + result[name] = ToolDescription(description=description, parameters=param_descs) + return result 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..0a9b92420573 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/_resolver.py @@ -0,0 +1,249 @@ +# --------------------------------------------------------- +# 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:: + + / + └── / + ├── config.json + └── skills/ + └── / + └── SKILL.md +""" + +from __future__ import annotations + +import json +import logging +import os +import pathlib +import shutil +import urllib.error +import urllib.parse +import urllib.request +from typing import Any +from azure.ai.agentserver.optimization._models import ( + OptimizationConfig) + +logger = logging.getLogger("azure.ai.agentserver.optimization") + +_downloaded: set[str] = set() + + +def resolve_candidate( + candidate_id: str, + job_id: str, + endpoint: str, + local_dir: pathlib.Path | None = None, +) -> dict[str, Any] | None: + """Resolve a candidate's full config from the optimization service. + + Downloads config and skills into ``local_dir//`` + following the standard local directory layout. + Returns ``None`` if the call fails. + """ + if candidate_id in _downloaded: + if local_dir is not None and (local_dir / candidate_id).is_dir(): + logger.debug("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) + + headers = _build_headers() + + # ── Step 1: Fetch config ───────────────────────────────────────── + config = _api_get_json(f"{endpoint}/agent_optimization_jobs/{job_id}/candidates/{candidate_id}/config", headers) + if config is None: + return None + + logger.info( + "Resolved candidate %s: model=%s, instructions=%d chars, skills=%d, tool_descriptions=%d", + candidate_id, + config.get("model", "?"), + len(config.get("instructions", "")), + len(config.get("skills", [])), + len(config.get("toolDescriptions", {}) or config.get("tool_descriptions", {})), + ) + + # ── 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(endpoint, job_id, candidate_id, headers, 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) + + _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 + + If the folder already exists it is removed and re-created. + """ + 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_descriptions / toolDescriptions as dict format + tool_descs = config.get("tool_descriptions") or config.get("toolDescriptions") + tools_list = config.get("tools") + if tool_descs and isinstance(tool_descs, dict): + tools_file = candidate_path / OptimizationConfig.TOOLS_FILE + tools_file.write_text(json.dumps(tool_descs, indent=2, ensure_ascii=False), encoding="utf-8") + elif 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") + + logger.info("Persisted config to local layout: %s", candidate_path) + + +def _download_skill_files( + endpoint: str, + job_id: str, + candidate_id: str, + headers: dict[str, str], + candidate_path: pathlib.Path, +) -> None: + """Fetch manifest and download skill files into candidate_path/skills//SKILL.md.""" + manifest = _api_get_json(f"{endpoint}/agent_optimization_jobs/{job_id}/candidates/{candidate_id}", headers) + 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( + f"{endpoint}/agent_optimization_jobs/{job_id}/candidates/{candidate_id}/files", + headers, + params={"path": file_path}, + ) + 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 + 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.""" + path = file_entry.get("path", "") + file_type = file_entry.get("type", "") + return file_type == "skill" or path.startswith("skills/") + + +# ── HTTP helpers ───────────────────────────────────────────────────── + + +def _build_headers() -> dict[str, str]: + headers: dict[str, str] = {"Accept": "application/json"} + token = _get_bearer_token() + if token: + headers["Authorization"] = f"Bearer {token}" + return headers + + +def _api_get_json(url: str, headers: dict[str, str]) -> dict[str, Any] | None: + """GET a JSON endpoint, return parsed dict or None on failure.""" + logger.debug("GET %s", url) + try: + req = urllib.request.Request(url, method="GET", headers=headers) + with urllib.request.urlopen(req, timeout=30) as resp: # noqa: S310 + return json.loads(resp.read().decode("utf-8")) + except (urllib.error.URLError, json.JSONDecodeError, OSError) as exc: + logger.error("GET %s failed: %s", url, exc) + return None + + +def _api_get_text( + url: str, headers: dict[str, str], params: dict[str, str] | None = None +) -> str | None: + """GET an endpoint, return response body as text or None on failure.""" + if params: + query = "&".join(f"{k}={urllib.parse.quote(v)}" for k, v in params.items()) + url = f"{url}?{query}" + logger.debug("GET %s", url) + try: + req = urllib.request.Request(url, method="GET", headers=headers) + with urllib.request.urlopen(req, timeout=60) as resp: # noqa: S310 + return resp.read().decode("utf-8") + except (urllib.error.URLError, OSError) as exc: + logger.error("GET %s failed: %s", url, exc) + return None + + +def _get_bearer_token() -> str | None: + """Acquire a bearer token for the resolver API. + + Uses ``azure-identity`` if available; returns ``None`` otherwise. + This keeps azure-identity as an optional dependency. + """ + try: + from azure.identity import DefaultAzureCredential # type: ignore[import-untyped] + + cred = DefaultAzureCredential() + token = cred.get_token("https://ai.azure.com/.default") + return token.token + except Exception: # noqa: BLE001 + 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/dev_requirements.txt b/sdk/agentserver/azure-ai-agentserver-optimization/dev_requirements.txt new file mode 100644 index 000000000000..190d2e5520b2 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/dev_requirements.txt @@ -0,0 +1,2 @@ +-e ../../../eng/tools/azure-sdk-tools +pytest 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..e11feb8656ba --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/pyproject.toml @@ -0,0 +1,61 @@ +[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-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..280cfc794766 --- /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_JOB_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) 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..fe4d7587b671 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py @@ -0,0 +1,1028 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for load_config — priority resolution, fallback, and edge cases.""" + +import json +from unittest.mock import patch + +import pytest + +from azure.ai.agentserver.optimization import ( + CandidateConfig, + OptimizationConfig, + Skill, + ToolDescription, + load_config, +) +from azure.ai.agentserver.optimization._models import ( + MetadataConfig, + _parse_skills, + _parse_tool_descriptions, + _parse_tools_list, +) +from azure.ai.agentserver.optimization._config import ( + _load_tool_descriptions, + _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 are set, load_config returns caller-supplied defaults.""" + + def test_returns_default_instructions(self): + config = load_config(default_instructions="Be helpful.") + assert config.instructions == "Be helpful." + assert config.source == "defaults" + + def test_returns_default_model(self): + config = load_config(default_model="gpt-4o") + assert config.model == "gpt-4o" + + def test_returns_default_temperature(self): + config = load_config(default_temperature=0.5) + assert config.temperature == 0.5 + + def test_returns_default_skills_dir(self): + config = load_config(default_skills_dir="/some/path") + assert config.skills_dir == "/some/path" + + def test_empty_skills_by_default(self): + config = load_config() + assert config.skills == [] + assert not config.has_skills + + def test_empty_tool_descriptions_by_default(self): + config = load_config() + assert config.tool_descriptions == {} + assert not config.has_tool_descriptions + + def test_falls_back_to_model_deployment_name_env(self, monkeypatch): + monkeypatch.setenv("MODEL_DEPLOYMENT_NAME", "gpt-4o-mini") + config = load_config() + assert config.model == "gpt-4o-mini" + + def test_explicit_model_overrides_env(self, monkeypatch): + monkeypatch.setenv("MODEL_DEPLOYMENT_NAME", "gpt-4o-mini") + config = load_config(default_model="gpt-4o") + assert config.model == "gpt-4o" + + def test_default_instructions_value(self): + config = load_config() + assert config.instructions == "You are a helpful assistant." + + +# ── 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(default_instructions="default") + 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.", + "tool_descriptions": { + "lookup_travel_policy": { + "description": "Look up the company travel policy.", + "parameters": {}, + }, + "check_department_budget": { + "description": "Check remaining budget.", + "parameters": {"dept": "Department name"}, + }, + }, + } + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) + config = load_config() + assert config.has_tool_descriptions + assert "lookup_travel_policy" in config.tool_descriptions + td = config.tool_descriptions["check_department_budget"] + assert isinstance(td, ToolDescription) + assert td.description == "Check remaining budget." + assert td.parameters == {"dept": "Department name"} + + def test_env_config_with_legacy_toolDescriptions(self, monkeypatch): + payload = { + "instructions": "With tools.", + "toolDescriptions": { + "search": {"description": "Search something.", "parameters": {}}, + }, + } + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) + config = load_config() + assert config.has_tool_descriptions + assert "search" in config.tool_descriptions + + 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 config.has_tool_descriptions + assert "lookup_policy" in config.tool_descriptions + td = config.tool_descriptions["lookup_policy"] + assert td.description == "Look up policy" + assert td.parameters == {"dept": "Department name"} + + def test_tool_descriptions_takes_priority_over_legacy(self, monkeypatch): + payload = { + "instructions": "Both.", + "tool_descriptions": {"new_tool": {"description": "New"}}, + "toolDescriptions": {"old_tool": {"description": "Old"}}, + } + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) + config = load_config() + assert "new_tool" in config.tool_descriptions + assert "old_tool" not in config.tool_descriptions + + def test_tool_descriptions_takes_priority_over_tools_list(self, monkeypatch): + payload = { + "instructions": "Both.", + "tool_descriptions": {"dict_tool": {"description": "Dict"}}, + "tools": [{"type": "function", "function": {"name": "list_tool", "description": "List"}}], + } + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) + config = load_config() + assert "dict_tool" in config.tool_descriptions + assert "list_tool" not in config.tool_descriptions + + def test_bad_json_falls_through(self, monkeypatch): + monkeypatch.setenv("OPTIMIZATION_CONFIG", "not-json{{{") + config = load_config(default_instructions="fallback") + assert config.instructions == "fallback" + assert config.source == "defaults" + + def test_empty_env_var_ignored(self, monkeypatch): + monkeypatch.setenv("OPTIMIZATION_CONFIG", " ") + config = load_config(default_instructions="fallback") + assert config.source == "defaults" + + def test_partial_config_uses_defaults(self, monkeypatch): + payload = {"model": "gpt-4o"} + monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) + config = load_config( + default_instructions="My default", + default_temperature=0.7, + ) + assert config.instructions == "My default" + assert config.model == "gpt-4o" + assert config.temperature == 0.7 + + 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_JOB_ID", "job-1") + 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 + JOB_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_JOB_ID", "job-42") + monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") + monkeypatch.setattr( + "azure.ai.agentserver.optimization._config.resolve_candidate", + lambda cid, job_id, 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 config.job_id == "job-42" + assert len(config.skills) == 1 + + def test_resolver_failure_falls_to_defaults(self, monkeypatch): + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "bad-id") + monkeypatch.setenv("OPTIMIZATION_JOB_ID", "job-1") + monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") + monkeypatch.setattr( + "azure.ai.agentserver.optimization._config.resolve_candidate", + lambda cid, job_id, endpoint, local_dir=None: None, + ) + config = load_config(default_instructions="fallback") + assert config.source == "defaults" + assert config.instructions == "fallback" + + def test_missing_job_id_skips_resolver(self, monkeypatch): + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-1") + monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") + # No JOB_ID set + config = load_config(default_instructions="default") + assert config.source == "defaults" + + def test_missing_endpoint_skips_resolver(self, monkeypatch): + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-1") + monkeypatch.setenv("OPTIMIZATION_JOB_ID", "job-1") + # No ENDPOINT set + config = load_config(default_instructions="default") + assert config.source == "defaults" + + 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_JOB_ID", "job-1") + 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, job_id, 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 len(config.skills) == 1 + assert config.skills[0].name == "math" + assert config.skills[0].description == "Do math" + assert config.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 = {"search": {"description": "Search stuff", "parameters": {"q": "query"}}} + (candidate_dir / "tools.json").write_text(json.dumps(tools)) + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert "search" in config.tool_descriptions + assert config.tool_descriptions["search"].description == "Search stuff" + + 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 "get_weather" in config.tool_descriptions + assert config.tool_descriptions["get_weather"].description == "Get the weather" + assert config.tool_descriptions["get_weather"].parameters == {"city": "City name"} + + def test_missing_instructions_uses_default(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(default_instructions="My default") + assert config.instructions == "My default" + + def test_nonexistent_local_dir_falls_to_defaults(self, monkeypatch): + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", "/nonexistent/path") + config = load_config(default_instructions="fallback") + assert config.source == "defaults" + + 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(default_instructions="default") + assert config.source == "defaults" + + +# ── _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 + + +# ── _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 + + +# ── Graceful error handling ───────────────────────────────────────── + + +class TestGracefulErrorHandling: + """load_config never crashes — always returns a valid config.""" + + def test_unexpected_exception_returns_defaults(self, monkeypatch): + """Any unexpected error in _load_config_inner returns defaults.""" + monkeypatch.setattr( + "azure.ai.agentserver.optimization._config._load_config_inner", + lambda **kwargs: (_ for _ in ()).throw(RuntimeError("boom")), + ) + config = load_config(default_instructions="safe") + assert config.source == "defaults" + assert config.instructions == "safe" + + def test_load_config_never_raises(self, monkeypatch): + """Even with corrupted env vars, load_config returns something.""" + monkeypatch.setenv("OPTIMIZATION_CONFIG", "{invalid") + monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "x") + monkeypatch.setenv("OPTIMIZATION_JOB_ID", "y") + monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://nope") + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", "/nonexistent") + config = load_config(default_instructions="fallback") + assert isinstance(config, OptimizationConfig) + assert config.instructions == "fallback" + + +# ── 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_has_tool_descriptions(self): + config = OptimizationConfig( + instructions="", model=None, temperature=None, + tool_descriptions={"t": ToolDescription(description="d")}, + ) + assert config.has_tool_descriptions + + def test_no_tool_descriptions(self): + config = OptimizationConfig( + instructions="", model=None, temperature=None, + ) + assert not config.has_tool_descriptions + + def test_get_tool_description(self): + td = ToolDescription(description="Search things", parameters={"q": "query"}) + config = OptimizationConfig( + instructions="", model=None, temperature=None, + tool_descriptions={"search": td}, + ) + assert config.get_tool_description("search") is td + assert config.get_tool_description("missing") is None + + def test_get_tool_param_description(self): + td = ToolDescription(description="Search", parameters={"q": "The query"}) + config = OptimizationConfig( + instructions="", model=None, temperature=None, + tool_descriptions={"search": td}, + ) + assert config.get_tool_param_description("search", "q") == "The query" + assert config.get_tool_param_description("search", "missing") is None + assert config.get_tool_param_description("missing", "q") is None + + 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" + + +# ── ToolDescription ────────────────────────────────────────────────── + + +class TestToolDescription: + """Tests for ToolDescription dataclass.""" + + def test_from_dict(self): + td = ToolDescription.from_dict({ + "description": "Search things", + "parameters": {"q": "The query", "limit": "Max results"}, + }) + assert td.description == "Search things" + assert td.parameters == {"q": "The query", "limit": "Max results"} + + def test_from_dict_defaults(self): + td = ToolDescription.from_dict({}) + assert td.description == "" + assert td.parameters == {} + + def test_from_dict_missing_parameters(self): + td = ToolDescription.from_dict({"description": "No params"}) + assert td.description == "No params" + assert td.parameters == {} + + +# ── 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"}, + ], + "tool_descriptions": { + "lookup_travel_policy": { + "description": "Look up travel policy.", + "parameters": {}, + }, + "get_flight_alternatives": { + "description": "Find cheaper flights.", + "parameters": {"destination": "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_descriptions) == 2 + td = candidate.tool_descriptions["get_flight_alternatives"] + assert td.description == "Find cheaper flights." + assert td.parameters["destination"] == "The destination city" + + 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_descriptions == {} + + def test_legacy_toolDescriptions_key(self): + payload = { + "toolDescriptions": { + "search": {"description": "Search", "parameters": {}}, + }, + } + candidate = CandidateConfig.from_dict(payload) + assert "search" in candidate.tool_descriptions + + 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 "get_weather" in candidate.tool_descriptions + assert candidate.tool_descriptions["get_weather"].description == "Get weather" + assert candidate.tool_descriptions["get_weather"].parameters == {"city": "City"} + + +# ── 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" + + +# ── _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 + + +# ── _parse_tool_descriptions ──────────────────────────────────────── + + +class TestParseToolDescriptions: + """Tests for _parse_tool_descriptions edge cases.""" + + def test_empty_data(self): + assert _parse_tool_descriptions({}) == {} + + def test_tool_descriptions_dict(self): + data = {"tool_descriptions": {"t1": {"description": "D1", "parameters": {}}}} + result = _parse_tool_descriptions(data) + assert "t1" in result + assert result["t1"].description == "D1" + + def test_toolDescriptions_camelCase(self): + data = {"toolDescriptions": {"t2": {"description": "D2"}}} + result = _parse_tool_descriptions(data) + assert "t2" in result + + def test_tool_descriptions_wins_over_toolDescriptions(self): + data = { + "tool_descriptions": {"winner": {"description": "W"}}, + "toolDescriptions": {"loser": {"description": "L"}}, + } + result = _parse_tool_descriptions(data) + assert "winner" in result + assert "loser" not in result + + def test_tools_list_fallback(self): + data = { + "tools": [ + {"type": "function", "function": {"name": "f1", "description": "Func"}}, + ] + } + result = _parse_tool_descriptions(data) + assert "f1" in result + assert result["f1"].description == "Func" + + def test_string_value_coerced(self): + data = {"tool_descriptions": {"t": "just a string"}} + result = _parse_tool_descriptions(data) + assert result["t"].description == "just a string" + + def test_non_dict_raw_ignored(self): + data = {"tool_descriptions": "not a dict"} + result = _parse_tool_descriptions(data) + assert result == {} + + +# ── _parse_tools_list ──────────────────────────────────────────────── + + +class TestParseToolsList: + """Tests for _parse_tools_list (OpenAI function-calling format).""" + + def test_basic(self): + tools = [ + { + "type": "function", + "function": { + "name": "search", + "description": "Search things", + "parameters": { + "type": "object", + "properties": { + "q": {"type": "string", "description": "Query"}, + }, + }, + }, + }, + ] + result = _parse_tools_list(tools) + assert "search" in result + assert result["search"].description == "Search things" + assert result["search"].parameters == {"q": "Query"} + + def test_no_parameters(self): + tools = [ + {"type": "function", "function": {"name": "noop", "description": "Do nothing"}}, + ] + result = _parse_tools_list(tools) + assert result["noop"].parameters == {} + + def test_skips_non_dict_items(self): + result = _parse_tools_list(["garbage", 42]) + assert result == {} + + def test_skips_items_without_function(self): + result = _parse_tools_list([{"type": "code_interpreter"}]) + assert result == {} + + def test_skips_items_without_name(self): + result = _parse_tools_list([ + {"type": "function", "function": {"description": "nameless"}}, + ]) + assert result == {} + + def test_empty_list(self): + assert _parse_tools_list([]) == {} + + def test_param_without_description_skipped(self): + tools = [ + { + "type": "function", + "function": { + "name": "f", + "description": "F", + "parameters": { + "type": "object", + "properties": { + "has_desc": {"type": "string", "description": "Yes"}, + "no_desc": {"type": "integer"}, + }, + }, + }, + }, + ] + result = _parse_tools_list(tools) + assert result["f"].parameters == {"has_desc": "Yes"} + + def test_multiple_functions(self): + tools = [ + {"type": "function", "function": {"name": "a", "description": "A"}}, + {"type": "function", "function": {"name": "b", "description": "B"}}, + ] + result = _parse_tools_list(tools) + assert len(result) == 2 + + +# ── _load_tool_descriptions (file loading) ────────────────────────── + + +class TestLoadToolDescriptions: + """Tests for _load_tool_descriptions from tools.json.""" + + def test_load_dict_format(self, tmp_path): + tool_file = tmp_path / "tools.json" + tools = {"my_tool": {"description": "My tool", "parameters": {"x": "input"}}} + tool_file.write_text(json.dumps(tools)) + result = _load_tool_descriptions(tool_file) + assert "my_tool" in result + assert isinstance(result["my_tool"], ToolDescription) + assert result["my_tool"].parameters == {"x": "input"} + + 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_descriptions(tool_file) + assert "f1" in result + assert result["f1"].description == "Func 1" + + def test_missing_file_returns_empty(self, tmp_path): + result = _load_tool_descriptions(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_descriptions(tool_file) + assert result == {} + + def test_non_dict_non_list_returns_empty(self, tmp_path): + tool_file = tmp_path / "tools.json" + tool_file.write_text('"just a string"') + result = _load_tool_descriptions(tool_file) + assert result == {} + + def test_string_value_in_dict(self, tmp_path): + tool_file = tmp_path / "tools.json" + tool_file.write_text(json.dumps({"t": "simple description"})) + result = _load_tool_descriptions(tool_file) + assert result["t"].description == "simple description" + + +# ── 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" + + 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" 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..e41751e45e87 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py @@ -0,0 +1,431 @@ +# --------------------------------------------------------- +# 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_headers, + _get_bearer_token, +) +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() + + +ENDPOINT = "http://fake-endpoint" +JOB_ID = "job-42" + + +# ── resolve_candidate ─────────────────────────────────────────────── + + +class TestResolveCandidate: + """Tests for resolve_candidate function.""" + + def test_returns_none_on_api_failure(self): + with patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value=None, + ): + result = resolve_candidate("cand-1", job_id=JOB_ID, 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._api_get_json", + return_value=config, + ): + result = resolve_candidate("cand-1", job_id=JOB_ID, endpoint=ENDPOINT) + assert result is not None + assert result["instructions"] == "Optimized." + assert result["model"] == "gpt-4o" + + def test_uses_correct_url(self): + """Verify the API route follows agent_optimization_jobs/{jobId}/candidates/{candidateId}/config.""" + called_urls: list[str] = [] + + def capture_url(url, headers): + called_urls.append(url) + return {"instructions": "ok"} + + with patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + side_effect=capture_url, + ): + resolve_candidate("cand-abc", job_id="job-xyz", endpoint="http://api.test") + assert called_urls[0] == "http://api.test/agent_optimization_jobs/job-xyz/candidates/cand-abc/config" + + def test_marks_downloaded_after_success(self): + with patch( + "azure.ai.agentserver.optimization._resolver._api_get_json", + return_value={"instructions": "ok"}, + ): + resolve_candidate("cand-mark", job_id=JOB_ID, 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", job_id=JOB_ID, 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._api_get_json", + return_value=config, + ): + result = resolve_candidate( + "cand-gone", job_id=JOB_ID, 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._api_get_json", + return_value=None, + ): + resolve_candidate("cand-fail", job_id=JOB_ID, 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_dict_format(self, tmp_path): + candidate_path = tmp_path / "cand-4" + config = { + "tool_descriptions": { + "search": {"description": "Search it", "parameters": {"q": "query"}}, + } + } + _persist_to_local_layout(candidate_path, config) + + tools = json.loads((candidate_path / "tools.json").read_text()) + assert tools["search"]["description"] == "Search it" + + def test_writes_tools_json_from_toolDescriptions(self, tmp_path): + candidate_path = tmp_path / "cand-5" + config = { + "toolDescriptions": { + "lookup": {"description": "Look up policy"}, + } + } + _persist_to_local_layout(candidate_path, config) + + tools = json.loads((candidate_path / "tools.json").read_text()) + assert tools["lookup"]["description"] == "Look up policy" + + 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_tool_descriptions_wins_over_tools_list(self, tmp_path): + """tool_descriptions dict takes priority over tools list.""" + candidate_path = tmp_path / "cand-7" + config = { + "tool_descriptions": {"search": {"description": "Dict format"}}, + "tools": [{"type": "function", "function": {"name": "f1"}}], + } + _persist_to_local_layout(candidate_path, config) + + tools = json.loads((candidate_path / "tools.json").read_text()) + assert isinstance(tools, dict) + assert "search" in tools + + 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 + + +# ── _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, + "tool_descriptions": { + "search": {"description": "Find things", "parameters": {"q": "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(default_instructions="unused") + assert loaded.instructions == "Round-trip test." + assert loaded.model == "gpt-4o" + assert loaded.temperature == 0.3 + assert "search" in loaded.tool_descriptions + assert loaded.source.startswith("local:") + + +# ── _download_skill_files ─────────────────────────────────────────── + + +class TestDownloadSkillFiles: + """Tests for _download_skill_files.""" + + def test_downloads_skill_files(self, tmp_path): + candidate_path = tmp_path / "cand-sk" + candidate_path.mkdir() + manifest = { + "files": [ + {"path": "skills/math/SKILL.md", "type": "skill"}, + ] + } + + def mock_get(url, headers, **kwargs): + if "config" in url: + return None + if "files" in url: + return None + return manifest + + def mock_text(url, headers, params=None): + return "# Math Skill\nDo math." + + with ( + patch("azure.ai.agentserver.optimization._resolver._api_get_json", side_effect=lambda u, h: manifest), + patch("azure.ai.agentserver.optimization._resolver._api_get_text", side_effect=mock_text), + ): + _download_skill_files(ENDPOINT, JOB_ID, "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): + 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(ENDPOINT, JOB_ID, "cand-no-manifest", {}, candidate_path) + assert not (candidate_path / "skills").exists() + + def test_skips_when_no_skill_files_in_manifest(self, tmp_path): + 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(ENDPOINT, JOB_ID, "cand-no-skills", {}, candidate_path) + assert not (candidate_path / "skills").exists() + + def test_skips_empty_path_entries(self, tmp_path): + 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(ENDPOINT, JOB_ID, "cand-empty-path", {}, candidate_path) + assert not (candidate_path / "skills").exists() + + def test_handles_download_failure(self, tmp_path): + 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(ENDPOINT, JOB_ID, "cand-dl-fail", {}, candidate_path) + # No crash, skill file simply not written + assert not (candidate_path / "skills" / "bad" / "SKILL.md").exists() + + +# ── _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._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", job_id=JOB_ID, 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 TestBuildHeaders: + """Tests for _build_headers.""" + + def test_includes_accept_header(self): + headers = _build_headers() + assert headers["Accept"] == "application/json" + + def test_includes_auth_when_token_available(self): + with patch( + "azure.ai.agentserver.optimization._resolver._get_bearer_token", + return_value="fake-token", + ): + headers = _build_headers() + assert headers["Authorization"] == "Bearer fake-token" + + def test_no_auth_when_no_token(self): + with patch( + "azure.ai.agentserver.optimization._resolver._get_bearer_token", + return_value=None, + ): + headers = _build_headers() + assert "Authorization" not in headers + + +class TestGetBearerToken: + """Tests for _get_bearer_token.""" + + def test_returns_none_without_azure_identity(self): + with patch.dict("sys.modules", {"azure.identity": None}): + token = _get_bearer_token() + assert token is None or isinstance(token, str) + + def test_returns_none_on_exception(self): + mock_identity = MagicMock() + mock_identity.DefaultAzureCredential.side_effect = Exception("No cred") + with patch.dict("sys.modules", {"azure.identity": mock_identity}): + token = _get_bearer_token() + assert token is None 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 From 45a426da60425b77e156d895cca8613fb5cc2c84 Mon Sep 17 00:00:00 2001 From: zyysurely Date: Fri, 22 May 2026 21:27:00 -0700 Subject: [PATCH 2/5] fix comments and guard path --- .../CHANGELOG.md | 19 ++-- .../README.md | 69 +++++++++----- .../ai/agentserver/optimization/__init__.py | 12 +-- .../ai/agentserver/optimization/_config.py | 88 +++++++++++++++--- .../ai/agentserver/optimization/_resolver.py | 28 ++++-- .../cspell.json | 11 +++ .../tests/conftest.py | 1 + .../tests/test_config.py | 91 ++++++++++++++++++- .../tests/test_resolver.py | 54 +++++++++-- 9 files changed, 309 insertions(+), 64 deletions(-) create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/cspell.json diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md index e93f27a39c91..bb104911e6f6 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md @@ -1,14 +1,19 @@ # Release History -## 0.1.0b1 (Unreleased) +## 1.0.0b1 ### Features Added - Initial beta release. -- `load_config()` — single-call config loader with graceful fallback. -- `OptimizationConfig` dataclass with instructions, model, temperature, skills, and tool definitions. -- Candidate resolution via `OPTIMIZATION_CANDIDATE_ID` env var and remote resolver API. -- Inline JSON config via `OPTIMIZATION_CONFIG` env var. -- Local directory layout support (`metadata.yaml` + `instructions.md` + `skills/`). +- `load_config()` — single-call config loader with 4-priority resolution and graceful fallback (never crashes). +- `OptimizationConfig` dataclass with instructions, model, temperature, skills, tool descriptions, source tracking, candidate_id, and job_id. +- 4-priority resolution order: + 1. Inline JSON via `OPTIMIZATION_CONFIG` env var. + 2. Resolver API via `OPTIMIZATION_CANDIDATE_ID` + `OPTIMIZATION_JOB_ID` + `OPTIMIZATION_RESOLVE_ENDPOINT`. + 3. Local directory layout (`OPTIMIZATION_LOCAL_DIR`, defaults to `.agent_configs/`). + 4. Caller-supplied defaults. +- Local directory layout: `metadata.yaml` + `instructions.md` + `tools.json` + `skills/` per candidate, with `baseline/` fallback. +- 3 tool description formats: `tool_descriptions` dict, `toolDescriptions` (legacy camelCase), and OpenAI function-calling `tools` list. - Skill loading from `SKILL.md` files with YAML frontmatter. -- Tool definition loading from `tools.json`. +- Resolver API persists fetched configs to local directory for offline use. +- Path traversal (zip-slip) protection on all untrusted path inputs. diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/README.md b/sdk/agentserver/azure-ai-agentserver-optimization/README.md index e3e1b173ec53..f84c5970886d 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/README.md +++ b/sdk/agentserver/azure-ai-agentserver-optimization/README.md @@ -1,32 +1,22 @@ -# azure-ai-agentserver-optimization +# Azure AI Agent Server Optimization client library for Python -Optimization config loader for Azure AI Hosted Agents. +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 descriptions) from multiple sources with graceful fallback — your agent works unchanged when not running under optimization. -Provides a single `load_config()` call that 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 -## Installation +### Install the package ```bash pip install azure-ai-agentserver-optimization ``` -## Quick Start +### Prerequisites -```python -from azure.ai.agentserver.optimization import load_config +- Python 3.10 or later -config = load_config(default_instructions="You are a helpful assistant.") +## Key concepts -# Use config in your agent -print(config.instructions) # optimized or default -print(config.model) # optimized or default -print(config.temperature) # optimized or default -print(config.skills) # learned skills (empty list if none) -print(config.tool_descriptions) # optimized tool descriptions (empty dict if none) -print(config.source) # "api:candidate:abc", "env:OPTIMIZATION_CONFIG", "local:...", or "defaults" -``` - -## Resolution Order +### Resolution Order `load_config()` resolves from four sources in order — first match wins: @@ -39,7 +29,7 @@ print(config.source) # "api:candidate:abc", "env:OPTIMIZATION_CONFIG Any unexpected error is caught and logged — `load_config()` always returns a valid `OptimizationConfig`. -## Environment Variables +### Environment Variables | Variable | Description | |----------|-------------| @@ -50,7 +40,7 @@ Any unexpected error is caught and logged — `load_config()` always returns a v | `OPTIMIZATION_LOCAL_DIR` | Path to the local config directory (default: `.agent_configs/`). | | `MODEL_DEPLOYMENT_NAME` | Fallback model name when no model is resolved or specified. | -## Local Directory Layout +### Local Directory Layout When using the local directory (Priority 3) or after the resolver API persists a candidate (Priority 2), the directory uses the following structure: @@ -72,7 +62,7 @@ When using the local directory (Priority 3) or after the resolver API persists a └── SKILL.md ``` -## Tool Description Formats +### Tool Description Formats `tools.json` and the inline JSON config support three formats: @@ -107,7 +97,7 @@ When using the local directory (Priority 3) or after the resolver API persists a ] ``` -## OptimizationConfig Properties +### OptimizationConfig Properties | Property | Type | Description | |----------|------|-------------| @@ -123,6 +113,41 @@ When using the local directory (Priority 3) or after the resolver API persists a | `has_skills` | `bool` | Whether skills are available. | | `has_tool_descriptions` | `bool` | Whether tool descriptions are available. | +## Examples + +```python +from azure.ai.agentserver.optimization import load_config + +config = load_config(default_instructions="You are a helpful assistant.") + +# Use config in your agent +print(config.instructions) # optimized or default +print(config.model) # optimized or default +print(config.temperature) # optimized or default +print(config.skills) # learned skills (empty list if none) +print(config.tool_descriptions) # optimized tool descriptions (empty dict if none) +print(config.source) # "env:OPTIMIZATION_CONFIG", "api:candidate:abc", "local:...", or "defaults" +``` + +## 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 all three env vars are set: `OPTIMIZATION_CANDIDATE_ID`, `OPTIMIZATION_JOB_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()` returns defaults unexpectedly** — check logs for warnings about path traversal, bad JSON, or missing directories. + +## 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/ai/agentserver/optimization/__init__.py b/sdk/agentserver/azure-ai-agentserver-optimization/azure/ai/agentserver/optimization/__init__.py index 0288fa3806f9..bf41412c0809 100644 --- 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 @@ -9,12 +9,12 @@ from azure.ai.agentserver.optimization import load_config config = load_config(default_instructions="You are a helpful assistant.") - # config.instructions — optimized or default - # config.model — optimized or default - # config.temperature — optimized or default - # config.skills — learned skills (empty if none) - # config.tool_definitions — optimized tool definitions (empty if none) - # config.source — "api:candidate:abc", "env:config", or "defaults" + # config.instructions — optimized or default + # config.model — optimized or default + # config.temperature — optimized or default + # config.skills — learned skills (empty list if none) + # config.tool_descriptions — optimized tool descriptions (empty dict if none) + # config.source — "env:OPTIMIZATION_CONFIG", "api:candidate:abc", "local:", or "defaults" Resolution order (first match wins): 1. OPTIMIZATION_CONFIG env var → inline JSON (used by temp agent versions) 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 index 332f4db1aaeb..7eed6c485b96 100644 --- 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 @@ -185,13 +185,29 @@ def _resolve_local_dir() -> Path: (``".agent_configs"``) when the env var is not set. """ local_dir_env = os.environ.get(OptimizationConfig.ENV_LOCAL_DIR, "").strip() - local_dir = Path(local_dir_env) if local_dir_env else Path(OptimizationConfig.DEFAULT_LOCAL_DIR) + explicitly_set = bool(local_dir_env) + local_dir = Path(local_dir_env) if explicitly_set else Path(OptimizationConfig.DEFAULT_LOCAL_DIR) + + # Guard: reject paths with ".." components (path traversal) + if ".." in local_dir.parts: + logger.warning( + "OPTIMIZATION_LOCAL_DIR contains '..' path traversal: %r — ignoring", + local_dir_env, + ) + local_dir = Path(OptimizationConfig.DEFAULT_LOCAL_DIR) + explicitly_set = False + 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 @@ -249,20 +265,33 @@ def _load_candidate_from_metadata( meta = MetadataConfig.from_dict(raw) - # Read instructions from the referenced file + # Read instructions from the referenced file (guard against traversal) instructions_path = candidate_path / meta.instruction_file - if instructions_path.is_file(): + if _is_safe_child(candidate_path, instructions_path) and instructions_path.is_file(): instructions = instructions_path.read_text(encoding="utf-8").strip() else: + if not _is_safe_child(candidate_path, instructions_path): + logger.warning("Path traversal in instruction_file: %r", meta.instruction_file) instructions = default_instructions - # Resolve skills directory - skills_path = (candidate_path / meta.skill_dir).resolve() - skills = _load_skills_from_dir(skills_path) if skills_path.is_dir() else [] - skills_dir = str(skills_path) if skills_path.is_dir() else default_skills_dir - - # Load tool descriptions - tool_descriptions = _load_tool_descriptions(candidate_path / meta.tool_file) + # Resolve skills directory (guard against traversal) + skills_path = candidate_path / meta.skill_dir + if _is_safe_child(candidate_path, skills_path) and skills_path.resolve().is_dir(): + skills = _load_skills_from_dir(skills_path.resolve()) + skills_dir = str(skills_path.resolve()) + else: + if not _is_safe_child(candidate_path, skills_path): + logger.warning("Path traversal in skill_dir: %r", meta.skill_dir) + skills = [] + skills_dir = default_skills_dir + + # Load tool descriptions (guard against traversal) + tool_file_path = candidate_path / meta.tool_file + if _is_safe_child(candidate_path, tool_file_path): + tool_descriptions = _load_tool_descriptions(tool_file_path) + else: + logger.warning("Path traversal in tool_file: %r", meta.tool_file) + tool_descriptions = {} return OptimizationConfig( instructions=instructions, @@ -301,7 +330,11 @@ def _load_tool_descriptions(tool_file: Path) -> dict[str, ToolDescription]: def _parse_simple_yaml(path: Path) -> dict: - """Minimal key: value parser for metadata.yaml when PyYAML is not installed.""" + """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. + """ result: dict = {} try: for line in path.read_text(encoding="utf-8").splitlines(): @@ -310,12 +343,31 @@ def _parse_simple_yaml(path: Path) -> dict: continue if ":" in line: key, _, value = line.partition(":") - result[key.strip()] = value.strip() + 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.""" + 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. @@ -324,12 +376,24 @@ def _resolve_candidate_folder(local_dir: Path, candidate_id: str | None) -> Path """ if candidate_id: exact = local_dir / candidate_id + if not _is_safe_child(local_dir, exact): + logger.warning("Path traversal detected in candidate_id: %r", candidate_id) + return None if exact.is_dir(): return exact baseline = local_dir / OptimizationConfig.BASELINE_DIR return baseline if baseline.is_dir() else None +def _is_safe_child(parent: Path, child: Path) -> bool: + """Return True if *child* is strictly inside *parent* (no traversal).""" + try: + child.resolve().relative_to(parent.resolve()) + return True + except ValueError: + return False + + def _load_skills_from_dir(skills_dir: Path) -> list[Skill]: """Load skills from a directory of skill folders. 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 index 0a9b92420573..d81008159136 100644 --- 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 @@ -9,7 +9,9 @@ / └── / - ├── config.json + ├── metadata.yaml + ├── instructions.md + ├── tools.json └── skills/ └── / └── SKILL.md @@ -19,15 +21,14 @@ import json import logging -import os import pathlib import shutil import urllib.error import urllib.parse import urllib.request from typing import Any -from azure.ai.agentserver.optimization._models import ( - OptimizationConfig) + +from azure.ai.agentserver.optimization._models import OptimizationConfig logger = logging.getLogger("azure.ai.agentserver.optimization") @@ -46,6 +47,13 @@ def resolve_candidate( following the standard local directory layout. Returns ``None`` if the call fails. """ + # Guard against path traversal in candidate_id + if local_dir is not None: + candidate_path_check = (local_dir / candidate_id).resolve() + if not str(candidate_path_check).startswith(str(local_dir.resolve())): + logger.error("Path traversal detected in candidate_id: %r — aborting", candidate_id) + return None + if candidate_id in _downloaded: if local_dir is not None and (local_dir / candidate_id).is_dir(): logger.debug("Candidate %s already downloaded — skipping", candidate_id) @@ -123,11 +131,11 @@ def _persist_to_local_layout(candidate_path: pathlib.Path, config: dict[str, Any instr_file.write_text(instructions, encoding="utf-8") # tools.json — write tool_descriptions / toolDescriptions as dict format - tool_descs = config.get("tool_descriptions") or config.get("toolDescriptions") + tool_data = config.get("tool_descriptions") or config.get("toolDescriptions") tools_list = config.get("tools") - if tool_descs and isinstance(tool_descs, dict): + if tool_data and isinstance(tool_data, dict): tools_file = candidate_path / OptimizationConfig.TOOLS_FILE - tools_file.write_text(json.dumps(tool_descs, indent=2, ensure_ascii=False), encoding="utf-8") + tools_file.write_text(json.dumps(tool_data, indent=2, ensure_ascii=False), encoding="utf-8") elif 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") @@ -180,7 +188,11 @@ def _download_skill_files( if rel_path.startswith(prefix): rel_path = rel_path[len(prefix):] - out_path = skills_dir / rel_path + 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)) 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..4fa70e4507c8 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/cspell.json @@ -0,0 +1,11 @@ +{ + "ignoreWords": [ + "agentserver", + "autouse", + "cand" + ], + "ignorePaths": [ + "*.json", + "*.rst" + ] +} diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/tests/conftest.py b/sdk/agentserver/azure-ai-agentserver-optimization/tests/conftest.py index 280cfc794766..b904d3574863 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/tests/conftest.py +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/conftest.py @@ -21,3 +21,4 @@ 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 index fe4d7587b671..527e139e5038 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py @@ -4,7 +4,6 @@ """Tests for load_config — priority resolution, fallback, and edge cases.""" import json -from unittest.mock import patch import pytest @@ -455,6 +454,48 @@ def test_no_candidate_no_baseline_falls_to_defaults(self, monkeypatch, tmp_path) config = load_config(default_instructions="default") assert config.source == "defaults" + def test_metadata_traversal_instruction_file(self, monkeypatch, tmp_path): + """instruction_file with '../' in metadata.yaml is rejected.""" + candidate_dir = tmp_path / "baseline" + candidate_dir.mkdir() + (candidate_dir / "metadata.yaml").write_text( + "instruction_file: ../../etc/passwd\n" + ) + # Create the traversal target to prove it's NOT read + secret = tmp_path / "secret.txt" + secret.write_text("SECRET DATA") + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config(default_instructions="safe default") + assert config.instructions == "safe default" + assert "SECRET" not in config.instructions + + def test_metadata_traversal_skill_dir(self, monkeypatch, tmp_path): + """skill_dir with '../' in metadata.yaml is rejected.""" + candidate_dir = tmp_path / "baseline" + candidate_dir.mkdir() + (candidate_dir / "metadata.yaml").write_text( + "skill_dir: ../../other_skills\n" + ) + (candidate_dir / "instructions.md").write_text("ok") + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert config.skills == [] + + def test_metadata_traversal_tool_file(self, monkeypatch, tmp_path): + """tool_file with '../' in metadata.yaml is rejected.""" + candidate_dir = tmp_path / "baseline" + candidate_dir.mkdir() + (candidate_dir / "metadata.yaml").write_text( + "tool_file: ../../secrets.json\n" + ) + (candidate_dir / "instructions.md").write_text("ok") + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config() + assert config.tool_descriptions == {} + # ── _resolve_local_dir ────────────────────────────────────────────── @@ -472,6 +513,19 @@ def test_uses_env_var(self, monkeypatch, tmp_path): local_dir = _resolve_local_dir() assert local_dir == tmp_path + def test_rejects_dotdot_traversal(self, monkeypatch): + """OPTIMIZATION_LOCAL_DIR with '..' is rejected, falls back to default.""" + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", "../../etc/sensitive") + local_dir = _resolve_local_dir() + assert local_dir.name == ".agent_configs" + + def test_rejects_dotdot_in_absolute_path(self, monkeypatch, tmp_path): + """Even absolute paths with '..' are rejected.""" + malicious = str(tmp_path / ".." / ".." / "etc") + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", malicious) + local_dir = _resolve_local_dir() + assert local_dir.name == ".agent_configs" + # ── _resolve_candidate_folder ─────────────────────────────────────── @@ -502,6 +556,16 @@ def test_returns_none_no_id_no_baseline(self, tmp_path): result = _resolve_candidate_folder(tmp_path, None) assert result is None + def test_rejects_traversal_candidate_id(self, tmp_path): + """candidate_id with '../' is rejected as path traversal.""" + result = _resolve_candidate_folder(tmp_path, "../../etc") + assert result is None + + def test_rejects_absolute_candidate_id(self, tmp_path): + """Absolute path in candidate_id is rejected.""" + result = _resolve_candidate_folder(tmp_path, "/etc/passwd") + assert result is None + # ── Graceful error handling ───────────────────────────────────────── @@ -1009,7 +1073,8 @@ def test_basic_parsing(self, tmp_path): 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 result["temperature"] == 0.5 + assert isinstance(result["temperature"], float) def test_skips_comments_and_blanks(self, tmp_path): f = tmp_path / "test.yaml" @@ -1026,3 +1091,25 @@ def test_colon_in_value(self, tmp_path): 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 diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py index e41751e45e87..f55ba2ebdfa5 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py @@ -277,13 +277,6 @@ def test_downloads_skill_files(self, tmp_path): ] } - def mock_get(url, headers, **kwargs): - if "config" in url: - return None - if "files" in url: - return None - return manifest - def mock_text(url, headers, params=None): return "# Math Skill\nDo math." @@ -341,6 +334,53 @@ def test_handles_download_failure(self, tmp_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): + """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(ENDPOINT, JOB_ID, "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() + + +# ── Path traversal in resolve_candidate ───────────────────────────── + + +class TestPathTraversalGuard: + """Tests for path traversal prevention in resolve_candidate.""" + + def test_rejects_traversal_candidate_id(self, tmp_path): + """candidate_id with '../' is rejected before any API call.""" + result = resolve_candidate( + "../../etc", job_id=JOB_ID, endpoint=ENDPOINT, local_dir=tmp_path + ) + assert result is None + + def test_rejects_absolute_candidate_id(self, tmp_path): + """Absolute path in candidate_id is rejected.""" + result = resolve_candidate( + "/etc/passwd", job_id=JOB_ID, endpoint=ENDPOINT, local_dir=tmp_path + ) + assert result is None + + 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._api_get_json", + return_value=config, + ): + result = resolve_candidate( + "valid-candidate-123", job_id=JOB_ID, endpoint=ENDPOINT, local_dir=tmp_path + ) + assert result is not None + # ── _is_skill_file ────────────────────────────────────────────────── From 7f5c298012aed11089e9235ebfd2985ad84ef7a3 Mon Sep 17 00:00:00 2001 From: zyysurely Date: Sat, 23 May 2026 00:13:18 -0700 Subject: [PATCH 3/5] more checks fix --- .../CHANGELOG.md | 11 +- .../README.md | 52 +++- .../ai/agentserver/optimization/__init__.py | 3 +- .../ai/agentserver/optimization/_config.py | 10 +- .../ai/agentserver/optimization/_models.py | 56 +++- .../ai/agentserver/optimization/_resolver.py | 135 +++++---- .../pyproject.toml | 1 + .../tests/test_config.py | 262 ++++++++++++++++- .../tests/test_integration.py | 274 ++++++++++++++++++ .../tests/test_resolver.py | 167 ++++++----- 10 files changed, 825 insertions(+), 146 deletions(-) create mode 100644 sdk/agentserver/azure-ai-agentserver-optimization/tests/test_integration.py diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md index bb104911e6f6..41d008fd7d11 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md @@ -1,12 +1,19 @@ # Release History -## 1.0.0b1 +## 1.0.0b1 (2026-05-22) ### Features Added - Initial beta release. - `load_config()` — single-call config loader with 4-priority resolution and graceful fallback (never crashes). -- `OptimizationConfig` dataclass with instructions, model, temperature, skills, tool descriptions, source tracking, candidate_id, and job_id. +- `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_dir, tool descriptions, source tracking, candidate_id, and job_id. +- `OptimizationConfig.apply_tool_descriptions(tools)` — patch `__doc__` on @tool-decorated functions from optimized descriptions. +- `OptimizationConfig.compose_instructions()` — append skill catalog to instructions. +- `OptimizationConfig.get_tool_description(name)` / `get_tool_param_description(name, param)` — look up individual optimized descriptions. +- `CandidateConfig` — typed representation of the resolver API payload. +- `Skill` — learned skill model (name, description, body). +- `ToolDescription` — optimized tool description model (description, parameters). - 4-priority resolution order: 1. Inline JSON via `OPTIMIZATION_CONFIG` env var. 2. Resolver API via `OPTIMIZATION_CANDIDATE_ID` + `OPTIMIZATION_JOB_ID` + `OPTIMIZATION_RESOLVE_ENDPOINT`. diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/README.md b/sdk/agentserver/azure-ai-agentserver-optimization/README.md index f84c5970886d..862e3343c72a 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/README.md +++ b/sdk/agentserver/azure-ai-agentserver-optimization/README.md @@ -113,22 +113,70 @@ When using the local directory (Priority 3) or after the resolver API persists a | `has_skills` | `bool` | Whether skills are available. | | `has_tool_descriptions` | `bool` | Whether tool descriptions are available. | +### Public API + +| Export | Type | Description | +|--------|------|-------------| +| `load_config(...)` | 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_descriptions. | +| `OptimizationConfig.apply_tool_descriptions(tools)` | method | Patch `__doc__` on tool functions from optimized descriptions. | +| `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). | +| `ToolDescription` | dataclass | Optimized tool description (description, parameters). | + ## Examples +### Basic usage + ```python from azure.ai.agentserver.optimization import load_config config = load_config(default_instructions="You are a helpful assistant.") -# Use config in your agent print(config.instructions) # optimized or default print(config.model) # optimized or default print(config.temperature) # optimized or default -print(config.skills) # learned skills (empty list if none) print(config.tool_descriptions) # optimized tool descriptions (empty dict if none) print(config.source) # "env:OPTIMIZATION_CONFIG", "api:candidate:abc", "local:...", or "defaults" ``` +### Apply optimized tool descriptions + +```python +from azure.ai.agentserver.optimization import load_config + +config = load_config(default_instructions="You are a travel agent.") + +# 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__ 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(default_instructions="You are a helpful assistant.") + +# 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: 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 index bf41412c0809..49c75199a7ce 100644 --- 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 @@ -23,7 +23,7 @@ 4. Defaults → your hardcoded values (agent works normally) """ -from azure.ai.agentserver.optimization._config import load_config +from azure.ai.agentserver.optimization._config import load_config, load_skills_from_dir from azure.ai.agentserver.optimization._models import ( CandidateConfig, OptimizationConfig, @@ -38,5 +38,6 @@ "Skill", "ToolDescription", "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 index 7eed6c485b96..e93fa9f36140 100644 --- 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 @@ -83,7 +83,7 @@ def load_config( default_temperature=default_temperature, default_skills_dir=default_skills_dir, ) - except Exception as exc: # noqa: BLE001 + except Exception as exc: # noqa: BLE001 # pylint: disable=broad-exception-caught logger.error("Unexpected error loading optimization config — returning defaults: %s", exc) model = default_model or os.environ.get("MODEL_DEPLOYMENT_NAME") return OptimizationConfig( @@ -181,7 +181,7 @@ def _load_config_inner( def _resolve_local_dir() -> Path: """Resolve the local optimization directory path. - Falls back to :pyattr:`OptimizationConfig.DEFAULT_LOCAL_DIR` + Falls back to ``OptimizationConfig.DEFAULT_LOCAL_DIR`` (``".agent_configs"``) when the env var is not set. """ local_dir_env = os.environ.get(OptimizationConfig.ENV_LOCAL_DIR, "").strip() @@ -275,14 +275,13 @@ def _load_candidate_from_metadata( instructions = default_instructions # Resolve skills directory (guard against traversal) + skills_dir: str | None skills_path = candidate_path / meta.skill_dir if _is_safe_child(candidate_path, skills_path) and skills_path.resolve().is_dir(): - skills = _load_skills_from_dir(skills_path.resolve()) skills_dir = str(skills_path.resolve()) else: if not _is_safe_child(candidate_path, skills_path): logger.warning("Path traversal in skill_dir: %r", meta.skill_dir) - skills = [] skills_dir = default_skills_dir # Load tool descriptions (guard against traversal) @@ -297,7 +296,6 @@ def _load_candidate_from_metadata( instructions=instructions, model=meta.model or default_model, temperature=meta.temperature if meta.temperature is not None else default_temperature, - skills=skills, skills_dir=skills_dir, tool_descriptions=tool_descriptions, source=f"local:{candidate_path}", @@ -394,7 +392,7 @@ def _is_safe_child(parent: Path, child: Path) -> bool: return False -def _load_skills_from_dir(skills_dir: Path) -> list[Skill]: +def load_skills_from_dir(skills_dir: Path) -> list[Skill]: """Load skills from a directory of skill folders. Expected layout:: 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 index f88b4eb82e1b..8ae4b612d601 100644 --- 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 @@ -6,9 +6,12 @@ from __future__ import annotations +import logging from dataclasses import dataclass, field, fields from typing import Any, ClassVar +logger = logging.getLogger(__name__) + @dataclass class Skill: @@ -116,7 +119,7 @@ def from_dict(cls, data: dict) -> MetadataConfig: @dataclass -class OptimizationConfig: +class OptimizationConfig: # pylint: disable=too-many-instance-attributes """Resolved optimization config. When not running under optimization, all fields contain the defaults @@ -166,6 +169,47 @@ def get_tool_param_description(self, tool_name: str, param_name: str) -> str | N return None return td.parameters.get(param_name) + def apply_tool_descriptions(self, tools: list) -> list: # pylint: disable=too-many-nested-blocks + """Apply optimized tool descriptions to a list of tool functions. + + Patches ``__doc__`` (used by the Agent Framework as the tool + description) on each tool function whose name appears in + :attr:`tool_descriptions`. + + Args: + tools: List of @tool-decorated functions. + + Returns: + The same list of tools (mutated in place). + """ + if not self.tool_descriptions: + return tools + for tool_fn in tools: + tool_name = getattr(tool_fn, "__name__", None) or getattr(tool_fn, "name", None) + if tool_name and tool_name in self.tool_descriptions: + overrides = self.tool_descriptions[tool_name] + if overrides.description: + # Patch .description (AIFunction/ToolProtocol) and __doc__ (plain functions) + try: + tool_fn.description = overrides.description + except AttributeError: + pass + tool_fn.__doc__ = overrides.description + logger.debug("Applied optimized description for tool '%s'", tool_name) + # Patch parameter descriptions on AIFunction's input_model + if overrides.parameters: + input_model = getattr(tool_fn, "input_model", None) + if input_model and hasattr(input_model, "model_fields"): + patched = False + for param_name, param_desc in overrides.parameters.items(): + if param_name in input_model.model_fields: + input_model.model_fields[param_name].description = param_desc + patched = True + if patched: + input_model.model_rebuild(force=True) + logger.debug("Applied optimized parameter descriptions for tool '%s'", tool_name) + return tools + def compose_instructions(self) -> str: """Return instructions with skill catalog appended (if any).""" if not self.skills: @@ -239,12 +283,12 @@ def _parse_tools_list(tools: list) -> dict[str, ToolDescription]: continue description = func.get("description", "") params_schema = func.get("parameters", {}) - param_descs: dict[str, str] = {} + param_descriptions: dict[str, str] = {} if isinstance(params_schema, dict): props = params_schema.get("properties", {}) if isinstance(props, dict): - for pname, pval in props.items(): - if isinstance(pval, dict) and "description" in pval: - param_descs[pname] = pval["description"] - result[name] = ToolDescription(description=description, parameters=param_descs) + for param_name, param_val in props.items(): + if isinstance(param_val, dict) and "description" in param_val: + param_descriptions[param_name] = param_val["description"] + result[name] = ToolDescription(description=description, parameters=param_descriptions) return result 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 index d81008159136..6aa120ba9109 100644 --- 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 @@ -23,17 +23,23 @@ import logging import pathlib import shutil -import urllib.error -import urllib.parse -import urllib.request 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" +_JOBS_PATH = "agent_optimization_jobs" +_AUTH_SCOPE = "https://ai.azure.com/.default" + def resolve_candidate( candidate_id: str, @@ -61,11 +67,16 @@ def resolve_candidate( logger.warning("Candidate %s was downloaded but folder is missing — re-downloading", candidate_id) _downloaded.discard(candidate_id) - headers = _build_headers() + client = _build_client(endpoint) # ── Step 1: Fetch config ───────────────────────────────────────── - config = _api_get_json(f"{endpoint}/agent_optimization_jobs/{job_id}/candidates/{candidate_id}/config", headers) + config = _api_get_json( + client, + f"/{_JOBS_PATH}/{job_id}/candidates/{candidate_id}/config", + params={"api-version": _API_VERSION}, + ) if config is None: + client.close() return None logger.info( @@ -82,7 +93,7 @@ def resolve_candidate( candidate_path = local_dir / candidate_id try: _persist_to_local_layout(candidate_path, config) - _download_skill_files(endpoint, job_id, candidate_id, headers, candidate_path) + _download_skill_files(client, job_id, 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 @@ -90,6 +101,7 @@ def resolve_candidate( if skills_path.is_dir(): config["skills_dir"] = str(skills_path) + client.close() _downloaded.add(candidate_id) return config @@ -102,7 +114,10 @@ def _persist_to_local_layout(candidate_path: pathlib.Path, config: dict[str, Any / ├── metadata.yaml ├── instructions.md - └── tools.json + ├── tools.json + └── skills/ + └── / + └── SKILL.md If the folder already exists it is removed and re-created. """ @@ -140,18 +155,44 @@ def _persist_to_local_layout(candidate_path: pathlib.Path, config: dict[str, Any 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( - endpoint: str, + client: PipelineClient, job_id: str, candidate_id: str, - headers: dict[str, str], candidate_path: pathlib.Path, ) -> None: """Fetch manifest and download skill files into candidate_path/skills//SKILL.md.""" - manifest = _api_get_json(f"{endpoint}/agent_optimization_jobs/{job_id}/candidates/{candidate_id}", headers) + manifest = _api_get_json( + client, + f"/{_JOBS_PATH}/{job_id}/candidates/{candidate_id}", + params={"api-version": _API_VERSION}, + ) if manifest is None: logger.debug("Could not fetch manifest for candidate %s", candidate_id) return @@ -174,9 +215,9 @@ def _download_skill_files( continue content = _api_get_text( - f"{endpoint}/agent_optimization_jobs/{job_id}/candidates/{candidate_id}/files", - headers, - params={"path": file_path}, + client, + f"/{_JOBS_PATH}/{job_id}/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) @@ -205,57 +246,53 @@ def _is_skill_file(file_entry: dict) -> bool: return file_type == "skill" or path.startswith("skills/") -# ── HTTP helpers ───────────────────────────────────────────────────── +# ── HTTP helpers (azure.core transport) ────────────────────────────── + +def _build_client(endpoint: str) -> PipelineClient: + """Create a PipelineClient with credential-based auth and retry.""" + policies: list = [RetryPolicy()] + try: + from azure.identity import DefaultAzureCredential -def _build_headers() -> dict[str, str]: - headers: dict[str, str] = {"Accept": "application/json"} - token = _get_bearer_token() - if token: - headers["Authorization"] = f"Bearer {token}" - return headers + 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(url: str, headers: dict[str, str]) -> dict[str, Any] | None: +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.""" + url = f"{client._base_url.rstrip('/')}{path}" # pylint: disable=protected-access + request = HttpRequest("GET", url, params=params) logger.debug("GET %s", url) try: - req = urllib.request.Request(url, method="GET", headers=headers) - with urllib.request.urlopen(req, timeout=30) as resp: # noqa: S310 - return json.loads(resp.read().decode("utf-8")) - except (urllib.error.URLError, json.JSONDecodeError, OSError) as exc: + 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( - url: str, headers: dict[str, str], params: dict[str, str] | None = None + client: PipelineClient, path: str, params: dict[str, str] | None = None ) -> str | None: """GET an endpoint, return response body as text or None on failure.""" - if params: - query = "&".join(f"{k}={urllib.parse.quote(v)}" for k, v in params.items()) - url = f"{url}?{query}" + url = f"{client._base_url.rstrip('/')}{path}" # pylint: disable=protected-access + request = HttpRequest("GET", url, params=params) logger.debug("GET %s", url) try: - req = urllib.request.Request(url, method="GET", headers=headers) - with urllib.request.urlopen(req, timeout=60) as resp: # noqa: S310 - return resp.read().decode("utf-8") - except (urllib.error.URLError, OSError) as exc: + 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 - - -def _get_bearer_token() -> str | None: - """Acquire a bearer token for the resolver API. - - Uses ``azure-identity`` if available; returns ``None`` otherwise. - This keeps azure-identity as an optional dependency. - """ - try: - from azure.identity import DefaultAzureCredential # type: ignore[import-untyped] - - cred = DefaultAzureCredential() - token = cred.get_token("https://ai.azure.com/.default") - return token.token - except Exception: # noqa: BLE001 - return None diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/pyproject.toml b/sdk/agentserver/azure-ai-agentserver-optimization/pyproject.toml index e11feb8656ba..13e7aaec3ca5 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/pyproject.toml +++ b/sdk/agentserver/azure-ai-agentserver-optimization/pyproject.toml @@ -22,6 +22,7 @@ keywords = ["azure", "azure sdk", "agent", "agentserver", "optimization"] dependencies = [ "pyyaml>=6.0", + "azure-core>=1.31.0", "azure-identity>=1.15.0", ] diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py index 527e139e5038..7623feffe816 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py @@ -388,10 +388,16 @@ def test_loads_skills_from_local_dir(self, monkeypatch, tmp_path): monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) config = load_config() - assert len(config.skills) == 1 - assert config.skills[0].name == "math" - assert config.skills[0].description == "Do math" - assert config.skills[0].body == "Body here." + 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" @@ -1113,3 +1119,251 @@ def test_coerces_bool(self, tmp_path): 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.""" + + def _make_config(self, tool_descriptions=None): + return OptimizationConfig( + instructions="test", + model=None, + temperature=None, + tool_descriptions=tool_descriptions or {}, + ) + + def test_patches_docstring(self): + def lookup_policy(): + """Original description.""" + + config = self._make_config( + {"lookup_policy": ToolDescription(description="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( + {"my_tool": ToolDescription(description="New.")} + ) + tools = [my_tool] + result = config.apply_tool_descriptions(tools) + assert result is tools + + def test_skips_tools_not_in_descriptions(self): + def unknown_tool(): + """Should not change.""" + + config = self._make_config( + {"other_tool": ToolDescription(description="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( + {"my_tool": ToolDescription(description="")} + ) + config.apply_tool_descriptions([my_tool]) + assert my_tool.__doc__ == "Original doc." + + def test_noop_when_no_tool_descriptions(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( + {"something": ToolDescription(description="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({ + "tool_a": ToolDescription(description="A optimized."), + "tool_c": ToolDescription(description="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( + {"my_tool": ToolDescription(description="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( + {"search_flights": ToolDescription(description="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( + {"my_tool": ToolDescription(description="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({ + "search_flights": ToolDescription( + description="Find flights.", + parameters={"destination": "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({ + "my_tool": ToolDescription( + description="New.", + parameters={"x": "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({ + "my_tool": ToolDescription( + description="New doc.", + parameters={"unknown_param": "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({ + "my_tool": ToolDescription(description="New.", parameters={}), + }) + 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( + {"tool": ToolDescription(description="X.")} + ) + result = config.apply_tool_descriptions([]) + assert result == [] 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..05dda5114af5 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_integration.py @@ -0,0 +1,274 @@ +# --------------------------------------------------------- +# 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, + ToolDescription, + 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, + "tool_descriptions": { + "search_flights": { + "description": "Find the cheapest flight options.", + "parameters": {"destination": "City name"}, + }, + "book_hotel": { + "description": "Reserve a hotel room.", + "parameters": {}, + }, + }, + } + 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(default_instructions="fallback") + assert config.source == "env:OPTIMIZATION_CONFIG" + assert config.instructions == "Optimized prompt." + assert config.has_tool_descriptions + + 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." + assert config.get_tool_param_description("lookup_policy", "dept") == "Department name" + + +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 = { + "search": {"description": "Search destinations.", "parameters": {"q": "Query"}}, + } + (candidate_dir / "tools.json").write_text(json.dumps(tools_data)) + + monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) + config = load_config(default_instructions="fallback") + + # 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 descriptions also loaded + assert config.has_tool_descriptions + assert config.tool_descriptions["search"].description == "Search destinations." + + 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 = { + "search_flights": {"description": "Find flights between cities.", "parameters": {}}, + "book_flight": {"description": "Book the selected flight.", "parameters": {}}, + } + (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( + default_instructions="Default prompt.", + default_model="gpt-3.5-turbo", + ) + 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_descriptions=config.tool_descriptions, + 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, everything works with defaults.""" + + def my_tool(): + """Original.""" + + config = load_config( + default_instructions="Be helpful.", + default_model="gpt-4o", + default_temperature=0.7, + ) + assert config.source == "defaults" + assert config.instructions == "Be helpful." + assert config.model == "gpt-4o" + assert config.temperature == 0.7 + assert not config.has_tool_descriptions + assert not config.has_skills + assert config.skills_dir is None + + # apply_tool_descriptions is a no-op + config.apply_tool_descriptions([my_tool]) + assert my_tool.__doc__ == "Original." + + # load_skills_from_dir with None skills_dir — user checks before calling + # This is the expected pattern: + if config.skills_dir: + skills = load_skills_from_dir(Path(config.skills_dir)) + else: + skills = [] + assert skills == [] diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py index f55ba2ebdfa5..778d4dcf487c 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py @@ -14,8 +14,7 @@ _persist_to_local_layout, _download_skill_files, _is_skill_file, - _build_headers, - _get_bearer_token, + _build_client, ) from azure.ai.agentserver.optimization._models import OptimizationConfig @@ -28,6 +27,14 @@ def clear_downloaded(): _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" JOB_ID = "job-42" @@ -39,9 +46,12 @@ class TestResolveCandidate: """Tests for resolve_candidate function.""" def test_returns_none_on_api_failure(self): - with patch( - "azure.ai.agentserver.optimization._resolver._api_get_json", - return_value=None, + 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", job_id=JOB_ID, endpoint=ENDPOINT) assert result is None @@ -53,34 +63,43 @@ def test_returns_config_on_success(self): "temperature": 0.2, "skills": [], } - with patch( - "azure.ai.agentserver.optimization._resolver._api_get_json", - return_value=config, + 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", job_id=JOB_ID, endpoint=ENDPOINT) assert result is not None assert result["instructions"] == "Optimized." assert result["model"] == "gpt-4o" - def test_uses_correct_url(self): - """Verify the API route follows agent_optimization_jobs/{jobId}/candidates/{candidateId}/config.""" - called_urls: list[str] = [] + def test_uses_correct_path(self): + """Verify the API route follows /agent_optimization_jobs/{jobId}/candidates/{candidateId}/config.""" + called_args: list = [] - def capture_url(url, headers): - called_urls.append(url) + def capture_call(client, path, params=None): + called_args.append((path, params)) return {"instructions": "ok"} - with patch( - "azure.ai.agentserver.optimization._resolver._api_get_json", - side_effect=capture_url, + 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", job_id="job-xyz", endpoint="http://api.test") - assert called_urls[0] == "http://api.test/agent_optimization_jobs/job-xyz/candidates/cand-abc/config" + assert called_args[0][0] == "/agent_optimization_jobs/job-xyz/candidates/cand-abc/config" def test_marks_downloaded_after_success(self): - with patch( - "azure.ai.agentserver.optimization._resolver._api_get_json", - return_value={"instructions": "ok"}, + 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", job_id=JOB_ID, endpoint=ENDPOINT) assert "cand-mark" in _downloaded @@ -99,9 +118,12 @@ 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._api_get_json", - return_value=config, + 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", job_id=JOB_ID, endpoint=ENDPOINT, local_dir=None, @@ -111,9 +133,12 @@ def test_redownloads_if_folder_missing(self): assert result["instructions"] == "re-downloaded" def test_does_not_mark_downloaded_on_api_failure(self): - with patch( - "azure.ai.agentserver.optimization._resolver._api_get_json", - return_value=None, + 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", job_id=JOB_ID, endpoint=ENDPOINT) assert "cand-fail" not in _downloaded @@ -268,7 +293,7 @@ def test_round_trip(self, monkeypatch, tmp_path): class TestDownloadSkillFiles: """Tests for _download_skill_files.""" - def test_downloads_skill_files(self, tmp_path): + def test_downloads_skill_files(self, tmp_path, mock_client): candidate_path = tmp_path / "cand-sk" candidate_path.mkdir() manifest = { @@ -277,30 +302,33 @@ def test_downloads_skill_files(self, tmp_path): ] } - def mock_text(url, headers, params=None): + 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=lambda u, h: manifest), + 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(ENDPOINT, JOB_ID, "cand-sk", {}, candidate_path) + _download_skill_files(mock_client, JOB_ID, "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): + 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(ENDPOINT, JOB_ID, "cand-no-manifest", {}, candidate_path) + _download_skill_files(mock_client, JOB_ID, "cand-no-manifest", candidate_path) assert not (candidate_path / "skills").exists() - def test_skips_when_no_skill_files_in_manifest(self, tmp_path): + 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"}]} @@ -308,10 +336,10 @@ def test_skips_when_no_skill_files_in_manifest(self, tmp_path): "azure.ai.agentserver.optimization._resolver._api_get_json", return_value=manifest, ): - _download_skill_files(ENDPOINT, JOB_ID, "cand-no-skills", {}, candidate_path) + _download_skill_files(mock_client, JOB_ID, "cand-no-skills", candidate_path) assert not (candidate_path / "skills").exists() - def test_skips_empty_path_entries(self, tmp_path): + 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"}]} @@ -319,10 +347,10 @@ def test_skips_empty_path_entries(self, tmp_path): "azure.ai.agentserver.optimization._resolver._api_get_json", return_value=manifest, ): - _download_skill_files(ENDPOINT, JOB_ID, "cand-empty-path", {}, candidate_path) + _download_skill_files(mock_client, JOB_ID, "cand-empty-path", candidate_path) assert not (candidate_path / "skills").exists() - def test_handles_download_failure(self, tmp_path): + 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"}]} @@ -330,11 +358,11 @@ def test_handles_download_failure(self, tmp_path): 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(ENDPOINT, JOB_ID, "cand-dl-fail", {}, candidate_path) + _download_skill_files(mock_client, JOB_ID, "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): + 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() @@ -343,7 +371,7 @@ def test_rejects_traversal_in_file_path(self, tmp_path): 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(ENDPOINT, JOB_ID, "cand-traversal", {}, candidate_path) + _download_skill_files(mock_client, JOB_ID, "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() @@ -372,9 +400,12 @@ def test_rejects_absolute_candidate_id(self, tmp_path): 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._api_get_json", - return_value=config, + 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", job_id=JOB_ID, endpoint=ENDPOINT, local_dir=tmp_path @@ -410,6 +441,7 @@ class TestPersistErrorHandling: 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, @@ -431,41 +463,24 @@ def test_persist_oserror_does_not_crash(self, tmp_path): # ── HTTP helpers ───────────────────────────────────────────────────── -class TestBuildHeaders: - """Tests for _build_headers.""" +class TestBuildClient: + """Tests for _build_client.""" - def test_includes_accept_header(self): - headers = _build_headers() - assert headers["Accept"] == "application/json" + def test_returns_pipeline_client(self): + from azure.core import PipelineClient - def test_includes_auth_when_token_available(self): - with patch( - "azure.ai.agentserver.optimization._resolver._get_bearer_token", - return_value="fake-token", - ): - headers = _build_headers() - assert headers["Authorization"] == "Bearer fake-token" + 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 - def test_no_auth_when_no_token(self): with patch( - "azure.ai.agentserver.optimization._resolver._get_bearer_token", - return_value=None, + "azure.identity.DefaultAzureCredential", + side_effect=Exception("No cred"), ): - headers = _build_headers() - assert "Authorization" not in headers - - -class TestGetBearerToken: - """Tests for _get_bearer_token.""" - - def test_returns_none_without_azure_identity(self): - with patch.dict("sys.modules", {"azure.identity": None}): - token = _get_bearer_token() - assert token is None or isinstance(token, str) - - def test_returns_none_on_exception(self): - mock_identity = MagicMock() - mock_identity.DefaultAzureCredential.side_effect = Exception("No cred") - with patch.dict("sys.modules", {"azure.identity": mock_identity}): - token = _get_bearer_token() - assert token is None + client = _build_client("http://example.com") + assert isinstance(client, PipelineClient) + client.close() From 9c078e0b6db09642b9428398f8644aa9331d2208 Mon Sep 17 00:00:00 2001 From: zyysurely Date: Sun, 24 May 2026 16:46:36 -0700 Subject: [PATCH 4/5] depend on scaffolder --- .../CHANGELOG.md | 22 +- .../README.md | 111 +- .../ai/agentserver/optimization/__init__.py | 18 +- .../ai/agentserver/optimization/_config.py | 326 +++--- .../ai/agentserver/optimization/_models.py | 258 ++--- .../ai/agentserver/optimization/_resolver.py | 133 ++- .../cspell.json | 7 +- .../dev_requirements.txt | 4 + .../tests/conftest.py | 1 - .../tests/test_config.py | 957 +++++++++--------- .../tests/test_integration.py | 306 +++++- .../tests/test_resolver.py | 311 +++++- 12 files changed, 1459 insertions(+), 995 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md index 41d008fd7d11..39e3ccde1c36 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-optimization/CHANGELOG.md @@ -1,26 +1,24 @@ # Release History -## 1.0.0b1 (2026-05-22) +## 1.0.0b1 (2026-05-24) ### Features Added - Initial beta release. -- `load_config()` — single-call config loader with 4-priority resolution and graceful fallback (never crashes). +- `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_dir, tool descriptions, source tracking, candidate_id, and job_id. -- `OptimizationConfig.apply_tool_descriptions(tools)` — patch `__doc__` on @tool-decorated functions from optimized descriptions. +- `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. -- `OptimizationConfig.get_tool_description(name)` / `get_tool_param_description(name, param)` — look up individual optimized descriptions. - `CandidateConfig` — typed representation of the resolver API payload. - `Skill` — learned skill model (name, description, body). -- `ToolDescription` — optimized tool description model (description, parameters). - 4-priority resolution order: 1. Inline JSON via `OPTIMIZATION_CONFIG` env var. - 2. Resolver API via `OPTIMIZATION_CANDIDATE_ID` + `OPTIMIZATION_JOB_ID` + `OPTIMIZATION_RESOLVE_ENDPOINT`. - 3. Local directory layout (`OPTIMIZATION_LOCAL_DIR`, defaults to `.agent_configs/`). - 4. Caller-supplied defaults. + 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. -- 3 tool description formats: `tool_descriptions` dict, `toolDescriptions` (legacy camelCase), and OpenAI function-calling `tools` list. +- Tool definitions use the OpenAI function-calling list format exclusively. - Skill loading from `SKILL.md` files with YAML frontmatter. -- Resolver API persists fetched configs to local directory for offline use. -- Path traversal (zip-slip) protection on all untrusted path inputs. +- 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/README.md b/sdk/agentserver/azure-ai-agentserver-optimization/README.md index 862e3343c72a..9d0204d30072 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/README.md +++ b/sdk/agentserver/azure-ai-agentserver-optimization/README.md @@ -1,6 +1,6 @@ # 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 descriptions) from multiple sources with graceful fallback — your agent works unchanged when not running under optimization. +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 @@ -23,11 +23,11 @@ pip install azure-ai-agentserver-optimization | 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_JOB_ID`, `OPTIMIZATION_RESOLVE_ENDPOINT` | Fetches the candidate config from the remote optimization service and persists it to the local directory. | +| 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 | **Defaults** | *(none)* | Returns the caller-supplied defaults unchanged — the agent works normally. | +| 4 | **No config** | *(none)* | `required=True` (default) raises `ValueError`; `required=False` returns `None`. | -Any unexpected error is caught and logged — `load_config()` always returns a valid `OptimizationConfig`. +Any unexpected error is caught and logged — `load_config()` never crashes (only `ValueError` from `required=True` propagates). ### Environment Variables @@ -35,25 +35,23 @@ Any unexpected error is caught and logged — `load_config()` always returns a v |----------|-------------| | `OPTIMIZATION_CONFIG` | Inline JSON config (Priority 1). | | `OPTIMIZATION_CANDIDATE_ID` | Candidate ID for resolver API or local folder lookup. | -| `OPTIMIZATION_JOB_ID` | Job ID for the resolver API. | -| `OPTIMIZATION_RESOLVE_ENDPOINT` | Base URL of the optimization service. | +| `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/`). | -| `MODEL_DEPLOYMENT_NAME` | Fallback model name when no model is resolved or specified. | ### Local Directory Layout -When using the local directory (Priority 3) or after the resolver API persists a candidate (Priority 2), the directory uses the following structure: +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/ # fallback candidate +├── baseline/ # required — default candidate used at startup │ ├── metadata.yaml # model, temperature, file pointers │ ├── instructions.md # system prompt -│ ├── tools.json # tool descriptions (dict or OpenAI list format) +│ ├── tools.json # tool definitions (OpenAI function-calling list format) │ └── skills/ # learned skills │ └── / │ └── SKILL.md -└── / # same layout as baseline/ +└── / # optional — created by resolver API, same layout as baseline/ ├── metadata.yaml ├── instructions.md ├── tools.json @@ -62,23 +60,31 @@ When using the local directory (Priority 3) or after the resolver API persists a └── SKILL.md ``` -### Tool Description Formats +#### `metadata.yaml` -`tools.json` and the inline JSON config support three formats: +Points to the other files and sets model parameters. All fields are optional: -**Dict format** (`tool_descriptions`): -```json -{ - "lookup_policy": { - "description": "Look up the company travel policy.", - "parameters": {"dept": "Department name"} - } -} +```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. ``` -**Legacy camelCase** (`toolDescriptions`) — same structure, different key. `tool_descriptions` takes priority when both are present. +#### `tools.json` + +Tool definitions in the OpenAI function-calling list format: -**OpenAI function-calling list** (`tools`): ```json [ { @@ -97,34 +103,45 @@ When using the local directory (Priority 3) or after the resolver API persists a ] ``` +#### `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` | System prompt (optimized or default). | +| `instructions` | `str \| None` | System prompt (optimized or default). | | `model` | `str \| None` | Model deployment name. | | `temperature` | `float \| None` | Sampling temperature. | -| `skills` | `list[Skill]` | Learned skills. | -| `skills_dir` | `str \| None` | Path to skills directory. | -| `tool_descriptions` | `dict[str, ToolDescription]` | Optimized tool descriptions. | +| `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). | -| `job_id` | `str \| None` | Job ID (when resolved via API). | -| `has_skills` | `bool` | Whether skills are available. | -| `has_tool_descriptions` | `bool` | Whether tool descriptions are available. | +| `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(...)` | function | Load optimization config with 4-priority resolution. | +| `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_descriptions. | -| `OptimizationConfig.apply_tool_descriptions(tools)` | method | Patch `__doc__` on tool functions from optimized descriptions. | +| `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). | -| `ToolDescription` | dataclass | Optimized tool description (description, parameters). | ## Examples @@ -133,13 +150,15 @@ When using the local directory (Priority 3) or after the resolver API persists a ```python from azure.ai.agentserver.optimization import load_config -config = load_config(default_instructions="You are a helpful assistant.") +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 or default -print(config.model) # optimized or default -print(config.temperature) # optimized or default -print(config.tool_descriptions) # optimized tool descriptions (empty dict if none) -print(config.source) # "env:OPTIMIZATION_CONFIG", "api:candidate:abc", "local:...", or "defaults" +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 @@ -147,7 +166,7 @@ print(config.source) # "env:OPTIMIZATION_CONFIG", "api:candidate:abc ```python from azure.ai.agentserver.optimization import load_config -config = load_config(default_instructions="You are a travel agent.") +config = load_config() # Your @tool-decorated functions def search_flights(origin: str, destination: str): @@ -158,7 +177,7 @@ def book_hotel(city: str): """Book a hotel room.""" ... -# Patches __doc__ on matching tools with optimized descriptions +# Patches __doc__ and .description on matching tools with optimized descriptions config.apply_tool_descriptions([search_flights, book_hotel]) ``` @@ -168,7 +187,7 @@ config.apply_tool_descriptions([search_flights, book_hotel]) from pathlib import Path from azure.ai.agentserver.optimization import load_config, load_skills_from_dir -config = load_config(default_instructions="You are a helpful assistant.") +config = load_config() # Skills are not loaded inline — load them when needed if config.skills_dir: @@ -187,9 +206,9 @@ logging.getLogger("azure.ai.agentserver.optimization").setLevel(logging.DEBUG) ``` Common issues: -- **Config not loading from resolver API** — ensure all three env vars are set: `OPTIMIZATION_CANDIDATE_ID`, `OPTIMIZATION_JOB_ID`, and `OPTIMIZATION_RESOLVE_ENDPOINT`. +- **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()` returns defaults unexpectedly** — check logs for warnings about path traversal, bad JSON, or missing directories. +- **`load_config()` raises ValueError** — no config source was found; either set up a baseline folder or pass `required=False`. ## Next steps 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 index 49c75199a7ce..4763a8c69568 100644 --- 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 @@ -8,19 +8,15 @@ from azure.ai.agentserver.optimization import load_config - config = load_config(default_instructions="You are a helpful assistant.") - # config.instructions — optimized or default - # config.model — optimized or default - # config.temperature — optimized or default - # config.skills — learned skills (empty list if none) - # config.tool_descriptions — optimized tool descriptions (empty dict if none) - # config.source — "env:OPTIMIZATION_CONFIG", "api:candidate:abc", "local:", or "defaults" + 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 + JOB_ID + ENDPOINT → resolver API → full config + skills - 3. Local directory (.agent_configs/) → metadata.yaml + instructions.md + tools.json + skills/ - 4. Defaults → your hardcoded values (agent works normally) + 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 @@ -28,7 +24,6 @@ CandidateConfig, OptimizationConfig, Skill, - ToolDescription, ) from azure.ai.agentserver.optimization._version import VERSION @@ -36,7 +31,6 @@ "CandidateConfig", "OptimizationConfig", "Skill", - "ToolDescription", "load_config", "load_skills_from_dir", ] 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 index e93fa9f36140..bebf4db973af 100644 --- 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 @@ -10,7 +10,7 @@ ├── baseline/ (fallback candidate) │ ├── metadata.yaml (model, temperature, file pointers) │ ├── instructions.md (system prompt) - │ ├── tools.json (tool descriptions — dict or list format) + │ ├── tools.json (tool definitions — list format) │ └── skills/ (learned skills) │ └── / │ └── SKILL.md @@ -40,8 +40,6 @@ MetadataConfig, OptimizationConfig, Skill, - ToolDescription, - _parse_tools_list, ) from azure.ai.agentserver.optimization._resolver import resolve_candidate @@ -50,11 +48,9 @@ def load_config( *, - default_instructions: str = "You are a helpful assistant.", - default_model: str | None = None, - default_temperature: float | None = None, - default_skills_dir: str | None = None, -) -> OptimizationConfig: + config_dir: str | Path | None = None, + required: bool = True, +) -> OptimizationConfig | None: """Load optimization config with graceful fallback. Resolution order (first match wins): @@ -62,47 +58,58 @@ def load_config( 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``, - ``OPTIMIZATION_JOB_ID``, and ``OPTIMIZATION_RESOLVE_ENDPOINT`` - are all set. Fetches the candidate config from the remote - optimization service and persists it to the local directory. + 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). - The local directory defaults to ``.agent_configs/`` relative to - the main script, overridable via ``OPTIMIZATION_LOCAL_DIR``. - 4. **Defaults** — returns the caller-supplied defaults unchanged. - The agent works exactly as if optimization were not installed. - - Safe to call at module load time. Any unexpected error is caught - and logged — the caller always gets a valid config back. + ``//`` (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( - default_instructions=default_instructions, - default_model=default_model, - default_temperature=default_temperature, - default_skills_dir=default_skills_dir, - ) + 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 — returning defaults: %s", exc) - model = default_model or os.environ.get("MODEL_DEPLOYMENT_NAME") - return OptimizationConfig( - instructions=default_instructions, - model=model, - temperature=default_temperature, - skills_dir=default_skills_dir, - source="defaults", - ) + logger.error("Unexpected error loading optimization config: %s", exc) + return None def _load_config_inner( *, - default_instructions: str, - default_model: str | None, - default_temperature: float | None, - default_skills_dir: str | None, -) -> OptimizationConfig: - """Internal config loader — may raise on unexpected errors.""" + 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() @@ -112,15 +119,15 @@ def _load_config_inner( candidate = CandidateConfig.from_dict(cfg) logger.warning( "Loaded optimization config from %s env var (%d chars instructions)", - env_var, len(candidate.instructions or ""), + env_var, + len(candidate.instructions or ""), ) return OptimizationConfig( - instructions=candidate.instructions or default_instructions, - model=candidate.model or default_model, - temperature=candidate.temperature if candidate.temperature is not None else default_temperature, + instructions=candidate.instructions, + model=candidate.model, + temperature=candidate.temperature, skills=candidate.skills, - skills_dir=cfg.get("skills_dir", default_skills_dir), - tool_descriptions=candidate.tool_descriptions, + tool_definitions=candidate.tool_definitions, source=f"env:{env_var}", ) except (json.JSONDecodeError, TypeError) as exc: @@ -128,11 +135,14 @@ def _load_config_inner( # ── Priority 2: Candidate ID → resolver API ────────────────────── candidate_id = os.environ.get(OptimizationConfig.ENV_CANDIDATE_ID, "").strip() - job_id = os.environ.get(OptimizationConfig.ENV_JOB_ID, "").strip() - endpoint = os.environ.get(OptimizationConfig.ENV_RESOLVE_ENDPOINT, "").strip().rstrip("/") - if candidate_id and job_id and endpoint: - local_dir = _resolve_local_dir() - resolved = resolve_candidate(candidate_id, job_id=job_id, endpoint=endpoint, local_dir=local_dir) + 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( @@ -140,15 +150,14 @@ def _load_config_inner( candidate_id, ) return OptimizationConfig( - instructions=candidate.instructions or default_instructions, - model=candidate.model or default_model, - temperature=candidate.temperature if candidate.temperature is not None else default_temperature, + instructions=candidate.instructions, + model=candidate.model, + temperature=candidate.temperature, skills=candidate.skills, - skills_dir=resolved.get("skills_dir", default_skills_dir), - tool_descriptions=candidate.tool_descriptions, + skills_dir=resolved.get("skills_dir"), + tool_definitions=candidate.tool_definitions, source=f"api:candidate:{candidate_id}", candidate_id=candidate_id, - job_id=job_id, ) logger.warning( "Failed to resolve candidate %s — falling through to local/defaults", @@ -156,49 +165,53 @@ def _load_config_inner( ) # ── Priority 3: Local directory (.agent_configs/) ────────── - local_config = _load_local_dir( - candidate_id or None, default_instructions, - default_model, default_temperature, default_skills_dir, - ) + 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, + local_config.source, + local_config.candidate_id, ) return local_config - # ── Priority 4: Defaults ───────────────────────────────────────── - model = default_model or os.environ.get("MODEL_DEPLOYMENT_NAME") - return OptimizationConfig( - instructions=default_instructions, - model=model, - temperature=default_temperature, - skills_dir=default_skills_dir, - source="defaults", - ) + # ── 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 at least " + "an instructions.md file, or pass required=False." + ) + logger.warning("No optimization config found — returning None") + return None -def _resolve_local_dir() -> Path: +def _resolve_local_dir(config_dir: str | Path | None = None) -> Path: """Resolve the local optimization directory path. - Falls back to ``OptimizationConfig.DEFAULT_LOCAL_DIR`` - (``".agent_configs"``) when the env var is not set. - """ - 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) + Priority: *config_dir* argument → ``OPTIMIZATION_LOCAL_DIR`` env + var → ``OptimizationConfig.DEFAULT_LOCAL_DIR`` (``.agent_configs``). - # Guard: reject paths with ".." components (path traversal) - if ".." in local_dir.parts: - logger.warning( - "OPTIMIZATION_LOCAL_DIR contains '..' path traversal: %r — ignoring", - local_dir_env, + :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) ) - local_dir = Path(OptimizationConfig.DEFAULT_LOCAL_DIR) - explicitly_set = False 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: @@ -213,13 +226,18 @@ def _resolve_local_dir() -> Path: def _load_local_dir( candidate_id: str | None, - default_instructions: str, - default_model: str | None, - default_temperature: float | None, - default_skills_dir: str | None, + config_dir: str | Path | None, ) -> OptimizationConfig | None: - """Load optimization config from a local directory.""" - local_dir = _resolve_local_dir() + """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 @@ -229,25 +247,27 @@ def _load_local_dir( metadata_file = candidate_path / OptimizationConfig.METADATA_FILE - return _load_candidate_from_metadata( - candidate_path, metadata_file, candidate_id, - default_instructions, default_model, default_temperature, default_skills_dir, - ) + 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, - default_instructions: str, - default_model: str | None, - default_temperature: float | None, - default_skills_dir: 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: @@ -265,66 +285,59 @@ def _load_candidate_from_metadata( meta = MetadataConfig.from_dict(raw) - # Read instructions from the referenced file (guard against traversal) + # Read instructions from the referenced file instructions_path = candidate_path / meta.instruction_file - if _is_safe_child(candidate_path, instructions_path) and instructions_path.is_file(): - instructions = instructions_path.read_text(encoding="utf-8").strip() + if instructions_path.is_file(): + instructions: str | None = instructions_path.read_text(encoding="utf-8").strip() else: - if not _is_safe_child(candidate_path, instructions_path): - logger.warning("Path traversal in instruction_file: %r", meta.instruction_file) - instructions = default_instructions + instructions = None - # Resolve skills directory (guard against traversal) + # Resolve skills directory skills_dir: str | None skills_path = candidate_path / meta.skill_dir - if _is_safe_child(candidate_path, skills_path) and skills_path.resolve().is_dir(): + if skills_path.resolve().is_dir(): skills_dir = str(skills_path.resolve()) else: - if not _is_safe_child(candidate_path, skills_path): - logger.warning("Path traversal in skill_dir: %r", meta.skill_dir) - skills_dir = default_skills_dir + skills_dir = None - # Load tool descriptions (guard against traversal) + # Load tool definitions tool_file_path = candidate_path / meta.tool_file - if _is_safe_child(candidate_path, tool_file_path): - tool_descriptions = _load_tool_descriptions(tool_file_path) - else: - logger.warning("Path traversal in tool_file: %r", meta.tool_file) - tool_descriptions = {} + tool_definitions = _load_tool_definitions(tool_file_path) return OptimizationConfig( instructions=instructions, - model=meta.model or default_model, - temperature=meta.temperature if meta.temperature is not None else default_temperature, + model=meta.model, + temperature=meta.temperature, skills_dir=skills_dir, - tool_descriptions=tool_descriptions, + tool_definitions=tool_definitions, source=f"local:{candidate_path}", candidate_id=candidate_id, ) -def _load_tool_descriptions(tool_file: Path) -> dict[str, ToolDescription]: - """Load tool descriptions from a tools.json file. +def _load_tool_definitions(tool_file: Path) -> list[dict]: + """Load tool definitions from a tools.json file. + + Expects the OpenAI function-calling list format:: - Supports both dict format ``{name: {description, parameters}}`` - and OpenAI function-calling list format ``[{type, function: {name, ...}}]``. + [{"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 {} + return [] try: raw = tool_file.read_text(encoding="utf-8") data = json.loads(raw) - if isinstance(data, dict): - return { - name: ToolDescription.from_dict(v) if isinstance(v, dict) else ToolDescription(description=str(v)) - for name, v in data.items() - } if isinstance(data, list): - return _parse_tools_list(data) - return {} + return data + return [] except (json.JSONDecodeError, OSError) as exc: logger.warning("Failed to read tools file %s: %s", tool_file, exc) - return {} + return [] def _parse_simple_yaml(path: Path) -> dict: @@ -332,6 +345,11 @@ def _parse_simple_yaml(path: Path) -> dict: 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: @@ -348,7 +366,13 @@ def _parse_simple_yaml(path: Path) -> dict: def _coerce_yaml_value(value: str) -> Any: - """Coerce a YAML scalar string to the appropriate Python type.""" + """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": @@ -371,27 +395,22 @@ def _resolve_candidate_folder(local_dir: Path, candidate_id: str | None) -> Path 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 not _is_safe_child(local_dir, exact): - logger.warning("Path traversal detected in candidate_id: %r", candidate_id) - return None if exact.is_dir(): return exact baseline = local_dir / OptimizationConfig.BASELINE_DIR return baseline if baseline.is_dir() else None -def _is_safe_child(parent: Path, child: Path) -> bool: - """Return True if *child* is strictly inside *parent* (no traversal).""" - try: - child.resolve().relative_to(parent.resolve()) - return True - except ValueError: - return False - - def load_skills_from_dir(skills_dir: Path) -> list[Skill]: """Load skills from a directory of skill folders. @@ -400,6 +419,11 @@ def load_skills_from_dir(skills_dir: Path) -> list[Skill]: 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 [] @@ -428,7 +452,13 @@ def load_skills_from_dir(skills_dir: Path) -> list[Skill]: def _parse_skill_frontmatter(content: str) -> tuple[dict, str]: - """Extract YAML frontmatter and body from a SKILL.md file.""" + """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 @@ -437,7 +467,7 @@ def _parse_skill_frontmatter(content: str) -> tuple[dict, str]: return {}, content fm_text = content[3:end].strip() - body = content[end + 3:].strip() + body = content[end + 3 :].strip() frontmatter: dict = {} for line in fm_text.splitlines(): 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 index 8ae4b612d601..4d2c204d8200 100644 --- 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 @@ -27,33 +27,6 @@ class Skill: body: str = "" -@dataclass -class ToolDescription: - """Description-only projection of a tool, optimized by the service. - - The optimizer patches *descriptions* (human-readable text) — it does - **not** change the tool's JSON-Schema (type, required, etc.) because - the hosted agent owns the static definition. - - Matches the API contract:: - - { - "description": "Find cheaper flight alternatives.", - "parameters": {"destination": "The travel destination city"} - } - """ - - description: str - parameters: dict[str, str] = field(default_factory=dict) - - @classmethod - def from_dict(cls, data: dict[str, Any]) -> ToolDescription: - return cls( - description=data.get("description", ""), - parameters=data.get("parameters", {}), - ) - - @dataclass class CandidateConfig: """Typed representation of the candidate config payload from the API. @@ -67,7 +40,7 @@ class CandidateConfig: "model": "gpt-4o", "temperature": 0.7, "skills": [{"name": "...", "description": "...", "body": "..."}], - "tool_descriptions": {"lookup_policy": {"description": "...", "parameters": {}}} + "tools": [{"type": "function", "function": {"name": "...", ...}}] } """ @@ -76,18 +49,25 @@ class CandidateConfig: model: str | None = None temperature: float | None = None skills: list[Skill] = field(default_factory=list) - tool_descriptions: dict[str, ToolDescription] = field(default_factory=dict) + 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.""" + """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_descriptions=_parse_tool_descriptions(data), + tool_definitions=tools if isinstance(tools, list) else [], ) @@ -112,7 +92,13 @@ class MetadataConfig: @classmethod def from_dict(cls, data: dict) -> MetadataConfig: - """Create from a parsed YAML dict, ignoring unknown keys.""" + """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) @@ -122,15 +108,14 @@ def from_dict(cls, data: dict) -> MetadataConfig: class OptimizationConfig: # pylint: disable=too-many-instance-attributes """Resolved optimization config. - When not running under optimization, all fields contain the defaults - you passed to :func:`load_config` — your agent works unchanged. + 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" - ENV_JOB_ID: ClassVar[str] = "OPTIMIZATION_JOB_ID" DEFAULT_LOCAL_DIR: ClassVar[str] = ".agent_configs" METADATA_FILE: ClassVar[str] = "metadata.yaml" @@ -140,92 +125,122 @@ class OptimizationConfig: # pylint: disable=too-many-instance-attributes SKILL_FILE: ClassVar[str] = "SKILL.md" BASELINE_DIR: ClassVar[str] = "baseline" - instructions: str - model: str | None - temperature: float | None + 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_descriptions: dict[str, ToolDescription] = field(default_factory=dict) + tool_definitions: list[dict] = field(default_factory=list) source: str = "defaults" candidate_id: str | None = None - job_id: str | None = None @property def has_skills(self) -> bool: return len(self.skills) > 0 or self.skills_dir is not None - @property - def has_tool_descriptions(self) -> bool: - return len(self.tool_descriptions) > 0 - - def get_tool_description(self, tool_name: str) -> ToolDescription | None: - """Look up the optimized description for a specific tool.""" - return self.tool_descriptions.get(tool_name) - - def get_tool_param_description(self, tool_name: str, param_name: str) -> str | None: - """Look up the optimized description for a specific tool parameter.""" - td = self.tool_descriptions.get(tool_name) - if td is None: - return None - return td.parameters.get(param_name) + def apply_tool_descriptions(self, tools: list) -> list: + """Apply optimized tool definitions to a list of tool functions. - def apply_tool_descriptions(self, tools: list) -> list: # pylint: disable=too-many-nested-blocks - """Apply optimized tool descriptions 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. - Patches ``__doc__`` (used by the Agent Framework as the tool - description) on each tool function whose name appears in - :attr:`tool_descriptions`. - - Args: - tools: List of @tool-decorated functions. - - Returns: - The same list of tools (mutated in place). + :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_descriptions: + 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 tool_name and tool_name in self.tool_descriptions: - overrides = self.tool_descriptions[tool_name] - if overrides.description: - # Patch .description (AIFunction/ToolProtocol) and __doc__ (plain functions) - try: - tool_fn.description = overrides.description - except AttributeError: - pass - tool_fn.__doc__ = overrides.description - logger.debug("Applied optimized description for tool '%s'", tool_name) - # Patch parameter descriptions on AIFunction's input_model - if overrides.parameters: - input_model = getattr(tool_fn, "input_model", None) - if input_model and hasattr(input_model, "model_fields"): - patched = False - for param_name, param_desc in overrides.parameters.items(): - if param_name in input_model.model_fields: - input_model.model_fields[param_name].description = param_desc - patched = True - if patched: - input_model.model_rebuild(force=True) - logger.debug("Applied optimized parameter descriptions for tool '%s'", tool_name) + 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 with skill catalog appended (if any). + + :return: Instructions text with skills appended. + :rtype: str + """ + base = self.instructions or "" if not self.skills: - return self.instructions + return base - lines = [self.instructions, "", "## Available Skills"] + 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 helpers (used by CandidateConfig.from_dict) ────────────── +# ── 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.""" + """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"): @@ -237,58 +252,3 @@ def _parse_skills(raw: list) -> list[Skill]: ) ) return skills - - -def _parse_tool_descriptions(data: dict[str, Any]) -> dict[str, ToolDescription]: - """Parse tool descriptions from an API response dict. - - Supports three formats: - - ``tool_descriptions`` / ``toolDescriptions``: ``{name: {description, parameters}}`` - - ``tools``: OpenAI function-calling list ``[{type: function, function: {name, description, parameters}}]`` - - ``tool_descriptions`` wins over ``toolDescriptions`` wins over ``tools``. - """ - raw = data.get("tool_descriptions") or data.get("toolDescriptions") - if isinstance(raw, dict): - return { - name: ToolDescription.from_dict(v) if isinstance(v, dict) else ToolDescription(description=str(v)) - for name, v in raw.items() - } - - tools_list = data.get("tools") - if isinstance(tools_list, list): - return _parse_tools_list(tools_list) - - return {} - - -def _parse_tools_list(tools: list) -> dict[str, ToolDescription]: - """Parse tool descriptions from OpenAI function-calling list format. - - Expected shape:: - - [{"type": "function", "function": {"name": "...", "description": "...", "parameters": {...}}}] - - Extracts per-parameter descriptions from ``parameters.properties..description``. - """ - result: dict[str, ToolDescription] = {} - for item in tools: - if not isinstance(item, dict): - continue - func = item.get("function", {}) - if not isinstance(func, dict): - continue - name = func.get("name") - if not name: - continue - description = func.get("description", "") - params_schema = func.get("parameters", {}) - param_descriptions: dict[str, str] = {} - if isinstance(params_schema, dict): - props = params_schema.get("properties", {}) - if isinstance(props, dict): - for param_name, param_val in props.items(): - if isinstance(param_val, dict) and "description" in param_val: - param_descriptions[param_name] = param_val["description"] - result[name] = ToolDescription(description=description, parameters=param_descriptions) - return result 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 index 6aa120ba9109..09cf08dfc4ba 100644 --- 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 @@ -37,34 +37,40 @@ # API path and version constants _API_VERSION = "2025-11-15-preview" -_JOBS_PATH = "agent_optimization_jobs" _AUTH_SCOPE = "https://ai.azure.com/.default" def resolve_candidate( candidate_id: str, - job_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. - """ - # Guard against path traversal in candidate_id - if local_dir is not None: - candidate_path_check = (local_dir / candidate_id).resolve() - if not str(candidate_path_check).startswith(str(local_dir.resolve())): - logger.error("Path traversal detected in candidate_id: %r — aborting", candidate_id) - return None + :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.debug("Candidate %s already downloaded — skipping", candidate_id) + 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) + logger.warning( + "Candidate %s was downloaded but folder is missing — re-downloading", + candidate_id, + ) _downloaded.discard(candidate_id) client = _build_client(endpoint) @@ -72,7 +78,7 @@ def resolve_candidate( # ── Step 1: Fetch config ───────────────────────────────────────── config = _api_get_json( client, - f"/{_JOBS_PATH}/{job_id}/candidates/{candidate_id}/config", + f"/candidates/{candidate_id}/config", params={"api-version": _API_VERSION}, ) if config is None: @@ -80,12 +86,12 @@ def resolve_candidate( return None logger.info( - "Resolved candidate %s: model=%s, instructions=%d chars, skills=%d, tool_descriptions=%d", + "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("toolDescriptions", {}) or config.get("tool_descriptions", {})), + len(config.get("tools", [])), ) # ── Step 2: Persist to local directory layout ──────────────────── @@ -93,9 +99,11 @@ def resolve_candidate( candidate_path = local_dir / candidate_id try: _persist_to_local_layout(candidate_path, config) - _download_skill_files(client, job_id, candidate_id, candidate_path) + _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) + 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(): @@ -106,7 +114,9 @@ def resolve_candidate( return config -def _persist_to_local_layout(candidate_path: pathlib.Path, config: dict[str, Any]) -> None: +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:: @@ -120,6 +130,11 @@ def _persist_to_local_layout(candidate_path: pathlib.Path, config: dict[str, Any └── 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) @@ -145,15 +160,13 @@ def _persist_to_local_layout(candidate_path: pathlib.Path, config: dict[str, Any instr_file = candidate_path / OptimizationConfig.INSTRUCTIONS_FILE instr_file.write_text(instructions, encoding="utf-8") - # tools.json — write tool_descriptions / toolDescriptions as dict format - tool_data = config.get("tool_descriptions") or config.get("toolDescriptions") + # tools.json — write tool definitions in list format tools_list = config.get("tools") - if tool_data and isinstance(tool_data, dict): + if tools_list and isinstance(tools_list, list): tools_file = candidate_path / OptimizationConfig.TOOLS_FILE - tools_file.write_text(json.dumps(tool_data, indent=2, ensure_ascii=False), encoding="utf-8") - elif 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") + 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", []) @@ -176,21 +189,30 @@ def _persist_to_local_layout(candidate_path: pathlib.Path, config: dict[str, Any 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 %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, - job_id: str, candidate_id: str, candidate_path: pathlib.Path, ) -> None: - """Fetch manifest and download skill files into candidate_path/skills//SKILL.md.""" + """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"/{_JOBS_PATH}/{job_id}/candidates/{candidate_id}", + f"/candidates/{candidate_id}", params={"api-version": _API_VERSION}, ) if manifest is None: @@ -205,7 +227,8 @@ def _download_skill_files( logger.info( "Downloading %d skill file(s) for candidate %s", - len(skill_files), candidate_id, + len(skill_files), + candidate_id, ) skills_dir = candidate_path / OptimizationConfig.SKILLS_DIR @@ -216,7 +239,7 @@ def _download_skill_files( content = _api_get_text( client, - f"/{_JOBS_PATH}/{job_id}/candidates/{candidate_id}/files", + f"/candidates/{candidate_id}/files", params={"path": file_path, "api-version": _API_VERSION}, ) if content is None: @@ -227,11 +250,13 @@ def _download_skill_files( rel_path = file_path prefix = OptimizationConfig.SKILLS_DIR + "/" if rel_path.startswith(prefix): - rel_path = rel_path[len(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) + logger.warning( + "Path traversal detected in skill file path: %r — skipping", file_path + ) continue out_path.parent.mkdir(parents=True, exist_ok=True) @@ -240,7 +265,13 @@ def _download_skill_files( def _is_skill_file(file_entry: dict) -> bool: - """Check if a manifest entry is a skill file.""" + """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/") @@ -250,7 +281,13 @@ def _is_skill_file(file_entry: dict) -> bool: def _build_client(endpoint: str) -> PipelineClient: - """Create a PipelineClient with credential-based auth and retry.""" + """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 @@ -258,14 +295,26 @@ def _build_client(endpoint: str) -> PipelineClient: 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") + 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.""" + """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) @@ -283,7 +332,17 @@ def _api_get_json( 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.""" + """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) diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/cspell.json b/sdk/agentserver/azure-ai-agentserver-optimization/cspell.json index 4fa70e4507c8..b3814147ea7e 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/cspell.json +++ b/sdk/agentserver/azure-ai-agentserver-optimization/cspell.json @@ -2,7 +2,12 @@ "ignoreWords": [ "agentserver", "autouse", - "cand" + "cand", + "dataclass", + "frontmatter", + "paramtype", + "pathlib", + "rtype" ], "ignorePaths": [ "*.json", diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/dev_requirements.txt b/sdk/agentserver/azure-ai-agentserver-optimization/dev_requirements.txt index 190d2e5520b2..c29511d76ecb 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/dev_requirements.txt +++ b/sdk/agentserver/azure-ai-agentserver-optimization/dev_requirements.txt @@ -1,2 +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/tests/conftest.py b/sdk/agentserver/azure-ai-agentserver-optimization/tests/conftest.py index b904d3574863..e806f9eaf652 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/tests/conftest.py +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/conftest.py @@ -9,7 +9,6 @@ ENV_VARS = [ "OPTIMIZATION_CONFIG", "OPTIMIZATION_CANDIDATE_ID", - "OPTIMIZATION_JOB_ID", "OPTIMIZATION_LOCAL_DIR", "OPTIMIZATION_RESOLVE_ENDPOINT", "MODEL_DEPLOYMENT_NAME", diff --git a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py index 7623feffe816..d212eafb82b3 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_config.py @@ -11,17 +11,14 @@ CandidateConfig, OptimizationConfig, Skill, - ToolDescription, load_config, ) from azure.ai.agentserver.optimization._models import ( MetadataConfig, _parse_skills, - _parse_tool_descriptions, - _parse_tools_list, ) from azure.ai.agentserver.optimization._config import ( - _load_tool_descriptions, + _load_tool_definitions, _parse_simple_yaml, _parse_skill_frontmatter, _resolve_candidate_folder, @@ -41,48 +38,32 @@ def clear_downloaded(): class TestDefaults: - """When no env vars are set, load_config returns caller-supplied defaults.""" + """When no env vars or config dir are set.""" - def test_returns_default_instructions(self): - config = load_config(default_instructions="Be helpful.") - assert config.instructions == "Be helpful." - assert config.source == "defaults" + def test_required_true_raises_by_default(self): + with pytest.raises(ValueError, match="No optimization config found"): + load_config() - def test_returns_default_model(self): - config = load_config(default_model="gpt-4o") - assert config.model == "gpt-4o" - - def test_returns_default_temperature(self): - config = load_config(default_temperature=0.5) - assert config.temperature == 0.5 - - def test_returns_default_skills_dir(self): - config = load_config(default_skills_dir="/some/path") - assert config.skills_dir == "/some/path" - - def test_empty_skills_by_default(self): - config = load_config() - assert config.skills == [] - assert not config.has_skills - - def test_empty_tool_descriptions_by_default(self): - config = load_config() - assert config.tool_descriptions == {} - assert not config.has_tool_descriptions + 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() - assert config.model == "gpt-4o-mini" + config = load_config(required=False) + assert config is None - def test_explicit_model_overrides_env(self, monkeypatch): - monkeypatch.setenv("MODEL_DEPLOYMENT_NAME", "gpt-4o-mini") - config = load_config(default_model="gpt-4o") + 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" - - def test_default_instructions_value(self): - config = load_config() - assert config.instructions == "You are a helpful assistant." + assert config.source.startswith("local:") # ── Inline JSON env var (Priority 1) ──────────────────────────────── @@ -98,7 +79,7 @@ def test_loads_from_env_config(self, monkeypatch): "temperature": 0.3, } monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) - config = load_config(default_instructions="default") + config = load_config() assert config.instructions == "Optimized prompt." assert config.model == "gpt-4o" assert config.temperature == 0.3 @@ -124,37 +105,35 @@ def test_env_config_with_skills(self, monkeypatch): def test_env_config_with_tool_descriptions(self, monkeypatch): payload = { "instructions": "With tools.", - "tool_descriptions": { - "lookup_travel_policy": { - "description": "Look up the company travel policy.", - "parameters": {}, + "tools": [ + { + "type": "function", + "function": { + "name": "lookup_travel_policy", + "description": "Look up the company travel policy.", + "parameters": {"type": "object", "properties": {}}, + }, }, - "check_department_budget": { - "description": "Check remaining budget.", - "parameters": {"dept": "Department name"}, + { + "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 config.has_tool_descriptions - assert "lookup_travel_policy" in config.tool_descriptions - td = config.tool_descriptions["check_department_budget"] - assert isinstance(td, ToolDescription) - assert td.description == "Check remaining budget." - assert td.parameters == {"dept": "Department name"} - - def test_env_config_with_legacy_toolDescriptions(self, monkeypatch): - payload = { - "instructions": "With tools.", - "toolDescriptions": { - "search": {"description": "Search something.", "parameters": {}}, - }, + ], } monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) config = load_config() - assert config.has_tool_descriptions - assert "search" in config.tool_descriptions + 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.""" @@ -178,61 +157,32 @@ def test_env_config_with_tools_list(self, monkeypatch): } monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) config = load_config() - assert config.has_tool_descriptions - assert "lookup_policy" in config.tool_descriptions - td = config.tool_descriptions["lookup_policy"] - assert td.description == "Look up policy" - assert td.parameters == {"dept": "Department name"} - - def test_tool_descriptions_takes_priority_over_legacy(self, monkeypatch): - payload = { - "instructions": "Both.", - "tool_descriptions": {"new_tool": {"description": "New"}}, - "toolDescriptions": {"old_tool": {"description": "Old"}}, - } - monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) - config = load_config() - assert "new_tool" in config.tool_descriptions - assert "old_tool" not in config.tool_descriptions + assert len(config.tool_definitions) == 1 + assert config.tool_definitions[0]["function"]["name"] == "lookup_policy" - def test_tool_descriptions_takes_priority_over_tools_list(self, monkeypatch): - payload = { - "instructions": "Both.", - "tool_descriptions": {"dict_tool": {"description": "Dict"}}, - "tools": [{"type": "function", "function": {"name": "list_tool", "description": "List"}}], - } - monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) - config = load_config() - assert "dict_tool" in config.tool_descriptions - assert "list_tool" not in config.tool_descriptions def test_bad_json_falls_through(self, monkeypatch): monkeypatch.setenv("OPTIMIZATION_CONFIG", "not-json{{{") - config = load_config(default_instructions="fallback") - assert config.instructions == "fallback" - assert config.source == "defaults" + config = load_config(required=False) + assert config is None def test_empty_env_var_ignored(self, monkeypatch): monkeypatch.setenv("OPTIMIZATION_CONFIG", " ") - config = load_config(default_instructions="fallback") - assert config.source == "defaults" + config = load_config(required=False) + assert config is None - def test_partial_config_uses_defaults(self, monkeypatch): + def test_partial_config_fills_none(self, monkeypatch): payload = {"model": "gpt-4o"} monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(payload)) - config = load_config( - default_instructions="My default", - default_temperature=0.7, - ) - assert config.instructions == "My default" + config = load_config() + assert config.instructions is None assert config.model == "gpt-4o" - assert config.temperature == 0.7 + 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_JOB_ID", "job-1") monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") config = load_config() assert config.source == "env:OPTIMIZATION_CONFIG" @@ -242,7 +192,7 @@ def test_env_config_takes_priority_over_candidate_id(self, monkeypatch): class TestCandidateResolver: - """OPTIMIZATION_CANDIDATE_ID + JOB_ID + ENDPOINT triggers resolver API.""" + """OPTIMIZATION_CANDIDATE_ID + ENDPOINT triggers resolver API.""" def test_candidate_id_calls_resolver(self, monkeypatch): resolved = { @@ -252,54 +202,41 @@ def test_candidate_id_calls_resolver(self, monkeypatch): "skills": [{"name": "s1", "description": "d1"}], } monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-123") - monkeypatch.setenv("OPTIMIZATION_JOB_ID", "job-42") monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") monkeypatch.setattr( "azure.ai.agentserver.optimization._config.resolve_candidate", - lambda cid, job_id, endpoint, local_dir=None: resolved, + 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 config.job_id == "job-42" assert len(config.skills) == 1 def test_resolver_failure_falls_to_defaults(self, monkeypatch): monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "bad-id") - monkeypatch.setenv("OPTIMIZATION_JOB_ID", "job-1") monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") monkeypatch.setattr( "azure.ai.agentserver.optimization._config.resolve_candidate", - lambda cid, job_id, endpoint, local_dir=None: None, + lambda cid, endpoint, local_dir=None: None, ) - config = load_config(default_instructions="fallback") - assert config.source == "defaults" - assert config.instructions == "fallback" - - def test_missing_job_id_skips_resolver(self, monkeypatch): - monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-1") - monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://fake") - # No JOB_ID set - config = load_config(default_instructions="default") - assert config.source == "defaults" + config = load_config(required=False) + assert config is None def test_missing_endpoint_skips_resolver(self, monkeypatch): monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-1") - monkeypatch.setenv("OPTIMIZATION_JOB_ID", "job-1") # No ENDPOINT set - config = load_config(default_instructions="default") - assert config.source == "defaults" + 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_JOB_ID", "job-1") 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, job_id, endpoint, local_dir=None: None, + lambda cid, endpoint, local_dir=None: None, ) # Set up local dir with this candidate candidate_dir = tmp_path / "cand-local" @@ -404,13 +341,15 @@ def test_loads_tools_dict_from_local_dir(self, monkeypatch, tmp_path): candidate_dir.mkdir() (candidate_dir / "metadata.yaml").write_text("tool_file: tools.json\n") (candidate_dir / "instructions.md").write_text("With tools.") - tools = {"search": {"description": "Search stuff", "parameters": {"q": "query"}}} + 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 "search" in config.tool_descriptions - assert config.tool_descriptions["search"].description == "Search stuff" + 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.""" @@ -436,71 +375,77 @@ def test_loads_tools_list_from_local_dir(self, monkeypatch, tmp_path): monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) config = load_config() - assert "get_weather" in config.tool_descriptions - assert config.tool_descriptions["get_weather"].description == "Get the weather" - assert config.tool_descriptions["get_weather"].parameters == {"city": "City name"} + assert len(config.tool_definitions) == 1 + assert config.tool_definitions[0]["function"]["name"] == "get_weather" - def test_missing_instructions_uses_default(self, monkeypatch, tmp_path): + 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(default_instructions="My default") - assert config.instructions == "My default" + 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(default_instructions="fallback") - assert config.source == "defaults" + 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(default_instructions="default") - assert config.source == "defaults" + config = load_config(required=False) + assert config is None - def test_metadata_traversal_instruction_file(self, monkeypatch, tmp_path): - """instruction_file with '../' in metadata.yaml is rejected.""" + 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: ../../etc/passwd\n" + "instruction_file: ../shared_instructions.md\n" ) - # Create the traversal target to prove it's NOT read - secret = tmp_path / "secret.txt" - secret.write_text("SECRET DATA") monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) - config = load_config(default_instructions="safe default") - assert config.instructions == "safe default" - assert "SECRET" not in config.instructions + config = load_config() + assert config.instructions == "Shared prompt." - def test_metadata_traversal_skill_dir(self, monkeypatch, tmp_path): - """skill_dir with '../' in metadata.yaml is rejected.""" + 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: ../../other_skills\n" + "skill_dir: ../shared_skills\n" ) - (candidate_dir / "instructions.md").write_text("ok") monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) config = load_config() - assert config.skills == [] + assert config.skills_dir is not None - def test_metadata_traversal_tool_file(self, monkeypatch, tmp_path): - """tool_file with '../' in metadata.yaml is rejected.""" + 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: ../../secrets.json\n" + "tool_file: ../shared_tools.json\n" ) - (candidate_dir / "instructions.md").write_text("ok") monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) config = load_config() - assert config.tool_descriptions == {} + assert len(config.tool_definitions) == 1 + assert config.tool_definitions[0]["function"]["name"] == "search" # ── _resolve_local_dir ────────────────────────────────────────────── @@ -519,16 +464,31 @@ def test_uses_env_var(self, monkeypatch, tmp_path): local_dir = _resolve_local_dir() assert local_dir == tmp_path - def test_rejects_dotdot_traversal(self, monkeypatch): - """OPTIMIZATION_LOCAL_DIR with '..' is rejected, falls back to default.""" - monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", "../../etc/sensitive") + 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.name == ".agent_configs" + assert local_dir == tmp_path - def test_rejects_dotdot_in_absolute_path(self, monkeypatch, tmp_path): - """Even absolute paths with '..' are rejected.""" - malicious = str(tmp_path / ".." / ".." / "etc") - monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", malicious) + 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" @@ -562,43 +522,47 @@ def test_returns_none_no_id_no_baseline(self, tmp_path): result = _resolve_candidate_folder(tmp_path, None) assert result is None - def test_rejects_traversal_candidate_id(self, tmp_path): - """candidate_id with '../' is rejected as path traversal.""" - result = _resolve_candidate_folder(tmp_path, "../../etc") - assert result is None - - def test_rejects_absolute_candidate_id(self, tmp_path): - """Absolute path in candidate_id is rejected.""" - result = _resolve_candidate_folder(tmp_path, "/etc/passwd") - 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.""" + """load_config never crashes — always returns a valid config (unless required ValueError).""" - def test_unexpected_exception_returns_defaults(self, monkeypatch): - """Any unexpected error in _load_config_inner returns defaults.""" + 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(default_instructions="safe") - assert config.source == "defaults" - assert config.instructions == "safe" + config = load_config(required=False) + assert config is None - def test_load_config_never_raises(self, monkeypatch): - """Even with corrupted env vars, load_config returns something.""" + 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_JOB_ID", "y") monkeypatch.setenv("OPTIMIZATION_RESOLVE_ENDPOINT", "http://nope") monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", "/nonexistent") - config = load_config(default_instructions="fallback") - assert isinstance(config, OptimizationConfig) - assert config.instructions == "fallback" + 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 ──────────────────────────────────── @@ -649,38 +613,6 @@ def test_no_skills(self): ) assert not config.has_skills - def test_has_tool_descriptions(self): - config = OptimizationConfig( - instructions="", model=None, temperature=None, - tool_descriptions={"t": ToolDescription(description="d")}, - ) - assert config.has_tool_descriptions - - def test_no_tool_descriptions(self): - config = OptimizationConfig( - instructions="", model=None, temperature=None, - ) - assert not config.has_tool_descriptions - - def test_get_tool_description(self): - td = ToolDescription(description="Search things", parameters={"q": "query"}) - config = OptimizationConfig( - instructions="", model=None, temperature=None, - tool_descriptions={"search": td}, - ) - assert config.get_tool_description("search") is td - assert config.get_tool_description("missing") is None - - def test_get_tool_param_description(self): - td = ToolDescription(description="Search", parameters={"q": "The query"}) - config = OptimizationConfig( - instructions="", model=None, temperature=None, - tool_descriptions={"search": td}, - ) - assert config.get_tool_param_description("search", "q") == "The query" - assert config.get_tool_param_description("search", "missing") is None - assert config.get_tool_param_description("missing", "q") is None - def test_constants(self): assert OptimizationConfig.DEFAULT_LOCAL_DIR == ".agent_configs" assert OptimizationConfig.METADATA_FILE == "metadata.yaml" @@ -691,31 +623,6 @@ def test_constants(self): assert OptimizationConfig.BASELINE_DIR == "baseline" -# ── ToolDescription ────────────────────────────────────────────────── - - -class TestToolDescription: - """Tests for ToolDescription dataclass.""" - - def test_from_dict(self): - td = ToolDescription.from_dict({ - "description": "Search things", - "parameters": {"q": "The query", "limit": "Max results"}, - }) - assert td.description == "Search things" - assert td.parameters == {"q": "The query", "limit": "Max results"} - - def test_from_dict_defaults(self): - td = ToolDescription.from_dict({}) - assert td.description == "" - assert td.parameters == {} - - def test_from_dict_missing_parameters(self): - td = ToolDescription.from_dict({"description": "No params"}) - assert td.description == "No params" - assert td.parameters == {} - - # ── CandidateConfig ───────────────────────────────────────────────── @@ -732,16 +639,10 @@ def test_full_payload(self): {"name": "budget-checker", "description": "Check budget", "body": "# Budget"}, {"name": "policy-reviewer", "description": "Review policy"}, ], - "tool_descriptions": { - "lookup_travel_policy": { - "description": "Look up travel policy.", - "parameters": {}, - }, - "get_flight_alternatives": { - "description": "Find cheaper flights.", - "parameters": {"destination": "The destination city"}, - }, - }, + "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" @@ -751,10 +652,8 @@ def test_full_payload(self): assert len(candidate.skills) == 2 assert candidate.skills[0].name == "budget-checker" assert candidate.skills[0].body == "# Budget" - assert len(candidate.tool_descriptions) == 2 - td = candidate.tool_descriptions["get_flight_alternatives"] - assert td.description == "Find cheaper flights." - assert td.parameters["destination"] == "The destination city" + 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({}) @@ -763,16 +662,7 @@ def test_minimal_payload(self): assert candidate.model is None assert candidate.temperature is None assert candidate.skills == [] - assert candidate.tool_descriptions == {} - - def test_legacy_toolDescriptions_key(self): - payload = { - "toolDescriptions": { - "search": {"description": "Search", "parameters": {}}, - }, - } - candidate = CandidateConfig.from_dict(payload) - assert "search" in candidate.tool_descriptions + assert candidate.tool_definitions == [] def test_tools_list_format(self): payload = { @@ -793,9 +683,28 @@ def test_tools_list_format(self): ], } candidate = CandidateConfig.from_dict(payload) - assert "get_weather" in candidate.tool_descriptions - assert candidate.tool_descriptions["get_weather"].description == "Get weather" - assert candidate.tool_descriptions["get_weather"].parameters == {"city": "City"} + 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 ────────────────────────────────────────────────── @@ -822,6 +731,16 @@ def test_from_dict_defaults(self): 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 ─────────────────────────────────────────────────── @@ -860,150 +779,11 @@ def test_mixed_valid_invalid(self): assert len(result) == 2 -# ── _parse_tool_descriptions ──────────────────────────────────────── - - -class TestParseToolDescriptions: - """Tests for _parse_tool_descriptions edge cases.""" - - def test_empty_data(self): - assert _parse_tool_descriptions({}) == {} +# ── _load_tool_definitions (file loading) ─────────────────────────── - def test_tool_descriptions_dict(self): - data = {"tool_descriptions": {"t1": {"description": "D1", "parameters": {}}}} - result = _parse_tool_descriptions(data) - assert "t1" in result - assert result["t1"].description == "D1" - def test_toolDescriptions_camelCase(self): - data = {"toolDescriptions": {"t2": {"description": "D2"}}} - result = _parse_tool_descriptions(data) - assert "t2" in result - - def test_tool_descriptions_wins_over_toolDescriptions(self): - data = { - "tool_descriptions": {"winner": {"description": "W"}}, - "toolDescriptions": {"loser": {"description": "L"}}, - } - result = _parse_tool_descriptions(data) - assert "winner" in result - assert "loser" not in result - - def test_tools_list_fallback(self): - data = { - "tools": [ - {"type": "function", "function": {"name": "f1", "description": "Func"}}, - ] - } - result = _parse_tool_descriptions(data) - assert "f1" in result - assert result["f1"].description == "Func" - - def test_string_value_coerced(self): - data = {"tool_descriptions": {"t": "just a string"}} - result = _parse_tool_descriptions(data) - assert result["t"].description == "just a string" - - def test_non_dict_raw_ignored(self): - data = {"tool_descriptions": "not a dict"} - result = _parse_tool_descriptions(data) - assert result == {} - - -# ── _parse_tools_list ──────────────────────────────────────────────── - - -class TestParseToolsList: - """Tests for _parse_tools_list (OpenAI function-calling format).""" - - def test_basic(self): - tools = [ - { - "type": "function", - "function": { - "name": "search", - "description": "Search things", - "parameters": { - "type": "object", - "properties": { - "q": {"type": "string", "description": "Query"}, - }, - }, - }, - }, - ] - result = _parse_tools_list(tools) - assert "search" in result - assert result["search"].description == "Search things" - assert result["search"].parameters == {"q": "Query"} - - def test_no_parameters(self): - tools = [ - {"type": "function", "function": {"name": "noop", "description": "Do nothing"}}, - ] - result = _parse_tools_list(tools) - assert result["noop"].parameters == {} - - def test_skips_non_dict_items(self): - result = _parse_tools_list(["garbage", 42]) - assert result == {} - - def test_skips_items_without_function(self): - result = _parse_tools_list([{"type": "code_interpreter"}]) - assert result == {} - - def test_skips_items_without_name(self): - result = _parse_tools_list([ - {"type": "function", "function": {"description": "nameless"}}, - ]) - assert result == {} - - def test_empty_list(self): - assert _parse_tools_list([]) == {} - - def test_param_without_description_skipped(self): - tools = [ - { - "type": "function", - "function": { - "name": "f", - "description": "F", - "parameters": { - "type": "object", - "properties": { - "has_desc": {"type": "string", "description": "Yes"}, - "no_desc": {"type": "integer"}, - }, - }, - }, - }, - ] - result = _parse_tools_list(tools) - assert result["f"].parameters == {"has_desc": "Yes"} - - def test_multiple_functions(self): - tools = [ - {"type": "function", "function": {"name": "a", "description": "A"}}, - {"type": "function", "function": {"name": "b", "description": "B"}}, - ] - result = _parse_tools_list(tools) - assert len(result) == 2 - - -# ── _load_tool_descriptions (file loading) ────────────────────────── - - -class TestLoadToolDescriptions: - """Tests for _load_tool_descriptions from tools.json.""" - - def test_load_dict_format(self, tmp_path): - tool_file = tmp_path / "tools.json" - tools = {"my_tool": {"description": "My tool", "parameters": {"x": "input"}}} - tool_file.write_text(json.dumps(tools)) - result = _load_tool_descriptions(tool_file) - assert "my_tool" in result - assert isinstance(result["my_tool"], ToolDescription) - assert result["my_tool"].parameters == {"x": "input"} +class TestLoadToolDefinitions: + """Tests for _load_tool_definitions from tools.json.""" def test_load_list_format(self, tmp_path): tool_file = tmp_path / "tools.json" @@ -1011,31 +791,49 @@ def test_load_list_format(self, tmp_path): {"type": "function", "function": {"name": "f1", "description": "Func 1"}}, ] tool_file.write_text(json.dumps(tools)) - result = _load_tool_descriptions(tool_file) - assert "f1" in result - assert result["f1"].description == "Func 1" + 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_descriptions(tmp_path / "nonexistent.json") - assert result == {} + 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_descriptions(tool_file) - assert result == {} + result = _load_tool_definitions(tool_file) + assert result == [] - def test_non_dict_non_list_returns_empty(self, tmp_path): + 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_descriptions(tool_file) - assert result == {} + result = _load_tool_definitions(tool_file) + assert result == [] - def test_string_value_in_dict(self, tmp_path): + 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_descriptions(tool_file) - assert result["t"].description == "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 ─────────────────────────────────────── @@ -1127,21 +925,29 @@ def test_coerces_bool(self, tmp_path): class TestApplyToolDescriptions: """Tests for OptimizationConfig.apply_tool_descriptions.""" - def _make_config(self, tool_descriptions=None): + @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_descriptions=tool_descriptions or {}, + tool_definitions=tool_definitions or [], ) def test_patches_docstring(self): def lookup_policy(): """Original description.""" - config = self._make_config( - {"lookup_policy": ToolDescription(description="Optimized 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." @@ -1150,20 +956,20 @@ def test_returns_same_list(self): def my_tool(): """Original.""" - config = self._make_config( - {"my_tool": ToolDescription(description="New.")} - ) + 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_descriptions(self): + def test_skips_tools_not_in_definitions(self): def unknown_tool(): """Should not change.""" - config = self._make_config( - {"other_tool": ToolDescription(description="Something.")} - ) + config = self._make_config([ + self._tool_def("other_tool", "Something."), + ]) config.apply_tool_descriptions([unknown_tool]) assert unknown_tool.__doc__ == "Should not change." @@ -1171,13 +977,13 @@ def test_skips_empty_description(self): def my_tool(): """Original doc.""" - config = self._make_config( - {"my_tool": ToolDescription(description="")} - ) + 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_descriptions(self): + def test_noop_when_no_tool_definitions(self): def my_tool(): """Keep me.""" @@ -1187,9 +993,9 @@ def my_tool(): def test_handles_objects_without_name(self): """Non-function items without __name__ are silently skipped.""" - config = self._make_config( - {"something": ToolDescription(description="X.")} - ) + config = self._make_config([ + self._tool_def("something", "X."), + ]) obj = object() # Should not raise config.apply_tool_descriptions([obj]) @@ -1204,10 +1010,10 @@ def tool_b(): def tool_c(): """C original.""" - config = self._make_config({ - "tool_a": ToolDescription(description="A optimized."), - "tool_c": ToolDescription(description="C optimized."), - }) + 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." @@ -1220,9 +1026,9 @@ class FakeTool: __doc__ = "Original." tool = FakeTool() - config = self._make_config( - {"my_tool": ToolDescription(description="Patched via .name attr.")} - ) + 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." @@ -1233,9 +1039,9 @@ def search_flights(): search_flights.description = "Original." # type: ignore[attr-defined] - config = self._make_config( - {"search_flights": ToolDescription(description="Optimized flights search.")} - ) + 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." @@ -1251,9 +1057,9 @@ def description(self): return "read-only" tool = ReadOnlyTool() - config = self._make_config( - {"my_tool": ToolDescription(description="Patched.")} - ) + 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 @@ -1280,12 +1086,14 @@ def search_flights(destination: str, date: str): search_flights.input_model = FakeInputModel # type: ignore[attr-defined] - config = self._make_config({ - "search_flights": ToolDescription( - description="Find flights.", - parameters={"destination": "The travel destination city"}, - ), - }) + 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" @@ -1297,12 +1105,12 @@ def test_skips_param_patch_when_no_input_model(self): def my_tool(): """Original.""" - config = self._make_config({ - "my_tool": ToolDescription( - description="New.", - parameters={"x": "Some param"}, - ), - }) + 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." @@ -1325,12 +1133,12 @@ def my_tool(): my_tool.input_model = FakeInputModel # type: ignore[attr-defined] - config = self._make_config({ - "my_tool": ToolDescription( - description="New doc.", - parameters={"unknown_param": "Should be ignored"}, - ), - }) + 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 @@ -1354,16 +1162,201 @@ def my_tool(): my_tool.input_model = FakeInputModel # type: ignore[attr-defined] - config = self._make_config({ - "my_tool": ToolDescription(description="New.", parameters={}), - }) + 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( - {"tool": ToolDescription(description="X.")} - ) + 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 index 05dda5114af5..276f8ea7d0ef 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_integration.py +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_integration.py @@ -11,7 +11,6 @@ from azure.ai.agentserver.optimization import ( OptimizationConfig, Skill, - ToolDescription, load_config, load_skills_from_dir, ) @@ -34,16 +33,29 @@ def test_env_config_apply_tool_descriptions(self, monkeypatch): "instructions": "Optimized prompt.", "model": "gpt-4o", "temperature": 0.5, - "tool_descriptions": { - "search_flights": { - "description": "Find the cheapest flight options.", - "parameters": {"destination": "City name"}, + "tools": [ + { + "type": "function", + "function": { + "name": "search_flights", + "description": "Find the cheapest flight options.", + "parameters": { + "type": "object", + "properties": { + "destination": {"type": "string", "description": "City name"}, + }, + }, + }, }, - "book_hotel": { - "description": "Reserve a hotel room.", - "parameters": {}, + { + "type": "function", + "function": { + "name": "book_hotel", + "description": "Reserve a hotel room.", + "parameters": {"type": "object", "properties": {}}, + }, }, - }, + ], } monkeypatch.setenv("OPTIMIZATION_CONFIG", json.dumps(cfg)) @@ -56,10 +68,10 @@ def book_hotel(city: str): def unrelated_tool(): """Should stay unchanged.""" - config = load_config(default_instructions="fallback") + config = load_config() assert config.source == "env:OPTIMIZATION_CONFIG" assert config.instructions == "Optimized prompt." - assert config.has_tool_descriptions + 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." @@ -95,7 +107,6 @@ def lookup_policy(dept: str): config = load_config() config.apply_tool_descriptions([lookup_policy]) assert lookup_policy.__doc__ == "Look up travel policy." - assert config.get_tool_param_description("lookup_policy", "dept") == "Department name" class TestLoadConfigAndLoadSkills: @@ -121,13 +132,13 @@ def test_local_dir_skills_workflow(self, monkeypatch, tmp_path): ) # Create tools.json - tools_data = { - "search": {"description": "Search destinations.", "parameters": {"q": "Query"}}, - } + 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(default_instructions="fallback") + config = load_config() # Verify config loaded assert config.instructions == "You are a travel agent." @@ -144,9 +155,9 @@ def test_local_dir_skills_workflow(self, monkeypatch, tmp_path): assert "budget-checker" in names assert "route-planner" in names - # Verify tool descriptions also loaded - assert config.has_tool_descriptions - assert config.tool_descriptions["search"].description == "Search destinations." + # 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.""" @@ -184,10 +195,10 @@ def test_complete_agent_setup(self, monkeypatch, tmp_path): (candidate_dir / "instructions.md").write_text( "You are a concise travel booking assistant." ) - tools_data = { - "search_flights": {"description": "Find flights between cities.", "parameters": {}}, - "book_flight": {"description": "Book the selected flight.", "parameters": {}}, - } + 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) @@ -200,10 +211,7 @@ def test_complete_agent_setup(self, monkeypatch, tmp_path): monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "candidate-v2") # Step 1: Load config - config = load_config( - default_instructions="Default prompt.", - default_model="gpt-3.5-turbo", - ) + config = load_config() assert config.instructions == "You are a concise travel booking assistant." assert config.model == "gpt-4o-mini" assert config.temperature == 0.3 @@ -233,7 +241,7 @@ def book_flight(flight_id): temperature=config.temperature, skills=skills, skills_dir=config.skills_dir, - tool_descriptions=config.tool_descriptions, + tool_definitions=config.tool_definitions, source=config.source, candidate_id=config.candidate_id, ) @@ -243,32 +251,226 @@ def book_flight(flight_id): assert "Handle rebooking requests" in composed def test_defaults_workflow_no_optimization(self): - """When no optimization is configured, everything works with defaults.""" + """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 my_tool(): + 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.""" - config = load_config( - default_instructions="Be helpful.", - default_model="gpt-4o", - default_temperature=0.7, + 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." ) - assert config.source == "defaults" - assert config.instructions == "Be helpful." + + config = load_config(config_dir=tmp_path) + assert config.instructions == "Custom agent." assert config.model == "gpt-4o" - assert config.temperature == 0.7 - assert not config.has_tool_descriptions - assert not config.has_skills - assert config.skills_dir is None - - # apply_tool_descriptions is a no-op - config.apply_tool_descriptions([my_tool]) - assert my_tool.__doc__ == "Original." - - # load_skills_from_dir with None skills_dir — user checks before calling - # This is the expected pattern: - if config.skills_dir: - skills = load_skills_from_dir(Path(config.skills_dir)) - else: - skills = [] - assert skills == [] + 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 index 778d4dcf487c..3e841958bf2a 100644 --- a/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py +++ b/sdk/agentserver/azure-ai-agentserver-optimization/tests/test_resolver.py @@ -36,7 +36,6 @@ def mock_client(): ENDPOINT = "http://fake-endpoint" -JOB_ID = "job-42" # ── resolve_candidate ─────────────────────────────────────────────── @@ -53,7 +52,7 @@ def test_returns_none_on_api_failure(self): return_value=None, ), ): - result = resolve_candidate("cand-1", job_id=JOB_ID, endpoint=ENDPOINT) + result = resolve_candidate("cand-1", endpoint=ENDPOINT) assert result is None def test_returns_config_on_success(self): @@ -70,13 +69,13 @@ def test_returns_config_on_success(self): return_value=config, ), ): - result = resolve_candidate("cand-1", job_id=JOB_ID, endpoint=ENDPOINT) + 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 /agent_optimization_jobs/{jobId}/candidates/{candidateId}/config.""" + """Verify the API route follows /candidates/{candidateId}/config.""" called_args: list = [] def capture_call(client, path, params=None): @@ -90,8 +89,8 @@ def capture_call(client, path, params=None): side_effect=capture_call, ), ): - resolve_candidate("cand-abc", job_id="job-xyz", endpoint="http://api.test") - assert called_args[0][0] == "/agent_optimization_jobs/job-xyz/candidates/cand-abc/config" + 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 ( @@ -101,7 +100,7 @@ def test_marks_downloaded_after_success(self): return_value={"instructions": "ok"}, ), ): - resolve_candidate("cand-mark", job_id=JOB_ID, endpoint=ENDPOINT) + resolve_candidate("cand-mark", endpoint=ENDPOINT) assert "cand-mark" in _downloaded def test_skips_if_already_downloaded_and_folder_exists(self, tmp_path): @@ -110,7 +109,7 @@ def test_skips_if_already_downloaded_and_folder_exists(self, tmp_path): _downloaded.add("cand-skip") result = resolve_candidate( - "cand-skip", job_id=JOB_ID, endpoint=ENDPOINT, local_dir=tmp_path, + "cand-skip", endpoint=ENDPOINT, local_dir=tmp_path, ) assert result is None @@ -126,7 +125,7 @@ def test_redownloads_if_folder_missing(self): ), ): result = resolve_candidate( - "cand-gone", job_id=JOB_ID, endpoint=ENDPOINT, local_dir=None, + "cand-gone", endpoint=ENDPOINT, local_dir=None, ) # local_dir is None → can't check folder → should re-download assert result is not None @@ -140,7 +139,7 @@ def test_does_not_mark_downloaded_on_api_failure(self): return_value=None, ), ): - resolve_candidate("cand-fail", job_id=JOB_ID, endpoint=ENDPOINT) + resolve_candidate("cand-fail", endpoint=ENDPOINT) assert "cand-fail" not in _downloaded @@ -177,29 +176,32 @@ def test_no_instructions_file_when_empty(self, tmp_path): assert not (candidate_path / "instructions.md").exists() - def test_writes_tools_json_dict_format(self, tmp_path): + def test_writes_tools_json_list_format_from_tools_key(self, tmp_path): candidate_path = tmp_path / "cand-4" config = { - "tool_descriptions": { - "search": {"description": "Search it", "parameters": {"q": "query"}}, - } + "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 tools["search"]["description"] == "Search it" + assert isinstance(tools, list) + assert tools[0]["function"]["name"] == "search" - def test_writes_tools_json_from_toolDescriptions(self, tmp_path): + def test_writes_tools_json_multiple_tools(self, tmp_path): candidate_path = tmp_path / "cand-5" config = { - "toolDescriptions": { - "lookup": {"description": "Look up policy"}, - } + "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 tools["lookup"]["description"] == "Look up policy" + 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" @@ -214,18 +216,15 @@ def test_writes_tools_json_list_format(self, tmp_path): assert isinstance(tools, list) assert tools[0]["function"]["name"] == "f1" - def test_tool_descriptions_wins_over_tools_list(self, tmp_path): - """tool_descriptions dict takes priority over tools list.""" + 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 = { - "tool_descriptions": {"search": {"description": "Dict format"}}, - "tools": [{"type": "function", "function": {"name": "f1"}}], + "tools": "not a list", } _persist_to_local_layout(candidate_path, config) - tools = json.loads((candidate_path / "tools.json").read_text()) - assert isinstance(tools, dict) - assert "search" in tools + assert not (candidate_path / "tools.json").exists() def test_no_tools_file_when_no_tools(self, tmp_path): candidate_path = tmp_path / "cand-8" @@ -255,6 +254,52 @@ def test_metadata_without_model_and_temperature(self, tmp_path): 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 ──────────────────────────────────── @@ -269,9 +314,9 @@ def test_round_trip(self, monkeypatch, tmp_path): "instructions": "Round-trip test.", "model": "gpt-4o", "temperature": 0.3, - "tool_descriptions": { - "search": {"description": "Find things", "parameters": {"q": "query"}}, - }, + "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) @@ -279,11 +324,12 @@ def test_round_trip(self, monkeypatch, tmp_path): # Now load via local dir monkeypatch.setenv("OPTIMIZATION_LOCAL_DIR", str(tmp_path)) monkeypatch.setenv("OPTIMIZATION_CANDIDATE_ID", "cand-rt") - loaded = load_config(default_instructions="unused") + loaded = load_config() assert loaded.instructions == "Round-trip test." assert loaded.model == "gpt-4o" assert loaded.temperature == 0.3 - assert "search" in loaded.tool_descriptions + assert len(loaded.tool_definitions) == 1 + assert loaded.tool_definitions[0]["function"]["name"] == "search" assert loaded.source.startswith("local:") @@ -312,7 +358,7 @@ def mock_text(client, path, params=None): 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, JOB_ID, "cand-sk", candidate_path) + _download_skill_files(mock_client, "cand-sk", candidate_path) skill_file = candidate_path / "skills" / "math" / "SKILL.md" assert skill_file.exists() @@ -325,7 +371,7 @@ def test_skips_when_no_manifest(self, tmp_path, mock_client): "azure.ai.agentserver.optimization._resolver._api_get_json", return_value=None, ): - _download_skill_files(mock_client, JOB_ID, "cand-no-manifest", candidate_path) + _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): @@ -336,7 +382,7 @@ def test_skips_when_no_skill_files_in_manifest(self, tmp_path, mock_client): "azure.ai.agentserver.optimization._resolver._api_get_json", return_value=manifest, ): - _download_skill_files(mock_client, JOB_ID, "cand-no-skills", candidate_path) + _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): @@ -347,7 +393,7 @@ def test_skips_empty_path_entries(self, tmp_path, mock_client): "azure.ai.agentserver.optimization._resolver._api_get_json", return_value=manifest, ): - _download_skill_files(mock_client, JOB_ID, "cand-empty-path", candidate_path) + _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): @@ -358,7 +404,7 @@ def test_handles_download_failure(self, tmp_path, mock_client): 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, JOB_ID, "cand-dl-fail", candidate_path) + _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() @@ -371,32 +417,78 @@ def test_rejects_traversal_in_file_path(self, tmp_path, mock_client): 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, JOB_ID, "cand-traversal", candidate_path) + _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) -# ── Path traversal in resolve_candidate ───────────────────────────── + 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 = [] -class TestPathTraversalGuard: - """Tests for path traversal prevention in resolve_candidate.""" + def mock_json(client, path, params=None): + called_paths.append(("json", path)) + return {"files": [{"path": "skills/s1/SKILL.md", "type": "skill"}]} - def test_rejects_traversal_candidate_id(self, tmp_path): - """candidate_id with '../' is rejected before any API call.""" - result = resolve_candidate( - "../../etc", job_id=JOB_ID, endpoint=ENDPOINT, local_dir=tmp_path - ) - assert result is None + def mock_text(client, path, params=None): + called_paths.append(("text", path)) + return "content" - def test_rejects_absolute_candidate_id(self, tmp_path): - """Absolute path in candidate_id is rejected.""" - result = resolve_candidate( - "/etc/passwd", job_id=JOB_ID, endpoint=ENDPOINT, local_dir=tmp_path - ) - assert result is None + 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"} @@ -408,11 +500,120 @@ def test_normal_candidate_id_allowed(self, tmp_path): ), ): result = resolve_candidate( - "valid-candidate-123", job_id=JOB_ID, endpoint=ENDPOINT, local_dir=tmp_path + "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 ────────────────────────────────────────────────── @@ -452,7 +653,7 @@ def test_persist_oserror_does_not_crash(self, tmp_path): ), ): result = resolve_candidate( - "cand-io", job_id=JOB_ID, endpoint=ENDPOINT, local_dir=tmp_path, + "cand-io", endpoint=ENDPOINT, local_dir=tmp_path, ) # Config is still returned from API even if persist fails assert result is not None From 4158e5072ca5fb4daaf47e6c62fe597ec4422f97 Mon Sep 17 00:00:00 2001 From: zyysurely Date: Sun, 24 May 2026 17:03:21 -0700 Subject: [PATCH 5/5] update warning --- .../azure/ai/agentserver/optimization/_config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 index bebf4db973af..dabc868a9b88 100644 --- 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 @@ -179,8 +179,8 @@ def _load_config_inner( 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 at least " - "an instructions.md file, or pass required=False." + 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