{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5437be1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load env variables and create client\n",
    "from dotenv import load_dotenv\n",
    "from anthropic import Anthropic\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "client = Anthropic()\n",
    "model = \"claude-3-5-haiku-latest\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3b0d8e9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper functions\n",
    "def add_user_message(messages, text):\n",
    "    user_message = {\"role\": \"user\", \"content\": text}\n",
    "    messages.append(user_message)\n",
    "\n",
    "\n",
    "def add_assistant_message(messages, text):\n",
    "    assistant_message = {\"role\": \"assistant\", \"content\": text}\n",
    "    messages.append(assistant_message)\n",
    "\n",
    "\n",
    "def chat(messages, system=None, temperature=1.0, stop_sequences=[]):\n",
    "    params = {\n",
    "        \"model\": model,\n",
    "        \"max_tokens\": 1000,\n",
    "        \"messages\": messages,\n",
    "        \"temperature\": temperature,\n",
    "        \"stop_sequences\": stop_sequences,\n",
    "    }\n",
    "\n",
    "    if system:\n",
    "        params[\"system\"] = system\n",
    "\n",
    "    message = client.messages.create(**params)\n",
    "    return message.content[0].text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e788701",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to generate a new dataset\n",
    "import json\n",
    "\n",
    "\n",
    "def generate_dataset():\n",
    "    prompt = \"\"\"\n",
    "Generate a evaluation dataset for a prompt evaluation. The dataset will be used to evaluate prompts\n",
    "that generate Python, JSON, or Regex specifically for AWS-related tasks. Generate an array of JSON objects,\n",
    "each representing task that requires Python, JSON, or a Regex to complete.\n",
    "\n",
    "Example output:\n",
    "```json\n",
    "[\n",
    "    {\n",
    "        \"task\": \"Description of task\",\n",
    "        \"format\": \"json\" or \"python\" or \"regex\"\n",
    "    },\n",
    "    ...additional\n",
    "]\n",
    "```\n",
    "\n",
    "* Focus on tasks that can be solved by writing a single Python function, a single JSON object, or a regular expression.\n",
    "* Focus on tasks that do not require writing much code\n",
    "\n",
    "Please generate 3 objects.\n",
    "\"\"\"\n",
    "\n",
    "    messages = []\n",
    "    add_user_message(messages, prompt)\n",
    "    add_assistant_message(messages, \"```json\")\n",
    "    text = chat(messages, stop_sequences=[\"```\"])\n",
    "    return json.loads(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "438ed743",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate the dataset and write it to 'dataset.json'\n",
    "dataset = generate_dataset()\n",
    "with open(\"dataset.json\", \"w\") as f:\n",
    "    json.dump(dataset, f, indent=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36b89174",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to grade a test case + output using a model\n",
    "def grade_by_model(test_case, output):\n",
    "    eval_prompt = f\"\"\"\n",
    "You are an expert AWS code reviewer. Your task is to evaluate the following AI-generated solution.\n",
    "\n",
    "Original Task:\n",
    "<task>\n",
    "{test_case[\"task\"]}\n",
    "</task>\n",
    "\n",
    "Solution to Evaluate:\n",
    "<solution>\n",
    "{output}\n",
    "</solution>\n",
    "\n",
    "Output Format\n",
    "Provide your evaluation as a structured JSON object with the following fields, in this specific order:\n",
    "- \"strengths\": An array of 1-3 key strengths\n",
    "- \"weaknesses\": An array of 1-3 key areas for improvement\n",
    "- \"reasoning\": A concise explanation of your overall assessment\n",
    "- \"score\": A number between 1-10\n",
    "\n",
    "Respond with JSON. Keep your response concise and direct.\n",
    "Example response shape:\n",
    "{{\n",
    "    \"strengths\": string[],\n",
    "    \"weaknesses\": string[],\n",
    "    \"reasoning\": string,\n",
    "    \"score\": number\n",
    "}}\n",
    "    \"\"\"\n",
    "\n",
    "    messages = []\n",
    "    add_user_message(messages, eval_prompt)\n",
    "    add_assistant_message(messages, \"```json\")\n",
    "    eval_text = chat(messages, stop_sequences=[\"```\"])\n",
    "    return json.loads(eval_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83809a7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Passes a test case into Claude\n",
    "def run_prompt(test_case):\n",
    "    prompt = f\"\"\"\n",
    "Please solve the following task:\n",
    "\n",
    "{test_case[\"task\"]}\n",
    "\"\"\"\n",
    "\n",
    "    messages = []\n",
    "    add_user_message(messages, prompt)\n",
    "    output = chat(messages)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9bcc4671",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to execute a single test case and grade the output\n",
    "def run_test_case(test_case):\n",
    "    \"\"\"Calls run_prompt, then grades the result\"\"\"\n",
    "    output = run_prompt(test_case)\n",
    "\n",
    "    model_grade = grade_by_model(test_case, output)\n",
    "    score = model_grade[\"score\"]\n",
    "    reasoning = model_grade[\"reasoning\"]\n",
    "\n",
    "    return {\n",
    "        \"output\": output,\n",
    "        \"test_case\": test_case,\n",
    "        \"score\": score,\n",
    "        \"reasoning\": reasoning,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5fa99d36",
   "metadata": {},
   "outputs": [],
   "source": [
    "from statistics import mean\n",
    "\n",
    "\n",
    "def run_eval(dataset):\n",
    "    \"\"\"Loads the dataset and calls run_test_case with each case\"\"\"\n",
    "    results = []\n",
    "\n",
    "    for test_case in dataset:\n",
    "        result = run_test_case(test_case)\n",
    "        results.append(result)\n",
    "\n",
    "    average_score = mean([result[\"score\"] for result in results])\n",
    "    print(f\"Average score: {average_score}\")\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "30fae983",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average score: 7.333333333333333\n"
     ]
    }
   ],
   "source": [
    "with open(\"dataset.json\", \"r\") as f:\n",
    "    dataset = json.load(f)\n",
    "\n",
    "results = run_eval(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "dbcc6111",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[\n",
      "  {\n",
      "    \"output\": \"\\n^[a-zA-Z0-9_-]{1,64}$\\n\",\n",
      "    \"test_case\": {\n",
      "      \"task\": \"Create a regular expression to validate an AWS IAM username, which must be 1-64 characters long, contain only alphanumeric characters, underscores, and hyphens\"\n",
      "    },\n",
      "    \"score\": 8,\n",
      "    \"reasoning\": \"The regex correctly captures most IAM username requirements, with strong character set validation and length constraints. Minor refinements could improve robustness against edge cases.\"\n",
      "  },\n",
      "  {\n",
      "    \"output\": \"\\ndef categorize_ec2_instance(instance_type):\\n    # Split the instance type into prefix and size\\n    try:\\n        prefix, size = instance_type.split('.')\\n    except ValueError:\\n        return \\\"Invalid instance type format\\\"\\n\\n    # Define compute categories based on instance type prefixes\\n    compute_categories = {\\n        't': 'Burstable',\\n        'm': 'General Purpose', \\n        'c': 'Compute Optimized',\\n        'r': 'Memory Optimized',\\n        'i': 'Storage Optimized',\\n        'g': 'GPU Instances',\\n        'p': 'GPU/Machine Learning'\\n    }\\n\\n    # Map size to descriptive categories \\n    size_categories = {\\n        'nano': 'Extra Small',\\n        'micro': 'Very Small', \\n        'small': 'Small',\\n        'medium': 'Medium',\\n        'large': 'Large', \\n        'xlarge': 'Extra Large',\\n        '2xlarge': '2X Large',\\n        '4xlarge': '4X Large', \\n        '8xlarge': '8X Large',\\n        '16xlarge': '16X Large'\\n    }\\n\\n    # Determine compute category\\n    compute_category = compute_categories.get(prefix[0], 'Unknown Category')\\n\\n    # Determine size category  \\n    size_category = size_categories.get(size, 'Unknown Size')\\n\\n    return {\\n        'Compute Category': compute_category, \\n        'Size Category': size_category\\n    }\\n\\n# Example usage\\nprint(categorize_ec2_instance('t2.micro'))\\nprint(categorize_ec2_instance('c5.large'))\\nprint(categorize_ec2_instance('r6g.xlarge'))\\n\",\n",
      "    \"test_case\": {\n",
      "      \"task\": \"Write a Python function that converts an EC2 instance type (e.g., 't2.micro') into its corresponding compute category and size\"\n",
      "    },\n",
      "    \"score\": 7,\n",
      "    \"reasoning\": \"The solution provides a solid initial approach to categorizing EC2 instance types with clear category mapping and error handling. It offers a flexible method for understanding instance type characteristics, though it would benefit from more comprehensive AWS instance type validation and expanded size mapping.\"\n",
      "  },\n",
      "  {\n",
      "    \"output\": \"\\n{\\n    \\\"FunctionName\\\": \\\"MyLambdaFunction\\\",\\n    \\\"Runtime\\\": \\\"python3.9\\\",\\n    \\\"Handler\\\": \\\"index.lambda_handler\\\",\\n    \\\"MemorySize\\\": 256,\\n    \\\"Timeout\\\": 30,\\n    \\\"Environment\\\": {\\n        \\\"Variables\\\": {\\n            \\\"LOGGING_LEVEL\\\": \\\"INFO\\\"\\n        }\\n    },\\n    \\\"Role\\\": \\\"arn:aws:iam::123456789012:role/lambda-execution-role\\\"\\n}\\n\",\n",
      "    \"test_case\": {\n",
      "      \"task\": \"Generate a JSON configuration that defines an AWS Lambda function's basic settings, including runtime, handler, and memory allocation\"\n",
      "    },\n",
      "    \"score\": 7,\n",
      "    \"reasoning\": \"The solution provides a basic, functional Lambda function configuration with key parameters. However, it lacks production-readiness and flexibility for different deployment environments.\"\n",
      "  }\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "print(json.dumps(results, indent=2))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}