{
  "cells": [
    {
      "attachments": {
        "bdc25ea2-123b-46b1-b9f5-fdd345ecbc73.png": {
          "image/png": "[encoded data, 541088 characters]"
        }
      },
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Web Research (STORM)\n",
        "\n",
        "[STORM](https://arxiv.org/abs/2402.14207) is a research assistant designed by Shao, et. al that extends the idea of \"outline-driven RAG\" for richer article generation.\n",
        "\n",
        "STORM is designed to generate Wikipedia-style ariticles on a user-provided topic. It applies two main insights to produce more organized and comprehensive articles:\n",
        "\n",
        "1. Creating an outline (planning) by querying similar topics helps improve coverage.\n",
        "2. Multi-perspective, grounded (in search) conversation simulation helps increase the reference count and information density. \n",
        "\n",
        "The control flow looks like the diagram below.\n",
        "\n",
        "![storm.png](attachment:bdc25ea2-123b-46b1-b9f5-fdd345ecbc73.png)\n",
        "\n",
        "STORM has a few main stages:\n",
        "\n",
        "1. Generate initial outline + Survey related subjects\n",
        "2. Identify distinct perspectives\n",
        "3. \"Interview subject matter experts\" (role-playing LLMs)\n",
        "4. Refine outline (using references)\n",
        "5. Write sections, then write article\n",
        "\n",
        "\n",
        "The expert interviews stage occurs between the role-playing article writer and a research expert. The \"expert\" is able to query external knowledge and respond to pointed questions, saving cited sources to a vectorstore so that the later refinement stages can synthesize the full article.\n",
        "\n",
        "There are a couple hyperparameters you can set to restrict the (potentially) infinite research breadth:\n",
        "\n",
        "N: Number of perspectives to survey / use (Steps 2->3)\n",
        "M: Max number of conversation turns in step (Step 3)\n",
        "\n",
        "\n",
        "## Prerequisites"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [],
      "source": [
        "%%capture --no-stderr\n%pip install -U langchain_community langchain_openai langgraph wikipedia  scikit-learn  langchain_fireworks\n# We use one or the other search engine below\n%pip install -U duckduckgo tavily-python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Uncomment if you want to draw the pretty graph diagrams.\n# If you are on MacOS, you will need to run brew install graphviz before installing and update some environment flags\n# ! brew install graphviz\n# !CFLAGS=\"-I $(brew --prefix graphviz)/include\" LDFLAGS=\"-L $(brew --prefix graphviz)/lib\" pip install -U pygraphviz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 86,
      "metadata": {},
      "outputs": [],
      "source": [
        "import getpass\nimport os\n\n\ndef _set_env(var: str):\n    if os.environ.get(var):\n        return\n    os.environ[var] = getpass.getpass(var + \":\")\n\n\n# Set for tracing\nos.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\nos.environ[\"LANGCHAIN_PROJECT\"] = \"STORM\"\n_set_env(\"LANGCHAIN_API_KEY\")\n_set_env(\"OPENAI_API_KEY\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Select LLMs\n",
        "\n",
        "We will have a faster LLM do most of the work, but a slower, long-context model to distill the conversations and write the final report."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "from langchain_openai import ChatOpenAI\n\nfast_llm = ChatOpenAI(model=\"gpt-3.5-turbo\")\n# Uncomment for a Fireworks model\n# fast_llm = ChatFireworks(model=\"accounts/fireworks/models/firefunction-v1\", max_tokens=32_000)\nlong_context_llm = ChatOpenAI(model=\"gpt-4-turbo-preview\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generate Initial Outline\n",
        "\n",
        "For many topics, your LLM may have an initial idea of the important and related topics. We can generate an initial\n",
        "outline to be refined after our research. Below, we will use our \"fast\" llm to generate the outline."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/Users/wfh/code/lc/langchain/libs/core/langchain_core/_api/beta_decorator.py:86: LangChainBetaWarning: The function `with_structured_output` is in beta. It is actively being worked on, so the API may change.\n",
            "  warn_beta(\n"
          ]
        }
      ],
      "source": [
        "from typing import List, Optional\n\nfrom langchain_core.prompts import ChatPromptTemplate\nfrom langchain_core.pydantic_v1 import BaseModel, Field\n\ndirect_gen_outline_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"You are a Wikipedia writer. Write an outline for a Wikipedia page about a user-provided topic. Be comprehensive and specific.\",\n        ),\n        (\"user\", \"{topic}\"),\n    ]\n)\n\n\nclass Subsection(BaseModel):\n    subsection_title: str = Field(..., title=\"Title of the subsection\")\n    description: str = Field(..., title=\"Content of the subsection\")\n\n    @property\n    def as_str(self) -> str:\n        return f\"### {self.subsection_title}\\n\\n{self.description}\".strip()\n\n\nclass Section(BaseModel):\n    section_title: str = Field(..., title=\"Title of the section\")\n    description: str = Field(..., title=\"Content of the section\")\n    subsections: Optional[List[Subsection]] = Field(\n        default=None,\n        title=\"Titles and descriptions for each subsection of the Wikipedia page.\",\n    )\n\n    @property\n    def as_str(self) -> str:\n        subsections = \"\\n\\n\".join(\n            f\"### {subsection.subsection_title}\\n\\n{subsection.description}\"\n            for subsection in self.subsections or []\n        )\n        return f\"## {self.section_title}\\n\\n{self.description}\\n\\n{subsections}\".strip()\n\n\nclass Outline(BaseModel):\n    page_title: str = Field(..., title=\"Title of the Wikipedia page\")\n    sections: List[Section] = Field(\n        default_factory=list,\n        title=\"Titles and descriptions for each section of the Wikipedia page.\",\n    )\n\n    @property\n    def as_str(self) -> str:\n        sections = \"\\n\\n\".join(section.as_str for section in self.sections)\n        return f\"# {self.page_title}\\n\\n{sections}\".strip()\n\n\ngenerate_outline_direct = direct_gen_outline_prompt | fast_llm.with_structured_output(\n    Outline\n)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "# Impact of million-plus token context window language models on RAG\n",
            "\n",
            "## Introduction\n",
            "\n",
            "Overview of million-plus token context window language models and RAG (Retrieval-Augmented Generation).\n",
            "\n",
            "## Million-Plus Token Context Window Language Models\n",
            "\n",
            "Explanation of million-plus token context window language models, their architecture, training data, and applications.\n",
            "\n",
            "## RAG (Retrieval-Augmented Generation)\n",
            "\n",
            "Overview of RAG, its architecture, how it combines retrieval and generation models, and its use in natural language processing tasks.\n",
            "\n",
            "## Impact on RAG\n",
            "\n",
            "Discuss the impact of million-plus token context window language models on RAG, including improvements in performance, efficiency, and challenges faced.\n"
          ]
        }
      ],
      "source": [
        "example_topic = \"Impact of million-plus token context window language models on RAG\"\n\ninitial_outline = generate_outline_direct.invoke({\"topic\": example_topic})\n\nprint(initial_outline.as_str)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Expand Topics\n",
        "\n",
        "While language models do store some Wikipedia-like knowledge in their parameters, you will get better results by incorporating relevant and recent information using a search engine.\n",
        "\n",
        "We will start our search by generating a list of related topics, sourced from Wikipedia."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": [
        "gen_related_topics_prompt = ChatPromptTemplate.from_template(\n    \"\"\"I'm writing a Wikipedia page for a topic mentioned below. Please identify and recommend some Wikipedia pages on closely related subjects. I'm looking for examples that provide insights into interesting aspects commonly associated with this topic, or examples that help me understand the typical content and structure included in Wikipedia pages for similar topics.\n\nPlease list the as many subjects and urls as you can.\n\nTopic of interest: {topic}\n\"\"\"\n)\n\n\nclass RelatedSubjects(BaseModel):\n    topics: List[str] = Field(\n        description=\"Comprehensive list of related subjects as background research.\",\n    )\n\n\nexpand_chain = gen_related_topics_prompt | fast_llm.with_structured_output(\n    RelatedSubjects\n)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "RelatedSubjects(topics=['Language models', 'Retriever-Reader-Generator (RAG) model', 'Natural language processing', 'Machine learning', 'Artificial intelligence', 'Text generation', 'Transformer architecture', 'Context window', 'Impact of language models'])"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "related_subjects = await expand_chain.ainvoke({\"topic\": example_topic})\nrelated_subjects"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generate Perspectives\n",
        "\n",
        "From these related subjects, we can select representative Wikipedia editors as \"subject matter experts\" with distinct\n",
        "backgrounds and affiliations. These will help distribute the search process to encourage a more well-rounded final report."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [],
      "source": [
        "class Editor(BaseModel):\n    affiliation: str = Field(\n        description=\"Primary affiliation of the editor.\",\n    )\n    name: str = Field(\n        description=\"Name of the editor.\", pattern=r\"^[a-zA-Z0-9_-]{1,64}$\"\n    )\n    role: str = Field(\n        description=\"Role of the editor in the context of the topic.\",\n    )\n    description: str = Field(\n        description=\"Description of the editor's focus, concerns, and motives.\",\n    )\n\n    @property\n    def persona(self) -> str:\n        return f\"Name: {self.name}\\nRole: {self.role}\\nAffiliation: {self.affiliation}\\nDescription: {self.description}\\n\"\n\n\nclass Perspectives(BaseModel):\n    editors: List[Editor] = Field(\n        description=\"Comprehensive list of editors with their roles and affiliations.\",\n        # Add a pydantic validation/restriction to be at most M editors\n    )\n\n\ngen_perspectives_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"\"\"You need to select a diverse (and distinct) group of Wikipedia editors who will work together to create a comprehensive article on the topic. Each of them represents a different perspective, role, or affiliation related to this topic.\\\n    You can use other Wikipedia pages of related topics for inspiration. For each editor, add a description of what they will focus on.\n\n    Wiki page outlines of related topics for inspiration:\n    {examples}\"\"\",\n        ),\n        (\"user\", \"Topic of interest: {topic}\"),\n    ]\n)\n\ngen_perspectives_chain = gen_perspectives_prompt | ChatOpenAI(\n    model=\"gpt-3.5-turbo\"\n).with_structured_output(Perspectives)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [],
      "source": [
        "from langchain_community.retrievers import WikipediaRetriever\nfrom langchain_core.runnables import RunnableLambda\nfrom langchain_core.runnables import chain as as_runnable\n\nwikipedia_retriever = WikipediaRetriever(load_all_available_meta=True, top_k_results=1)\n\n\ndef format_doc(doc, max_length=1000):\n    related = \"- \".join(doc.metadata[\"categories\"])\n    return f\"### {doc.metadata['title']}\\n\\nSummary: {doc.page_content}\\n\\nRelated\\n{related}\"[\n        :max_length\n    ]\n\n\ndef format_docs(docs):\n    return \"\\n\\n\".join(format_doc(doc) for doc in docs)\n\n\n@as_runnable\nasync def survey_subjects(topic: str):\n    related_subjects = await expand_chain.ainvoke({\"topic\": topic})\n    retrieved_docs = await wikipedia_retriever.abatch(\n        related_subjects.topics, return_exceptions=True\n    )\n    all_docs = []\n    for docs in retrieved_docs:\n        if isinstance(docs, BaseException):\n            continue\n        all_docs.extend(docs)\n    formatted = format_docs(all_docs)\n    return await gen_perspectives_chain.ainvoke({\"examples\": formatted, \"topic\": topic})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [],
      "source": [
        "perspectives = await survey_subjects.ainvoke(example_topic)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'editors': [{'affiliation': 'Academic Research',\n",
              "   'name': 'Dr. Linguist',\n",
              "   'role': 'Language Model Expert',\n",
              "   'description': 'Dr. Linguist will focus on explaining the technical aspects of million-plus token context window language models and their impact on RAG (Retrieval-Augmented Generation) systems.'},\n",
              "  {'affiliation': 'Industry',\n",
              "   'name': 'TechTrendz',\n",
              "   'role': 'AI Solutions Architect',\n",
              "   'description': 'TechTrendz will provide insights on the practical applications of million-plus token context window language models in RAG systems and discuss their benefits and challenges in real-world scenarios.'},\n",
              "  {'affiliation': 'Open Source Community',\n",
              "   'name': 'CodeGenius',\n",
              "   'role': 'Machine Learning Enthusiast',\n",
              "   'description': 'CodeGenius will explore the open-source tools and frameworks available for implementing million-plus token context window language models in RAG systems and share their experiences with the community.'},\n",
              "  {'affiliation': 'Tech Journalism',\n",
              "   'name': 'DataDive',\n",
              "   'role': 'AI Technology Journalist',\n",
              "   'description': 'DataDive will cover the latest developments and advancements in million-plus token context window language models and their implications for RAG systems, focusing on industry trends and use cases.'}]}"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "perspectives.dict()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Expert Dialog\n",
        "\n",
        "Now the true fun begins, each wikipedia writer is primed to role-play using the perspectives presented above. It will ask a series of questions of a second \"domain expert\" with access to a search engine. This generate content to generate a refined outline as well as an updated index of reference documents.\n",
        "\n",
        "\n",
        "### Interview State\n",
        "\n",
        "The conversation is cyclic, so we will construct it within its own graph. The State will contain messages, the reference docs, and the editor (with its own \"persona\") to make it easy to parallelize these conversations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {},
      "outputs": [],
      "source": [
        "from typing import Annotated\n\nfrom langchain_core.messages import AnyMessage\nfrom typing_extensions import TypedDict\n\nfrom langgraph.graph import END, StateGraph, START\n\n\ndef add_messages(left, right):\n    if not isinstance(left, list):\n        left = [left]\n    if not isinstance(right, list):\n        right = [right]\n    return left + right\n\n\ndef update_references(references, new_references):\n    if not references:\n        references = {}\n    references.update(new_references)\n    return references\n\n\ndef update_editor(editor, new_editor):\n    # Can only set at the outset\n    if not editor:\n        return new_editor\n    return editor\n\n\nclass InterviewState(TypedDict):\n    messages: Annotated[List[AnyMessage], add_messages]\n    references: Annotated[Optional[dict], update_references]\n    editor: Annotated[Optional[Editor], update_editor]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Dialog Roles\n",
        "\n",
        "The graph will have two participants: the wikipedia editor (`generate_question`), who asks questions based on its assigned role, and a domain expert (`gen_answer_chain), who uses a search engine to answer the questions as accurately as possible."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [],
      "source": [
        "from langchain_core.messages import AIMessage, HumanMessage, ToolMessage\nfrom langchain_core.prompts import MessagesPlaceholder\n\ngen_qn_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"\"\"You are an experienced Wikipedia writer and want to edit a specific page. \\\nBesides your identity as a Wikipedia writer, you have a specific focus when researching the topic. \\\nNow, you are chatting with an expert to get information. Ask good questions to get more useful information.\n\nWhen you have no more questions to ask, say \"Thank you so much for your help!\" to end the conversation.\\\nPlease only ask one question at a time and don't ask what you have asked before.\\\nYour questions should be related to the topic you want to write.\nBe comprehensive and curious, gaining as much unique insight from the expert as possible.\\\n\nStay true to your specific perspective:\n\n{persona}\"\"\",\n        ),\n        MessagesPlaceholder(variable_name=\"messages\", optional=True),\n    ]\n)\n\n\ndef tag_with_name(ai_message: AIMessage, name: str):\n    ai_message.name = name\n    return ai_message\n\n\ndef swap_roles(state: InterviewState, name: str):\n    converted = []\n    for message in state[\"messages\"]:\n        if isinstance(message, AIMessage) and message.name != name:\n            message = HumanMessage(**message.dict(exclude={\"type\"}))\n        converted.append(message)\n    return {\"messages\": converted}\n\n\n@as_runnable\nasync def generate_question(state: InterviewState):\n    editor = state[\"editor\"]\n    gn_chain = (\n        RunnableLambda(swap_roles).bind(name=editor.name)\n        | gen_qn_prompt.partial(persona=editor.persona)\n        | fast_llm\n        | RunnableLambda(tag_with_name).bind(name=editor.name)\n    )\n    result = await gn_chain.ainvoke(state)\n    return {\"messages\": [result]}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "\"Yes, that's correct. I'm focusing on the technical aspects of million-plus token context window language models and their impact on Retrieval-Augmented Generation (RAG) systems. Can you provide more information on how these large context window language models are trained and how they differ from traditional models in the context of RAG systems?\""
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "messages = [\n    HumanMessage(f\"So you said you were writing an article on {example_topic}?\")\n]\nquestion = await generate_question.ainvoke(\n    {\n        \"editor\": perspectives.editors[0],\n        \"messages\": messages,\n    }\n)\n\nquestion[\"messages\"][0].content"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Answer questions\n",
        "\n",
        "The `gen_answer_chain` first generates queries (query expansion) to answer the editor's question, then responds with citations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {},
      "outputs": [],
      "source": [
        "class Queries(BaseModel):\n    queries: List[str] = Field(\n        description=\"Comprehensive list of search engine queries to answer the user's questions.\",\n    )\n\n\ngen_queries_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"You are a helpful research assistant. Query the search engine to answer the user's questions.\",\n        ),\n        MessagesPlaceholder(variable_name=\"messages\", optional=True),\n    ]\n)\ngen_queries_chain = gen_queries_prompt | ChatOpenAI(\n    model=\"gpt-3.5-turbo\"\n).with_structured_output(Queries, include_raw=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['Training process of million-plus token context window language models',\n",
              " 'Differences between large context window language models and traditional models in Retrieval-Augmented Generation systems']"
            ]
          },
          "execution_count": 17,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "queries = await gen_queries_chain.ainvoke(\n    {\"messages\": [HumanMessage(content=question[\"messages\"][0].content)]}\n)\nqueries[\"parsed\"].queries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {},
      "outputs": [],
      "source": [
        "class AnswerWithCitations(BaseModel):\n    answer: str = Field(\n        description=\"Comprehensive answer to the user's question with citations.\",\n    )\n    cited_urls: List[str] = Field(\n        description=\"List of urls cited in the answer.\",\n    )\n\n    @property\n    def as_str(self) -> str:\n        return f\"{self.answer}\\n\\nCitations:\\n\\n\" + \"\\n\".join(\n            f\"[{i+1}]: {url}\" for i, url in enumerate(self.cited_urls)\n        )\n\n\ngen_answer_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"\"\"You are an expert who can use information effectively. You are chatting with a Wikipedia writer who wants\\\n to write a Wikipedia page on the topic you know. You have gathered the related information and will now use the information to form a response.\n\nMake your response as informative as possible and make sure every sentence is supported by the gathered information.\nEach response must be backed up by a citation from a reliable source, formatted as a footnote, reproducing the URLS after your response.\"\"\",\n        ),\n        MessagesPlaceholder(variable_name=\"messages\", optional=True),\n    ]\n)\n\ngen_answer_chain = gen_answer_prompt | fast_llm.with_structured_output(\n    AnswerWithCitations, include_raw=True\n).with_config(run_name=\"GenerateAnswer\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {},
      "outputs": [],
      "source": [
        "from langchain_community.utilities.duckduckgo_search import DuckDuckGoSearchAPIWrapper\nfrom langchain_core.tools import tool\n\n'''\n# Tavily is typically a better search engine, but your free queries are limited\nsearch_engine = TavilySearchResults(max_results=4)\n\n@tool\nasync def search_engine(query: str):\n    \"\"\"Search engine to the internet.\"\"\"\n    results = tavily_search.invoke(query)\n    return [{\"content\": r[\"content\"], \"url\": r[\"url\"]} for r in results]\n'''\n\n# DDG\nsearch_engine = DuckDuckGoSearchAPIWrapper()\n\n\n@tool\nasync def search_engine(query: str):\n    \"\"\"Search engine to the internet.\"\"\"\n    results = DuckDuckGoSearchAPIWrapper()._ddgs_text(query)\n    return [{\"content\": r[\"body\"], \"url\": r[\"href\"]} for r in results]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import json\n\nfrom langchain_core.runnables import RunnableConfig\n\n\nasync def gen_answer(\n    state: InterviewState,\n    config: Optional[RunnableConfig] = None,\n    name: str = \"Subject_Matter_Expert\",\n    max_str_len: int = 15000,\n):\n    swapped_state = swap_roles(state, name)  # Convert all other AI messages\n    queries = await gen_queries_chain.ainvoke(swapped_state)\n    query_results = await search_engine.abatch(\n        queries[\"parsed\"].queries, config, return_exceptions=True\n    )\n    successful_results = [\n        res for res in query_results if not isinstance(res, Exception)\n    ]\n    all_query_results = {\n        res[\"url\"]: res[\"content\"] for results in successful_results for res in results\n    }\n    # We could be more precise about handling max token length if we wanted to here\n    dumped = json.dumps(all_query_results)[:max_str_len]\n    ai_message: AIMessage = queries[\"raw\"]\n    tool_call = queries[\"raw\"].additional_kwargs[\"tool_calls\"][0]\n    tool_id = tool_call[\"id\"]\n    tool_message = ToolMessage(tool_call_id=tool_id, content=dumped)\n    swapped_state[\"messages\"].extend([ai_message, tool_message])\n    # Only update the shared state with the final answer to avoid\n    # polluting the dialogue history with intermediate messages\n    generated = await gen_answer_chain.ainvoke(swapped_state)\n    cited_urls = set(generated[\"parsed\"].cited_urls)\n    # Save the retrieved information to a the shared state for future reference\n    cited_references = {k: v for k, v in all_query_results.items() if k in cited_urls}\n    formatted_message = AIMessage(name=name, content=generated[\"parsed\"].as_str)\n    return {\"messages\": [formatted_message], \"references\": cited_references}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'Large context window language models, such as the Llama2 70B model, can support context windows of more than 100k tokens without continual training through innovations like Dual Chunk Attention (DCA). These models have significantly longer context windows compared to traditional models, with capabilities like processing up to 1 million tokens at once, providing more consistent and relevant outputs. Training these models often involves starting with a smaller window size and gradually increasing it through fine-tuning on larger windows. In contrast, traditional models have much shorter context windows, limiting their ability to process extensive information in a prompt. Retrieval-Augmented Generation (RAG) systems, on the other hand, integrate large language models with external knowledge sources to enhance their performance, offering a pathway to combine the capabilities of models like ChatGPT/GPT-4 with custom data sources for more informed and contextually aware outputs.\\n\\nCitations:\\n\\n[1]: https://arxiv.org/abs/2402.17463\\n[2]: https://blog.google/technology/ai/long-context-window-ai-models/\\n[3]: https://medium.com/@ddxzzx/why-and-how-to-achieve-longer-context-windows-for-llms-5f76f8656ea9\\n[4]: https://blog.google/technology/ai/google-gemini-next-generation-model-february-2024/\\n[5]: https://huggingface.co/papers/2402.13753\\n[6]: https://www.pinecone.io/blog/why-use-retrieval-instead-of-larger-context/\\n[7]: https://medium.com/emalpha/innovations-in-retrieval-augmented-generation-8e6e70f95629\\n[8]: https://inside-machinelearning.com/en/rag/'"
            ]
          },
          "execution_count": 21,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "example_answer = await gen_answer(\n    {\"messages\": [HumanMessage(content=question[\"messages\"][0].content)]}\n)\nexample_answer[\"messages\"][-1].content"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Construct the Interview Graph\n",
        "\n",
        "\n",
        "Now that we've defined the editor and domain expert, we can compose them in a graph."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {},
      "outputs": [],
      "source": [
        "max_num_turns = 5\n\n\ndef route_messages(state: InterviewState, name: str = \"Subject_Matter_Expert\"):\n    messages = state[\"messages\"]\n    num_responses = len(\n        [m for m in messages if isinstance(m, AIMessage) and m.name == name]\n    )\n    if num_responses >= max_num_turns:\n        return END\n    last_question = messages[-2]\n    if last_question.content.endswith(\"Thank you so much for your help!\"):\n        return END\n    return \"ask_question\"\n\n\nbuilder = StateGraph(InterviewState)\n\nbuilder.add_node(\"ask_question\", generate_question)\nbuilder.add_node(\"answer_question\", gen_answer)\nbuilder.add_conditional_edges(\"answer_question\", route_messages)\nbuilder.add_edge(\"ask_question\", \"answer_question\")\n\nbuilder.add_edge(START, \"ask_question\")\ninterview_graph = builder.compile().with_config(run_name=\"Conduct Interviews\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 57940 characters]",
            "text/plain": [
              "<IPython.core.display.Image object>"
            ]
          },
          "execution_count": 46,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from IPython.display import Image\n\n# Feel free to comment out if you have\n# not installed pygraphviz\nImage(interview_graph.get_graph().draw_png())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "ask_question\n",
            "--  [AIMessage(content=\"Yes, that's correct. I am focusing on the technical aspects of million-plus token context window language models and their impact on RAG systems. Can you provide more insight into how these large context window models affect the performance and capabilities of RAG systems?\", name\n",
            "answer_question\n",
            "--  [AIMessage(content='The introduction of large context window language models, such as Gemini 1.5 with a 1 million token context window, has raised concerns in the AI community regarding its impact on Retrieval-Augmented Generation (RAG) systems. RAG systems represent a significant advancement over t\n",
            "ask_question\n",
            "--  [AIMessage(content='Thank you for the detailed explanation and resources. Could you elaborate on the specific challenges and opportunities that million-plus token context window language models present for RAG systems in terms of improving generation quality, addressing data biases, and the potentia\n",
            "answer_question\n",
            "--  [AIMessage(content='Million-plus token context window language models present both challenges and opportunities for RAG systems. Challenges include the increased computational cost and complexity associated with processing larger context windows, potential issues with retaining factual accuracy when\n",
            "ask_question\n",
            "--  [AIMessage(content='Thank you for the detailed information and references provided. It has been insightful to understand both the challenges and opportunities that million-plus token context window language models bring to RAG systems. I appreciate your assistance in shedding light on this complex t\n",
            "answer_question\n",
            "--  [AIMessage(content=\"You're welcome! If you have any more questions or need further assistance in the future, feel free to reach out. Good luck with your article on RAG systems and million-plus token context window language models!\\n\\nCitations:\\n\\n[1]: https://www.nerdwallet.com/article/finance/exam\n",
            "__end__\n",
            "--  [AIMessage(content='So you said you were writing an article on Impact of million-plus token context window language models on RAG?', name='Subject Matter Expert'), AIMessage(content=\"Yes, that's correct. I am focusing on the technical aspects of million-plus token context window language models and \n"
          ]
        }
      ],
      "source": [
        "final_step = None\n\ninitial_state = {\n    \"editor\": perspectives.editors[0],\n    \"messages\": [\n        AIMessage(\n            content=f\"So you said you were writing an article on {example_topic}?\",\n            name=\"Subject_Matter_Expert\",\n        )\n    ],\n}\nasync for step in interview_graph.astream(initial_state):\n    name = next(iter(step))\n    print(name)\n    print(\"-- \", str(step[name][\"messages\"])[:300])\n    if END in step:\n        final_step = step"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {},
      "outputs": [],
      "source": [
        "final_state = next(iter(final_step.values()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Refine Outline\n",
        "\n",
        "At this point in STORM, we've conducted a large amount of research from different perspectives. It's time to refine the original outline based on these investigations. Below, create a chain using the LLM with a long context window to update the original outline."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 53,
      "metadata": {},
      "outputs": [],
      "source": [
        "refine_outline_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"\"\"You are a Wikipedia writer. You have gathered information from experts and search engines. Now, you are refining the outline of the Wikipedia page. \\\nYou need to make sure that the outline is comprehensive and specific. \\\nTopic you are writing about: {topic} \n\nOld outline:\n\n{old_outline}\"\"\",\n        ),\n        (\n            \"user\",\n            \"Refine the outline based on your conversations with subject-matter experts:\\n\\nConversations:\\n\\n{conversations}\\n\\nWrite the refined Wikipedia outline:\",\n        ),\n    ]\n)\n\n# Using turbo preview since the context can get quite long\nrefine_outline_chain = refine_outline_prompt | long_context_llm.with_structured_output(\n    Outline\n)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {},
      "outputs": [],
      "source": [
        "refined_outline = refine_outline_chain.invoke(\n    {\n        \"topic\": example_topic,\n        \"old_outline\": initial_outline.as_str,\n        \"conversations\": \"\\n\\n\".join(\n            f\"### {m.name}\\n\\n{m.content}\" for m in final_state[\"messages\"]\n        ),\n    }\n)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "# Impact of million-plus token context window language models on RAG\n",
            "\n",
            "## Introduction\n",
            "\n",
            "Provides a brief overview of million-plus token context window language models and their relevance to Retrieval-Augmented Generation (RAG) systems, setting the stage for a deeper exploration of their impact.\n",
            "\n",
            "## Background\n",
            "\n",
            "A foundational section to understand the core concepts involved.\n",
            "\n",
            "### Million-Plus Token Context Window Language Models\n",
            "\n",
            "Explains what million-plus token context window language models are, including notable examples like Gemini 1.5, focusing on their architecture, training data, and the evolution of their applications.\n",
            "\n",
            "### Retrieval-Augmented Generation (RAG)\n",
            "\n",
            "Describes the RAG framework, its unique approach of combining retrieval and generation models for enhanced natural language processing, and its significance in the AI landscape.\n",
            "\n",
            "## Impact on RAG Systems\n",
            "\n",
            "Delves into the effects of million-plus token context window language models on RAG, highlighting both the challenges and opportunities presented.\n",
            "\n",
            "### Performance and Efficiency\n",
            "\n",
            "Discusses how large context window models influence RAG performance, including aspects of latency, computational demands, and overall efficiency.\n",
            "\n",
            "### Generation Quality and Diversity\n",
            "\n",
            "Explores the impact on generation quality, the potential for more accurate and diverse outputs, and how these models address data biases and factual accuracy.\n",
            "\n",
            "### Technical Challenges\n",
            "\n",
            "Identifies specific technical hurdles such as prompt template design, context length limitations, and similarity searches in vector databases, and how they affect RAG systems.\n",
            "\n",
            "### Opportunities and Advancements\n",
            "\n",
            "Outlines the new capabilities and improvements in agent interaction, information retrieval, and response relevance that these models bring to RAG systems.\n",
            "\n",
            "## Future Directions\n",
            "\n",
            "Considers ongoing research and potential future developments in the integration of million-plus token context window language models with RAG systems, including speculation on emerging trends and technologies.\n",
            "\n",
            "## Conclusion\n",
            "\n",
            "Summarizes the key points discussed in the article, reaffirming the significant impact of million-plus token context window language models on RAG systems.\n"
          ]
        }
      ],
      "source": [
        "print(refined_outline.as_str)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generate Article\n",
        "\n",
        "Now it's time to generate the full article. We will first divide-and-conquer, so that each section can be tackled by an individual llm. Then we will prompt the long-form LLM to refine the finished article (since each section may use an inconsistent voice).\n",
        "\n",
        "#### Create Retriever\n",
        "\n",
        "The research process uncovers a large number of reference documents that we may want to query during the final article-writing process.\n",
        "\n",
        "First, create the retriever:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {},
      "outputs": [],
      "source": [
        "from langchain_community.vectorstores import SKLearnVectorStore\nfrom langchain_core.documents import Document\nfrom langchain_openai import OpenAIEmbeddings\n\nembeddings = OpenAIEmbeddings(model=\"text-embedding-3-small\")\nreference_docs = [\n    Document(page_content=v, metadata={\"source\": k})\n    for k, v in final_state[\"references\"].items()\n]\n# This really doesn't need to be a vectorstore for this size of data.\n# It could just be a numpy matrix. Or you could store documents\n# across requests if you want.\nvectorstore = SKLearnVectorStore.from_documents(\n    reference_docs,\n    embedding=embeddings,\n)\nretriever = vectorstore.as_retriever(k=10)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[Document(page_content='In Retrieval Augmented Generation (RAG), a longer context augments our model with more information. For LLMs that power agents, such as chatbots, longer context means more tools and capabilities. When summarizing, longer context means more comprehensive summaries. There exist plenty of use-cases for LLMs that are unlocked by longer context lengths.', metadata={'id': '20454848-23ac-4649-b083-81980532a77b', 'source': 'https://www.anyscale.com/blog/fine-tuning-llms-for-longer-context-and-better-rag-systems'}),\n",
              " Document(page_content='By the way, the context limits differ among models: two Claude models offer a 100K token context window, which works out to about 75,000 words, which is much higher than most other LLMs. The ...', metadata={'id': '1ee2d2bb-8f8e-4a7e-b45e-608b0804fe4c', 'source': 'https://www.infoworld.com/article/3712227/what-is-rag-more-accurate-and-reliable-llms.html'}),\n",
              " Document(page_content='Figure 1: LLM response accuracy goes down when context needed to answer correctly is found in the middle of the context window. The problem gets worse with larger context models. The problem gets ...', metadata={'id': 'a41d69e6-62eb-4abd-90ad-0892a2836cba', 'source': 'https://medium.com/@jm_51428/long-context-window-models-vs-rag-a73c35a763f2'}),\n",
              " Document(page_content='To improve performance, we used retrieval-augmented generation (RAG) to prompt an LLM with accurate up-to-date information. As a result of using RAG, the writing quality of the LLM improves substantially, which has implications for the practical usability of LLMs in clinical trial-related writing.', metadata={'id': 'e1af6e30-8c2b-495b-b572-ac6a29067a94', 'source': 'https://arxiv.org/abs/2402.16406'})]"
            ]
          },
          "execution_count": 29,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "retriever.invoke(\"What's a long context LLM anyway?\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Generate Sections\n",
        "\n",
        "Now you can generate the sections using the indexed docs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {},
      "outputs": [],
      "source": [
        "class SubSection(BaseModel):\n    subsection_title: str = Field(..., title=\"Title of the subsection\")\n    content: str = Field(\n        ...,\n        title=\"Full content of the subsection. Include [#] citations to the cited sources where relevant.\",\n    )\n\n    @property\n    def as_str(self) -> str:\n        return f\"### {self.subsection_title}\\n\\n{self.content}\".strip()\n\n\nclass WikiSection(BaseModel):\n    section_title: str = Field(..., title=\"Title of the section\")\n    content: str = Field(..., title=\"Full content of the section\")\n    subsections: Optional[List[Subsection]] = Field(\n        default=None,\n        title=\"Titles and descriptions for each subsection of the Wikipedia page.\",\n    )\n    citations: List[str] = Field(default_factory=list)\n\n    @property\n    def as_str(self) -> str:\n        subsections = \"\\n\\n\".join(\n            subsection.as_str for subsection in self.subsections or []\n        )\n        citations = \"\\n\".join([f\" [{i}] {cit}\" for i, cit in enumerate(self.citations)])\n        return (\n            f\"## {self.section_title}\\n\\n{self.content}\\n\\n{subsections}\".strip()\n            + f\"\\n\\n{citations}\".strip()\n        )\n\n\nsection_writer_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"You are an expert Wikipedia writer. Complete your assigned WikiSection from the following outline:\\n\\n\"\n            \"{outline}\\n\\nCite your sources, using the following references:\\n\\n<Documents>\\n{docs}\\n<Documents>\",\n        ),\n        (\"user\", \"Write the full WikiSection for the {section} section.\"),\n    ]\n)\n\n\nasync def retrieve(inputs: dict):\n    docs = await retriever.ainvoke(inputs[\"topic\"] + \": \" + inputs[\"section\"])\n    formatted = \"\\n\".join(\n        [\n            f'<Document href=\"{doc.metadata[\"source\"]}\"/>\\n{doc.page_content}\\n</Document>'\n            for doc in docs\n        ]\n    )\n    return {\"docs\": formatted, **inputs}\n\n\nsection_writer = (\n    retrieve\n    | section_writer_prompt\n    | long_context_llm.with_structured_output(WikiSection)\n)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "## Background\n",
            "\n",
            "To fully appreciate the impact of million-plus token context window language models on Retrieval-Augmented Generation (RAG) systems, it's essential to first understand the foundational concepts that underpin these technologies. This background section provides a comprehensive overview of both million-plus token context window language models and RAG, setting the stage for a deeper exploration of their integration and subsequent impacts on artificial intelligence and natural language processing.\n",
            "\n",
            "### Million-Plus Token Context Window Language Models\n",
            "\n",
            "Million-plus token context window language models, such as Gemini 1.5, represent a significant leap forward in the field of language modeling. These models are designed to process and understand large swathes of text, sometimes exceeding a million tokens in a single pass. The ability to handle such vast amounts of information at once allows for a deeper understanding of context and nuance, which is crucial for generating coherent and relevant text outputs. The development of these models involves sophisticated architecture and extensive training data, pushing the boundaries of what's possible in natural language processing. Over time, the applications of these models have evolved, extending their utility beyond mere text generation to complex tasks like sentiment analysis, language translation, and more.\n",
            "\n",
            "### Retrieval-Augmented Generation (RAG)\n",
            "\n",
            "The Retrieval-Augmented Generation framework represents a novel approach in the realm of artificial intelligence, blending the strengths of both retrieval and generation models to enhance natural language processing capabilities. At its core, RAG leverages a two-step process: initially, it uses a query to retrieve relevant documents or data from a knowledge base; this information is then utilized to inform and guide the generation of responses by a language model. This method addresses the limitations of fixed context windows by converting text to vector embeddings, facilitating a dynamic and flexible interaction with a vast array of information. RAG's unique approach has cemented its significance in the AI landscape, offering a pathway to more accurate, informative, and contextually relevant text generation.\n"
          ]
        }
      ],
      "source": [
        "section = await section_writer.ainvoke(\n    {\n        \"outline\": refined_outline.as_str,\n        \"section\": refined_outline.sections[1].section_title,\n        \"topic\": example_topic,\n    }\n)\nprint(section.as_str)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Generate final article\n",
        "\n",
        "Now we can rewrite the draft to appropriately group all the citations and maintain a consistent voice."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {},
      "outputs": [],
      "source": [
        "from langchain_core.output_parsers import StrOutputParser\n\nwriter_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"You are an expert Wikipedia author. Write the complete wiki article on {topic} using the following section drafts:\\n\\n\"\n            \"{draft}\\n\\nStrictly follow Wikipedia format guidelines.\",\n        ),\n        (\n            \"user\",\n            'Write the complete Wiki article using markdown format. Organize citations using footnotes like \"[1]\",'\n            \" avoiding duplicates in the footer. Include URLs in the footer.\",\n        ),\n    ]\n)\n\nwriter = writer_prompt | long_context_llm | StrOutputParser()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "# Impact of Million-Plus Token Context Window Language Models on Retrieval-Augmented Generation (RAG)\n",
            "\n",
            "The integration of million-plus token context window language models into Retrieval-Augmented Generation (RAG) systems marks a pivotal advancement in the field of artificial intelligence (AI) and natural language processing (NLP). This article delves into the background of both technologies, explores their convergence, and examines the profound effects of this integration on the capabilities and applications of AI-driven language models.\n",
            "\n",
            "## Contents\n",
            "\n",
            "1. [Background](#Background)\n",
            "    1. [Million-Plus Token Context Window Language Models](#Million-Plus-Token-Context-Window-Language-Models)\n",
            "    2. [Retrieval-Augmented Generation (RAG)](#Retrieval-Augmented-Generation-(RAG))\n",
            "2. [Integration of Million-Plus Token Context Window Models and RAG](#Integration-of-Million-Plus-Token-Context-Window-Models-and-RAG)\n",
            "3. [Impact on Natural Language Processing](#Impact-on-Natural-Language-Processing)\n",
            "4. [Applications](#Applications)\n",
            "5. [Challenges and Limitations](#Challenges-and-Limitations)\n",
            "6. [Future Directions](#Future-Directions)\n",
            "7. [Conclusion](#Conclusion)\n",
            "8. [References](#References)\n",
            "\n",
            "## Background\n",
            "\n",
            "### Million-Plus Token Context Window Language Models\n",
            "\n",
            "Million-plus token context window language models, exemplified by systems like Gemini 1.5, have revolutionized language modeling by their ability to process and interpret extensive texts, potentially exceeding a million tokens in a single analysis[1]. The capacity to manage such large volumes of data enables these models to grasp context and subtlety to a degree previously unattainable, enhancing their effectiveness in generating text that is coherent, relevant, and nuanced. The development of these models has been characterized by innovative architecture and the utilization of vast training datasets, pushing the envelope of natural language processing capabilities[2].\n",
            "\n",
            "### Retrieval-Augmented Generation (RAG)\n",
            "\n",
            "RAG systems represent an innovative paradigm in AI, merging the strengths of retrieval-based and generative models to improve the quality and relevance of text generation[3]. By initially retrieving related documents or data in response to a query, and subsequently using this information to guide the generation process, RAG overcomes the limitations inherent in fixed context windows. This methodology allows for dynamic access to a broad range of information, significantly enhancing the model's ability to generate accurate, informative, and contextually appropriate responses[4].\n",
            "\n",
            "## Integration of Million-Plus Token Context Window Models and RAG\n",
            "\n",
            "The integration of million-plus token context window models with RAG systems has been a natural progression in the quest for more sophisticated NLP solutions. By combining the extensive contextual understanding afforded by large context window models with the dynamic, information-rich capabilities of RAG, researchers and developers have been able to create AI systems that exhibit unprecedented levels of understanding, coherence, and relevance in text generation[5].\n",
            "\n",
            "## Impact on Natural Language Processing\n",
            "\n",
            "The fusion of these technologies has had a significant impact on the field of NLP, leading to advancements in several key areas:\n",
            "- **Enhanced Understanding**: The combined system exhibits a deeper comprehension of both the immediate context and broader subject matter[6].\n",
            "- **Improved Coherence**: Generated text is more coherent over longer passages, maintaining consistency and relevance[7].\n",
            "- **Increased Relevance**: Outputs are more contextually relevant, drawing accurately from a wider range of sources[8].\n",
            "\n",
            "## Applications\n",
            "\n",
            "This technological convergence has broadened the applicability of NLP systems in numerous fields, including but not limited to:\n",
            "- **Automated Content Creation**: Generating written content that is both informative and contextually appropriate for various platforms[9].\n",
            "- **Customer Support**: Providing answers that are not only accurate but also tailored to the specific context of user inquiries[10].\n",
            "- **Research Assistance**: Assisting in literature review and data analysis by retrieving and synthesizing relevant information from vast databases[11].\n",
            "\n",
            "## Challenges and Limitations\n",
            "\n",
            "Despite their advancements, the integration of these technologies faces several challenges:\n",
            "- **Computational Resources**: The processing of million-plus tokens and the dynamic retrieval of relevant information require significant computational power[12].\n",
            "- **Data Privacy and Security**: Ensuring the confidentiality and integrity of the data accessed by these systems poses ongoing concerns[13].\n",
            "- **Bias and Fairness**: The potential for inheriting and amplifying biases from training data remains a critical issue to address[14].\n",
            "\n",
            "## Future Directions\n",
            "\n",
            "Future research is likely to focus on optimizing computational efficiency, enhancing the models' ability to understand and generate more diverse and nuanced text, and addressing ethical considerations associated with AI and NLP technologies[15].\n",
            "\n",
            "## Conclusion\n",
            "\n",
            "The integration of million-plus token context window language models with RAG systems represents a milestone in the evolution of natural language processing, offering enhanced capabilities that have significant implications across various applications. As these technologies continue to evolve, they promise to further transform the landscape of AI-driven language models.\n",
            "\n",
            "## References\n",
            "\n",
            "1. Gemini 1.5 Documentation. (n.d.).\n",
            "2. The Evolution of Language Models. (2022).\n",
            "3. Introduction to Retrieval-Augmented Generation. (2021).\n",
            "4. Leveraging Large Context Windows for NLP. (2023).\n",
            "5. Integrating Context Window Models with RAG. (2023).\n",
            "6. Deep Learning in NLP. (2020).\n",
            "7. Coherence in Text Generation. (2019).\n",
            "8. Contextual Relevance in AI. (2021).\n",
            "9. Applications of NLP in Content Creation. (2022).\n",
            "10. AI in Customer Support. (2023).\n",
            "11. NLP for Research Assistance. (2021).\n",
            "12. Computational Challenges in NLP. (2022).\n",
            "13. Data Privacy in AI Systems. (2020).\n",
            "14. Addressing Bias in AI. (2021).\n",
            "15. Future of NLP Technologies. (2023)."
          ]
        }
      ],
      "source": [
        "for tok in writer.stream({\"topic\": example_topic, \"draft\": section.as_str}):\n    print(tok, end=\"\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Final Flow\n",
        "\n",
        "Now it's time to string everything together. We will have 6 main stages in sequence:\n",
        ".\n",
        "1. Generate the initial outline + perspectives\n",
        "2. Batch converse with each perspective to expand the content for the article\n",
        "3. Refine the outline based on the conversations\n",
        "4. Index the reference docs from the conversations\n",
        "5. Write the individual sections of the article\n",
        "6. Write the final wiki\n",
        "\n",
        "The state tracks the outputs of each stage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 55,
      "metadata": {},
      "outputs": [],
      "source": [
        "class ResearchState(TypedDict):\n    topic: str\n    outline: Outline\n    editors: List[Editor]\n    interview_results: List[InterviewState]\n    # The final sections output\n    sections: List[WikiSection]\n    article: str"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 80,
      "metadata": {},
      "outputs": [],
      "source": [
        "import asyncio\n\n\nasync def initialize_research(state: ResearchState):\n    topic = state[\"topic\"]\n    coros = (\n        generate_outline_direct.ainvoke({\"topic\": topic}),\n        survey_subjects.ainvoke(topic),\n    )\n    results = await asyncio.gather(*coros)\n    return {\n        **state,\n        \"outline\": results[0],\n        \"editors\": results[1].editors,\n    }\n\n\nasync def conduct_interviews(state: ResearchState):\n    topic = state[\"topic\"]\n    initial_states = [\n        {\n            \"editor\": editor,\n            \"messages\": [\n                AIMessage(\n                    content=f\"So you said you were writing an article on {topic}?\",\n                    name=\"Subject_Matter_Expert\",\n                )\n            ],\n        }\n        for editor in state[\"editors\"]\n    ]\n    # We call in to the sub-graph here to parallelize the interviews\n    interview_results = await interview_graph.abatch(initial_states)\n\n    return {\n        **state,\n        \"interview_results\": interview_results,\n    }\n\n\ndef format_conversation(interview_state):\n    messages = interview_state[\"messages\"]\n    convo = \"\\n\".join(f\"{m.name}: {m.content}\" for m in messages)\n    return f'Conversation with {interview_state[\"editor\"].name}\\n\\n' + convo\n\n\nasync def refine_outline(state: ResearchState):\n    convos = \"\\n\\n\".join(\n        [\n            format_conversation(interview_state)\n            for interview_state in state[\"interview_results\"]\n        ]\n    )\n\n    updated_outline = await refine_outline_chain.ainvoke(\n        {\n            \"topic\": state[\"topic\"],\n            \"old_outline\": state[\"outline\"].as_str,\n            \"conversations\": convos,\n        }\n    )\n    return {**state, \"outline\": updated_outline}\n\n\nasync def index_references(state: ResearchState):\n    all_docs = []\n    for interview_state in state[\"interview_results\"]:\n        reference_docs = [\n            Document(page_content=v, metadata={\"source\": k})\n            for k, v in interview_state[\"references\"].items()\n        ]\n        all_docs.extend(reference_docs)\n    await vectorstore.aadd_documents(all_docs)\n    return state\n\n\nasync def write_sections(state: ResearchState):\n    outline = state[\"outline\"]\n    sections = await section_writer.abatch(\n        [\n            {\n                \"outline\": refined_outline.as_str,\n                \"section\": section.section_title,\n                \"topic\": state[\"topic\"],\n            }\n            for section in outline.sections\n        ]\n    )\n    return {\n        **state,\n        \"sections\": sections,\n    }\n\n\nasync def write_article(state: ResearchState):\n    topic = state[\"topic\"]\n    sections = state[\"sections\"]\n    draft = \"\\n\\n\".join([section.as_str for section in sections])\n    article = await writer.ainvoke({\"topic\": topic, \"draft\": draft})\n    return {\n        **state,\n        \"article\": article,\n    }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Create the graph"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 73,
      "metadata": {},
      "outputs": [],
      "source": [
        "from langgraph.checkpoint.memory import MemorySaver\n\nbuilder_of_storm = StateGraph(ResearchState)\n\nnodes = [\n    (\"init_research\", initialize_research),\n    (\"conduct_interviews\", conduct_interviews),\n    (\"refine_outline\", refine_outline),\n    (\"index_references\", index_references),\n    (\"write_sections\", write_sections),\n    (\"write_article\", write_article),\n]\nfor i in range(len(nodes)):\n    name, node = nodes[i]\n    builder_of_storm.add_node(name, node)\n    if i > 0:\n        builder_of_storm.add_edge(nodes[i - 1][0], name)\n\nbuilder_of_storm.add_edge(START, nodes[0][0])\nbuilder_of_storm.add_edge(nodes[-1][0], END)\nstorm = builder_of_storm.compile(checkpointer=MemorySaver())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 74,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 64392 characters]",
            "text/plain": [
              "<IPython.core.display.Image object>"
            ]
          },
          "execution_count": 74,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "Image(storm.get_graph().draw_png())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 75,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "init_research\n",
            "--  {'topic': 'Groq, NVIDIA, Llamma.cpp and the future of LLM Inference', 'outline': Outline(page_title='Groq, NVIDIA, Llamma.cpp and the future of LLM Inference', sections=[Section(section_title='Introduction', description='Overview of Groq, NVIDIA, Llamma.cpp, and their significance in the field of La\n",
            "conduct_interviews\n",
            "--  {'topic': 'Groq, NVIDIA, Llamma.cpp and the future of LLM Inference', 'outline': Outline(page_title='Groq, NVIDIA, Llamma.cpp and the future of LLM Inference', sections=[Section(section_title='Introduction', description='Overview of Groq, NVIDIA, Llamma.cpp, and their significance in the field of La\n",
            "refine_outline\n",
            "--  {'topic': 'Groq, NVIDIA, Llamma.cpp and the future of LLM Inference', 'outline': Outline(page_title='Groq, NVIDIA, Llamma.cpp and the Future of LLM Inference', sections=[Section(section_title='Introduction', description='An overview of the significance and roles of Groq, NVIDIA, and Llamma.cpp in th\n",
            "index_references\n",
            "--  {'topic': 'Groq, NVIDIA, Llamma.cpp and the future of LLM Inference', 'outline': Outline(page_title='Groq, NVIDIA, Llamma.cpp and the Future of LLM Inference', sections=[Section(section_title='Introduction', description='An overview of the significance and roles of Groq, NVIDIA, and Llamma.cpp in th\n",
            "write_sections\n",
            "--  {'topic': 'Groq, NVIDIA, Llamma.cpp and the future of LLM Inference', 'outline': Outline(page_title='Groq, NVIDIA, Llamma.cpp and the Future of LLM Inference', sections=[Section(section_title='Introduction', description='An overview of the significance and roles of Groq, NVIDIA, and Llamma.cpp in th\n",
            "write_article\n",
            "--  {'topic': 'Groq, NVIDIA, Llamma.cpp and the future of LLM Inference', 'outline': Outline(page_title='Groq, NVIDIA, Llamma.cpp and the Future of LLM Inference', sections=[Section(section_title='Introduction', description='An overview of the significance and roles of Groq, NVIDIA, and Llamma.cpp in th\n",
            "__end__\n",
            "--  {'topic': 'Groq, NVIDIA, Llamma.cpp and the future of LLM Inference', 'outline': Outline(page_title='Groq, NVIDIA, Llamma.cpp and the Future of LLM Inference', sections=[Section(section_title='Introduction', description='An overview of the significance and roles of Groq, NVIDIA, and Llamma.cpp in th\n"
          ]
        }
      ],
      "source": [
        "config = {\"configurable\": {\"thread_id\": \"my-thread\"}}\nasync for step in storm.astream(\n    {\n        \"topic\": \"Groq, NVIDIA, Llamma.cpp and the future of LLM Inference\",\n    },\n    config,\n):\n    name = next(iter(step))\n    print(name)\n    print(\"-- \", str(step[name])[:300])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 82,
      "metadata": {},
      "outputs": [],
      "source": [
        "checkpoint = storm.get_state(config)\narticle = checkpoint.values[\"article\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Render the Wiki\n",
        "\n",
        "Now we can render the final wiki page!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 83,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "# Large Language Model (LLM) Inference Technologies\n",
              "\n",
              "### Contents\n",
              "1. [Introduction](#Introduction)\n",
              "2. [Groq's Advancements in LLM Inference](#Groqs-Advancements-in-LLM-Inference)\n",
              "3. [NVIDIA's Contributions to LLM Inference](#NVIDIAs-Contributions-to-LLM-Inference)\n",
              "   1. [Hardware Innovations](#Hardware-Innovations)\n",
              "   2. [Software Solutions](#Software-Solutions)\n",
              "   3. [Research and Development](#Research-and-Development)\n",
              "4. [Llamma.cpp: Accelerating LLM Inference](#Llammacpp-Accelerating-LLM-Inference)\n",
              "5. [The Future of LLM Inference](#The-Future-of-LLM-Inference)\n",
              "6. [References](#References)\n",
              "\n",
              "### Introduction\n",
              "\n",
              "The advent of million-plus token context window language models, such as Gemini 1.5, has significantly advanced the field of artificial intelligence, particularly in natural language processing (NLP). These models have expanded the capabilities of machine learning in understanding and generating text over vastly larger contexts than previously possible. This leap in technology has paved the way for transformative applications across various domains, including the integration into Retrieval-Augmented Generation (RAG) systems to produce more accurate and contextually rich responses. \n",
              "\n",
              "### Groq's Advancements in LLM Inference\n",
              "\n",
              "Groq has introduced the Groq Linear Processor Unit (LPU), a purpose-built hardware architecture for LLM inference. This innovation positions Groq as a leader in efficient and high-performance LLM processing by optimizing the hardware specifically for LLM tasks. The Groq LPU dramatically reduces latency and increases the throughput of LLM inferences, facilitating advancements in a wide range of applications, from natural language processing to broader artificial intelligence technologies[1].\n",
              "\n",
              "### NVIDIA's Contributions to LLM Inference\n",
              "\n",
              "NVIDIA has played a pivotal role in advancing LLM inference through its GPUs, optimized for AI and machine learning workloads, and specialized software frameworks. The company's GPU architecture and software solutions, such as the CUDA Deep Neural Network library (cuDNN) and the TensorRT inference optimizer, are designed to accelerate computational processes and improve LLM performance. NVIDIA's active participation in research and development further underscores its commitment to enhancing the capabilities of LLMs[1].\n",
              "\n",
              "#### Hardware Innovations\n",
              "\n",
              "NVIDIA's GPU architecture facilitates high throughput and parallel processing for LLM inference tasks, significantly reducing inference time and enabling complex models to be used in real-time applications.\n",
              "\n",
              "#### Software Solutions\n",
              "\n",
              "NVIDIA's suite of software tools, including cuDNN and TensorRT, optimizes LLM performance on its hardware, streamlining the deployment of LLMs by improving their efficiency and reducing latency.\n",
              "\n",
              "#### Research and Development\n",
              "\n",
              "NVIDIA collaborates with academic and industry partners to develop new techniques and models that push the boundaries of LLM technology, aiming to make LLMs more powerful and applicable across a broader range of tasks.\n",
              "\n",
              "### Llamma.cpp: Accelerating LLM Inference\n",
              "\n",
              "Llamma.cpp is a framework developed to enhance the speed and efficiency of LLM inference. By integrating specialized hardware, such as Groq's LPU, and optimizing for parallel processing, Llamma.cpp significantly accelerates computation times and reduces energy consumption. The framework supports million-plus token context window models, enabling applications requiring deep contextual understanding and extensive knowledge retrieval[1][2].\n",
              "\n",
              "### The Future of LLM Inference\n",
              "\n",
              "The future of LLM inference is poised for transformative changes with advances in purpose-built hardware architectures like Groq's LPU. These innovations promise to enhance the speed and efficiency of LLM processing, leading to more interactive, capable, and integrated AI applications. The potential for advanced hardware and sophisticated LLMs to enable near-instantaneous processing of complex queries and interactions opens new avenues for research and application in various fields, suggesting a future where AI is seamlessly integrated into society[1][2].\n",
              "\n",
              "### References\n",
              "\n",
              "[1] \"Groq's LPU: Advancing LLM Inference Efficiency,\" Prompt Engineering. https://promptengineering.org/groqs-lpu-advancing-llm-inference-efficiency/\n",
              "\n",
              "[2] \"The Speed of Thought: Harnessing the Fastest LLM with Groq's LPU,\" Medium. https://medium.com/@anasdavoodtk1/the-speed-of-thought-harnessing-the-fastest-llm-with-groqs-lpu-11bb00864e9c"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "execution_count": 83,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from IPython.display import Markdown\n\n# We will down-header the sections to create less confusion in this notebook\nMarkdown(article.replace(\"\\n#\", \"\\n##\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.12.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}
