{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import PromptTemplate, ChatPromptTemplate, HumanMessagePromptTemplate\n",
    "from langchain.llms import OpenAI\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "\n",
    "from langchain.output_parsers import PydanticOutputParser\n",
    "from pydantic import BaseModel, Field, validator\n",
    "from typing import List"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'text-davinci-003'\n",
    "temperature = 0.7\n",
    "model = OpenAI(model_name=model_name, temperature=temperature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define your desired data structure.\n",
    "class Joke(BaseModel):\n",
    "    setup: str = Field(description=\"question to set up a joke\")\n",
    "    punchline: str = Field(description=\"answer to resolve the joke\")\n",
    "    \n",
    "    # You can add custom validation logic easily with Pydantic.\n",
    "    @validator('setup')\n",
    "    def question_ends_with_question_mark(cls, field):\n",
    "        # 如果是中文，这里需要判断中文问号和英文问号。否则这里会抛出ValueError。\n",
    "        # if field[-1] != '?' and field[-1] != '？':\n",
    "        if field[-1] != '?' and field[-1] != '？':\n",
    "            raise ValueError(\"Badly formed question!\")\n",
    "        return field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up a parser + inject instructions into the prompt template.\n",
    "parser = PydanticOutputParser(pydantic_object=Joke)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = PromptTemplate(\n",
    "    template=\"Answer the user query.\\n{format_instructions}\\n{query}\\n\",\n",
    "    input_variables=[\"query\"],\n",
    "    partial_variables={\"format_instructions\": parser.get_format_instructions()}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "prompt: text='Answer the user query.\\nThe output should be formatted as a JSON instance that conforms to the JSON schema below.\\n\\nAs an example, for the schema {\"properties\": {\"foo\": {\"title\": \"Foo\", \"description\": \"a list of strings\", \"type\": \"array\", \"items\": {\"type\": \"string\"}}}, \"required\": [\"foo\"]}}\\nthe object {\"foo\": [\"bar\", \"baz\"]} is a well-formatted instance of the schema. The object {\"properties\": {\"foo\": [\"bar\", \"baz\"]}} is not well-formatted.\\n\\nHere is the output schema:\\n```\\n{\"properties\": {\"setup\": {\"title\": \"Setup\", \"description\": \"question to set up a joke\", \"type\": \"string\"}, \"punchline\": {\"title\": \"Punchline\", \"description\": \"answer to resolve the joke\", \"type\": \"string\"}}, \"required\": [\"setup\", \"punchline\"]}\\n```\\n用中文回答，讲个笑话。\\n'\n"
     ]
    }
   ],
   "source": [
    "# And a query intented to prompt a language model to populate the data structure.\n",
    "joke_query = \"用中文回答，讲个笑话。\"\n",
    "_input = prompt.format_prompt(query=joke_query)\n",
    "print(\"prompt:\", _input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = model(_input.to_string())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 setup='为什么山羊爱上蜜蜂？' punchline='因为它们都喜欢花蜜！'\n",
      "1 setup='为什么鸡蛋不能上树？' punchline='因为它们太害怕摔下来会摔碎！'\n",
      "2 setup='为什么海洋里的鱼不会发烧？' punchline='因为它们住在温度常定的水里！'\n",
      "3 setup='为什么老鼠喜欢住在洞里？' punchline=\"因为它们可以说：'我家有钥匙！'\"\n",
      "4 setup='为什么火车总是不让狗上车？' punchline='因为它们没有票！'\n",
      "5 setup='为什么小明总是迟到？' punchline='因为他睡得太晚。'\n",
      "6 setup='为什么狐狸不喜欢吃草？' punchline='因为它害怕被牛发现！'\n",
      "7 setup='什么东西很长却没有腿？' punchline='一条彩虹！'\n",
      "8 setup='为什么老虎不吃素食？' punchline='因为他不喜欢蔬菜！'\n",
      "9 setup='为什么小明每次考试都不及格？' punchline='因为他总是把“考”理解成“糊弄”。'\n"
     ]
    }
   ],
   "source": [
    "for i in range(10):\n",
    "    output = model(_input.to_string())\n",
    "    print(i, parser.parse(output))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.10.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}