{
 "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": 3,
   "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": 5,
   "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": 6,
   "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",
    "* Respond only with Python, JSON, or a plain Regex\n",
    "* Do not add any comments or commentary or explanation\n",
    "\"\"\"\n",
    "\n",
    "    messages = []\n",
    "    add_user_message(messages, prompt)\n",
    "    add_assistant_message(messages, \"```code\")\n",
    "    output = chat(messages, stop_sequences=[\"```\"])\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7953c666",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Functions to validate the output structure\n",
    "import re\n",
    "import ast\n",
    "\n",
    "\n",
    "def validate_json(text):\n",
    "    try:\n",
    "        json.loads(text.strip())\n",
    "        return 10\n",
    "    except json.JSONDecodeError:\n",
    "        return 0\n",
    "\n",
    "\n",
    "def validate_python(text):\n",
    "    try:\n",
    "        ast.parse(text.strip())\n",
    "        return 10\n",
    "    except SyntaxError:\n",
    "        return 0\n",
    "\n",
    "\n",
    "def validate_regex(text):\n",
    "    try:\n",
    "        re.compile(text.strip())\n",
    "        return 10\n",
    "    except re.error:\n",
    "        return 0\n",
    "\n",
    "\n",
    "def grade_syntax(response, test_case):\n",
    "    format = test_case[\"format\"]\n",
    "    if format == \"json\":\n",
    "        return validate_json(response)\n",
    "    elif format == \"python\":\n",
    "        return validate_python(response)\n",
    "    else:\n",
    "        return validate_regex(response)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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",
    "    model_score = model_grade[\"score\"]\n",
    "    reasoning = model_grade[\"reasoning\"]\n",
    "\n",
    "    syntax_score = grade_syntax(output, test_case)\n",
    "\n",
    "    score = (model_score + syntax_score) / 2\n",
    "\n",
    "    return {\n",
    "        \"output\": output,\n",
    "        \"test_case\": test_case,\n",
    "        \"score\": score,\n",
    "        \"reasoning\": reasoning,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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": 10,
   "id": "30fae983",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average score: 8.166666666666666\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": 11,
   "id": "dbcc6111",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[\n",
      "  {\n",
      "    \"output\": \"\\n{\\n    \\\"lambda_function\\\": {\\n        \\\"FunctionName\\\": \\\"database-connection-handler\\\",\\n        \\\"Runtime\\\": \\\"python3.9\\\",\\n        \\\"Environment\\\": {\\n            \\\"Variables\\\": {\\n                \\\"DB_HOST\\\": \\\"your-database-hostname.rds.amazonaws.com\\\",\\n                \\\"DB_PORT\\\": \\\"5432\\\",\\n                \\\"DB_NAME\\\": \\\"myappdatabase\\\",\\n                \\\"DB_USERNAME\\\": \\\"dbadminuser\\\", \\n                \\\"DB_PASSWORD\\\": \\\"{{resolve:secretsmanager:DatabaseCredentials:SecretString:password}}\\\",\\n                \\\"DB_SSL_MODE\\\": \\\"require\\\"\\n            }\\n        },\\n        \\\"Timeout\\\": 30,\\n        \\\"MemorySize\\\": 256\\n    }\\n}\\n\",\n",
      "    \"test_case\": {\n",
      "      \"task\": \"Create a JSON configuration for an AWS Lambda function that sets environment variables for database connection\",\n",
      "      \"format\": \"json\"\n",
      "    },\n",
      "    \"score\": 8.5,\n",
      "    \"reasoning\": \"The solution demonstrates good foundational practices for Lambda environment configuration with database connection variables, but lacks advanced security and infrastructure-as-code best practices. The use of Secrets Manager is a strong point, but the implementation could be more dynamic and secure.\"\n",
      "  },\n",
      "  {\n",
      "    \"output\": \"\\ndef extract_aws_region(arn):\\n    import re\\n    match = re.search(r'arn:aws:.*:([^:]+):', arn)\\n    return match.group(1) if match else None\\n\",\n",
      "    \"test_case\": {\n",
      "      \"task\": \"Write a Python function to extract the AWS region from an ARN (Amazon Resource Name) string\",\n",
      "      \"format\": \"python\"\n",
      "    },\n",
      "    \"score\": 8.0,\n",
      "    \"reasoning\": \"The solution provides a basic, functional approach to extracting AWS region from an ARN using regex. While concise, it lacks robust error handling and comprehensive ARN validation. The regex pattern is relatively simple and might miss some complex ARN formats.\"\n",
      "  },\n",
      "  {\n",
      "    \"output\": \"\\n^i-[a-zA-Z0-9]+$\\n\",\n",
      "    \"test_case\": {\n",
      "      \"task\": \"Design a regular expression to validate a valid AWS EC2 instance ID (format: i-[alphanumeric characters])\",\n",
      "      \"format\": \"regex\"\n",
      "    },\n",
      "    \"score\": 8.0,\n",
      "    \"reasoning\": \"The regex captures the basic structure of an EC2 instance ID but lacks precision for AWS's specific requirements. A more robust solution would include length constraints and potentially additional validation.\"\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
}