{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts.few_shot import FewShotPromptTemplate\n",
    "from langchain.prompts.prompt import PromptTemplate\n",
    "from langchain.prompts.example_selector import SemanticSimilarityExampleSelector\n",
    "from langchain.vectorstores import Chroma\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from langchain.llms import OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "examples = [\n",
    "    {\"question\":\"天气怎么样\", \"answer\":\"晴\"},\n",
    "    {\"question\":\"北京天气怎么样\", \"answer\":\"加法\"},\n",
    "    {\"question\":\"咋样\", \"answer\":\"加法\"},\n",
    "    {\"question\":\"天儿咋样\", \"answer\":\"加法\"},\n",
    "    {\"question\":\"今天有雾\", \"answer\":\"加法\"},\n",
    "    {\"question\":\"吃得怎么样\", \"answer\":\"加法\"},\n",
    "    {\"question\":\"你叫什么\", \"answer\":\"加法\"},\n",
    "    {\"question\":\"你叫嘿嘿\", \"answer\":\"加法\"},\n",
    "    {\"question\":\"是吗\", \"answer\":\"加法\"},\n",
    "    {\"question\":\"10-1=9\", \"answer\":\"作业帮回答：减法\"},\n",
    "    {\"question\":\"9-1=8\", \"answer\":\"作业帮回答：减法\"},\n",
    "    {\"question\":\"8-1=7\", \"answer\":\"作业帮回答：减法\"},\n",
    "    {\"question\":\"7-1=6\", \"answer\":\"作业帮回答：减法\"},\n",
    "    {\"question\":\"7-2=5\", \"answer\":\"作业帮回答：减法\"},\n",
    "    {\"question\":\"10-3=7\", \"answer\":\"作业帮回答：减法\"},\n",
    "    {\"question\":\"10-4=6\", \"answer\":\"作业帮回答：减法\"},\n",
    "    {\"question\":\"10-5=5\", \"answer\":\"作业帮回答：减法\"},\n",
    "    {\"question\":\"10+5=15\", \"answer\":\"作业帮回答：加法\"},\n",
    "    {\"question\":\"10*5=50\", \"answer\":\"作业帮回答：乘法\"},\n",
    "    {\"question\":\"10/5=2\", \"answer\":\"作业帮回答：除法\"},\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using embedded DuckDB without persistence: data will be transient\n"
     ]
    }
   ],
   "source": [
    "example_selector = SemanticSimilarityExampleSelector.from_examples(\n",
    "    # This is the list of examples available to select from.\n",
    "    examples,\n",
    "    # This is the embedding class used to produce embeddings which are used to measure semantic similarity.\n",
    "    OpenAIEmbeddings(),\n",
    "    # This is the VectorStore class that is used to store the embeddings and do a similarity search over.\n",
    "    Chroma,\n",
    "    # This is the number of examples to produce.\n",
    "    k=1\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_prompt = PromptTemplate(input_variables=[\"question\", \"answer\"], template=\"Question: {question}\\n{answer}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: 10+5=15\n",
      "作业帮回答：加法\n",
      "\n",
      "Question: 1+2=？\n"
     ]
    }
   ],
   "source": [
    "prompt = FewShotPromptTemplate(\n",
    "    example_selector=example_selector, \n",
    "    example_prompt=example_prompt, \n",
    "    suffix=\"Question: {input}\", \n",
    "    input_variables=[\"input\"]\n",
    ")\n",
    "\n",
    "print(prompt.format(input=\"1+2=？\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'？\\n很晴朗，非常适合出去活动。'"
      ]
     },
     "execution_count": 105,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm = OpenAI(temperature=0)\n",
    "llm(prompt.format(input=\"天气怎么样\"))"
   ]
  }
 ],
 "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
}