{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "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": 5,
   "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": [
    "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",
    "    },\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",
    "\"\"\""
   ]
  }
 ],
 "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
}