{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "^C\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "grpcio-tools 1.66.1 requires protobuf<6.0dev,>=5.26.1, but you have protobuf 4.25.5 which is incompatible.\n",
      "langchain-chroma 0.1.3 requires langchain-core<0.3,>=0.1.40, but you have langchain-core 0.3.5 which is incompatible.\n",
      "langchain-huggingface 0.0.3 requires langchain-core<0.3,>=0.1.52, but you have langchain-core 0.3.5 which is incompatible.\n",
      "ragas 0.1.20 requires langchain-core<0.3, but you have langchain-core 0.3.5 which is incompatible.\n"
     ]
    }
   ],
   "source": [
    "%pip install -qU langchain-community tiktoken langchain-openai langchainhub langchain langgraph langchain-text-splitters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install qdrant-client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import getpass\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API Key:\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import PyMuPDFLoader\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_community.vectorstores import Qdrant\n",
    "\n",
    "pdfs = [\n",
    "    \"C:/Users/andre/OneDrive/Documents/AIE4/AIE4/Midterm/Blueprint-for-an-AI-Bill-of-Rights.pdf\",\n",
    "    \"C:/Users/andre/OneDrive/Documents/AIE4/AIE4/Midterm/NIST_report.pdf\",\n",
    "]\n",
    "\n",
    "docs = [PyMuPDFLoader(pdf).load() for pdf in pdfs]\n",
    "\n",
    "docs_list = [item for sublist in docs for item in sublist]\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=500, chunk_overlap=50\n",
    ")\n",
    "\n",
    "doc_splits = text_splitter.split_documents(docs_list)\n",
    "\n",
    "embeddings = OpenAIEmbeddings(model=\"text-embedding-3-small\")\n",
    "\n",
    "vectorstore = Qdrant.from_documents(\n",
    "    documents=doc_splits,\n",
    "    embedding=embeddings,\n",
    "    location=\":memory:\",\n",
    "    collection_name=\"rag-agentic\"\n",
    ")\n",
    "\n",
    "retriever = vectorstore.as_retriever()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.tools.retriever import create_retriever_tool\n",
    "\n",
    "retriever_tool = create_retriever_tool(\n",
    "    retriever,\n",
    "    \"retrieve_blog_posts\",\n",
    "    \"Search and return information about the responsible and ethical use of AI along with the development of policies and practices to protect civil rights and promote democratic values in the building, deployment, and government of automated systems.\",\n",
    ")\n",
    "\n",
    "tools = [retriever_tool]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated, Sequence, TypedDict\n",
    "\n",
    "from langchain_core.messages import BaseMessage\n",
    "\n",
    "from langgraph.graph.message import add_messages\n",
    "\n",
    "\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": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated, Literal, Sequence, TypedDict\n",
    "\n",
    "from langchain import hub\n",
    "from langchain_core.messages import BaseMessage, HumanMessage\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "# NOTE: you must use langchain-core >= 0.3 with Pydantic v2\n",
    "from pydantic import BaseModel, Field\n",
    "from langgraph.prebuilt import tools_condition\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "### Edges\n",
    "\n",
    "\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(temperature=0, model=\"gpt-4o-mini\", streaming=True)\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": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"LANGCHAIN_API_KEY\"] = getpass.getpass(\"LangChain API Key:\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "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(temperature=0, streaming=True, model=\"gpt-4o-mini\")\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]}\n",
    "\n",
    "\n",
    "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(temperature=0, model=\"gpt-4o-mini\", streaming=True)\n",
    "    response = model.invoke(msg)\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "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",
    "\n",
    "    # Prompt\n",
    "    prompt = hub.pull(\"rlm/rag-prompt\")\n",
    "\n",
    "    # LLM\n",
    "    llm = ChatOpenAI(model_name=\"gpt-4o-mini\", temperature=0, streaming=True)\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 | StrOutputParser()\n",
    "\n",
    "    # Run\n",
    "    response = rag_chain.invoke({\"context\": docs, \"question\": question})\n",
    "    return {\"messages\": [response]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import END, StateGraph, START\n",
    "from langgraph.prebuilt import ToolNode\n",
    "\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",
    "workflow.add_edge(\"generate\", END)\n",
    "workflow.add_edge(\"rewrite\", \"agent\")\n",
    "\n",
    "# Compile\n",
    "graph = workflow.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---CALL AGENT---\n",
      "\"Output from node 'agent':\"\n",
      "'---'\n",
      "('Some problems with AI include biases in algorithms that can lead to unfair '\n",
      " 'treatment of individuals and the potential for job displacement as '\n",
      " 'automation increases. Additionally, concerns about privacy, security, and '\n",
      " 'the ethical implications of decision-making by AI systems pose significant '\n",
      " 'challenges.')\n",
      "'\\n---\\n'\n"
     ]
    }
   ],
   "source": [
    "import pprint\n",
    "\n",
    "inputs = {\n",
    "    \"messages\": [\n",
    "        (\"user\", \"What are some problems with AI? Give me a response in two sentences or less\"),\n",
    "    ]\n",
    "}\n",
    "\n",
    "\n",
    "for output in graph.stream(inputs):\n",
    "    for key, value in output.items():\n",
    "        pprint.pprint(f\"Output from node '{key}':\")\n",
    "        pprint.pprint(\"---\")\n",
    "        pprint.pprint(value['messages'][0].content, indent=2, width=80, depth=None)\n",
    "    pprint.pprint(\"\\n---\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm-ops",
   "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
}