diff --git a/packages/gen/docs/gen_ai_hub/examples/prompt-optimization.ipynb b/packages/gen/docs/gen_ai_hub/examples/prompt-optimization.ipynb new file mode 100644 index 00000000..94197368 --- /dev/null +++ b/packages/gen/docs/gen_ai_hub/examples/prompt-optimization.ipynb @@ -0,0 +1,645 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "intro", + "metadata": {}, + "source": [ + "# Prompt Optimization\n", + "\n", + "This notebook shows how to use the **Prompt Optimization** feature to improve a prompt template for one or more target models, based on a labelled dataset and an optimization metric.\n", + "\n", + "The SDK handles:\n", + "- Uploading your dataset to object storage\n", + "- Registering the AI Core configuration and execution\n", + "- Polling for completion\n", + "- Fetching the optimized prompt templates and evaluation scores" + ] + }, + { + "cell_type": "markdown", + "id": "setup-header", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "setup-client", + "metadata": {}, + "outputs": [], + "source": [ + "from gen_ai_hub.optimizations.client import OptimizationClient\n", + "from gen_ai_hub.optimizations.models import PromptOptimizationConfig\n", + "from dotenv import load_dotenv\n", + "import os\n", + "\n", + "load_dotenv(override=True)\n", + "\n", + "client = OptimizationClient(\n", + " base_url=os.getenv(\"AICORE_BASE_URL\"),\n", + " auth_url=os.getenv(\"AICORE_AUTH_URL\"),\n", + " client_id=os.getenv(\"AICORE_CLIENT_ID\"),\n", + " client_secret=os.getenv(\"AICORE_CLIENT_SECRET\"),\n", + " resource_group=os.getenv(\"AICORE_RESOURCE_GROUP\", \"default\"),\n", + " aws_access_key_id=os.getenv(\"AWS_ACCESS_KEY_ID\"),\n", + " aws_secret_access_key=os.getenv(\"AWS_SECRET_ACCESS_KEY\"),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "setup-secret-header", + "metadata": {}, + "source": [ + "## One-Time Setup: Object Store Secret\n", + "\n", + "The SDK needs an object store secret to upload your dataset and store output artifacts. Run this once per resource group — you can skip it on subsequent runs." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "setup-secret", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "\n", + "AWS_S3_ENDPOINT = os.getenv(\"AWS_S3_ENDPOINT\", \"s3-eu-central-1.amazonaws.com\")\n", + "AWS_BUCKET_ID = os.getenv(\"AWS_BUCKET_ID\")\n", + "AWS_REGION = os.getenv(\"AWS_REGION\", \"eu-central-1\")\n", + "\n", + "default_secret_creds = {\n", + " \"data\": {},\n", + " \"type\": \"S3\",\n", + " \"pathPrefix\": \"sdkOutputFiles\",\n", + " \"endpoint\": AWS_S3_ENDPOINT,\n", + " \"bucket\": AWS_BUCKET_ID,\n", + " \"region\": AWS_REGION,\n", + " \"usehttps\": \"1\",\n", + "}\n", + "\n", + "response = client.setup(default_secret_body=default_secret_creds, replace_existing=True)" + ] + }, + { + "cell_type": "markdown", + "id": "dataset-header", + "metadata": {}, + "source": [ + "## Dataset Format\n", + "\n", + "Your dataset must be a JSON file — a list of objects where each entry has exactly two top-level keys:\n", + "\n", + "- **`fields`**: a dictionary of input variables. The keys must match the placeholder variables defined in your base prompt template (e.g. `{{?input}}`, `{{?page}}`, `{{?format_instructions}}`). Any number of fields is supported. Keys and values must be strings and must be consistent across all entries.\n", + "- **`answer`**: the expected output string (can be plain text or a JSON string)\n", + "\n", + "```json\n", + "[\n", + " {\n", + " \"fields\": {\n", + " \"input\": \"Your input text here\"\n", + " },\n", + " \"answer\": \"Expected output\"\n", + " }\n", + "]\n", + "```\n", + "\n", + "For prompts with multiple placeholders:\n", + "\n", + "```json\n", + "[\n", + " {\n", + " \"fields\": {\n", + " \"page\": \"document text...\",\n", + " \"format_instructions\": \"Return JSON with fields X, Y, Z\"\n", + " },\n", + " \"answer\": \"{\\\"x\\\": \\\"...\\\", \\\"y\\\": \\\"...\\\", \\\"z\\\": \\\"...\\\"}\"\n", + " }\n", + "]\n", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "config-header", + "metadata": {}, + "source": [ + "## Defining the Optimization Config\n", + "\n", + "`PromptOptimizationConfig` specifies everything the optimizer needs:\n", + "\n", + "| Parameter | Required | Description |\n", + "|---|---|---|\n", + "| `dataset_path` | Yes* | Local path to your labelled JSON dataset. The SDK uploads it automatically. |\n", + "| `base_prompt` | Yes | Starting prompt template in `/:` format |\n", + "| `target_models` | Yes | Models to optimize for, e.g. `[\"gpt-4o:2024-11-20\"]` |\n", + "| `target_prompt_mapping` | Yes | Maps each target model to the output prompt name in the Prompt Registry |\n", + "| `optimization_metric` | Yes* | System metric to optimize for, e.g. `\"JSON_Match\"` |\n", + "| `custom_metric_id` | Yes* | Alternative to `optimization_metric` — use the ID of a custom metric |\n", + "| `base_model` | No | Reference model used internally during optimization |\n", + "| `include_few_shot_examples` | No | Whether to include few-shot examples (default: `False`) |\n", + "| `maximize` | No | Whether higher metric scores are better (default: `True`) |\n", + "| `correctness_cutoff` | No | Score threshold for correctness classification |\n", + "| `prompt_template_scope` | No | `\"tenant\"` or `\"resourcegroup\"` (default: `\"tenant\"`) |\n", + "| `prototype_mode` | No | Use as few as 3 samples for quick prototyping (default: `False`) |\n", + "\n", + "*Either `dataset_path` or `artifact_id`+`dataset` must be provided. Either `optimization_metric` or `custom_metric_id` must be provided." + ] + }, + { + "cell_type": "markdown", + "id": "c05a506f", + "metadata": {}, + "source": [ + "## Creating a Base Prompt Template\n", + "\n", + "Before running optimization, you need a base prompt template in the Prompt Registry. The example below creates one — skip this if you already have one and just set `base_prompt` to your existing `/:`." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "f0bb77fd", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Base prompt created: 252f28cf-1c04-4438-b131-095dde3a6a5c\n", + "Use as base_prompt: genai-optimizations/my-base-prompt:0.0.1\n" + ] + } + ], + "source": [ + "from gen_ai_hub.prompt_registry.client import PromptTemplateClient\n", + "from gen_ai_hub.prompt_registry.models.prompt_template import PromptTemplate, PromptTemplateSpec\n", + "\n", + "BASE_PROMPT_NAME = \"my-base-prompt\"\n", + "BASE_PROMPT_VERSION = \"0.0.1\"\n", + "BASE_PROMPT_SCENARIO = \"genai-optimizations\"\n", + "\n", + "prompt_client = PromptTemplateClient(proxy_client=client._gen_ai_hub_proxy_client)\n", + "\n", + "spec = PromptTemplateSpec(\n", + " template=[\n", + " PromptTemplate(role=\"system\", content=\"You are a helpful assistant\"),\n", + " PromptTemplate(\n", + " role=\"user\",\n", + " content=(\n", + " \"Giving the following message --- {{?input}} --- \"\n", + " \"Extract and return a json with the following keys and values: \"\n", + " \"- 'urgency' as one of `high`, `medium`, `low` \"\n", + " \"- 'sentiment' as one of `negative`, `neutral`, `positive` \"\n", + " \"- 'categories' Create a dictionary with categories as keys and boolean values (True/False), \"\n", + " \"where the value indicates whether the category is one of the best matching support category tags from: \"\n", + " \"`emergency_repair_services`, `routine_maintenance_requests`, `quality_and_safety_concerns`, \"\n", + " \"`specialized_cleaning_services`, `general_inquiries`, `sustainability_and_environmental_practices`, \"\n", + " \"`training_and_support_requests`, `cleaning_services_scheduling`, `customer_feedback_and_complaints`, \"\n", + " \"`facility_management_issues` \"\n", + " \"Your complete message should be a valid json string that can be read directly and only contain \"\n", + " \"the keys mentioned in the list above. Never enclose it in ```json...```, no newlines, no unnecessary whitespaces.\"\n", + " ),\n", + " ),\n", + " ],\n", + " defaults={\"input\": \"\"},\n", + " additional_fields={\n", + " \"modelParams\": {\"temperature\": 0.7, \"max_tokens\": 100},\n", + " \"modelGroup\": \"chat\",\n", + " },\n", + ")\n", + "\n", + "response = prompt_client.create_prompt_template(\n", + " name=BASE_PROMPT_NAME,\n", + " version=BASE_PROMPT_VERSION,\n", + " scenario=BASE_PROMPT_SCENARIO,\n", + " prompt_template_spec=spec,\n", + ")\n", + "base_prompt_id = response.id\n", + "base_prompt_ref = f\"{BASE_PROMPT_SCENARIO}/{BASE_PROMPT_NAME}:{BASE_PROMPT_VERSION}\"\n", + "print(f\"Base prompt created: {base_prompt_id}\")\n", + "print(f\"Use as base_prompt: {base_prompt_ref}\")" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "config-code", + "metadata": {}, + "outputs": [], + "source": [ + "# Option 1: local dataset file — SDK uploads it to S3 automatically.\n", + "PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(\"__file__\"), \"../../..\"))\n", + "dataset_path = os.path.join(PROJECT_ROOT, \"integration_tests/optimizations/po_dataset.json\")\n", + "\n", + "optimization_config = PromptOptimizationConfig(\n", + " dataset_path=dataset_path,\n", + " base_prompt=base_prompt_ref,\n", + " target_models=[\"gemini-2.5-pro:001\"],\n", + " target_prompt_mapping={\n", + " \"gemini-2.5-pro:001\": \"base-prompt-custom-gemini-25-pro:0.0.1\",\n", + " },\n", + " optimization_metric=\"JSON_Match\",\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "534bcca8", + "metadata": {}, + "outputs": [], + "source": [ + "# Option 2: reuse an already-uploaded artifact — skips the S3 upload.\n", + "#\n", + "# optimization_config = PromptOptimizationConfig(\n", + "# artifact_id=\"\",\n", + "# dataset=\"testdata/.json\",\n", + "# base_prompt=\"genai-optimizations/my-base-prompt:1.0.0\",\n", + "# target_models=[\"gpt-4o:2024-11-20\"],\n", + "# target_prompt_mapping={\n", + "# \"gpt-4o:2024-11-20\": \"my-optimized-prompt:0.0.1\",\n", + "# },\n", + "# optimization_metric=\"JSON_Match\",\n", + "# )" + ] + }, + { + "cell_type": "markdown", + "id": "run-header", + "metadata": {}, + "source": [ + "## Running the Optimization\n", + "\n", + "`client.optimize()` validates the config, uploads the dataset, registers the AI Core execution, and returns an `OptimizationRun` object immediately. The job runs asynchronously in the background." + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "run-code", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Optimization started. Execution ID: e001e2b8bdc9d575\n" + ] + } + ], + "source": [ + "optimization_run = client.optimize(optimization_config)\n", + "print(f\"Optimization started. Execution ID: {optimization_run._run_context.execution_id}\")" + ] + }, + { + "cell_type": "markdown", + "id": "wait-header", + "metadata": {}, + "source": [ + "## Wait for Completion\n", + "\n", + "Use `wait_for_completion()` to block until the job finishes. You can also call `get_current_status()` at any point to check progress without blocking." + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "wait-code", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Status: Status.COMPLETED\n" + ] + } + ], + "source": [ + "optimization_run.wait_for_completion(timeout=3600) # timeout in seconds, default is 1 hour\n", + "print(f\"Status: {optimization_run.get_current_status()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "debug-header", + "metadata": {}, + "source": [ + "## Debugging\n", + "\n", + "If the job fails, use these helpers to investigate." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "debug-code", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "\n", + "# Structured summary of which step failed and why\n", + "debug_info = optimization_run.get_debug_info()\n", + "print(\"Debug info:\", json.dumps(debug_info, indent=4, default=str))\n", + "\n", + "# Full execution logs\n", + "debug_logs = optimization_run.get_debug_logs()\n", + "print(\"Logs:\", json.dumps(debug_logs, indent=4, default=str))" + ] + }, + { + "cell_type": "markdown", + "id": "results-header", + "metadata": {}, + "source": [ + "## Viewing Results\n", + "\n", + "`results()` returns an `OptimizationResults` object with:\n", + "\n", + "- **`metrics`** — pre/post evaluation scores per model from the tracking service\n", + "- **`prompts`** — the optimized prompt template from the Prompt Registry for each target model\n", + "\n", + "Printing the object shows a formatted summary of both.\n", + "\n", + "> **Note:** If the optimized prompt could not be fetched from the Prompt Registry, the results will still contain the metrics but the prompt section will be omitted. In that case, you can retrieve the optimized prompt from the execution logs using `get_debug_logs()`" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "id": "results-code", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "OptimizationResults\n", + "══════════════════════════════════════════════════════════════════════\n", + "\n", + " Score Summary\n", + " ┌────────────────────┬──────────┬───────┬───────┬─────────────┐\n", + " │ Model │ Baseline │ Pre │ Post │ Improvement │\n", + " ├────────────────────┼──────────┼───────┼───────┼─────────────┤\n", + " │ gemini-2.5-pro:001 │ – │ 0.910 │ 0.963 │ ▲ +5.8% │\n", + " └────────────────────┴──────────┴───────┴───────┴─────────────┘\n", + "\n", + " Evaluation Details ─ gemini-2.5-pro:001\n", + " ┌────────────┬───────┬───────┐\n", + " │ Metric │ Pre │ Post │\n", + " ├────────────┼───────┼───────┤\n", + " │ f1 │ 0.910 │ 0.963 │\n", + " │ tp │ 273 │ 289 │\n", + " │ llm_p │ 300 │ 300 │\n", + " │ recall │ 0.910 │ 0.963 │\n", + " │ ground_p │ 300 │ 300 │\n", + " │ precision │ 0.910 │ 0.963 │\n", + " │ is_correct │ 23 │ 25 │\n", + " └────────────┴───────┴───────┘\n", + "\n", + " Custom Info ─ gemini-2.5-pro:001\n", + " llm_request_metrics: {'provider_request_counts': [{'provider': {'model_name': 'gpt-5', 'model_params': {'reasoning_effort': 'medium'}, 'model_version': '2025-08-07'}, 'num_requests': 26, 'purpose': 'optimization', 'input_tokens': 171485, 'output_tokens': 116063}, {'provider': {'model_name': 'gemini-2.5-pro', 'model_params': {}, 'model_version': '001'}, 'num_requests': 470, 'purpose': 'target', 'input_tokens': 624256, 'output_tokens': 583104}, {'provider': {'model_name': 'gemini-2.5-pro', 'model_params': {}, 'model_version': '001'}, 'num_requests': 22, 'purpose': 'optimization', 'input_tokens': 161751, 'output_tokens': 69707}]}\n", + "\n", + " Optimised Prompt ─ gemini-2.5-pro:001\n", + " ┌──────────────────────────────────────────────────────────────────────┐\n", + " │ [system] You are a structured classification assistant that │\n", + " │ extracts metadata from support messages and outputs │\n", + " │ strict, valid JSON only. Your job: given one message, │\n", + " │ determine (a) urgency, (b) sentiment, and (c) category │\n", + " │ flags. Output rules (must follow exactly): - Produce one │\n", + " │ minified JSON object with exactly these top-level keys: │\n", + " │ \"urgency\", \"sentiment\", \"categories\". - Use double quotes │\n", + " │ for all keys and string values. - Use lowercase JSON │\n", + " │ booleans true/false inside \"categories\". - Single line │\n", + " │ only (no newlines), no extra text, no code fences, no │\n", + " │ trailing commas, no additional keys. - Multi-label is │\n", + " │ allowed and expected: evaluate each category independently │\n", + " │ and set true for every applicable category; set false │\n", + " │ otherwise. Include all category keys in the \"categories\" │\n", + " │ object, spelled exactly. Categories (evaluate │\n", + " │ independently): - emergency_repair_services: Urgent fixes │\n", + " │ to broken systems/equipment or utilities (not cleaning │\n", + " │ tasks). - routine_maintenance_requests: Regular or │\n", + " │ scheduled upkeep/repairs without urgency. - │\n", + " │ quality_and_safety_concerns: Safety, quality, sanitation, │\n", + " │ hazards, or regulatory/compliance risks explicitly │\n", + " │ mentioned (e.g., spill, contamination, biohazard, │\n", + " │ chemical, hazmat, mold, unsafe, violation). - │\n", + " │ specialized_cleaning_services: Specialized or hazardous │\n", + " │ cleaning (e.g., biohazard, chemical, post-construction, │\n", + " │ hazmat). - general_inquiries: Requests for information, │\n", + " │ brochures, pricing, or service details. - │\n", + " │ sustainability_and_environmental_practices: Eco-friendly │\n", + " │ products, recycling, sustainability policies/practices, │\n", + " │ energy efficiency. - training_and_support_requests: │\n", + " │ Training sessions, documentation, how-to support. - │\n", + " │ cleaning_services_scheduling: Explicit │\n", + " │ scheduling/rescheduling of cleaning dates/times or │\n", + " │ availability checks. - customer_feedback_and_complaints: │\n", + " │ Feedback, praise, or complaints about service quality. - │\n", + " │ facility_management_issues: Facility operations topics │\n", + " │ beyond cleaning services, including (a) │\n", + " │ malfunctions/outages of building systems/utilities │\n", + " │ (elevators, HVAC, plumbing, power), (b) contracts and │\n", + " │ renewals, (c) coordination/communication or service-level │\n", + " │ discussions, (d) strategic facility planning/operations, │\n", + " │ and (e) sustainability or energy-efficiency consultations. │\n", + " │ Think step by step internally to apply the rules in the │\n", + " │ prompt template, but output only the final JSON object. │\n", + " │ │\n", + " │ [user] # Role and Objective Classify the message and return │\n", + " │ strict, minified JSON with: urgency, sentiment, and a │\n", + " │ categories object containing all category flags. # │\n", + " │ Instructions - Determine and return: - \"urgency\": one of │\n", + " │ \"high\", \"medium\", or \"low\" using the rules below. - │\n", + " │ \"sentiment\": one of \"negative\", \"neutral\", or \"positive\" │\n", + " │ using the rules below. - \"categories\": an object with │\n", + " │ all category keys present and boolean values (true/false) │\n", + " │ indicating applicability. - Multi-label is allowed: set │\n", + " │ true for every applicable category; set false otherwise. │\n", + " │ Evaluate each category independently. - Formatting must be │\n", + " │ a single-line, minified JSON object with exactly the keys │\n", + " │ \"urgency\", \"sentiment\", and \"categories\". Use double │\n", + " │ quotes for keys/strings and lowercase true/false. No extra │\n", + " │ keys, text, or code fences. - Use the category definitions │\n", + " │ provided in the system prompt. Apply the decision rules │\n", + " │ below for edge cases. - Canonical JSON skeleton (use exact │\n", + " │ key spelling and structure; values should reflect your │\n", + " │ classification): {\"urgency\":\"\",\"sentiment │\n", + " │ \":\"\",\"categories\":{\"emergency_r │\n", + " │ epair_services\":false,\"routine_maintenance_requests\":false │\n", + " │ ,\"quality_and_safety_concerns\":false,\"specialized_cleaning │\n", + " │ _services\":false,\"general_inquiries\":false,\"sustainability │\n", + " │ _and_environmental_practices\":false,\"training_and_support_ │\n", + " │ requests\":false,\"cleaning_services_scheduling\":false,\"cust │\n", + " │ omer_feedback_and_complaints\":false,\"facility_management_i │\n", + " │ ssues\":false}} ## Decision and Tie-breaker Rules - │\n", + " │ Scheduling vs general inquiries: If the message is │\n", + " │ booking/rescheduling cleaning or asking │\n", + " │ availability/timing, set cleaning_services_scheduling=true │\n", + " │ and general_inquiries=false unless there is a separate │\n", + " │ request for brochures/pricing/policies (then also set │\n", + " │ general_inquiries=true). Requests like \"send personnel\" or │\n", + " │ \"dispatch a team\" without dates/availability do NOT │\n", + " │ trigger cleaning_services_scheduling. - Specialized │\n", + " │ cleaning vs quality_and_safety_concerns: Set │\n", + " │ quality_and_safety_concerns only when explicit │\n", + " │ hazards/compliance/sanitation risks are mentioned (e.g., │\n", + " │ spill, contamination, biohazard, chemical, hazmat, mold, │\n", + " │ violation). Do not infer hazards from specialized cleaning │\n", + " │ alone. - Facility management scope: Set │\n", + " │ facility_management_issues=true for (a) │\n", + " │ malfunctions/outages of building systems/utilities, (b) │\n", + " │ contracts/renewals and service-level terms, (c) │\n", + " │ communication/coordination issues, (d) strategic facility │\n", + " │ planning/operations topics, and (e) sustainability or │\n", + " │ energy-efficiency consultations. If an urgent repair is │\n", + " │ requested for building systems, also set │\n", + " │ emergency_repair_services=true. - Training and │\n", + " │ safety/compliance: Set training_and_support_requests=true │\n", + " │ for requests for training/materials. Also set │\n", + " │ quality_and_safety_concerns=true only when there are │\n", + " │ explicit hazard/compliance/sanitation risks or │\n", + " │ regulatory/compliance requirements (e.g., facility safety │\n", + " │ compliance documentation). Generic training on cleaning │\n", + " │ protocols or safety-themed training without explicit │\n", + " │ risk/violation does not require │\n", + " │ quality_and_safety_concerns. - Customer feedback: Set │\n", + " │ customer_feedback_and_complaints=true for complaints or │\n", + " │ praise about service quality; co-tag │\n", + " │ quality_and_safety_concerns only when sanitation/safety │\n", + " │ risks are explicitly cited. - Do not infer hazards unless │\n", + " │ explicitly stated. ## Urgency Rules - high: Explicit │\n", + " │ urgency cues (\"urgent\", \"immediate\", \"ASAP\", \"emergency\"), │\n", + " │ explicit hazards/spills, or active service interruptions │\n", + " │ requiring immediate action. - medium: Explicit time bounds │\n", + " │ or requests that need attention soon but are not │\n", + " │ emergencies (e.g., scheduling or training with a specific │\n", + " │ timeframe like \"next week\" or \"by Friday\"). - low: General │\n", + " │ information or routine maintenance scheduling without │\n", + " │ explicit time pressure. - Tie-breaker: If no clear urgency │\n", + " │ cues are present, default to low. ## Sentiment Rules and │\n", + " │ Cues - negative: Complaints/dissatisfaction or disruption │\n", + " │ from malfunctions/outages (e.g., \"unacceptable\", │\n", + " │ \"disappointed\", \"frustrated\", \"broken\", \"stuck\", \"outage\", │\n", + " │ \"not working\", \"need immediate repair\"). - positive: │\n", + " │ Explicit praise/appreciation (e.g., \"thank you\", │\n", + " │ \"appreciate\", \"great\", \"excellent\"); also consider │\n", + " │ proactive improvement or sustainability/efficiency │\n", + " │ consultation requests as positive when phrased │\n", + " │ constructively and with no negative cues. - neutral: │\n", + " │ Purely factual requests or hazard notices without │\n", + " │ expressed dissatisfaction. - Tie-breaker: If positive vs │\n", + " │ neutral is ambiguous and there are no explicit cues, │\n", + " │ default to neutral. # Reasoning Steps 1) Extract key │\n", + " │ cues: urgency terms, hazard mentions, │\n", + " │ scheduling/availability, malfunction/outage, │\n", + " │ training/safety/compliance, info requests, and feedback │\n", + " │ tone. 2) Assign \"urgency\" using the urgency rules. 3) │\n", + " │ Assign \"sentiment\" using the sentiment rules and tie- │\n", + " │ breakers. 4) Evaluate each category independently; apply │\n", + " │ the decision rules and set all applicable categories to │\n", + " │ true; otherwise false. Include all category keys. 5) │\n", + " │ Validate JSON: exactly the three top-level keys; double │\n", + " │ quotes for keys/strings; lowercase true/false; single-line │\n", + " │ minified; no extra text or trailing commas. # Context │\n", + " │ {{?input}} # Final instructions - │\n", + " │ Include all category keys with correct spelling; set true │\n", + " │ for every applicable category (multi-label allowed). - │\n", + " │ Prefer cleaning_services_scheduling over general_inquiries │\n", + " │ for booking/availability; only add general_inquiries if │\n", + " │ separate general info is requested. - Only set │\n", + " │ quality_and_safety_concerns when │\n", + " │ hazards/compliance/sanitation risks are explicit or when │\n", + " │ compliance documentation is requested; do not infer │\n", + " │ hazards. - Co-tag facility_management_issues with │\n", + " │ emergency_repair_services for urgent building system │\n", + " │ malfunctions. - When cues are unclear: default urgency to │\n", + " │ low and sentiment to neutral. - Output only the JSON │\n", + " │ object, single line, minified, with the exact keys │\n", + " │ \"urgency\", \"sentiment\", and \"categories\"; use double │\n", + " │ quotes and lowercase booleans. │\n", + " └──────────────────────────────────────────────────────────────────────┘\n" + ] + } + ], + "source": [ + "results = optimization_run.results()\n", + "print(results)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "da95e246-2aba-4793-8170-3671391f78bc", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "markdown", + "id": "e64fccfe", + "metadata": {}, + "source": [ + "## Cleanup\n", + "\n", + "Delete the base prompt template created earlier." + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "fd671332", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Base prompt deleted: 252f28cf-1c04-4438-b131-095dde3a6a5c\n" + ] + } + ], + "source": [ + "prompt_client.delete_prompt_template_by_id(base_prompt_id)\n", + "print(f\"Base prompt deleted: {base_prompt_id}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "da70481a-52a1-4fd5-96b2-1d75434b6bb8", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.7" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/packages/gen/gen_ai_hub/evaluations/client.py b/packages/gen/gen_ai_hub/evaluations/client.py index 59b399bf..49ac1e6e 100644 --- a/packages/gen/gen_ai_hub/evaluations/client.py +++ b/packages/gen/gen_ai_hub/evaluations/client.py @@ -181,6 +181,10 @@ def __init__( self._validate_provider_params() logger.info("Initialization of the client completed!") + @property + def _gen_ai_hub_proxy_client(self): + return self.__gen_ai_hub_proxy_client + def __repr__(self): attrs = ", ".join(f"{k}={v!r}" for k, v in vars(self).items()) return f"{self.__class__.__name__}({attrs})" diff --git a/packages/gen/gen_ai_hub/optimizations/__init__.py b/packages/gen/gen_ai_hub/optimizations/__init__.py new file mode 100644 index 00000000..e2cf27b6 --- /dev/null +++ b/packages/gen/gen_ai_hub/optimizations/__init__.py @@ -0,0 +1,4 @@ +from .client import OptimizationClient +from .models import OptimizationRun, PromptOptimizationConfig, OptimizationResults + +__all__ = ["OptimizationClient", "OptimizationRun", "PromptOptimizationConfig", "OptimizationResults"] diff --git a/packages/gen/gen_ai_hub/optimizations/client.py b/packages/gen/gen_ai_hub/optimizations/client.py new file mode 100644 index 00000000..f6b7ddf5 --- /dev/null +++ b/packages/gen/gen_ai_hub/optimizations/client.py @@ -0,0 +1,59 @@ +"""Client for submitting and managing prompt optimization jobs on generative AI Hub.""" +from gen_ai_hub.evaluations._internal._models import _AWSObjectStoreData +from gen_ai_hub.evaluations.client import EvaluationClient +from gen_ai_hub.evaluations.constants import DEFAULT_KEY +from gen_ai_hub.evaluations.exceptions.error_codes import ErrorCode +from gen_ai_hub.evaluations.helpers.collector import ValidationCollector +from gen_ai_hub.evaluations.utils.oss_secret_utils import fetch_object_store_secret_by_name +from gen_ai_hub.optimizations.models.optimization_config import PromptOptimizationConfig +from gen_ai_hub.optimizations.models.optimization_run import OptimizationRun +from gen_ai_hub.optimizations.optimization_flow import optimization_job_flow +from gen_ai_hub.optimizations.utils import validate_optimization_config + + +class OptimizationClient(EvaluationClient): + """Client for running prompt optimization jobs against a target metric on generative AI Hub.""" + + def optimize(self, optimization_config: PromptOptimizationConfig) -> OptimizationRun: + """Submit a prompt optimization job and return an OptimizationRun to track its progress.""" + error_collector = ValidationCollector() + try: + if self.default_object_store_secret_name is None: + response = fetch_object_store_secret_by_name( + self.ai_core_client, + DEFAULT_KEY, + self.resource_group, + error_collector, + ) + if response is None: + error_collector.add_error( + ErrorCode.MISSING_DEFAULT_OBJECT_STORE_SECRET_ERROR.value, + "Default Object Store secret is required to run optimize function. " + "Please use setup() function to create one!", + ) + + error_collector.raise_if_errors() + + validate_optimization_config( + optimization_config, + self.ai_core_client, + self.resource_group, + error_collector, + ) + error_collector.raise_if_errors() + + object_store_credentials = _AWSObjectStoreData( + aws_access_key_id=self.aws_access_key_id, + aws_secret_access_key=self.aws_secret_access_key, + ) + return optimization_job_flow( + optimization_config, + object_store_credentials, + self.ai_core_client, + self.resource_group, + error_collector, + proxy_client=self._gen_ai_hub_proxy_client, + ) + except Exception as exc: + error_collector.raise_if_errors() + raise RuntimeError("Optimize function failed!") from exc diff --git a/packages/gen/gen_ai_hub/optimizations/constants.py b/packages/gen/gen_ai_hub/optimizations/constants.py new file mode 100644 index 00000000..f7d6a646 --- /dev/null +++ b/packages/gen/gen_ai_hub/optimizations/constants.py @@ -0,0 +1,3 @@ +OPTIMIZATIONS_SCENARIO_ID = "genai-optimizations" +OPTIMIZATIONS_CONFIG_PREFIX_KEY = "optimization-config-" +OPTIMIZATIONS_ARTIFACT_KEY = "prompt-data" diff --git a/packages/gen/gen_ai_hub/optimizations/models/__init__.py b/packages/gen/gen_ai_hub/optimizations/models/__init__.py new file mode 100644 index 00000000..fb01c493 --- /dev/null +++ b/packages/gen/gen_ai_hub/optimizations/models/__init__.py @@ -0,0 +1,5 @@ +from .optimization_config import PromptOptimizationConfig +from .optimization_run import OptimizationRun +from .optimization_results import OptimizationResults + +__all__ = ["PromptOptimizationConfig", "OptimizationRun", "OptimizationResults"] diff --git a/packages/gen/gen_ai_hub/optimizations/models/optimization_config.py b/packages/gen/gen_ai_hub/optimizations/models/optimization_config.py new file mode 100644 index 00000000..636d30a7 --- /dev/null +++ b/packages/gen/gen_ai_hub/optimizations/models/optimization_config.py @@ -0,0 +1,102 @@ +"""PromptOptimizationConfig: configuration model for prompt optimization jobs.""" +from typing import Dict, List, Optional + + +class PromptOptimizationConfig: + """Configuration for a prompt optimization job. + + :param target_prompt_mapping: Maps each target model (``name:version``) to the output + prompt name in the Prompt Registry. + :param target_models: List of models to optimize for, e.g. ``["gpt-4o:2024-11-20"]``. + :param base_prompt: Starting prompt template referenced as ``/:``. + :param optimization_metric: System-defined metric to optimize for (e.g. ``"JSON_Match"``). + Mutually exclusive with ``custom_metric_id``. + :param artifact_id: AI Core artifact ID of a pre-uploaded dataset. Mutually exclusive with ``dataset_path``. + :param dataset: Relative file key within the artifact (required when using ``artifact_id``). + :param dataset_path: Local path to a labelled JSON dataset. The SDK uploads it automatically. + Mutually exclusive with ``artifact_id``. + :param base_model: Reference model used internally during optimization. + :param include_few_shot_examples: Whether to include few-shot examples in the optimized prompt + (default: ``False``). + :param custom_metric_id: ID of a custom metric to optimize for. Mutually exclusive with + ``optimization_metric``. + :param maximize: Whether higher metric scores are better (default: ``True``). + :param correctness_cutoff: Score threshold for correctness classification. + :param prompt_template_scope: Scope for the output prompt template — ``"tenant"`` or + ``"resourcegroup"`` (default: ``"tenant"``). + :param prototype_mode: Use as few as 3 samples for quick prototyping (default: ``False``). + :param train_dataset_config: Optional training dataset config. Requires ``test_dataset_config``. + :param test_dataset_config: Optional test dataset config. Required when ``train_dataset_config`` is provided. + :param model_params: JSON string mapping model IDs to parameter dicts (e.g. ``temperature``, ``max_tokens``). + :param variable_mapping: JSON string mapping prompt template variable names to dataset field names. + :param field_evaluation_metrics: JSON string mapping response format field names to their evaluation metrics + (e.g. ``'{"urgency": "ExactMatch", "sentiment": "LLMaaJ:Sem_Sim_1"}'``). Requires a ``response_format`` + to be defined in the base prompt template. Mutually exclusive with ``optimization_metric`` and ``custom_metric_id``. + """ + + def __init__( + self, + target_prompt_mapping: Dict[str, str], + target_models: List[str], + base_prompt: str, + optimization_metric: Optional[str] = None, + artifact_id: Optional[str] = None, + dataset: Optional[str] = None, + dataset_path: Optional[str] = None, + base_model: Optional[str] = "none", + include_few_shot_examples: Optional[bool] = False, + custom_metric_id: Optional[str] = None, + maximize: Optional[bool] = True, + correctness_cutoff: Optional[float] = None, + prompt_template_scope: Optional[str] = "tenant", + prototype_mode: Optional[bool] = False, + train_dataset_config=None, + test_dataset_config=None, + model_params: Optional[str] = None, + variable_mapping: Optional[str] = None, + field_evaluation_metrics: Optional[str] = None, + ): + self.artifact_id = artifact_id + self.dataset = dataset + self.dataset_path = dataset_path + self.target_prompt_mapping = target_prompt_mapping + self.target_models = target_models + self.base_prompt = base_prompt + self.optimization_metric = optimization_metric + self.base_model = base_model + self.include_few_shot_examples = include_few_shot_examples + self.custom_metric_id = custom_metric_id + self.maximize = maximize + self.correctness_cutoff = correctness_cutoff + self.prompt_template_scope = prompt_template_scope + self.prototype_mode = prototype_mode + self.train_dataset_config = train_dataset_config + self.test_dataset_config = test_dataset_config + self.model_params = model_params + self.variable_mapping = variable_mapping + self.field_evaluation_metrics = field_evaluation_metrics + self._validate( + artifact_id, dataset, dataset_path, + optimization_metric, custom_metric_id, + train_dataset_config, test_dataset_config, + field_evaluation_metrics, + ) + + def _validate(self, artifact_id, dataset, dataset_path, + optimization_metric, custom_metric_id, + train_dataset_config, test_dataset_config, + field_evaluation_metrics=None): + if artifact_id is None and dataset_path is None: + raise ValueError("Either artifact_id or dataset_path must be provided.") + if artifact_id is not None and dataset_path is not None: + raise ValueError("Only one of artifact_id or dataset_path must be provided, not both.") + if artifact_id is not None and dataset is None: + raise ValueError("dataset (filename) must be provided when using artifact_id.") + if train_dataset_config is not None and test_dataset_config is None: + raise ValueError("test_dataset_config must be provided when train_dataset_config is provided.") + if optimization_metric is None and custom_metric_id is None and field_evaluation_metrics is None: + raise ValueError( + "At least one of optimization_metric, custom_metric_id, or field_evaluation_metrics must be provided." + ) + if optimization_metric is not None and custom_metric_id is not None: + raise ValueError("Only one of optimization_metric or custom_metric_id must be provided, not both.") diff --git a/packages/gen/gen_ai_hub/optimizations/models/optimization_results.py b/packages/gen/gen_ai_hub/optimizations/models/optimization_results.py new file mode 100644 index 00000000..79536c83 --- /dev/null +++ b/packages/gen/gen_ai_hub/optimizations/models/optimization_results.py @@ -0,0 +1,237 @@ +"""OptimizationResults: model for retrieving and displaying prompt optimization job outcomes.""" +import json +import textwrap + + +def _fmt_num(value): + if value is None: + return "–" + if isinstance(value, float): + return f"{value:.3f}" + if isinstance(value, int): + return f"{value:,}" + return str(value) + + +def _parse_custom_value(custom_val): + try: + return json.loads(custom_val.value) + except (json.JSONDecodeError, TypeError): + return custom_val.value + + +class OptimizationResults: + """Holds and displays the evaluation results produced by a completed prompt optimization job.""" + + def __init__(self, metrics, prompts: dict): + self.metrics = metrics + self.prompts = prompts + + @staticmethod + def _process_origin_resource(resource, entry, custom): + for metric in (resource.metrics or []): + entry.setdefault("baseline", metric.value) + for custom_val in (resource.custom_info or []): + custom[custom_val.name] = _parse_custom_value(custom_val) + + @staticmethod + def _process_target_metrics(resource, entry): + for metric in (resource.metrics or []): + label = next( + (lbl.value for lbl in (metric.labels or []) if lbl.name == "optimizer_metric_type"), + None, + ) + if not label: + continue + label_lower = label.lower() + if "pre" in label_lower: + entry["pre"] = metric.value + elif "post" in label_lower: + entry["post"] = metric.value + + @staticmethod + def _process_target_custom_info(resource, entry, custom): + for custom_val in (resource.custom_info or []): + val = _parse_custom_value(custom_val) + if custom_val.name == "pre_optimization_evaluation": + entry["pre_eval"] = val + elif custom_val.name == "post_optimization_evaluation": + entry["post_eval"] = val + else: + custom[custom_val.name] = val + + @staticmethod + def _extract_prompt(tmpl): + if tmpl is None: + return None + try: + spec = getattr(tmpl, "spec", None) + if spec and spec.template: + return [{"role": msg.role, "content": msg.content} for msg in spec.template] + except AttributeError: + pass + return None + + def _collect_data(self): + models = {} + + for resource in ((self.metrics and self.metrics.resources) or []): + tags = {tag.name: tag.value for tag in (resource.tags or [])} + model = tags.get("evaluation.ai.sap.com/model") or resource.execution_id + purpose = tags.get("evaluation.ai.sap.com/purpose", "unknown") + entry = models.setdefault(model, {}) + custom = entry.setdefault("custom", {}) + + if purpose == "origin": + self._process_origin_resource(resource, entry, custom) + elif purpose == "target": + self._process_target_metrics(resource, entry) + self._process_target_custom_info(resource, entry, custom) + + for model, tmpl in (self.prompts or {}).items(): + prompt = self._extract_prompt(tmpl) + if prompt is not None: + models.setdefault(model, {})["prompt"] = prompt + + return models + + @staticmethod + def _table(headers, rows, col_align=None): + num_cols = len(headers) + if col_align is None: + col_align = ["<"] * num_cols + + str_rows = [[str(cell) for cell in row] for row in rows] + widths = [len(h) for h in headers] + for row in str_rows: + for idx, cell in enumerate(row): + widths[idx] = max(widths[idx], len(cell)) + + def _row(cells, aligns): + parts = [] + for cell, wid, align in zip(cells, widths, aligns): + if align == ">": + formatted = cell.rjust(wid) + elif align == "^": + formatted = cell.center(wid) + else: + formatted = cell.ljust(wid) + parts.append(formatted) + return "│ " + " │ ".join(parts) + " │" + + top = "┌─" + "─┬─".join("─" * wid for wid in widths) + "─┐" + div = "├─" + "─┼─".join("─" * wid for wid in widths) + "─┤" + bot = "└─" + "─┴─".join("─" * wid for wid in widths) + "─┘" + lines = [top, _row(headers, ["^"] * num_cols), div] + lines += [_row(r, col_align) for r in str_rows] + lines.append(bot) + return "\n".join(lines) + + @staticmethod + def _prompt_box(messages, width=72): + prefix_w = 10 + text_w = width - 4 - prefix_w + lines = [] + for idx, msg in enumerate(messages): + if idx > 0: + lines.append("") + role_tag = f"[{msg['role']}]" + prefix = f"{role_tag:<{prefix_w}}" + wrapped = textwrap.wrap(msg["content"], width=text_w) or [""] + lines.append(prefix + wrapped[0]) + for extra in wrapped[1:]: + lines.append(" " * prefix_w + extra) + + inner = max((len(line) for line in lines), default=20) + top = "┌" + "─" * (inner + 2) + "┐" + bot = "└" + "─" * (inner + 2) + "┘" + body = ["│ " + line.ljust(inner) + " │" for line in lines] + return "\n".join([top] + body + [bot]) + + @staticmethod + def _compute_improvement(pre, post): + if pre is None or post is None: + return "─" + if pre != 0: + delta = (post - pre) / pre * 100 + return f"{'▲' if delta >= 0 else '▼'} {delta:+.1f}%" + diff = post - pre + return f"{'▲' if diff >= 0 else '▼'} {diff:+.3f}" + + def _build_score_rows(self, data): + rows = [] + for model, model_data in data.items(): + baseline, pre, post = model_data.get("baseline"), model_data.get("pre"), model_data.get("post") + if baseline is None and pre is None and post is None: + continue + rows.append([ + model, + f"{baseline:.3f}" if baseline is not None else "–", + f"{pre:.3f}" if pre is not None else "–", + f"{post:.3f}" if post is not None else "–", + self._compute_improvement(pre, post), + ]) + return rows + + def _render_eval_detail(self, model, model_data, out): + pre_eval = model_data.get("pre_eval") if isinstance(model_data.get("pre_eval"), dict) else None + post_eval = model_data.get("post_eval") if isinstance(model_data.get("post_eval"), dict) else None + if not (pre_eval or post_eval): + return + out.append(f" Evaluation Details ─ {model}") + keys = list(dict.fromkeys( + list(pre_eval.keys() if pre_eval else []) + + list(post_eval.keys() if post_eval else []) + )) + rows = [ + [key, + _fmt_num(pre_eval.get(key) if pre_eval else None), + _fmt_num(post_eval.get(key) if post_eval else None)] + for key in keys + ] + out.append(textwrap.indent( + self._table(["Metric", "Pre", "Post"], rows, ["<", ">", ">"]), + " " + )) + out.append("") + + def _render_custom_info(self, model, model_data, out): + custom = {key: val for key, val in model_data.get("custom", {}).items() if val is not None} + if not custom: + return + out.append(f" Custom Info ─ {model}") + for key, val in custom.items(): + out.append(f" {key}: {val}") + out.append("") + + def _render_prompt_section(self, model, model_data, out): + prompt = model_data.get("prompt") + if not prompt: + return + out.append(f" Optimised Prompt ─ {model}") + out.append(textwrap.indent(self._prompt_box(prompt), " ")) + out.append("") + + def __str__(self): + data = self._collect_data() + if not data: + return "OptimizationResults(no results available)" + + out = ["OptimizationResults", "═" * 70, ""] + + score_rows = self._build_score_rows(data) + if score_rows: + out.append(" Score Summary") + out.append(textwrap.indent( + self._table(["Model", "Baseline", "Pre", "Post", "Improvement"], + score_rows, ["<", "^", "^", "^", "^"]), + " " + )) + out.append("") + + for model, model_data in data.items(): + self._render_eval_detail(model, model_data, out) + self._render_custom_info(model, model_data, out) + self._render_prompt_section(model, model_data, out) + + return "\n".join(out).rstrip() diff --git a/packages/gen/gen_ai_hub/optimizations/models/optimization_run.py b/packages/gen/gen_ai_hub/optimizations/models/optimization_run.py new file mode 100644 index 00000000..89e1cf6a --- /dev/null +++ b/packages/gen/gen_ai_hub/optimizations/models/optimization_run.py @@ -0,0 +1,163 @@ +"""OptimizationRun: tracks the lifecycle and results of a prompt optimization job.""" +import logging + +from ai_api_client_sdk.models.metric_resource import MetricResource +from ai_api_client_sdk.models.status import Status +from ai_core_sdk.ai_core_v2_client import AICoreV2Client +from ai_core_sdk.tracking import Tracking + +from gen_ai_hub.evaluations.constants import ADDITIONAL_INFO_KEY +from gen_ai_hub.evaluations.models.evaluation_run import EvaluationRun +from gen_ai_hub.optimizations.models.optimization_results import OptimizationResults +from gen_ai_hub.prompt_registry.client import PromptTemplateClient + +logger = logging.getLogger(__name__) + + +class _OptimizationRunContext: + def __init__(self, execution_id, configuration_id, artifact_id, + ai_core_client, resource_group, proxy_client, target_prompt_mapping): + self.execution_id = execution_id + self.configuration_id = configuration_id + self.artifact_id = artifact_id + self.ai_core_client = ai_core_client + self.resource_group = resource_group + self.proxy_client = proxy_client + self.target_prompt_mapping = target_prompt_mapping + + +class OptimizationRun(EvaluationRun): + """Tracks execution state and exposes results for a submitted prompt optimization job.""" + + _STEP_ERROR_MESSAGES = { + "optimize": "Optimization job failed in the optimization step.", + "config": "Optimization job failed in Config Validation step.", + } + + def __init__( + self, + run_id: str, + execution_id: str, + ai_core_client: AICoreV2Client, + configuration_id: str = None, + artifact_id: str = None, + resource_group: str = None, + proxy_client=None, + target_prompt_mapping: dict = None, + ): + self.id = run_id + self.status = Status.UNKNOWN + self._run_context = _OptimizationRunContext( + execution_id=execution_id, + configuration_id=configuration_id, + artifact_id=artifact_id, + ai_core_client=ai_core_client, + resource_group=resource_group, + proxy_client=proxy_client, + target_prompt_mapping=target_prompt_mapping or {}, + ) + + def _enrich_failed_pods(self, failed_pods: list, workflow_lookup: dict) -> None: + for pod in failed_pods: + pod_name = pod.get("name", "") + if not self._apply_step_level_message(pod, pod_name): + self._apply_fallback_message(pod, pod_name, workflow_lookup) + + def _apply_fallback_message(self, pod: dict, pod_name: str, workflow_lookup: dict) -> None: + suffix = pod_name.split("-")[-1] + workflow = workflow_lookup.get(suffix) + if not workflow: + return + message = workflow.get("message", "Unknown error") + pod[ADDITIONAL_INFO_KEY] = f"Optimization job failed with error: {message}" + + def get_debug_info(self): + """Return structured debug information about the execution status.""" + execution_status_response = self._execution_status_fetcher() + current_status = execution_status_response.status + status_details = getattr(execution_status_response, "status_details", None) + + if not status_details: + return { + "status": current_status, + "details": ( + "No specific details found. Please use get_debug_logs() " + "to get more information." + ), + } + + failed_pod_details = self._extract_failed_pods(status_details) + workflow_lookup = self._build_workflow_lookup(status_details) + self._enrich_failed_pods(failed_pod_details, workflow_lookup) + + return {"status": current_status, "details": failed_pod_details} + + def aggregations(self): + """Fetch and return metric aggregations for the optimization run from the tracking service.""" + try: + tracking_client = Tracking( + base_url=self._run_context.ai_core_client.base_url, + token_creator=self._run_context.ai_core_client.rest_client.get_token, + resource_group=self._run_context.resource_group, + ) + result = tracking_client.query( + execution_ids=[self._run_context.execution_id], + resource_group=self._run_context.resource_group, + ) + path = f"/lm/metrics?tagFilters=evaluation.ai.sap.com/child-of={self._run_context.execution_id}" + child_response = self._run_context.ai_core_client.rest_client.get( + path=path, + resource_group=self._run_context.resource_group, + ) + child_resources = [ + MetricResource.from_dict(r) + for r in child_response.get("resources", []) + ] + result.resources = (result.resources or []) + child_resources + return result + except Exception as err: + logger.warning("Could not fetch aggregations for run %s: %s", self.id, err) + return None + + def results(self) -> OptimizationResults: + """Fetch and return the optimization results including metrics and optimized prompts.""" + execution_status = self._execution_status_fetcher().status + if execution_status != Status.COMPLETED: + if execution_status == Status.RUNNING: + raise ValueError( + "Status of the run is Running. Use wait_for_completion() first." + ) + raise ValueError( + f"Cannot fetch results — run is not completed. Current status: {execution_status}" + ) + + metrics = self.aggregations() + + prompt_client = PromptTemplateClient(proxy_client=self._run_context.proxy_client) + prompts = {} + for resource in (metrics.resources if metrics is not None else []): + tags = {t.name: t.value for t in (resource.tags or [])} + if tags.get("evaluation.ai.sap.com/purpose") != "target": + continue + model = tags.get("evaluation.ai.sap.com/model") + prompt_id = tags.get("evaluation.ai.sap.com/promptTemplateId") + if not model or not prompt_id: + continue + try: + prompts[model] = prompt_client.get_prompt_template_by_id(prompt_id) + except Exception as err: + logger.warning("Could not fetch prompt for model %s (id=%s): %s", model, prompt_id, err) + + if not prompts: + for model, prompt_ref in self._run_context.target_prompt_mapping.items(): + try: + name, version = prompt_ref.rsplit(":", 1) + response = prompt_client.get_prompt_templates( + scenario="genai-optimizations", name=name, version=version + ) + if response.resources: + prompts[model] = response.resources[0] + except Exception as err: + logger.warning("Could not fetch prompt for model %s (ref=%s): %s", model, prompt_ref, err) + + return OptimizationResults(metrics=metrics, prompts=prompts) diff --git a/packages/gen/gen_ai_hub/optimizations/optimization_flow.py b/packages/gen/gen_ai_hub/optimizations/optimization_flow.py new file mode 100644 index 00000000..bf2f4b54 --- /dev/null +++ b/packages/gen/gen_ai_hub/optimizations/optimization_flow.py @@ -0,0 +1,164 @@ +"""Orchestration flow for uploading datasets, registering AI Core configurations, and launching optimization jobs.""" +import json +import os +import uuid +from pathlib import Path + +from ai_api_client_sdk.models.artifact import Artifact +from ai_core_sdk.ai_core_v2_client import AICoreV2Client + +from gen_ai_hub.evaluations._internal._models import _AWSObjectStoreData +from gen_ai_hub.evaluations.constants import ( + AI_PROTOCOL_PREFIX, + AWS_OSS_PATH_PREFIX_URL_KEY, + DATASET_FOLDER_KEY, + DEFAULT_KEY, +) +from gen_ai_hub.evaluations.exceptions.error_codes import ErrorCode +from gen_ai_hub.evaluations.helpers.collector import ValidationCollector +from gen_ai_hub.evaluations.helpers.logging import get_logger +from gen_ai_hub.evaluations.utils.aicore_utils import ( + generate_random_id, + register_aicore_execution, + upload_file_to_aws_s3, +) +from gen_ai_hub.evaluations.utils.oss_secret_utils import fetch_object_store_secret_by_name +from gen_ai_hub.optimizations.constants import OPTIMIZATIONS_ARTIFACT_KEY, OPTIMIZATIONS_SCENARIO_ID +from gen_ai_hub.optimizations.models.optimization_config import PromptOptimizationConfig +from gen_ai_hub.optimizations.models.optimization_run import OptimizationRun +from gen_ai_hub.optimizations.utils import register_optimization_aicore_configuration + +logger = get_logger() + + +def _upload_optimization_dataset( + dataset_data, + dataset_type: str, + object_store_credentials: _AWSObjectStoreData, + ai_core_client: AICoreV2Client, + resource_group: str, + error_collector: ValidationCollector, +): + """Upload dataset to S3 and register an optimization artifact. Returns (artifact_id, dataset_file_key).""" + object_store_secret = fetch_object_store_secret_by_name( + ai_core_client, DEFAULT_KEY, resource_group, error_collector + ) + error_collector.raise_if_errors() + + bare_folder_id = generate_random_id() + path_prefix = object_store_secret.metadata.get(AWS_OSS_PATH_PREFIX_URL_KEY) + s3_root = os.path.join(path_prefix, bare_folder_id) if path_prefix else bare_folder_id + + dataset_file_name = f"{generate_random_id()[:7]}.{dataset_type}" + dataset_file_key = os.path.join(DATASET_FOLDER_KEY, dataset_file_name) + s3_file_key = os.path.join(s3_root, dataset_file_key) + + uploaded = upload_file_to_aws_s3( + object_store_credentials, + object_store_secret.metadata, + dataset_data, + s3_file_key, + dataset_type, + error_collector, + ) + if not uploaded: + error_collector.add_error( + ErrorCode.FILE_UPLOAD_ERROR, + f"Error uploading optimization dataset to object store at {s3_file_key}", + ) + error_collector.raise_if_errors() + + artifact_url = os.path.join(AI_PROTOCOL_PREFIX, DEFAULT_KEY, bare_folder_id) + try: + response = ai_core_client.artifact.create( + name=OPTIMIZATIONS_ARTIFACT_KEY + "-" + generate_random_id()[:7], + kind=Artifact.Kind.OTHER, + url=artifact_url, + scenario_id=OPTIMIZATIONS_SCENARIO_ID, + resource_group=resource_group, + ) + logger.info("Optimization artifact created: %s", response.id) + return response.id, dataset_file_key + except Exception as err: + error_collector.add_error( + ErrorCode.ARTIFACT_CREATION_FAILURE, + f"Error creating optimization artifact: {err}", + ) + error_collector.raise_if_errors() + return None, None + + +def optimization_job_flow( + optimization_config: PromptOptimizationConfig, + object_store_credentials: _AWSObjectStoreData, + ai_core_client: AICoreV2Client, + resource_group: str, + error_collector: ValidationCollector, + proxy_client=None, +) -> OptimizationRun: + """Upload the dataset, register AI Core configuration and execution, and return an OptimizationRun.""" + if optimization_config.artifact_id: + aicore_artifact_id = optimization_config.artifact_id + dataset_file_key = optimization_config.dataset + logger.info("Reusing existing artifact %s for optimization", aicore_artifact_id) + else: + dataset_path = optimization_config.dataset_path + dataset_type = Path(dataset_path).suffix.lstrip(".") + with open(dataset_path, "r", encoding="utf-8") as file: + dataset_data = json.load(file) + + aicore_artifact_id, dataset_file_key = _upload_optimization_dataset( + dataset_data, + dataset_type, + object_store_credentials, + ai_core_client, + resource_group, + error_collector, + ) + logger.info("Uploaded dataset and registered artifact %s", aicore_artifact_id) + + aicore_configuration_id = register_optimization_aicore_configuration( + aicore_artifact_id, + ai_core_client, + resource_group, + optimization_config, + dataset_file_key, + error_collector, + ) + logger.info("AI Core configuration ID: %s", aicore_configuration_id) + + if not aicore_configuration_id: + error_collector.add_error( + ErrorCode.CONFIGURATION_CREATION_FAILURE, + "Error while creating the optimization aicore configuration, so terminating the optimize function. " + "Please look into the error and try again", + ) + error_collector.raise_if_errors() + + aicore_execution_id = register_aicore_execution( + ai_core_client, + aicore_configuration_id, + resource_group, + error_collector, + ) + logger.info("AI Core execution ID: %s", aicore_execution_id) + logger.info("Dataset file key: %s", dataset_file_key) + + if not aicore_execution_id: + error_collector.add_error( + ErrorCode.EXECUTION_CREATION_FAILURE, + "Error while creating the optimization aicore execution, so terminating the optimize function. " + "Please look into the error and try again", + ) + error_collector.raise_if_errors() + + return OptimizationRun( + run_id=uuid.uuid4().hex, + execution_id=aicore_execution_id, + configuration_id=aicore_configuration_id, + artifact_id=aicore_artifact_id, + ai_core_client=ai_core_client, + resource_group=resource_group, + proxy_client=proxy_client, + target_prompt_mapping=optimization_config.target_prompt_mapping, + ) diff --git a/packages/gen/gen_ai_hub/optimizations/utils.py b/packages/gen/gen_ai_hub/optimizations/utils.py new file mode 100644 index 00000000..5dc729d8 --- /dev/null +++ b/packages/gen/gen_ai_hub/optimizations/utils.py @@ -0,0 +1,173 @@ +"""Optimization-specific helpers: AI Core configuration registration and config validation.""" +import json +import os + +from ai_api_client_sdk.models.input_artifact_binding import InputArtifactBinding +from ai_api_client_sdk.models.parameter_binding import ParameterBinding +from ai_core_sdk.ai_core_v2_client import AICoreV2Client + +from gen_ai_hub.evaluations.exceptions.error_codes import ErrorCode +from gen_ai_hub.evaluations.helpers.collector import ValidationCollector +from gen_ai_hub.evaluations.helpers.logging import get_logger +from gen_ai_hub.evaluations.utils.aicore_utils import generate_random_id +from gen_ai_hub.evaluations.utils.metric_client_utils import get_custom_metric_by_id +from gen_ai_hub.optimizations.constants import ( + OPTIMIZATIONS_CONFIG_PREFIX_KEY, + OPTIMIZATIONS_SCENARIO_ID, +) + +logger = get_logger() + + +def register_optimization_aicore_configuration( + aicore_artifact_id: str, + ai_core_client: AICoreV2Client, + resource_group: str, + optimization_config, + dataset_file_key: str, + error_collector: ValidationCollector, +): + """Register an AI Core configuration for an optimization run and return the configuration ID.""" + try: + target_models_str = ",".join(optimization_config.target_models) + target_prompt_mapping_str = ",".join( + f"{k}={v}" for k, v in optimization_config.target_prompt_mapping.items() + ) + + parameter_bindings_list = [ + ParameterBinding(key="basePrompt", value=optimization_config.base_prompt), + ParameterBinding(key="baseModel", value=optimization_config.base_model or "none"), + ParameterBinding(key="dataset", value=dataset_file_key), + ParameterBinding(key="targetModels", value=target_models_str), + ParameterBinding(key="targetPromptMapping", value=target_prompt_mapping_str), + ParameterBinding(key="optimizationMetric", value=optimization_config.optimization_metric or "none"), + ParameterBinding(key="customMetricId", value=optimization_config.custom_metric_id or "none"), + ParameterBinding(key="includeFewShotExamples", + value=str(optimization_config.include_few_shot_examples).lower()), + ParameterBinding(key="maximize", value=str(optimization_config.maximize).lower()), + ParameterBinding( + key="correctnessCutoff", + value=str(optimization_config.correctness_cutoff) if optimization_config.correctness_cutoff else "none", + ), + ParameterBinding(key="promptTemplateScope", value=optimization_config.prompt_template_scope or "tenant"), + ParameterBinding(key="prototypeMode", value=str(optimization_config.prototype_mode).lower()), + ] + train_cfg = optimization_config.train_dataset_config + test_cfg = optimization_config.test_dataset_config + train_path = (train_cfg.source.path if train_cfg and train_cfg.source else None) or "none" + test_path = (test_cfg.source.path if test_cfg and test_cfg.source else None) or "none" + parameter_bindings_list += [ + ParameterBinding(key="trainDataset", value=train_path), + ParameterBinding(key="testDataset", value=test_path), + ParameterBinding(key="fieldEvaluationMetrics", value=optimization_config.field_evaluation_metrics or "none"), + ParameterBinding(key="modelParams", value=optimization_config.model_params or "none"), + ParameterBinding(key="variableMapping", value=optimization_config.variable_mapping or "none"), + ] + + configuration_name = OPTIMIZATIONS_CONFIG_PREFIX_KEY + generate_random_id()[:7] + + response = ai_core_client.configuration.create( + name=configuration_name, + scenario_id=OPTIMIZATIONS_SCENARIO_ID, + executable_id="genai-optimizations", + parameter_bindings=parameter_bindings_list, + input_artifact_bindings=[ + InputArtifactBinding(key="prompt-data", artifact_id=aicore_artifact_id) + ], + resource_group=resource_group, + ) + configuration_id = response.id + logger.info("Optimization configuration id created is %s", configuration_id) + return configuration_id + except Exception as err: + error_collector.add_error( + ErrorCode.CONFIGURATION_CREATION_FAILURE, + f"Error occurred while attempting to create optimization aicore configuration with error of {err}", + ) + return None + + +def _validate_json_file(dataset_path: str, error_collector: ValidationCollector): + try: + with open(dataset_path, "r", encoding="utf-8") as file: + data = json.load(file) + if not data: + error_collector.add_error( + ErrorCode.EMPTY_FILE_DATA_ERROR, + f"Dataset file is empty: {dataset_path}", + ) + except ValueError as err: + error_collector.add_error( + ErrorCode.INVALID_JSON_DECODING_ERROR, + f"Dataset file is not valid JSON: {dataset_path} — {err}", + ) + + +def _validate_dataset_path(dataset_path: str, error_collector: ValidationCollector): + if not os.path.exists(dataset_path): + error_collector.add_error( + ErrorCode.INVALID_FILE_PATH_ERROR, + f"Dataset file not found: {dataset_path}", + ) + return + _, ext = os.path.splitext(dataset_path) + if ext.lower() != ".json": + error_collector.add_error( + ErrorCode.UNSUPPORTED_FILE_TYPE_ERROR, + f"Unsupported dataset file type '{ext}'. Only .json is supported for optimization.", + ) + else: + _validate_json_file(dataset_path, error_collector) + + +def _validate_target_prompt_mapping(optimization_config, error_collector: ValidationCollector): + missing = set(optimization_config.target_models) - set(optimization_config.target_prompt_mapping.keys()) + if missing: + error_collector.add_error( + ErrorCode.INVALID_PARAMETER_VALUE_ERROR, + f"target_prompt_mapping is missing entries for target_models: {sorted(missing)}", + ) + + +def _validate_custom_metric(optimization_config, ai_core_client, resource_group, error_collector): + metric_info = get_custom_metric_by_id( + optimization_config.custom_metric_id, + ai_core_client, + resource_group, + error_collector, + ) + if not metric_info: + error_collector.add_error( + ErrorCode.METRIC_SERVER_RESOLVE_ERROR, + f"custom_metric_id '{optimization_config.custom_metric_id}' could not be resolved " + "on the Metric Management Service.", + ) + + +def validate_optimization_config( + optimization_config, + ai_core_client: AICoreV2Client, + resource_group: str, + error_collector: ValidationCollector, +): + """Validate the optimization configuration, including dataset path, target models, base prompt, and custom metric.""" + if optimization_config.dataset_path is not None: + _validate_dataset_path(optimization_config.dataset_path, error_collector) + + if not optimization_config.target_models: + error_collector.add_error( + ErrorCode.INVALID_PARAMETER_VALUE_ERROR, + "target_models must be a non-empty list.", + ) + + if not optimization_config.base_prompt or not optimization_config.base_prompt.strip(): + error_collector.add_error( + ErrorCode.INVALID_PARAMETER_VALUE_ERROR, + "base_prompt must be a non-empty string.", + ) + + if optimization_config.target_models and optimization_config.target_prompt_mapping: + _validate_target_prompt_mapping(optimization_config, error_collector) + + if optimization_config.custom_metric_id: + _validate_custom_metric(optimization_config, ai_core_client, resource_group, error_collector) diff --git a/packages/gen/integration_tests/optimizations/__init__.py b/packages/gen/integration_tests/optimizations/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/packages/gen/integration_tests/optimizations/po_dataset_small.json b/packages/gen/integration_tests/optimizations/po_dataset_small.json new file mode 100644 index 00000000..e027309e --- /dev/null +++ b/packages/gen/integration_tests/optimizations/po_dataset_small.json @@ -0,0 +1,20 @@ +[ + { + "fields": { + "input": "Subject: Urgent Assistance Required for Specialized Cleaning Services\n\nDear ProCare Facility Solutions Support Team,\n..." + }, + "answer": "{\"categories\": {\"routine_maintenance_requests\": false, \"customer_feedback_and_complaints\": false, \"training_and_support_requests\": false, \"quality_and_safety_concerns\": false, \"sustainability_and_environmental_practices\": false, \"cleaning_services_scheduling\": false, \"specialized_cleaning_services\": true, \"emergency_repair_services\": false, \"facility_management_issues\": false, \"general_inquiries\": false}, \"sentiment\": \"neutral\", \"urgency\": \"high\"}" + }, + { + "fields": { + "input": "Subject: Weekly Office Cleaning Schedule\n\nHello,\n\nWe need to schedule our regular weekly office cleaning for next month. Please let me know your available time slots for Mondays and Wednesdays." + }, + "answer": "{\"categories\": {\"routine_maintenance_requests\": false, \"customer_feedback_and_complaints\": false, \"training_and_support_requests\": false, \"quality_and_safety_concerns\": false, \"sustainability_and_environmental_practices\": false, \"cleaning_services_scheduling\": true, \"specialized_cleaning_services\": false, \"emergency_repair_services\": false, \"facility_management_issues\": false, \"general_inquiries\": false}, \"sentiment\": \"neutral\", \"urgency\": \"low\"}" + }, + { + "fields": { + "input": "Subject: EMERGENCY - Water Leak in Main Conference Room\n\nURGENT: We have a major water leak in our main conference room. Water is spreading to adjacent offices. Need immediate assistance!" + }, + "answer": "{\"categories\": {\"routine_maintenance_requests\": false, \"customer_feedback_and_complaints\": false, \"training_and_support_requests\": false, \"quality_and_safety_concerns\": true, \"sustainability_and_environmental_practices\": false, \"cleaning_services_scheduling\": false, \"specialized_cleaning_services\": false, \"emergency_repair_services\": true, \"facility_management_issues\": true, \"general_inquiries\": false}, \"sentiment\": \"negative\", \"urgency\": \"high\"}" + } +] diff --git a/packages/gen/integration_tests/optimizations/test_base.py b/packages/gen/integration_tests/optimizations/test_base.py new file mode 100644 index 00000000..ad42310e --- /dev/null +++ b/packages/gen/integration_tests/optimizations/test_base.py @@ -0,0 +1,65 @@ +"""Base test class for optimization integration tests.""" +import os +import unittest + +from gen_ai_hub.optimizations.client import OptimizationClient + + +class OptimizationClientTestBase(unittest.TestCase): + """Base class for optimization client integration tests.""" + + @classmethod + def setUpClass(cls): + """Set up the test class with credentials.""" + cls.base_url = os.getenv("AICORE_BASE_URL") + cls.auth_url = os.getenv("AICORE_AUTH_URL") + cls.client_id = os.getenv("AICORE_CLIENT_ID") + cls.client_secret = os.getenv("AICORE_CLIENT_SECRET") + cls.resource_group = "default" + cls.aws_access_key_id = os.getenv("AWS_ACCESS_KEY_ID") + cls.aws_secret_access_key = os.getenv("AWS_SECRET_ACCESS_KEY") + current_dir = os.path.dirname(os.path.abspath(__file__)) + cls.dataset_path = os.path.abspath(os.path.join(current_dir, "po_dataset_small.json")) + + def setUp(self): + """Set up each test with a fresh optimization client instance and object store secrets.""" + self.client = OptimizationClient( + base_url=self.base_url, + auth_url=self.auth_url, + client_id=self.client_id, + client_secret=self.client_secret, + resource_group=self.resource_group, + aws_access_key_id=self.aws_access_key_id, + aws_secret_access_key=self.aws_secret_access_key, + ) + + AWS_S3_ENDPOINT = "s3-eu-central-1.amazonaws.com" + AWS_BUCKET_ID = "hcp-e597ff51-40f5-42c9-a75a-744281742e61" + AWS_REGION = "eu-central-1" + + default_secret_creds = { + "data": {}, + "type": "S3", + "pathPrefix": "sdkOutputFiles", + "endpoint": AWS_S3_ENDPOINT, + "bucket": AWS_BUCKET_ID, + "region": AWS_REGION, + "usehttps": "1", + } + + input_secret_creds = { + "data": {}, + "name": "sdk-data", + "type": "S3", + "pathPrefix": "sdk_input_files/data", + "endpoint": AWS_S3_ENDPOINT, + "bucket": AWS_BUCKET_ID, + "region": AWS_REGION, + "usehttps": "1", + } + + self.client.setup( + default_secret_body=default_secret_creds, + input_secret_body=input_secret_creds, + replace_existing=True, + ) diff --git a/packages/gen/integration_tests/optimizations/test_optimization_flow.py b/packages/gen/integration_tests/optimizations/test_optimization_flow.py new file mode 100644 index 00000000..600edd2d --- /dev/null +++ b/packages/gen/integration_tests/optimizations/test_optimization_flow.py @@ -0,0 +1,120 @@ +"""Integration tests for the prompt optimization flow.""" +import unittest + +from gen_ai_hub.optimizations.models.optimization_config import PromptOptimizationConfig +from gen_ai_hub.optimizations.models.optimization_results import OptimizationResults +from gen_ai_hub.prompt_registry.client import PromptTemplateClient +from gen_ai_hub.prompt_registry.models.prompt_template import PromptTemplate, PromptTemplateSpec +from integration_tests.optimizations.test_base import OptimizationClientTestBase + +BASE_PROMPT_NAME = "sdktest-opt-base-prompt" +BASE_PROMPT_VERSION = "0.0.1" +BASE_PROMPT_SCENARIO = "genai-optimizations" +TARGET_MODEL = "gemini-2.5-pro:001" +TARGET_PROMPT_NAME = "sdktest-opt-target-prompt" +TARGET_PROMPT_MAPPING = {TARGET_MODEL: f"{TARGET_PROMPT_NAME}:0.0.1"} +OPTIMIZATION_METRIC = "JSON_Match" + + +class TestOptimizationFlow(OptimizationClientTestBase): + """Integration tests for the optimization job flow.""" + + @classmethod + def setUpClass(cls): + """Create base prompt template for tests.""" + super().setUpClass() + cls.base_prompt_id = None + cls.base_prompt_ref = None + + def setUp(self): + """Set up client and create base prompt template before each test class run.""" + super().setUp() + if self.__class__.base_prompt_id is None: + try: + prompt_client = PromptTemplateClient( + proxy_client=self.client._gen_ai_hub_proxy_client + ) + spec = PromptTemplateSpec( + template=[ + PromptTemplate(role="system", content="You are a helpful assistant"), + PromptTemplate( + role="user", + content=( + "Giving the following message --- {{?input}} --- " + "Extract and return a json with the following keys and values: " + "- 'urgency' as one of `high`, `medium`, `low` " + "- 'sentiment' as one of `negative`, `neutral`, `positive` " + "- 'categories' Create a dictionary with categories as keys and boolean values (True/False), " + "where the value indicates whether the category is one of the best matching support category tags from: " + "`emergency_repair_services`, `routine_maintenance_requests`, `quality_and_safety_concerns`, " + "`specialized_cleaning_services`, `general_inquiries`, `sustainability_and_environmental_practices`, " + "`training_and_support_requests`, `cleaning_services_scheduling`, `customer_feedback_and_complaints`, " + "`facility_management_issues` " + "Your complete message should be a valid json string that can be read directly and only contain " + "the keys mentioned in the list above. Never enclose it in ```json...```, no newlines, no unnecessary whitespaces." + ), + ), + ], + defaults={"input": ""}, + additional_fields={ + "modelParams": {"temperature": 0.7, "max_tokens": 100}, + "modelGroup": "chat", + }, + ) + response = prompt_client.create_prompt_template( + name=BASE_PROMPT_NAME, + version=BASE_PROMPT_VERSION, + scenario=BASE_PROMPT_SCENARIO, + prompt_template_spec=spec, + ) + self.__class__.base_prompt_id = response.id + self.__class__.base_prompt_ref = f"{BASE_PROMPT_SCENARIO}/{BASE_PROMPT_NAME}:{BASE_PROMPT_VERSION}" + self.__class__.prompt_client = prompt_client + print(f"Base prompt template created: {self.__class__.base_prompt_id}") + except Exception as err: + self.fail(f"Failed to create base prompt template: {err}") + + @classmethod + def tearDownClass(cls): + """Delete the base prompt template created for tests.""" + if hasattr(cls, "prompt_client") and hasattr(cls, "base_prompt_id") and cls.base_prompt_id: + try: + cls.prompt_client.delete_prompt_template_by_id(cls.base_prompt_id) + print(f"Base prompt template deleted: {cls.base_prompt_id}") + except Exception as err: + print(f"Warning: Could not delete prompt template {cls.base_prompt_id}: {err}") + + def test_optimize_wait_for_completion(self): + """Test optimization job completes successfully and returns results.""" + config = PromptOptimizationConfig( + dataset_path=self.dataset_path, + target_prompt_mapping=TARGET_PROMPT_MAPPING, + target_models=[TARGET_MODEL], + base_prompt=self.base_prompt_ref, + optimization_metric=OPTIMIZATION_METRIC, + prototype_mode=True, + ) + + run = self.client.optimize(config) + self.assertIsNotNone(run) + + run.wait_for_completion() + + debug_info = run.get_debug_info() + current_status = run.get_current_status() + self.assertEqual( + current_status.name, "COMPLETED", + f"Run did not complete. Status: {current_status}, Debug info: {debug_info}", + ) + + results = run.results() + self.assertIsInstance(results, OptimizationResults) + self.assertIsNotNone(results.metrics) + self.assertTrue( + len(results.metrics.resources or []) > 0, + "Expected at least one metric resource in results", + ) + + +if __name__ == "__main__": + unittest.main()