Prompt optimization

# Loading the credentials from the env file
from gen_ai_hub.proxy.gen_ai_hub_proxy import GenAIHubProxyClient
from dotenv import load_dotenv
import os

load_dotenv(override=True)

# Fetching environment variables
AICORE_BASE_URL = os.getenv("AICORE_BASE_URL")
AICORE_RESOURCE_GROUP = os.getenv("AICORE_RESOURCE_GROUP")
AICORE_AUTH_URL = os.getenv("AICORE_AUTH_URL")
AICORE_CLIENT_ID = os.getenv("AICORE_CLIENT_ID")
AICORE_CLIENT_SECRET = os.getenv("AICORE_CLIENT_SECRET")

AWS_ACCESS_KEY = os.getenv("AWS_ACCESS_KEY")
AWS_BUCKET_ID = os.getenv("AWS_BUCKET_ID")
AWS_REGION = os.getenv("AWS_REGION")
AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")

# Initializing the GenAIHubProxyClient
client = GenAIHubProxyClient(
    base_url=AICORE_BASE_URL,
    auth_url=AICORE_AUTH_URL,
    client_id=AICORE_CLIENT_ID,
    client_secret=AICORE_CLIENT_SECRET,
    resource_group=AICORE_RESOURCE_GROUP
)

# Register S3 secret with AI Core which will be used an input source 
import requests

def register_oss_secret():
    headers = _get_headers()
    
    POST_SECRETS_ENDPOINT = '/v2/admin/objectStoreSecrets'
    request_url = f"{AICORE_BASE_URL}{POST_SECRETS_ENDPOINT}"
    
    request_body = {
        "name": "default",
        "data": {
            "AWS_ACCESS_KEY_ID": AWS_ACCESS_KEY,
            "AWS_SECRET_ACCESS_KEY": AWS_SECRET_ACCESS_KEY
        },
        "type": "S3",
        "bucket": AWS_BUCKET_ID,
        "endpoint": "s3-eu-central-1.amazonaws.com",
        "region": AWS_REGION,
        "pathPrefix": ""    
    }
    try:
        response = requests.post(
            request_url, headers=headers, data=json.dumps(request_body), timeout=120
        )
        result = response.json()
        print(result)
        return result
    except:
        logging.error("Error occurred while attempting to create object store secret")
        raise

register_oss_secret()

from typing import List
import requests
import mimetypes
from urllib.parse import quote
import pathlib
import json


def validate_dataset(dataset: str | pathlib.Path | list, expected_keys: None | List[str] = None) -> bool:
    if isinstance(dataset, (str, pathlib.Path)):
        with open(dataset, "r") as f:
            try:
                dataset = json.load(f)
            except json.JSONDecodeError as e:
                raise ValueError(f"Invalid JSON in file: {e}")
    if not isinstance(dataset, list):
        raise ValueError("Dataset must be a list of dictionaries.")

def upload_dataset(secret: str,
                   local_path: str | pathlib.Path,
                   remote_path: str,
                   scenario: str,
                   description: str | None = None,
                   overwrite: bool = False,
                   expected_keys: None | List[str] = None,

                   allow_bucket_root: bool = False) -> str:
    # Validate dataset
    validate_dataset(local_path, expected_keys)
    # check if secret exists
    secrets = [r.name for r in client.ai_core_client.object_store_secrets.query().resources]
    if secret not in secrets:
        raise ValueError(f"Secret '{secret}' not found in object store secrets. Known secrets: {secrets}")

    # Check if local path exists
    remote_path = remote_path.lstrip("/")
    if "/" not in remote_path and not allow_bucket_root:
        raise ValueError(
            "Remote path must use subdirectories. Otherwise the whole bucket will be used as an input artifact. Set allow_bucket_root=True to allow this."
        )

    # URL-encode the path parameter
    path = f"{secret}/" + remote_path.lstrip("/")
    encoded_path = quote(path, safe="")
    url = f"{client.ai_core_client.base_url}/lm/dataset/files/{encoded_path}"
    params = {"overwrite": str(overwrite).lower()}

    # Prepare headers
    headers = {
        **client.request_header,
        "Content-Type": "application/octet-stream",
    }
    # Guess MIME type
    guessed_type, _ = mimetypes.guess_type(local_path)
    if guessed_type:
        headers["Content-Type"] = guessed_type

    with open(local_path, "rb") as f:
        response = requests.put(url, params=params, headers=headers, data=f)

    # Handle response
    if response.status_code == 201:
        response = response.json()
    elif response.status_code in (400, 409, 413):
        # Return error details
        raise requests.HTTPError(f"Upload failed ({response.status_code}): {response.text}")
    else:
        response.raise_for_status()
    artifact_url = "/".join(response["url"].split("/")[:-1])
    for artifact in client.ai_core_client.artifact.query().resources:
        if response["url"].startswith(artifact.url + "/"):
            return artifact, response["url"].removeprefix(artifact.url).lstrip("/")

    # Create new artifact
    path = response["url"].split("/")[-1]
    new_artifact = client.ai_core_client.artifact.create(
        name=f"{scenario}-prompt-optimization-data",
        kind=Artifact.Kind.DATASET,
        url=artifact_url,
        scenario_id=scenario,
        description="Datasets for prompt optimization" if description is None else description,
        resource_group=headers[client.ai_core_client.rest_client.resource_group_header]
    )
    return new_artifact, path

artifact, dataset_path = upload_dataset(
    secret=dataset_secret,
    local_path=dataset_local_path,
    remote_path=dataset_remote_path,
    expected_keys=base_template.placeholders,
    scenario=scenario,
    overwrite=True
)

print(f"Dataset uploaded to {artifact.url}/{dataset_path} -> Artifact ID: {artifact.id}")

system: |-
  You are a helpful assistant.

user: |-
  Giving the following message:
  ---
  {{?input}}
  ---
  Extract and return a JSON object with the following structure:
  {
    "urgency": "<one of: high, medium, low>",
    "sentiment": "<one of: positive, neutral, negative>",
    "categories": {
      "emergency_repair_services": <true/false>,
      "routine_maintenance_requests": <true/false>,
      "quality_and_safety_concerns": <true/false>,
      "specialized_cleaning_services": <true/false>,
      "general_inquiries": <true/false>,
      "sustainability_and_environmental_practices": <true/false>,
      "training_and_support_requests": <true/false>,
      "cleaning_services_scheduling": <true/false>,
      "customer_feedback_and_complaints": <true/false>,
      "facility_management_issues": <true/false>
    }
  }

  Your response must:
  - Contain only this JSON structure (no extra text).
  - Be valid JSON (parsable without errors).
  - Match the keys and value types exactly.

from gen_ai_hub.prompt_registry.client import PromptTemplateClient
from gen_ai_hub.prompt_registry.models.prompt_template import PromptTemplateSpec, PromptTemplate

# Initialize Prompt Registry Client
prompt_registry_client = PromptTemplateClient(proxy_client=client)

prompt_template_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 follwoing 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 unnessacary whitespaces."""
            )
        )
    ]
)

# Create prompt template in registry
template = prompt_registry_client.create_prompt_template(
    scenario="genai-optimizations",
    name="facility-json-template",
    version="1.0.0",
    prompt_template_spec=prompt_template_spec
)

print(f"✅ Created Prompt Template with ID: {template.id}")

old_new_name_mapping = {
    "gemini-2.5-pro:001": "gemini-2.5-pro--001",
    "gpt-4o:2024-08-06": "openai/gpt-4o-2024-08-06"
}

old_new_name_mapping.update({old_new_name_mapping[k]: k for k, v in old_new_name_mapping.items()})


def create_config(metric: str,
                  reference_model: str,
                  targets: dict,
                  dataset_path: str,
                  scenario: str,
                  prompt: PromptTemplateSpec) -> str:
    assert metric in SUPPORTED_METRICS, f"Unsupported metric: {metric}. Supported metrics: {SUPPORTED_METRICS}"
    assert reference_model in SUPPORTED_MODELS, f"Unsupported reference model: {reference_model}. Supported models: {SUPPORTED_MODELS}"
    assert all(model in SUPPORTED_MODELS for model in targets.keys()), f"Unsupported target models: {targets}. Supported models: {SUPPORTED_MODELS}"
    input_parameters = [
        ParameterBinding(key="dataset", value=dataset_path),
        ParameterBinding(key="optimizationMetric", value=metric),
        ParameterBinding(key="basePrompt", value=f'{scenario}/{prompt["name"]}:{prompt["version"]}'),
        ParameterBinding(key="baseModel", value=reference_model),
        ParameterBinding(key="targetModels", value=','.join(targets.keys())),
        ParameterBinding(key="targetPromptMapping", value=",".join([f"{old_new_name_mapping[k]}={v}" for k, v in targets.items()]))
    ]
    existing_configs = client.ai_core_client.configuration.query(scenario_id='genai-optimizations', executable_ids=['genai-optimizations'])
    params = {par.key: par.value for par in input_parameters}
    for conf in existing_configs.resources:
        if {par.key: par.value for par in conf.parameter_bindings} == params:
            return conf.id

    input_artifacts = [InputArtifactBinding(key="prompt-data", artifact_id=artifact.id)]

    response = client.ai_core_client.configuration.create(
        name = "prompt-optimization-configuration", # custom name of configuration
        scenario_id = "genai-optimizations", # value from workflow
        executable_id = "genai-optimizations", # value from workflow
        resource_group = resource_group,
        parameter_bindings = input_parameters,
        input_artifact_bindings = input_artifacts
    )

    return response.id

# Create the configuration
configuration_id = create_config(
    metric=metric,
    reference_model=reference_model,
    targets=targets,
    dataset_path=dataset_path,
    scenario=scenario,
    prompt=prompt
)
print("Optimization Configuration ID:", configuration_id)