From 24954c24e0020f5d67674f5ae32b8170a1775853 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Thu, 29 Feb 2024 16:28:40 -0800 Subject: [PATCH] Storm Draft --- examples/storm/storm.ipynb | 1145 ++++++++++++++++++++++++++++++++++++ 1 file changed, 1145 insertions(+) create mode 100644 examples/storm/storm.ipynb diff --git a/examples/storm/storm.ipynb b/examples/storm/storm.ipynb new file mode 100644 index 000000000..c5b294ac8 --- /dev/null +++ b/examples/storm/storm.ipynb @@ -0,0 +1,1145 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# STORM\n", + "\n", + "[STORM](https://arxiv.org/abs/2402.14207) is a research assistant by Shao, et. al that extends the idea of \"outline-driven RAG\" for richer article generation.\n", + "\n", + "It is tasked with generating Wikipedia-like ariticles on a user-provided topic. It has a few main stages:\n", + "\n", + "1. Survey related subjects\n", + "2. Identify perspectives\n", + "3. \"Expert Interviews\" (between the writer and an agent role-playing as a perspective)\n", + "4. Refine article \n", + "\n", + "The expert interviews stage ocurrs between the article writer and each role-playing agent and itself is a loop, where the \"expert\" is able to query external knowledge and respond to pointed questions.\n", + "\n", + "Couple hyperparameters to restrict the infinite research breadth:\n", + "\n", + "N: Number of perspectives to survey / use (2->3)\n", + "M: Max number of conversation turns in step (3)\n", + "\n", + "The paper uses DSPY and few-shot examples to adapt but we'll just use functioncalling here." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# %pip install langchain_community langchain_openai langgraph wikipedia tavily-python scikit-learn" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/Users/wfh/code/lc/community/langgraph-engineer/.venv/lib/python3.11/site-packages/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 langchain_openai import ChatOpenAI\n", + "from langchain_core.pydantic_v1 import BaseModel, Field\n", + "from typing import List, Optional\n", + "from langchain_core.prompts import ChatPromptTemplate\n", + "\n", + "direct_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", + "\n", + "class 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", + "\n", + "class 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", + "\n", + "class Outline(BaseModel):\n", + " page_title: str = Field(..., title=\"Title of the Wikipedia page\")\n", + " sections: List[Section] = Field(\n", + " default_factory=list, 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", + "\n", + "generate_outline_direct = direct_gen_outline_prompt | ChatOpenAI(\n", + " model=\"gpt-3.5-turbo\"\n", + ").with_structured_output(Outline)" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "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, including architecture, training data, benefits, and challenges.\n", + "\n", + "## RAG (Retrieval-Augmented Generation)\n", + "\n", + "Explanation of RAG, its components, and how it integrates with million-plus token context window language models.\n", + "\n", + "## Impact on RAG\n", + "\n", + "Discuss the implications of using million-plus token context window language models with RAG, including improvements in performance, challenges, and future research directions.\n" + ] + } + ], + "source": [ + "example_topic = \"Impact of million-plus token context window language models on RAG\"\n", + "\n", + "initial_outline = generate_outline_direct.invoke({\"topic\": example_topic})\n", + "\n", + "print(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." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "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", + "\n", + "Please list the as many subjects and urls as you can.\n", + "\n", + "Topic of interest: {topic}\n", + "\"\"\"\n", + ")\n", + "\n", + "class RelatedSubjects(BaseModel):\n", + " topics: List[str] = Field(\n", + " description=\"Comprehensive list of related subjects as background research.\",\n", + " )\n", + "\n", + "\n", + "expand_chain = (\n", + " gen_related_topics_prompt \n", + " | ChatOpenAI(model=\"gpt-3.5-turbo\").with_structured_output(RelatedSubjects)\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "RelatedSubjects(topics=['Impact of million-plus token context window language models', 'RAG'])" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "related_subjects = await expand_chain.ainvoke({\"topic\": example_topic})\n", + "related_subjects" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "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.\",\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", + "\n", + "class 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", + "\n", + "gen_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", + "\n", + "gen_perspectives_chain = gen_perspectives_prompt | ChatOpenAI(\n", + " model=\"gpt-3.5-turbo\"\n", + ").with_structured_output(Perspectives)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_community.retrievers import WikipediaRetriever\n", + "from langchain_core.runnables import RunnableLambda, chain as as_runnable\n", + "\n", + "wikipedia_retriever = WikipediaRetriever(load_all_available_meta=True, top_k_results=1)\n", + "\n", + "\n", + "def 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", + "\n", + "def format_docs(docs):\n", + " return \"\\n\\n\".join(format_doc(doc) for doc in docs)\n", + "\n", + "\n", + "@as_runnable\n", + "async def survey_subjects(topic: str):\n", + " related_subjects = await expand_chain.ainvoke({\"topic\": topic})\n", + " retrieved_docs = await wikipedia_retriever.abatch(related_subjects.topics, return_exceptions=True)\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": 7, + "metadata": {}, + "outputs": [], + "source": [ + "perspectives = await survey_subjects.ainvoke(example_topic)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'editors': [{'affiliation': 'Language model research institute',\n", + " 'name': 'Alice',\n", + " 'role': 'Language model researcher',\n", + " 'description': 'Alice is a language model researcher focusing on the impact of million-plus token context window language models on RAG. She is interested in analyzing the advancements in language models and their implications on the RAG (Retrieve, Analyze, Generate) framework.'},\n", + " {'affiliation': 'Vector database company',\n", + " 'name': 'Bob',\n", + " 'role': 'Vector database expert',\n", + " 'description': 'Bob is a vector database expert who will provide insights on how vector databases can support the storage and retrieval of large language models like the million-plus token context window language models used in RAG systems.'},\n", + " {'affiliation': 'Natural language processing organization',\n", + " 'name': 'Charlie',\n", + " 'role': 'NLP specialist',\n", + " 'description': 'Charlie is an NLP specialist who will focus on the application of million-plus token context window language models in natural language processing tasks, particularly in the context of the RAG framework.'},\n", + " {'affiliation': 'Machine learning consultancy',\n", + " 'name': 'David',\n", + " 'role': 'Machine learning consultant',\n", + " 'description': 'David is a machine learning consultant with expertise in training and optimizing large language models. He will discuss the machine learning techniques involved in developing million-plus token context window models for RAG.'},\n", + " {'affiliation': 'Artificial intelligence company',\n", + " 'name': 'Eve',\n", + " 'role': 'AI expert',\n", + " 'description': 'Eve is an AI expert who will provide insights on the broader implications of million-plus token context window language models in artificial intelligence applications, including their impact on the RAG framework.'}]}" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "perspectives.dict()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Expert Dialog\n", + "\n", + "Now the true fun begins, the wikipedia writer will \"talk\" with expert agents primed to role-play using the perspectives presented above." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "from langgraph.graph import StateGraph, END\n", + "from typing_extensions import TypedDict\n", + "from langchain_core.messages import AnyMessage\n", + "from typing import Annotated, Sequence\n", + "\n", + "\n", + "def 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", + "\n", + "def update_references(references, new_references):\n", + " if not references:\n", + " references = {}\n", + " references.update(new_references)\n", + " return references\n", + "\n", + "\n", + "def 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", + "\n", + "class 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": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_core.prompts import MessagesPlaceholder\n", + "from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage\n", + "\n", + "\n", + "gen_qn_prompt = ChatPromptTemplate.from_messages(\n", + " [\n", + " (\n", + " \"system\",\n", + " \"\"\"You are an experienced Wikipedia writer and want to edit a specific page. \\\n", + "Besides your identity as a Wikipedia writer, you have a specific focus when researching the topic. \\\n", + "Now, you are chatting with an expert to get information. Ask good questions to get more useful information.\n", + "\n", + "When you have no more questions to ask, say \"Thank you so much for your help!\" to end the conversation.\\\n", + "Please only ask one question at a time and don't ask what you have asked before.\\\n", + "Your questions should be related to the topic you want to write.\n", + "Be comprehensive and curious, gaining as much unique insight from the expert as possible.\\\n", + "\n", + "Stay true to your specific perspective:\n", + "\n", + "{persona}\"\"\",\n", + " ),\n", + " MessagesPlaceholder(variable_name=\"messages\", optional=True),\n", + " ]\n", + ")\n", + "\n", + "\n", + "def tag_with_name(ai_message: AIMessage, name: str):\n", + " ai_message.name = name\n", + " return ai_message\n", + "\n", + "\n", + "def 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\n", + "async 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", + " | ChatOpenAI(model=\"gpt-3.5-turbo\")\n", + " | RunnableLambda(tag_with_name).bind(name=editor.name)\n", + " )\n", + " result = await gn_chain.ainvoke(state)\n", + " return {\n", + " \"messages\": [result]\n", + " }" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "\"Yes, that's correct. I am focusing on the impact of million-plus token context window language models on the RAG framework. These language models have significantly expanded the scope of information that can be considered when retrieving, analyzing, and generating content. Is there a specific aspect of this topic that you would like to explore further?\"" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "messages = [\n", + " HumanMessage(f\"So you said you were writing an article on {example_topic}?\")\n", + "]\n", + "question = await generate_question.ainvoke(\n", + " {\n", + " \"editor\": perspectives.editors[0],\n", + " \"messages\": messages,\n", + " }\n", + ")\n", + "\n", + "question[\"messages\"][0].content" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "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", + "\n", + "gen_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", + ")\n", + "gen_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": 13, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "['Impact of million-plus token context window language models on the RAG framework',\n", + " 'Benefits of using million-plus token context window language models in the RAG framework',\n", + " 'Challenges of integrating million-plus token context window language models with the RAG framework',\n", + " 'Comparison of different million-plus token context window language models in the RAG framework']" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "queries = await gen_queries_chain.ainvoke({\"messages\": [HumanMessage(content=question[\"messages\"][0].content)]})\n", + "queries['parsed'].queries" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "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(f\"[{i+1}]: {url}\" for i, url in enumerate(self.cited_urls))\n", + "\n", + "gen_answer_prompt = ChatPromptTemplate.from_messages(\n", + " [(\"system\", \"\"\"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", + "\n", + "Make your response as informative as possible and make sure every sentence is supported by the gathered information.\n", + "Each response must be backed up by a citation from a reliable source, formatted as a footnote, reproducing the URLS after your response.\"\"\"),\n", + " MessagesPlaceholder(variable_name=\"messages\", optional=True),\n", + " ]\n", + ")\n", + "\n", + "gen_answer_chain = gen_answer_prompt | ChatOpenAI(model=\"gpt-3.5-turbo\").with_structured_output(AnswerWithCitations, include_raw=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Reference Store\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", + "Here, we will createa multi-vector retriever and store all the searched documents inline." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_community.utilities.duckduckgo_search import DuckDuckGoSearchAPIWrapper\n", + "search_engine = DuckDuckGoSearchAPIWrapper()\n", + "from langchain_core.tools import tool\n", + "# TODO: remove when i get my api limit bumped\n", + "@tool\n", + "async def search_engine(query: str):\n", + " \"\"\"Search engine to the internet.\"\"\"\n", + " results = DuckDuckGoSearchAPIWrapper()._ddgs_text(\"beijing olympics\")\n", + " return [\n", + " {\"content\": r[\"body\"],\n", + " \"url\": r[\"href\"]} for r in results]\n" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_community.tools.tavily_search import TavilySearchResults\n", + "from langchain_core.runnables import RunnableConfig\n", + "import json\n", + "\n", + "# search_engine = TavilySearchResults(max_results=4)\n", + "\n", + "\n", + "async def gen_answer(\n", + " state: InterviewState,\n", + " config: RunnableConfig | None = 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(queries[\"parsed\"].queries, config, return_exceptions=True)\n", + " successful_results = [res for res in query_results if not isinstance(res, Exception)]\n", + " all_query_results = {res[\"url\"]: res[\"content\"] for results in successful_results for res in results}\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(\n", + " tool_call_id=tool_id,\n", + " content=dumped\n", + " )\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 {\n", + " \"messages\": [formatted_message],\n", + " \"references\": cited_references\n", + " }" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'Million-plus token context window language models have revolutionized the RAG (Retrieval-Augmented Generation) framework by significantly expanding the scope of information considered during content retrieval, analysis, and generation. These advanced language models, with their extensive context windows, allow for more nuanced understanding of text and context, leading to more accurate and contextually relevant content generation within the RAG framework.\\n\\nCitations:\\n\\n[1]: https://www.teamusa.com/olympic-games-beijing-2022\\n[2]: https://apnews.com/article/beijing-olympics-nhl-d32ef9ddd57b6be68f3ae5b55b47c3c4\\n[3]: https://www.britannica.com/topic/2008-Beijing-Olympic-Games-1702245\\n[4]: https://olympics.com/ioc/news/final-report-highlights-legacy-of-olympic-winter-games-beijing-2022'" + ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "example_answer = await gen_answer(\n", + " {\"messages\": [HumanMessage(content=question[\"messages\"][0].content)]}\n", + ")\n", + "example_answer['messages'][-1].content" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": {}, + "outputs": [], + "source": [ + "max_num_turns = 5\n", + "\n", + "\n", + "def 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", + "builder = StateGraph(InterviewState)\n", + "\n", + "builder.add_node(\"ask_question\", generate_question)\n", + "builder.add_node(\"answer_question\", gen_answer)\n", + "builder.add_conditional_edges(\"answer_question\", route_messages)\n", + "builder.add_edge(\"ask_question\", \"answer_question\")\n", + "\n", + "builder.set_entry_point(\"ask_question\")\n", + "interview_graph = builder.compile().with_config(run_name=\"Conduct Interviews\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "ask_question\n", + "-- [AIMessage(content=\"Yes, that's correct. I am interested in understanding how million-plus token context window language models are influencing the RAG (Retrieve, Analyze, Generate) framework. Do you have insights on how these advanced language models are changing the way information is retrieved, a\n", + "answer_question\n", + "-- [AIMessage(content='Million-plus token context window language models have significantly impacted the RAG (Retrieve, Analyze, Generate) framework in natural language processing tasks. These advanced language models, such as GPT-3 with 175 billion parameters, have revolutionized the way information i\n" + ] + } + ], + "source": [ + "final_step = None\n", + "\n", + "initial_state = {\n", + " \"editor\": perspectives.editors[0],\n", + " \"messages\": [AIMessage(content=f\"So you said you were writing an article on {example_topic}?\", name=\"Subject Matter Expert\")]\n", + "}\n", + "async 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": null, + "metadata": {}, + "outputs": [], + "source": [ + "final_state = next(iter(final_step.values()))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Refine Outline\n", + "\n", + "Now that we have all this cool stuff, let's distill it into a refined outline." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "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.\\\n", + "You need to make sure that the outline is comprehensive and specific.\\\n", + "Topic you are writing about: {topic}\\\n", + "Old outline: {old_outline}\"\"\",\n", + " ),\n", + " (\"user\", \"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", + "# Using turbo preview since the context can get quite long\n", + "refine_outline_chain = refine_outline_prompt | ChatOpenAI(model=\"gpt-4-turbo-preview\").with_structured_output(Outline)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "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": null, + "metadata": {}, + "outputs": [], + "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 divide-and-conquer, so that each section can be tackled by an individual llm." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class SubSection(BaseModel):\n", + " subsection_title: str = Field(..., title=\"Title of the subsection\")\n", + " content: str = Field(..., title=\"Full content of the subsection. Include [#] citations to the cited sources where relevant.\")\n", + "\n", + " @property\n", + " def as_str(self) -> str:\n", + " \n", + " return f\"### {self.subsection_title}\\n\\n{self.content}\".strip()\n", + "\n", + "\n", + "class 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\n", + " for subsection in self.subsections or []\n", + " )\n", + " citations = \"\\n\".join([f\" [{i}] {cit}\" for i, cit in enumerate(self.citations)])\n", + " return f\"## {self.section_title}\\n\\n{self.content}\\n\\n{subsections}\".strip() + f\"\\n\\n{citations}\".strip()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_core.documents import Document\n", + "\n", + "from langchain_community.vectorstores import SKLearnVectorStore\n", + "from langchain_openai import OpenAIEmbeddings\n", + "\n", + "embeddings = OpenAIEmbeddings(model=\"text-embedding-3-small\")\n", + "reference_docs = [Document(page_content=v, metadata={\"source\": k}) for k, v in final_state[\"references\"].items()]\n", + "# This really doesn't need to be a vectorstore.\n", + "# could just be a numpy matrix\n", + "vectorstore = SKLearnVectorStore.from_documents(\n", + " reference_docs,\n", + " embedding=embeddings,\n", + ")\n", + "retriever = vectorstore.as_retriever(k=10)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "section_writer_prompt = ChatPromptTemplate.from_messages(\n", + " [\n", + " (\"system\", \"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\\n{docs}\\n\"),\n", + " (\"user\", \"Write the full WikiSection for the {section} section.\")\n", + " ]\n", + ")\n", + "\n", + "\n", + "async def retrieve(inputs: dict):\n", + " docs = await retriever.ainvoke(inputs['topic'] + \": \" + inputs[\"section\"])\n", + " formatted = \"\\n\".join([f'\\n{doc.page_content}\\n' for doc in docs])\n", + " return {\n", + " \"docs\": formatted,\n", + " **inputs\n", + " }\n", + "\n", + "section_writer = (\n", + " retrieve\n", + " | section_writer_prompt\n", + " | ChatOpenAI(model=\"gpt-4-turbo-preview\").with_structured_output(WikiSection)\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "section = await section_writer.abatch(\n", + " {\"outline\": refined_outline.as_str,\n", + " \"section\" : refined_outline.sections[1].section_title,\n", + " \"topic\": example_topic,\n", + "}\n", + ")\n", + "print(section.as_str)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(section.as_str)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Generate final article" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_core.output_parsers import StrOutputParser\n", + "writer_prompt = ChatPromptTemplate.from_messages(\n", + " [\n", + " (\"system\", \"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", + " (\"user\", 'Write the complete Wiki article using markdown format. Organize citations using footnotes like \"[1]\", avoiding duplicates in the footer.')\n", + " ]\n", + ")\n", + "\n", + "writer = writer_prompt | ChatOpenAI(model=\"gpt-4-turbo-preview\") | StrOutputParser()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "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 3 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. Write the final wiki." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "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": null, + "metadata": {}, + "outputs": [], + "source": [ + "import asyncio\n", + "\n", + "async 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", + "async def conduct_interviews(state: ResearchState):\n", + " topic = state[\"topic\"]\n", + " initial_states = [{\n", + " \"editor\": editor,\n", + " \"messages\": [AIMessage(content=f\"So you said you were writing an article on {topic}?\", name=\"Subject Matter Expert\")]\n", + " } for editor in state[\"editors\"]]\n", + " # We call in to the sub-graph here\n", + " interview_results = await interview_graph.abatch(initial_states)\n", + " \n", + " return {\n", + " **state,\n", + " \"interview_results\": interview_results,\n", + " }\n", + "\n", + "\n", + "def format_conversation(interview_state):\n", + " messages = interview_state[\"messages\"]\n", + " convo = \"\\n\".join(\n", + " f\"{m.name}: {m.content}\" for m in final_state[\"messages\"]\n", + " )\n", + " return f'Conversation with {interview_state[\"editor\"].name}\\n\\n' + convo\n", + "\n", + "\n", + "async def refine_outline(state: ResearchState):\n", + " convos = \"\\n\\n\".join([format_conversation(interview_state) for interview_state in state[\"interview_results\"]])\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 {\n", + " **state,\n", + " \"outline\": updated_outline\n", + " }\n", + "\n", + "async def index_references(state: ResearchState):\n", + " all_docs = []\n", + " for interview_state in state[\"interview_results\"]:\n", + " reference_docs = [Document(page_content=v, metadata={\"source\": k}) for k, v in interview_state[\"references\"].items()]\n", + " all_docs.extend(reference_docs)\n", + " await vectorstore.aadd_documents(all_docs)\n", + " return state\n", + "\n", + "async 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", + "async 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 = writer.ainvoke({\"topic\": example_topic, \"draft\": draft})\n", + " return {\n", + " **state,\n", + " \"article\": article,\n", + " }" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "builder_of_storm = StateGraph(ResearchState)\n", + "\n", + "\n", + "nodes = [\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", + "]\n", + "for 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", + "\n", + "builder_of_storm.set_entry_point(nodes[0][0])\n", + "builder_of_storm.set_finish_point(nodes[-1][0])\n", + "storm = builder_of_storm.compile()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "async for step in storm.astream({\n", + " \"topic\": \"NVIDIA 2024 Q1 earnings report\",\n", + " }):\n", + " name = next(iter(step))\n", + " print(name)\n", + " print(\"-- \", str(step[name])[:300])\n", + " if END in step:\n", + " results = step" + ] + }, + { + "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.11.2" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +}