notebook fmt (#12498)

cookbook/LLaMA2_sql_chat.ipynb

@@ -60,22 +60,21 @@
     }
    ],
    "source": [
-    "# Local \n",
+    "# Local\n",
     "from langchain.chat_models import ChatOllama\n",
+    "\n",
     "llama2_chat = ChatOllama(model=\"llama2:13b-chat\")\n",
     "llama2_code = ChatOllama(model=\"codellama:7b-instruct\")\n",
     "\n",
     "# API\n",
     "from getpass import getpass\n",
     "from langchain.llms import Replicate\n",
+    "\n",
     "# REPLICATE_API_TOKEN = getpass()\n",
     "# os.environ[\"REPLICATE_API_TOKEN\"] = REPLICATE_API_TOKEN\n",
     "replicate_id = \"meta/llama-2-13b-chat:f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d\"\n",
     "llama2_chat_replicate = Replicate(\n",
-    "    model=replicate_id,\n",
-    "    input={\"temperature\": 0.01, \n",
-    "           \"max_length\": 500, \n",
-    "           \"top_p\": 1}\n",
+    "    model=replicate_id, input={\"temperature\": 0.01, \"max_length\": 500, \"top_p\": 1}\n",
     ")"
    ]
   },
@@ -110,11 +109,14 @@
    "outputs": [],
    "source": [
     "from langchain.utilities import SQLDatabase\n",
-    "db = SQLDatabase.from_uri(\"sqlite:///nba_roster.db\", sample_rows_in_table_info= 0)\n",
+    "\n",
+    "db = SQLDatabase.from_uri(\"sqlite:///nba_roster.db\", sample_rows_in_table_info=0)\n",
+    "\n",
     "\n",
     "def get_schema(_):\n",
     "    return db.get_table_info()\n",
     "\n",
+    "\n",
     "def run_query(query):\n",
     "    return db.run(query)"
    ]
@@ -149,26 +151,29 @@
    "source": [
     "# Prompt\n",
     "from langchain.prompts import ChatPromptTemplate\n",
+    "\n",
     "template = \"\"\"Based on the table schema below, write a SQL query that would answer the user's question:\n",
     "{schema}\n",
     "\n",
     "Question: {question}\n",
     "SQL Query:\"\"\"\n",
-    "prompt = ChatPromptTemplate.from_messages([\n",
-    "    (\"system\", \"Given an input question, convert it to a SQL query. No pre-amble.\"),\n",
-    "    (\"human\", template)\n",
-    "])\n",
+    "prompt = ChatPromptTemplate.from_messages(\n",
+    "    [\n",
+    "        (\"system\", \"Given an input question, convert it to a SQL query. No pre-amble.\"),\n",
+    "        (\"human\", template),\n",
+    "    ]\n",
+    ")\n",
     "\n",
     "# Chain to query\n",
     "from langchain.schema.output_parser import StrOutputParser\n",
     "from langchain.schema.runnable import RunnablePassthrough\n",
     "\n",
     "sql_response = (\n",
-    "        RunnablePassthrough.assign(schema=get_schema)\n",
-    "        | prompt\n",
-    "        | llm.bind(stop=[\"\\nSQLResult:\"])\n",
-    "        | StrOutputParser()\n",
-    "    )\n",
+    "    RunnablePassthrough.assign(schema=get_schema)\n",
+    "    | prompt\n",
+    "    | llm.bind(stop=[\"\\nSQLResult:\"])\n",
+    "    | StrOutputParser()\n",
+    ")\n",
     "\n",
     "sql_response.invoke({\"question\": \"What team is Klay Thompson on?\"})"
    ]
@@ -209,18 +214,23 @@
     "Question: {question}\n",
     "SQL Query: {query}\n",
     "SQL Response: {response}\"\"\"\n",
-    "prompt_response = ChatPromptTemplate.from_messages([\n",
-    "    (\"system\", \"Given an input question and SQL response, convert it to a natural langugae answer. No pre-amble.\"),\n",
-    "    (\"human\", template)\n",
-    "])\n",
+    "prompt_response = ChatPromptTemplate.from_messages(\n",
+    "    [\n",
+    "        (\n",
+    "            \"system\",\n",
+    "            \"Given an input question and SQL response, convert it to a natural langugae answer. No pre-amble.\",\n",
+    "        ),\n",
+    "        (\"human\", template),\n",
+    "    ]\n",
+    ")\n",
     "\n",
     "full_chain = (\n",
-    "    RunnablePassthrough.assign(query=sql_response) \n",
+    "    RunnablePassthrough.assign(query=sql_response)\n",
     "    | RunnablePassthrough.assign(\n",
     "        schema=get_schema,\n",
     "        response=lambda x: db.run(x[\"query\"]),\n",
     "    )\n",
-    "    | prompt_response \n",
+    "    | prompt_response\n",
     "    | llm\n",
     ")\n",
     "\n",
@@ -269,36 +279,42 @@
     "# Prompt\n",
     "from langchain.memory import ConversationBufferMemory\n",
     "from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
+    "\n",
     "template = \"\"\"Based on the table schema below, write a SQL query that would answer the user's question:\n",
     "{schema}\n",
     "\n",
     "Question: {question}\n",
     "SQL Query:\"\"\"\n",
-    "prompt = ChatPromptTemplate.from_messages([\n",
-    "    (\"system\", \"Given an input question, convert it to a SQL query. No pre-amble.\"),\n",
-    "    MessagesPlaceholder(variable_name=\"history\"),\n",
-    "    (\"human\", template)\n",
-    "])\n",
+    "prompt = ChatPromptTemplate.from_messages(\n",
+    "    [\n",
+    "        (\"system\", \"Given an input question, convert it to a SQL query. No pre-amble.\"),\n",
+    "        MessagesPlaceholder(variable_name=\"history\"),\n",
+    "        (\"human\", template),\n",
+    "    ]\n",
+    ")\n",
     "\n",
     "memory = ConversationBufferMemory(return_messages=True)\n",
     "\n",
-    "# Chain to query with memory \n",
+    "# Chain to query with memory\n",
     "from langchain.schema.runnable import RunnableLambda\n",
     "\n",
     "sql_chain = (\n",
     "    RunnablePassthrough.assign(\n",
-    "       schema=get_schema,\n",
-    "       history=RunnableLambda(lambda x: memory.load_memory_variables(x)[\"history\"])\n",
-    "    )| prompt\n",
+    "        schema=get_schema,\n",
+    "        history=RunnableLambda(lambda x: memory.load_memory_variables(x)[\"history\"]),\n",
+    "    )\n",
+    "    | prompt\n",
     "    | llm.bind(stop=[\"\\nSQLResult:\"])\n",
     "    | StrOutputParser()\n",
     ")\n",
     "\n",
+    "\n",
     "def save(input_output):\n",
     "    output = {\"output\": input_output.pop(\"output\")}\n",
     "    memory.save_context(input_output, output)\n",
-    "    return output['output']\n",
-    "    \n",
+    "    return output[\"output\"]\n",
+    "\n",
+    "\n",
     "sql_response_memory = RunnablePassthrough.assign(output=sql_chain) | save\n",
     "sql_response_memory.invoke({\"question\": \"What team is Klay Thompson on?\"})"
    ]
@@ -349,18 +365,23 @@
     "Question: {question}\n",
     "SQL Query: {query}\n",
     "SQL Response: {response}\"\"\"\n",
-    "prompt_response = ChatPromptTemplate.from_messages([\n",
-    "    (\"system\", \"Given an input question and SQL response, convert it to a natural langugae answer. No pre-amble.\"),\n",
-    "    (\"human\", template)\n",
-    "])\n",
+    "prompt_response = ChatPromptTemplate.from_messages(\n",
+    "    [\n",
+    "        (\n",
+    "            \"system\",\n",
+    "            \"Given an input question and SQL response, convert it to a natural langugae answer. No pre-amble.\",\n",
+    "        ),\n",
+    "        (\"human\", template),\n",
+    "    ]\n",
+    ")\n",
     "\n",
     "full_chain = (\n",
-    "    RunnablePassthrough.assign(query=sql_response_memory) \n",
+    "    RunnablePassthrough.assign(query=sql_response_memory)\n",
     "    | RunnablePassthrough.assign(\n",
     "        schema=get_schema,\n",
     "        response=lambda x: db.run(x[\"query\"]),\n",
     "    )\n",
-    "    | prompt_response \n",
+    "    | prompt_response\n",
     "    | llm\n",
     ")\n",
     "\n",