Prompt Optimization with Tool Calling and Response Formatting

from collections import defaultdict
from gen_ai_hub.proxy.gen_ai_hub_proxy import GenAIHubProxyClient
from dotenv import load_dotenv
import os
import json
import requests
import random
from urllib.parse import quote
from pathlib import Path
from typing import List, Tuple
import time
from ai_api_client_sdk.models.parameter_binding import ParameterBinding
from ai_api_client_sdk.models.input_artifact_binding import InputArtifactBinding
from pydantic import BaseModel
from ai_api_client_sdk.models.artifact import Artifact

load_dotenv(override=True)

# ── SAP AI Core client ────────────────────────────────────────────────────────
client = GenAIHubProxyClient(
    base_url=os.getenv("AICORE_BASE_URL"),
    auth_url=os.getenv("AICORE_AUTH_URL"),
    client_id=os.getenv("AICORE_CLIENT_ID"),
    client_secret=os.getenv("AICORE_CLIENT_SECRET"),
    resource_group=os.getenv("AICORE_RESOURCE_GROUP")
)
resource_group = client.request_header[
    client.ai_core_client.rest_client.resource_group_header
]

print("✅ Connected to AI Core")
print(f"   Resource group: {resource_group}")

def read_bfcl_file(file_path: Path) -> list:
    """Robust reader for BFCL v3 files (concatenated JSON objects)."""
    with open(file_path, "r") as f:
        content = f.read().strip()

    print(f"File size: {len(content):,} bytes")

    # Format 1: JSON array [ {...}, {...} ]
    if content.startswith("["):
        try:
            result = json.loads(content)
            if result and isinstance(result[0], str):
                print("Detected double-encoded strings — decoding...")
                result = [json.loads(item) for item in result]
            print(f"Loaded as JSON array: {len(result)} records")
            return result
        except json.JSONDecodeError:
            pass

    # Format 2: Standard JSONL — one complete object per line
    if "\n" in content:
        objects = []
        for line in content.split("\n"):
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
                if isinstance(obj, dict):
                    objects.append(obj)
            except json.JSONDecodeError:
                pass
        if objects:
            print(f"Loaded as JSONL: {len(objects)} records")
            return objects

    # Format 3: Concatenated JSON objects — use raw_decode to scan through
    objects = []
    decoder = json.JSONDecoder()
    idx = 0
    while idx < len(content):
        while idx < len(content) and content[idx] in " \t\n\r":
            idx += 1
        if idx >= len(content):
            break
        if content[idx] != "{":
            idx += 1
            continue
        try:
            obj, end_idx = decoder.raw_decode(content, idx)
            if isinstance(obj, dict):
                objects.append(obj)
            idx = end_idx
        except json.JSONDecodeError:
            idx += 1
            continue

    print(f"Loaded as concatenated JSON: {len(objects)} records")
    return objects

def dedupe_tool_name(tool_name: str, answers: list, occurrences: int):
    new_name = f"{tool_name}_n{occurrences}"
    new_answers = []
    for answer in answers:
        new_answer = {}
        for k, v in answer.items():
            new_answer[reformat_tool_name(new_name if k == tool_name else k)] = v
        new_answers.append(new_answer)
    return new_name, new_answers

def union_bfcl_tools(samples: list, tool_key: str = "function") -> list:
    """Union and deduplicate all tools across samples into one normalized list."""
    tool_map = {}
    tool_name_occurrences = defaultdict(int)
    tool_name_to_definition = {}
    for sample in samples:
        for tool in sample[tool_key]:
            tool_name = tool["name"]
            tool_name_occurrences[tool_name] += 1
            is_duplicate = (
                tool_name in tool_name_to_definition and
                tool != tool_name_to_definition[tool_name]
            )
            if is_duplicate:
                new_name, new_answers = dedupe_tool_name(
                    tool_name, sample["answer"], tool_name_occurrences[tool_name]
                )
                print(f"WARNING: duplicate tool {tool_name!r} → renamed to {new_name!r}")
                tool["name"]     = new_name
                sample["answer"] = new_answers
            tool_name_to_definition[tool["name"]] = tool
            normalized = normalize_bfcl_tool(tool)
            tool_map[normalized["function"]["name"]] = normalized
    return list(tool_map.values())

def detect_tool_key(sample: dict) -> str:
    """Detect which key holds the tool definitions in a BFCL sample."""
    for candidate in ["function", "functions", "tools", "tool"]:
        if candidate in sample:
            return candidate
    raise KeyError(
        f"Cannot find tool key in sample. Available keys: {list(sample.keys())}"
    )

def detect_question_key(sample: dict) -> str:
    """Detect which key holds the question/messages in a BFCL sample."""
    for candidate in ["question", "messages", "turns", "prompt"]:
        if candidate in sample:
            return candidate
    raise KeyError(
        f"Cannot find question key in sample. Available keys: {list(sample.keys())}"
    )

def build_golden(sample: dict, question_key: str = "question") -> dict:
    """Convert one BFCL sample to a SAP optimizer golden record."""
    raw_question = sample[question_key]
    if isinstance(raw_question[0], list):
        question = "\n".join(q["content"] for q in raw_question[0])
    elif isinstance(raw_question[0], dict):
        question = "\n".join(q["content"] for q in raw_question)
    else:
        question = str(raw_question)

    # Merge all tool calls into a single flat dict
    # e.g. [{"weather": {...}}, {"hotel": {...}}] → {"weather": {...}, "hotel": {...}}
    merged_answer = {}
    for answer in sample["answer"]:
        for key, value in answer.items():
            key = reformat_tool_name(key)
            # Flatten nested single-element arrays (BFCL quirk)
            if isinstance(value, list) and len(value) == 1 and isinstance(value[0], list):
                value = value[0]
            # Float → string to match normalized tool schema
            if isinstance(value, float):
                value = str(value)
            elif isinstance(value, list):
                value = [str(v) if isinstance(v, float) else v for v in value]
            merged_answer[key] = value

    # answer must be a JSON object string, not a JSON array string
    return {
        "fields": {"question": question},
        "answer": json.dumps(merged_answer)   # "{...}" not "[{...}]"
    }

def load_bfcl_dataset(dataset_path: Path, n_train: int, n_test: int) -> Tuple[list, list, list]:
    """Load, sample, normalise, and split BFCL v3 data."""
    samples = read_bfcl_file(dataset_path)
    print(f"Total records in file: {len(samples)}")

    n_total = min(n_train + n_test, len(samples))
    random.seed(42)
    sampled = random.sample(samples, n_total)

    tool_key     = detect_tool_key(sampled[0])
    question_key = detect_question_key(sampled[0])
    print(f"Tool key: '{tool_key}' | Question key: '{question_key}'")

    tools         = union_bfcl_tools(sampled, tool_key=tool_key)
    train_goldens = [build_golden(s, question_key=question_key) for s in sampled[:n_train]]
    test_goldens  = [build_golden(s, question_key=question_key) for s in sampled[n_train:]]
    return train_goldens, test_goldens, tools

dataset_path = Path(BFCL_DATASET)
train_goldens, test_goldens, tools = load_bfcl_dataset(
    dataset_path, N_TRAIN_SAMPLES, N_TEST_SAMPLES
)
print(f"Train goldens : {len(train_goldens)}")
print(f"Test goldens  : {len(test_goldens)}")
print(f"Unioned tools : {len(tools)}")
print(f"\nSample golden:\n{json.dumps(train_goldens[0], indent=2)}")

def build_tool_call_json_schema(tools: list) -> dict:
    """
    Build a JSON Schema for ResponseFormatJsonSchema from the unioned BFCL tool
    definitions. Mirrors the golden-answer shape used throughout this notebook:
    a JSON object whose keys are tool names, and whose values are that tool's
    arguments object.
    """
    properties = {}
    for entry in tools:
        fn = entry["function"]
        properties[fn["name"]] = {
            "type": "object",
            "description": fn.get("description", ""),
            "properties": fn["parameters"].get("properties", {}),
        }

    return {
        "title": "ToolCallResponse",
        "type": "object",
        "description": "One or more tool calls, keyed by tool name.",
        "properties": properties,
        "minProperties": 1,
    }

tool_call_schema = build_tool_call_json_schema(tools)
print(f"Built JSON schema with {len(tool_call_schema['properties'])} tool properties.")

def create_config(
    metric: str,
    reference_model: str,
    targets: dict,
    train_filename: str,
    test_filename: str,
    prompt_artifact_id: str,
    prompt_name: str,
    prompt_version: str,
    scenario: str,
) -> str:
    base_prompt = f"{scenario}/{prompt_name}:{prompt_version}"

    input_parameters = [
        ParameterBinding(key="optimizationMetric",     value=metric),
        ParameterBinding(key="basePrompt",             value=base_prompt),
        ParameterBinding(key="baseModel",              value=reference_model),
        ParameterBinding(key="targetModels",           value=",".join(targets.keys())),
        ParameterBinding(
            key="targetPromptMapping",
            value=",".join(f"{k}={v}" for k, v in targets.items())
        ),
        ParameterBinding(key="trainDataset",           value=train_filename),
        ParameterBinding(key="testDataset",            value=test_filename),
        ParameterBinding(key="maximize",               value="true"),
        ParameterBinding(key="correctnessCutoff",      value="none"),
        ParameterBinding(key="includeFewShotExamples", value="false"),
        ParameterBinding(key="promptTemplateScope",    value="tenant"),
        ParameterBinding(key="prototypeMode",          value="false"),
        ParameterBinding(key="fieldEvaluationMetrics", value="none"),
        ParameterBinding(key="modelParams",            value="none"),
        ParameterBinding(key="customMetricId",         value="none"),
    ]

    input_artifacts = [
        InputArtifactBinding(key="prompt-data", artifact_id=prompt_artifact_id)
    ]

    params_dict = {p.key: p.value for p in input_parameters}

    try:
        existing = client.ai_core_client.configuration.query(
            scenario_id=SCENARIO, resource_group=resource_group
        )
        for conf in existing.resources:
            if {p.key: p.value for p in conf.parameter_bindings} == params_dict:
                print(f"Reusing configuration: {conf.id}")
                return conf.id
    except Exception as e:
        print(f"Could not query configs: {e}")

    resp = client.ai_core_client.configuration.create(
        name="bfcl-tool-config",
        scenario_id=SCENARIO,
        executable_id=SCENARIO,
        resource_group=resource_group,
        parameter_bindings=input_parameters,
        input_artifact_bindings=input_artifacts,
    )
    print(f"Created configuration: {resp.id}")
    return resp.id

configuration_id = create_config(
    metric=METRIC,
    reference_model=REFERENCE_MODEL,
    targets=TARGET_MODELS,
    train_filename="bfcl_train.json",
    test_filename="bfcl_test.json",
    prompt_artifact_id=optimizer_artifact_id,
    prompt_name=PROMPT_NAME,
    prompt_version=PROMPT_VERSION,
    scenario=SCENARIO,
)
print(f"Configuration ID: {configuration_id}")

TERMINAL_STATES = {"COMPLETED", "FAILED", "DEAD", "STOPPED"}

while True:
    status = client.ai_core_client.execution.get(
        execution_id=execution_id, resource_group=resource_group
    )
    status_str = status.status.value if hasattr(status.status, "value") else str(status.status)
    print(f"[{time.strftime('%H:%M:%S')}] {status_str}", end="")

    if hasattr(status, "status_details") and status.status_details:
        progress = status.status_details.get("progress", "")
        print(f"  progress={progress}", end="")
    print()

    if status_str in TERMINAL_STATES:
        if status_str != "COMPLETED":
            try:
                logs = client.ai_core_client.execution.get_logs(
                    execution_id=execution_id, resource_group=resource_group
                )
                print("── Execution logs ──")
                for log in logs.data:
                    print(log.msg)
            except Exception as e:
                print(f"Could not fetch logs: {e}")
        break

    time.sleep(30)

print(f"\nFinal status: {status_str}")

from gen_ai_hub.orchestration.models.llm import LLM
from gen_ai_hub.orchestration.models.message import SystemMessage, UserMessage
from gen_ai_hub.orchestration.models.template import Template, TemplateValue
from gen_ai_hub.orchestration.models.config import OrchestrationConfig
from gen_ai_hub.orchestration.service import OrchestrationService
from gen_ai_hub.orchestration.models.response_format import ResponseFormatJsonSchema

# ── Load base + optimized templates from the registry ─────────────────────────
def get_prompt_template(template_id: str) -> dict:
    url = f"{client.ai_core_client.base_url}/lm/promptTemplates/{template_id}"
    res = requests.get(url, headers=client.request_header)
    res.raise_for_status()
    return res.json()

def extract_messages(template: dict) -> dict:
    messages = {}
    for msg in template.get("spec", {}).get("template", []):
        messages[msg["role"]] = msg["content"]
    return messages

url = f"{client.ai_core_client.base_url}/lm/promptTemplates"
res = requests.get(url, headers=client.request_header)
all_templates = {
    f"{t['name']}:{t['version']}": t["id"]
    for t in res.json().get("resources", [])
}

base_id      = all_templates.get("bfcl-tool-base:0.0.1")
optimized_id = all_templates.get("bfcl-tool-optimized-gemini25:0.0.1")

base_messages      = extract_messages(get_prompt_template(base_id))
optimized_messages = extract_messages(get_prompt_template(optimized_id))
print("✅ Loaded base and optimized prompt templates.")

# ── Inference helper (optionally with response_format) ────────────────────────
# Paste your orchestration deployment URL here
ORCHESTRATION_DEPLOYMENT_URL = "<YOUR_ORCHESTRATION_DEPLOYMENT_URL>"

def run_inference(system_prompt: str, user_template: str, question: str, model_name: str,
                  use_response_format: bool = False) -> str:
    user_content = user_template.replace("{{?question}}", question)

    template_kwargs = {
        "messages": [
            SystemMessage(system_prompt),
            UserMessage(user_content),
        ]
    }

    if use_response_format:
        template_kwargs["response_format"] = ResponseFormatJsonSchema(
            name="tool_call_response",
            description="Structured tool call output, keyed by tool name.",
            schema=tool_call_schema,
        )

    config = OrchestrationConfig(
        llm=LLM(name=model_name),
        template=Template(**template_kwargs),
    )

    service  = OrchestrationService(api_url=ORCHESTRATION_DEPLOYMENT_URL, config=config)
    response = service.run()
    return response.module_results.llm.choices[0].message.content

def clean_json_output(text: str) -> str:
    """Strip markdown code fences that models sometimes wrap around JSON."""
    text = text.strip()
    if text.startswith("```"):
        lines = text.split("\n")
        lines = [l for l in lines if not l.strip().startswith("```")]
        text = "\n".join(lines).strip()
    return text


def compare_prompts(question: str, model_name: str = "gemini-2.5-pro"):
    print("\n" + "=" * 70)
    print(f"QUESTION:\n{question}")
    print("=" * 70)

    print("\n📌 BASE PROMPT OUTPUT:")
    base_output = run_inference(
        system_prompt=base_messages["system"],
        user_template=base_messages["user"],
        question=question,
        model_name=model_name,
    )
    print(base_output)

    print("\n✅ OPTIMIZED PROMPT OUTPUT:")
    optimized_output = run_inference(
        system_prompt=optimized_messages["system"],
        user_template=optimized_messages["user"],
        question=question,
        model_name=model_name,
    )
    print(optimized_output)

    print("\n📊 COMPARISON:")
    for label, output in [("BASE", base_output), ("OPTIMIZED", optimized_output)]:
        cleaned = clean_json_output(output)
        try:
            parsed = json.loads(cleaned)
            print(f"{label:12s} → valid JSON ✅ | tools called: {list(parsed.keys())}")
        except json.JSONDecodeError:
            print(f"{label:12s} → invalid JSON ❌ | raw: {output[:120]}")

    return base_output, optimized_output

compare_prompts("What is the weather in Tokyo for the next 3 days in celsius?")