17_constitutional.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Constitutional AI\n",
    "\n",
    "Wir möchten, dass die Ausgabe von unserem Agenten sich immer an bestimmte Richtlinien hält.\n",
    "Wir können ein LLM als Revisor dahinter schalten. Das Revisor-LLM kennt die Richtlinien kennt und zwingt den Agenten dazu, bei einer Verletzung der Richtlinien die Ausgabe neu zu generieren.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import (\n",
    "    ChatPromptTemplate,\n",
    "    HumanMessagePromptTemplate,\n",
    "    MessagesPlaceholder,\n",
    "    SystemMessagePromptTemplate,\n",
    ")\n",
    "\n",
    "generator_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        SystemMessagePromptTemplate.from_template(\n",
    "            \"\"\"You are a helpful assistant. Answer this question as good as you can. You are being revised.\n",
    "            If you are being revised, just reformulate your original answer. Original Question:\"\"\"\n",
    "        ),\n",
    "        HumanMessagePromptTemplate.from_template(\"{query}\"),\n",
    "        MessagesPlaceholder(\"revision_dialogue\"),\n",
    "        SystemMessagePromptTemplate.from_template(\n",
    "            \"If the revisor asked you to reformulate your answer, please do so. Do NOT apologize. JUST REFORMULATE.\"\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "\n",
    "revisor_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        SystemMessagePromptTemplate.from_template(\n",
    "            \"\"\"You are a Revisor. You are revising the answer of a chatbot. Your principles are as follows:\n",
    "                {principles}\n",
    "                You can either return the answer from the chatbot to the user or tell the chatbot to reformulate his answer if any principles are violated.\n",
    "                Original Question:\"\"\"\n",
    "        ),\n",
    "        HumanMessagePromptTemplate.from_template(\"{query}\"),\n",
    "        MessagesPlaceholder(\"revision_dialogue\"),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import Annotated, Optional, Sequence, TypedDict\n",
    "\n",
    "from langchain_core.messages import BaseMessage\n",
    "\n",
    "\n",
    "class AgentState(TypedDict):\n",
    "    query: str\n",
    "    principles: str\n",
    "    revision_dialogue: Annotated[Sequence[BaseMessage], operator.add]\n",
    "    answer: Optional[str]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helpers import llm\n",
    "\n",
    "model = llm(model=\"gpt-4o\", temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "from langgraph.graph import END\n",
    "\n",
    "\n",
    "def generator(state):\n",
    "    generator_chain = generator_prompt | model\n",
    "\n",
    "    proposal = generator_chain.invoke(\n",
    "        {\"query\": state[\"query\"], \"revision_dialogue\": state[\"revision_dialogue\"] or []}\n",
    "    )\n",
    "\n",
    "    return {\"revision_dialogue\": [proposal]}\n",
    "\n",
    "\n",
    "def revisor(state):\n",
    "    class Revision(BaseModel):\n",
    "        acceptable: str = Field(\n",
    "            description=\"Is the answer of the chatbot compliant with the principles? answer is 'yes' or 'no'\"\n",
    "        )\n",
    "        critique: Optional[str] = Field(\n",
    "            description=\"Critique of the document. Leave out if answer is acceptable.\",\n",
    "            default=None,\n",
    "        )\n",
    "\n",
    "    revisor_model = model.with_structured_output(Revision)\n",
    "\n",
    "    revisor_chain = revisor_prompt | revisor_model\n",
    "\n",
    "    revision: Revision = revisor_chain.invoke(\n",
    "        {\n",
    "            \"query\": state[\"query\"],\n",
    "            \"revision_dialogue\": state[\"revision_dialogue\"],\n",
    "            \"principles\": state[\"principles\"],\n",
    "        }\n",
    "    )\n",
    "\n",
    "    accepted = revision.acceptable == \"yes\" and revision.critique is None\n",
    "\n",
    "    if accepted:\n",
    "        return {\"answer\": state[\"revision_dialogue\"][-1].content}\n",
    "    else:\n",
    "        return {\"revision_dialogue\": [HumanMessage(content=revision.critique)]}\n",
    "\n",
    "\n",
    "def is_accepted(state):\n",
    "    if state.get(\"answer\", None):\n",
    "        return END\n",
    "    else:\n",
    "        return \"generator\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import StateGraph\n",
    "\n",
    "workflow = StateGraph(AgentState)\n",
    "workflow.add_node(\"generator\", generator)\n",
    "workflow.add_node(\"revisor\", revisor)\n",
    "workflow.set_entry_point(\"generator\")\n",
    "workflow.add_edge(\"generator\", \"revisor\")\n",
    "workflow.add_conditional_edges(\"revisor\", is_accepted)\n",
    "\n",
    "app = workflow.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = {\n",
    "    \"query\": \"Tell me a joke about star trek (or any other sci-fi series).\",\n",
    "    \"principles\": \"The answer must not contain any jokes about star trek.\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for chunk in app.stream(input=inputs):\n",
    "    for k, v in chunk.items():\n",
    "        print(f\"answer from node: {k}\")\n",
    "        print(v)\n",
    "        if k == \"__end__\":\n",
    "            print(\"\\n--- Final answer ---\\n\")\n",
    "            print(v[\"answer\"])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}