rag.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a40a1b69-a23d-4b64-ab0e-4c37ded0e6d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pprint\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_chroma import Chroma\n",
    "from langchain_huggingface import HuggingFaceEmbeddings\n",
    "from langchain_community.document_loaders import PDFMinerLoader\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from typing import Annotated, Literal, Sequence, TypedDict\n",
    "\n",
    "from langchain import hub\n",
    "from langchain_core.messages import BaseMessage, HumanMessage,ToolMessage\n",
    "from langchain_core.prompts import PromptTemplate,ChatPromptTemplate\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "\n",
    "\n",
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "os.environ[\"LANGCHAIN_API_KEY\"] = \"LANGCHAIN_API_KEY\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6242db40-c7a2-427d-961f-5bc76725e852",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/alex/anaconda3/envs/rag/lib/python3.9/site-packages/sentence_transformers/cross_encoder/CrossEncoder.py:11: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
      "  from tqdm.autonotebook import tqdm, trange\n"
     ]
    }
   ],
   "source": [
    "# model\n",
    "llm = ChatOpenAI(\n",
    "    temperature=0.95,\n",
    "    model=\"glm-4\",\n",
    "    openai_api_key=\"openai_api_key\",\n",
    "    openai_api_base=\"https://open.bigmodel.cn/api/paas/v4/\"\n",
    ")\n",
    "embeddings = HuggingFaceEmbeddings(model_name='./bge-m3')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ac86f686-ba70-456f-b636-7aaabc2a6f48",
   "metadata": {},
   "outputs": [],
   "source": [
    "# document\n",
    "file_path = (\n",
    "    \"./data/建筑设计常用规范速查手册第四版.pdf\"\n",
    ")\n",
    "loader = PDFMinerLoader(file_path,extract_images=True)\n",
    "pages = loader.load()\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)\n",
    "splits = text_splitter.split_documents(pages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3d4d745-634e-4d53-87de-acd3de01524f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add to vectorstore\n",
    "vectorstore = Chroma.from_documents(\n",
    "    documents=splits,\n",
    "    collection_name=\"rag-chroma\",\n",
    "    embedding=embeddings,\n",
    ")\n",
    "retriever = vectorstore.as_retriever()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20ec12ab-244f-412b-8c45-41bcc89db0ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.tools.retriever import create_retriever_tool\n",
    "\n",
    "retriever_tool = create_retriever_tool(\n",
    "    retriever,\n",
    "    name=\"retrieve_Building_Rules\",\n",
    "    description=\"Search and return information about building design codes\"\n",
    ")\n",
    "\n",
    "tools = [retriever_tool]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa724b82-239c-4a0f-b870-5c7697c1eac7",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def grade_documents(state) -> Literal[\"generate\", \"rewrite\"]:\n",
    "    \"\"\"\n",
    "    Determines whether the retrieved documents are relevant to the question.\n",
    "\n",
    "    Args:\n",
    "        state (messages): The current state\n",
    "\n",
    "    Returns:\n",
    "        str: A decision for whether the documents are relevant or not\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---CHECK RELEVANCE---\")\n",
    "\n",
    "    # Data model\n",
    "    class grade(BaseModel):\n",
    "        \"\"\"Binary score for relevance check.\"\"\"\n",
    "\n",
    "        binary_score: str = Field(description=\"Relevance score 'yes' or 'no'\")\n",
    "\n",
    "    # LLM\n",
    "    model = ChatOpenAI(\n",
    "    temperature=0.95,\n",
    "    model=\"glm-4\",\n",
    "    openai_api_key=\"fcab620c96c888c22770f19ed85e107e.Uhcc2XfuZnbQlf7C\",\n",
    "    openai_api_base=\"https://open.bigmodel.cn/api/paas/v4/\"\n",
    ")\n",
    "\n",
    "    # LLM with tool and validation\n",
    "    llm_with_tool = model.with_structured_output(grade)\n",
    "\n",
    "    # Prompt\n",
    "    prompt = PromptTemplate(\n",
    "        template=\"\"\"You are a grader assessing relevance of a retrieved document to a user question. \\n \n",
    "        Here is the retrieved document: \\n\\n {context} \\n\\n\n",
    "        Here is the user question: {question} \\n\n",
    "        If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \\n\n",
    "        Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.\"\"\",\n",
    "        input_variables=[\"context\", \"question\"],\n",
    "    )\n",
    "\n",
    "    # Chain\n",
    "    chain = prompt | llm_with_tool\n",
    "\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "\n",
    "    question = messages[0].content\n",
    "    docs = last_message.content\n",
    "\n",
    "    scored_result = chain.invoke({\"question\": question, \"context\": docs})\n",
    "\n",
    "    score = scored_result.binary_score\n",
    "\n",
    "    if score == \"yes\":\n",
    "        print(\"---DECISION: DOCS RELEVANT---\")\n",
    "        return \"generate\"\n",
    "\n",
    "    else:\n",
    "        print(\"---DECISION: DOCS NOT RELEVANT---\")\n",
    "        print(score)\n",
    "        return \"rewrite\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a13ddff-4880-4772-a6e5-e3d271ffea15",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Nodes\n",
    "\n",
    "\n",
    "def agent(state):\n",
    "    \"\"\"\n",
    "    Invokes the agent model to generate a response based on the current state. Given\n",
    "    the question, it will decide to retrieve using the retriever tool, or simply end.\n",
    "\n",
    "    Args:\n",
    "        state (messages): The current state\n",
    "\n",
    "    Returns:\n",
    "        dict: The updated state with the agent response appended to messages\n",
    "    \"\"\"\n",
    "    print(\"---CALL AGENT---\")\n",
    "    messages = state[\"messages\"]\n",
    "    model = ChatOpenAI(\n",
    "    temperature=0.95,\n",
    "    model=\"glm-4\",\n",
    "    openai_api_key=\"fcab620c96c888c22770f19ed85e107e.Uhcc2XfuZnbQlf7C\",\n",
    "    openai_api_base=\"https://open.bigmodel.cn/api/paas/v4/\"\n",
    ")\n",
    "    model = model.bind_tools(tools)\n",
    "    response = model.invoke(messages)\n",
    "    # We return a list, because this will get added to the existing list\n",
    "    return {\"messages\": [response]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0417ab3a-eb53-43de-9773-3efaa0a752f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rewrite(state):\n",
    "    \"\"\"\n",
    "    Transform the query to produce a better question.\n",
    "\n",
    "    Args:\n",
    "        state (messages): The current state\n",
    "\n",
    "    Returns:\n",
    "        dict: The updated state with re-phrased question\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---TRANSFORM QUERY---\")\n",
    "    messages = state[\"messages\"]\n",
    "    question = messages[0].content\n",
    "\n",
    "    msg = [\n",
    "        HumanMessage(\n",
    "            content=f\"\"\" \\n \n",
    "    Look at the input and try to reason about the underlying semantic intent / meaning. \\n \n",
    "    Here is the initial question:\n",
    "    \\n ------- \\n\n",
    "    {question} \n",
    "    \\n ------- \\n\n",
    "    Formulate an improved question: \"\"\",\n",
    "        )\n",
    "    ]\n",
    "\n",
    "    # Grader\n",
    "    model = ChatOpenAI(\n",
    "    temperature=0.95,\n",
    "    model=\"glm-4\",\n",
    "    openai_api_key=\"fcab620c96c888c22770f19ed85e107e.Uhcc2XfuZnbQlf7C\",\n",
    "    openai_api_base=\"https://open.bigmodel.cn/api/paas/v4/\"\n",
    ")\n",
    "    response = model.invoke(msg)\n",
    "    return {\"messages\": [response]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eea9a77d-894e-415f-a8ca-911d4ee72ccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(state):\n",
    "    \"\"\"\n",
    "    Generate answer\n",
    "\n",
    "    Args:\n",
    "        state (messages): The current state\n",
    "\n",
    "    Returns:\n",
    "         dict: The updated state with re-phrased question\n",
    "    \"\"\"\n",
    "    print(\"---GENERATE---\")\n",
    "    messages = state[\"messages\"]\n",
    "    question = messages[0].content\n",
    "    last_message = messages[-1]\n",
    "\n",
    "    docs = last_message.content\n",
    "    name='retrieve_Building_Rules'\n",
    "    question = messages[0].content\n",
    "    docs = last_message.content\n",
    "    # Prompt\n",
    "    prompt = hub.pull(\"rlm/rag-prompt\")\n",
    "    \n",
    "    # LLM\n",
    "    llm = ChatOpenAI(\n",
    "    temperature=0.95,\n",
    "    model=\"glm-4\",\n",
    "    openai_api_key=\"fcab620c96c888c22770f19ed85e107e.Uhcc2XfuZnbQlf7C\",\n",
    "    openai_api_base=\"https://open.bigmodel.cn/api/paas/v4/\"\n",
    ")\n",
    "    # Post-processing\n",
    "    def format_docs(docs):\n",
    "        return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
    "\n",
    "    # Chain\n",
    "    rag_chain = prompt | llm \n",
    "\n",
    "    # Run\n",
    "    response = rag_chain.invoke({\"context\": docs, \"question\": question})\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "print(\"*\" * 20 + \"Prompt[rlm/rag-prompt]\" + \"*\" * 20)\n",
    "prompt = hub.pull(\"rlm/rag-prompt\").pretty_print()  # Show what the prompt looks like"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52efdf45-71e9-4ddb-828f-6f9e67af9c91",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45ccd399-00e2-451b-ad81-a52bf96dc862",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "### Hallucination Grader\n",
    "\n",
    "\n",
    "# Data model\n",
    "class GradeHallucinations(BaseModel):\n",
    "    \"\"\"Binary score for hallucination present in generation answer.\"\"\"\n",
    "\n",
    "    binary_score: str = Field(\n",
    "        description=\"Answer is grounded in the facts, 'yes' or 'no'\"\n",
    "    )\n",
    "\n",
    "\n",
    "# LLM\n",
    "llm = ChatOpenAI(\n",
    "temperature=0.95,\n",
    "model=\"glm-4\",\n",
    "openai_api_key=\"fcab620c96c888c22770f19ed85e107e.Uhcc2XfuZnbQlf7C\",\n",
    "openai_api_base=\"https://open.bigmodel.cn/api/paas/v4/\"\n",
    ")\n",
    "structured_llm_grader = llm.with_structured_output(GradeHallucinations)\n",
    "\n",
    "# Prompt\n",
    "system = \"\"\"You are a grader assessing whether an LLM generation is grounded in / supported by a set of retrieved facts. \\n \n",
    "     Give a binary score 'yes' or 'no'. 'Yes' means that the answer is grounded in / supported by the set of facts.\"\"\"\n",
    "hallucination_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"Set of facts: \\n\\n {documents} \\n\\n LLM generation: {generation}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "hallucination_grader = hallucination_prompt | structured_llm_grader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35db1967-f8ee-4556-b21f-7c3040acd8a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Answer Grader\n",
    "\n",
    "\n",
    "# Data model\n",
    "class GradeAnswer(BaseModel):\n",
    "    \"\"\"Binary score to assess answer addresses question.\"\"\"\n",
    "\n",
    "    binary_score: str = Field(\n",
    "        description=\"Answer addresses the question, 'yes' or 'no'\"\n",
    "    )\n",
    "\n",
    "\n",
    "# LLM with function call\n",
    "# LLM\n",
    "llm = ChatOpenAI(\n",
    "temperature=0.95,\n",
    "model=\"glm-4\",\n",
    "openai_api_key=\"fcab620c96c888c22770f19ed85e107e.Uhcc2XfuZnbQlf7C\",\n",
    "openai_api_base=\"https://open.bigmodel.cn/api/paas/v4/\"\n",
    ")\n",
    "structured_llm_grader = llm.with_structured_output(GradeAnswer)\n",
    "\n",
    "# Prompt\n",
    "system = \"\"\"You are a grader assessing whether an answer addresses / resolves a question \\n \n",
    "     Give a binary score 'yes' or 'no'. Yes' means that the answer resolves the question.\"\"\"\n",
    "answer_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"User question: \\n\\n {question} \\n\\n LLM generation: {generation}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "answer_grader = answer_prompt | structured_llm_grader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "642b21cd-6680-4c3f-9090-5bdbe4af8000",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "463cfabe-c126-450c-820a-5680ee4d3b47",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import filter_messages\n",
    "\n",
    "def grade_generation_v_documents_and_question(state):\n",
    "    \"\"\"\n",
    "    Determines whether the generation is grounded in the document and answers question.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        str: Decision for next node to call\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---CHECK HALLUCINATIONS---\")\n",
    "    messages = state[\"messages\"]\n",
    "    \n",
    "    last_message = messages[-1]\n",
    "    last_toolmessage = filter_messages(messages, include_types=[ToolMessage],include_names=[\"retrieve_Building_Rules\"])[-1]\n",
    "    documents = {last_toolmessage.content}\n",
    "    generation = {last_message.content}\n",
    "    question = {messages[0].content}\n",
    "    score = hallucination_grader.invoke(\n",
    "        {\"documents\": documents, \"generation\": generation}\n",
    "    )\n",
    "    grade = score.binary_score\n",
    "    # Check hallucination\n",
    "    if grade == \"yes\":\n",
    "        print(\"---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---\")\n",
    "        # Check question-answering\n",
    "        print(\"---GRADE GENERATION vs QUESTION---\")\n",
    "        score = answer_grader.invoke({\"question\": question, \"generation\": generation})\n",
    "        grade = score.binary_score\n",
    "        if grade == \"yes\":\n",
    "            print(\"---DECISION: GENERATION ADDRESSES QUESTION---\")\n",
    "            return \"useful\"\n",
    "        else:\n",
    "            print(\"---DECISION: GENERATION DOES NOT ADDRESS QUESTION---\")\n",
    "            return \"not useful\"\n",
    "    else:\n",
    "        print(\"---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---\")\n",
    "        return \"not supported\"\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce2b6996-a726-412c-acc4-bec7f6526d0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph.message import add_messages\n",
    "class AgentState(TypedDict):\n",
    "    # The add_messages function defines how an update should be processed\n",
    "    # Default is to replace. add_messages says \"append\"\n",
    "    messages: Annotated[Sequence[BaseMessage], add_messages]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3396fc9c-577a-4ec8-97af-ff4eee9ea41f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import END, StateGraph, START\n",
    "from langgraph.prebuilt import ToolNode\n",
    "from langgraph.prebuilt import tools_condition\n",
    "# Define a new graph\n",
    "workflow = StateGraph(AgentState)\n",
    "\n",
    "# Define the nodes we will cycle between\n",
    "workflow.add_node(\"agent\", agent)  # agent\n",
    "retrieve = ToolNode([retriever_tool])\n",
    "workflow.add_node(\"retrieve\", retrieve)  # retrieval\n",
    "workflow.add_node(\"rewrite\", rewrite)  # Re-writing the question\n",
    "workflow.add_node(\n",
    "    \"generate\", generate\n",
    ")  # Generating a response after we know the documents are relevant\n",
    "# Call agent node to decide to retrieve or not\n",
    "workflow.add_edge(START, \"agent\")\n",
    "\n",
    "# Decide whether to retrieve\n",
    "workflow.add_conditional_edges(\n",
    "    \"agent\",\n",
    "    # Assess agent decision\n",
    "    tools_condition,\n",
    "    {\n",
    "        # Translate the condition outputs to nodes in our graph\n",
    "        \"tools\": \"retrieve\",\n",
    "        END: END,\n",
    "    },\n",
    ")\n",
    "\n",
    "# Edges taken after the `action` node is called.\n",
    "workflow.add_conditional_edges(\n",
    "    \"retrieve\",\n",
    "    # Assess agent decision\n",
    "    grade_documents,\n",
    ")\n",
    "\n",
    "workflow.add_conditional_edges(\n",
    "    \"generate\",\n",
    "    grade_generation_v_documents_and_question,\n",
    "    {\n",
    "        \"not supported\": \"generate\",\n",
    "        \"useful\": END,\n",
    "        \"not useful\": \"rewrite\",\n",
    "    },\n",
    ")\n",
    "\n",
    "workflow.add_edge(\"rewrite\", \"agent\")\n",
    "\n",
    "memory = MemorySaver()\n",
    "# Compile\n",
    "graph = workflow.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29d62916-5262-49be-96a2-69496ef66672",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "config = {\n",
    "    \"configurable\": {\n",
    "        # Checkpoints are accessed by thread_id\n",
    "        \"thread_id\": 1,\n",
    "    }\n",
    "}\n",
    "\n",
    "for output in graph.stream({\"messages\": (\"user\",  \"医院建筑\")}, config, stream_mode=\"values\"\n",
    "    ):\n",
    "    for key, value in output.items():\n",
    "        pprint.pprint(f\"Output from node '{key}':\")\n",
    "        pprint.pprint(\"---\")\n",
    "        pprint.pprint(value, indent=2, width=80, depth=None)\n",
    "    pprint.pprint(\"\\n---\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd8ffd2b-db95-4e4b-b2af-d3ee1c5bc826",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a74a720-df1c-4b3c-84d8-be703c9ca427",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "try:\n",
    "    display(Image(graph.get_graph(xray=True).draw_mermaid_png()))\n",
    "except Exception:\n",
    "    # This requires some extra dependencies and is optional\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c41bf1ab-6f52-409c-a14e-19ddb3892391",
   "metadata": {},
   "outputs": [],
   "source": [
    "for output in graph.stream({\"messages\": (\"user\",  \"我之前的问题什么\")}, config, stream_mode=\"values\"\n",
    "    ):\n",
    "    for key, value in output.items():\n",
    "        pprint.pprint(f\"Output from node '{key}':\")\n",
    "        pprint.pprint(\"---\")\n",
    "        pprint.pprint(value, indent=2, width=80, depth=None)\n",
    "    pprint.pprint(\"\\n---\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d697744c-1ad8-434c-a514-ded58afd4807",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07f2c106-ac9c-453b-85dc-b2a7c0584558",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc403607-2571-4ab8-b01d-6a83d66e4ce9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e78bf736-a201-4136-9315-d1ed2ac138d3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed597c77-4902-4642-a0ad-e87be10cdddf",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [Conda Environment: rag]",
   "language": "python",
   "name": "rag-kernel"
  },
  "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.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}