code: get_ipython().run_line_magic('pip', 'install llama-index-embeddings-openai')
  get_ipython().run_line_magic('pip', 'install llama-index-postprocessor-cohere-rerank')
  get_ipython().run_line_magic('pip', 'install llama-index-llms-openai')
  
  
  import phoenix as px
  
  px.launch_app()
  import llama_index.core
  
  llama_index.core.set_global_handler("arize_phoenix")
  
  
  from llama_index.llms.openai import OpenAI
  from llama_index.embeddings.openai import OpenAIEmbedding
  from llama_index.core import Settings
  
  Settings.llm = OpenAI(model="gpt-3.5-turbo")
  Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
  
  
  from llama_index.core import SimpleDirectoryReader
  
  reader = SimpleDirectoryReader("../data/paul_graham")
  
  
  docs = reader.load_data()
  
  
  import os
  from llama_index.core import (
      StorageContext,
      VectorStoreIndex,
      load_index_from_storage,
  )
  
  if not os.path.exists("storage"):
      index = VectorStoreIndex.from_documents(docs)
      index.set_index_id("vector_index")
      index.storage_context.persist("./storage")
  else:
      storage_context = StorageContext.from_defaults(persist_dir="storage")
      index = load_index_from_storage(storage_context, index_id="vector_index")
  
  
  
  from llama_index.core.query_pipeline import QueryPipeline
  from llama_index.core import PromptTemplate
  
  prompt_str = "Please generate related movies to {movie_name}"
  prompt_tmpl = PromptTemplate(prompt_str)
  llm = OpenAI(model="gpt-3.5-turbo")
  
  p = QueryPipeline(chain=[prompt_tmpl, llm], verbose=True)
  
  
  output = p.run(movie_name="The Departed")
  
  
  print(str(output))
  
  
  
  from typing import List
  from pydantic import BaseModel, Field
  from llama_index.core.output_parsers import PydanticOutputParser
  
  
  class Movie(BaseModel):
      """Object representing a single movie."""
  
      name: str = Field(..., description="Name of the movie.")
      year: int = Field(..., description="Year of the movie.")
  
  
  class Movies(BaseModel):
      """Object representing a list of movies."""
  
      movies: List[Movie] = Field(..., description="List of movies.")
  
  
  llm = OpenAI(model="gpt-3.5-turbo")
  output_parser = PydanticOutputParser(Movies)
  json_prompt_str = """\
  Please generate related movies to {movie_name}. Output with the following JSON format: 
  """
  json_prompt_str = output_parser.format(json_prompt_str)
  
  
  json_prompt_tmpl = PromptTemplate(json_prompt_str)
  
  p = QueryPipeline(chain=[json_prompt_tmpl, llm, output_parser], verbose=True)
  output = p.run(movie_name="Toy Story")
  
  
  output
  
  
  
  
  prompt_str = "Please generate related movies to {movie_name}"
  prompt_tmpl = PromptTemplate(prompt_str)
  prompt_str2 = """\
  Here's some text:
  
  {text}
  
  Can you rewrite this with a summary of each movie?
  """
  prompt_tmpl2 = PromptTemplate(prompt_str2)
  llm = OpenAI(model="gpt-3.5-turbo")
  llm_c = llm.as_query_component(streaming=True)
  
  p = QueryPipeline(
      chain=[prompt_tmpl, llm_c, prompt_tmpl2, llm_c], verbose=True
  )
  
  
  output = p.run(movie_name="The Dark Knight")
  for o in output:
      print(o.delta, end="")
  
  
  
  p = QueryPipeline(
      chain=[
          json_prompt_tmpl,
          llm.as_query_component(streaming=True),
          output_parser,
      ],
      verbose=True,
  )
  output = p.run(movie_name="Toy Story")
  print(output)
  
  
  
  
  
  from llama_index.postprocessor.cohere_rerank import CohereRerank
  
  prompt_str1 = "Please generate a concise question about Paul Graham's life regarding the following topic {topic}"
  prompt_tmpl1 = PromptTemplate(prompt_str1)
  prompt_str2 = (
      "Please write a passage to answer the question\n"
      "Try to include as many key details as possible.\n"
      "\n"
      "\n"
      "{query_str}\n"
      "\n"
      "\n"
      'Passage:"""\n'
  )
  prompt_tmpl2 = PromptTemplate(prompt_str2)
  
  llm = OpenAI(model="gpt-3.5-turbo")
  retriever = index.as_retriever(similarity_top_k=5)
  p = QueryPipeline(
      chain=[prompt_tmpl1, llm, prompt_tmpl2, llm, retriever], verbose=True
  )
  
  
  nodes = p.run(topic="college")
  len(nodes)
  
  
  
  
  from llama_index.postprocessor.cohere_rerank import CohereRerank
  from llama_index.core.response_synthesizers import TreeSummarize
  
  
  prompt_str = "Please generate a question about Paul Graham's life regarding the following topic {topic}"
  prompt_tmpl = PromptTemplate(prompt_str)
  llm = OpenAI(model="gpt-3.5-turbo")
  retriever = index.as_retriever(similarity_top_k=3)
  reranker = CohereRerank()
  summarizer = TreeSummarize(llm=llm)
  
  
  p = QueryPipeline(verbose=True)
  p.add_modules(
      {
          "llm": llm,
          "prompt_tmpl": prompt_tmpl,
          "retriever": retriever,
          "summarizer": summarizer,
          "reranker": reranker,
      }
  )
  
  
  
  p.add_link("prompt_tmpl", "llm")
  p.add_link("llm", "retriever")
  p.add_link("retriever", "reranker", dest_key="nodes")
  p.add_link("llm", "reranker", dest_key="query_str")
  p.add_link("reranker", "summarizer", dest_key="nodes")
  p.add_link("llm", "summarizer", dest_key="query_str")
  
  print(summarizer.as_query_component().input_keys)
  
  
  
  from pyvis.network import Network
  
  net = Network(notebook=True, cdn_resources="in_line", directed=True)
  net.from_nx(p.dag)
  net.show("rag_dag.html")
  
  
  
  response = p.run(topic="YC")
  
  
  print(str(response))
  
  
  response = await p.arun(topic="YC")
  print(str(response))
  
  
  
  from llama_index.postprocessor.cohere_rerank import CohereRerank
  from llama_index.core.response_synthesizers import TreeSummarize
  from llama_index.core.query_pipeline import InputComponent
  
  retriever = index.as_retriever(similarity_top_k=5)
  summarizer = TreeSummarize(llm=OpenAI(model="gpt-3.5-turbo"))
  reranker = CohereRerank()
  
  
  p = QueryPipeline(verbose=True)
  p.add_modules(
      {
          "input": InputComponent(),
          "retriever": retriever,
          "summarizer": summarizer,
      }
  )
  p.add_link("input", "retriever")
  p.add_link("input", "summarizer", dest_key="query_str")
  p.add_link("retriever", "summarizer", dest_key="nodes")
  
  
  output = p.run(input="what did the author do in YC")
  
  
  print(str(output))
  
  
  
  from llama_index.core.query_pipeline import (
      CustomQueryComponent,
      InputKeys,
      OutputKeys,
  )
  from typing import Dict, Any
  from llama_index.core.llms.llm import LLM
  from pydantic import Field
  
  
  class RelatedMovieComponent(CustomQueryComponent):
      """Related movie component."""
  
      llm: LLM = Field(..., description="OpenAI LLM")
  
      def _validate_component_inputs(
          self, input: Dict[str, Any]
      ) -> Dict[str, Any]:
          """Validate component inputs during run_component."""
          return input
  
      @property
      def _input_keys(self) -> set:
          """Input keys dict."""
          return {"movie"}
  
      @property
      def _output_keys(self) -> set:
          return {"output"}
  
      def _run_component(self, **kwargs) -> Dict[str, Any]:
          """Run the component."""
          prompt_str = "Please generate related movies to {movie_name}"
          prompt_tmpl = PromptTemplate(prompt_str)
          p = QueryPipeline(chain=[prompt_tmpl, llm])
          return {"output": p.run(movie_name=kwargs["movie"])}
  
  
  
  llm = OpenAI(model="gpt-3.5-turbo")
  component = RelatedMovieComponent(llm=llm)
  
  prompt_str = """\
  Here's some text:
  
  {text}
  
  Can you rewrite this in the voice of Shakespeare?
  """
  prompt_tmpl = PromptTemplate(prompt_str)
  
  p = QueryPipeline(chain=[component, prompt_tmpl, llm], verbose=True)
  
  
  output = p.run(movie="Love Actually")
  
  
  print(str(output))
apis: llama_index.postprocessor.cohere_rerank.CohereRerank llama_index.core.StorageContext.from_defaults llama_index.core.output_parsers.PydanticOutputParser llama_index.core.PromptTemplate llama_index.core.response_synthesizers.TreeSummarize llama_index.core.query_pipeline.QueryPipeline llama_index.core.query_pipeline.InputComponent llama_index.core.VectorStoreIndex.from_documents llama_index.core.SimpleDirectoryReader llama_index.llms.openai.OpenAI llama_index.core.load_index_from_storage llama_index.embeddings.openai.OpenAIEmbedding
extract_api: [((259, 274), 'phoenix.launch_app', 'px.launch_app', ([], {}), '()\n', (272, 274), True, 'import phoenix as px\n'), ((510, 539), 'llama_index.llms.openai.OpenAI', 'OpenAI', ([], {'model': '"""gpt-3.5-turbo"""'}), "(model='gpt-3.5-turbo')\n", (516, 539), False, 'from llama_index.llms.openai import OpenAI\n'), ((563, 610), 'llama_index.embeddings.openai.OpenAIEmbedding', 'OpenAIEmbedding', ([], {'model': '"""text-embedding-3-small"""'}), "(model='text-embedding-3-small')\n", (578, 610), False, 'from llama_index.embeddings.openai import OpenAIEmbedding\n'), ((674, 718), 'llama_index.core.SimpleDirectoryReader', 'SimpleDirectoryReader', (['"""../data/paul_graham"""'], {}), "('../data/paul_graham')\n", (695, 718), False, 'from llama_index.core import SimpleDirectoryReader\n'), ((1374, 1400), 'llama_index.core.PromptTemplate', 'PromptTemplate', (['prompt_str'], {}), '(prompt_str)\n', (1388, 1400), False, 'from llama_index.core import PromptTemplate\n'), ((1407, 1436), 'llama_index.llms.openai.OpenAI', 'OpenAI', ([], {'model': '"""gpt-3.5-turbo"""'}), "(model='gpt-3.5-turbo')\n", (1413, 1436), False, 'from llama_index.llms.openai import OpenAI\n'), ((1442, 1495), 'llama_index.core.query_pipeline.QueryPipeline', 'QueryPipeline', ([], {'chain': '[prompt_tmpl, llm]', 'verbose': '(True)'}), '(chain=[prompt_tmpl, llm], verbose=True)\n', (1455, 1495), False, 'from llama_index.core.query_pipeline import QueryPipeline\n'), ((2038, 2067), 'llama_index.llms.openai.OpenAI', 'OpenAI', ([], {'model': '"""gpt-3.5-turbo"""'}), "(model='gpt-3.5-turbo')\n", (2044, 2067), False, 'from llama_index.llms.openai import OpenAI\n'), ((2084, 2112), 'llama_index.core.output_parsers.PydanticOutputParser', 'PydanticOutputParser', (['Movies'], {}), '(Movies)\n', (2104, 2112), False, 'from llama_index.core.output_parsers import PydanticOutputParser\n'), ((2305, 2336), 'llama_index.core.PromptTemplate', 'PromptTemplate', (['json_prompt_str'], {}), '(json_prompt_str)\n', (2319, 2336), False, 'from llama_index.core import PromptTemplate\n'), ((2342, 2415), 'llama_index.core.query_pipeline.QueryPipeline', 'QueryPipeline', ([], {'chain': '[json_prompt_tmpl, llm, output_parser]', 'verbose': '(True)'}), '(chain=[json_prompt_tmpl, llm, output_parser], verbose=True)\n', (2355, 2415), False, 'from llama_index.core.query_pipeline import QueryPipeline\n'), ((2544, 2570), 'llama_index.core.PromptTemplate', 'PromptTemplate', (['prompt_str'], {}), '(prompt_str)\n', (2558, 2570), False, 'from llama_index.core import PromptTemplate\n'), ((2687, 2714), 'llama_index.core.PromptTemplate', 'PromptTemplate', (['prompt_str2'], {}), '(prompt_str2)\n', (2701, 2714), False, 'from llama_index.core import PromptTemplate\n'), ((2721, 2750), 'llama_index.llms.openai.OpenAI', 'OpenAI', ([], {'model': '"""gpt-3.5-turbo"""'}), "(model='gpt-3.5-turbo')\n", (2727, 2750), False, 'from llama_index.llms.openai import OpenAI\n'), ((2803, 2879), 'llama_index.core.query_pipeline.QueryPipeline', 'QueryPipeline', ([], {'chain': '[prompt_tmpl, llm_c, prompt_tmpl2, llm_c]', 'verbose': '(True)'}), '(chain=[prompt_tmpl, llm_c, prompt_tmpl2, llm_c], verbose=True)\n', (2816, 2879), False, 'from llama_index.core.query_pipeline import QueryPipeline\n'), ((3388, 3415), 'llama_index.core.PromptTemplate', 'PromptTemplate', (['prompt_str1'], {}), '(prompt_str1)\n', (3402, 3415), False, 'from llama_index.core import PromptTemplate\n'), ((3635, 3662), 'llama_index.core.PromptTemplate', 'PromptTemplate', (['prompt_str2'], {}), '(prompt_str2)\n', (3649, 3662), False, 'from llama_index.core import PromptTemplate\n'), ((3670, 3699), 'llama_index.llms.openai.OpenAI', 'OpenAI', ([], {'model': '"""gpt-3.5-turbo"""'}), "(model='gpt-3.5-turbo')\n", (3676, 3699), False, 'from llama_index.llms.openai import OpenAI\n'), ((3755, 3843), 'llama_index.core.query_pipeline.QueryPipeline', 'QueryPipeline', ([], {'chain': '[prompt_tmpl1, llm, prompt_tmpl2, llm, retriever]', 'verbose': '(True)'}), '(chain=[prompt_tmpl1, llm, prompt_tmpl2, llm, retriever],\n    verbose=True)\n', (3768, 3843), False, 'from llama_index.core.query_pipeline import QueryPipeline\n'), ((4145, 4171), 'llama_index.core.PromptTemplate', 'PromptTemplate', (['prompt_str'], {}), '(prompt_str)\n', (4159, 4171), False, 'from llama_index.core import PromptTemplate\n'), ((4178, 4207), 'llama_index.llms.openai.OpenAI', 'OpenAI', ([], {'model': '"""gpt-3.5-turbo"""'}), "(model='gpt-3.5-turbo')\n", (4184, 4207), False, 'from llama_index.llms.openai import OpenAI\n'), ((4270, 4284), 'llama_index.postprocessor.cohere_rerank.CohereRerank', 'CohereRerank', ([], {}), '()\n', (4282, 4284), False, 'from llama_index.postprocessor.cohere_rerank import CohereRerank\n'), ((4298, 4320), 'llama_index.core.response_synthesizers.TreeSummarize', 'TreeSummarize', ([], {'llm': 'llm'}), '(llm=llm)\n', (4311, 4320), False, 'from llama_index.core.response_synthesizers import TreeSummarize\n'), ((4327, 4354), 'llama_index.core.query_pipeline.QueryPipeline', 'QueryPipeline', ([], {'verbose': '(True)'}), '(verbose=True)\n', (4340, 4354), False, 'from llama_index.core.query_pipeline import QueryPipeline\n'), ((4913, 4975), 'pyvis.network.Network', 'Network', ([], {'notebook': '(True)', 'cdn_resources': '"""in_line"""', 'directed': '(True)'}), "(notebook=True, cdn_resources='in_line', directed=True)\n", (4920, 4975), False, 'from pyvis.network import Network\n'), ((5451, 5465), 'llama_index.postprocessor.cohere_rerank.CohereRerank', 'CohereRerank', ([], {}), '()\n', (5463, 5465), False, 'from llama_index.postprocessor.cohere_rerank import CohereRerank\n'), ((5472, 5499), 'llama_index.core.query_pipeline.QueryPipeline', 'QueryPipeline', ([], {'verbose': '(True)'}), '(verbose=True)\n', (5485, 5499), False, 'from llama_index.core.query_pipeline import QueryPipeline\n'), ((6893, 6922), 'llama_index.llms.openai.OpenAI', 'OpenAI', ([], {'model': '"""gpt-3.5-turbo"""'}), "(model='gpt-3.5-turbo')\n", (6899, 6922), False, 'from llama_index.llms.openai import OpenAI\n'), ((7080, 7106), 'llama_index.core.PromptTemplate', 'PromptTemplate', (['prompt_str'], {}), '(prompt_str)\n', (7094, 7106), False, 'from llama_index.core import PromptTemplate\n'), ((7112, 7176), 'llama_index.core.query_pipeline.QueryPipeline', 'QueryPipeline', ([], {'chain': '[component, prompt_tmpl, llm]', 'verbose': '(True)'}), '(chain=[component, prompt_tmpl, llm], verbose=True)\n', (7125, 7176), False, 'from llama_index.core.query_pipeline import QueryPipeline\n'), ((871, 896), 'os.path.exists', 'os.path.exists', (['"""storage"""'], {}), "('storage')\n", (885, 896), False, 'import os\n'), ((910, 947), 'llama_index.core.VectorStoreIndex.from_documents', 'VectorStoreIndex.from_documents', (['docs'], {}), '(docs)\n', (941, 947), False, 'from llama_index.core import StorageContext, VectorStoreIndex, load_index_from_storage\n'), ((1062, 1113), 'llama_index.core.StorageContext.from_defaults', 'StorageContext.from_defaults', ([], {'persist_dir': '"""storage"""'}), "(persist_dir='storage')\n", (1090, 1113), False, 'from llama_index.core import StorageContext, VectorStoreIndex, load_index_from_storage\n'), ((1126, 1191), 'llama_index.core.load_index_from_storage', 'load_index_from_storage', (['storage_context'], {'index_id': '"""vector_index"""'}), "(storage_context, index_id='vector_index')\n", (1149, 1191), False, 'from llama_index.core import StorageContext, VectorStoreIndex, load_index_from_storage\n'), ((1780, 1824), 'pydantic.Field', 'Field', (['...'], {'description': '"""Name of the movie."""'}), "(..., description='Name of the movie.')\n", (1785, 1824), False, 'from pydantic import Field\n'), ((1841, 1885), 'pydantic.Field', 'Field', (['...'], {'description': '"""Year of the movie."""'}), "(..., description='Year of the movie.')\n", (1846, 1885), False, 'from pydantic import Field\n'), ((1988, 2029), 'pydantic.Field', 'Field', (['...'], {'description': '"""List of movies."""'}), "(..., description='List of movies.')\n", (1993, 2029), False, 'from pydantic import Field\n'), ((6161, 6197), 'pydantic.Field', 'Field', (['...'], {'description': '"""OpenAI LLM"""'}), "(..., description='OpenAI LLM')\n", (6166, 6197), False, 'from pydantic import Field\n'), ((5409, 5438), 'llama_index.llms.openai.OpenAI', 'OpenAI', ([], {'model': '"""gpt-3.5-turbo"""'}), "(model='gpt-3.5-turbo')\n", (5415, 5438), False, 'from llama_index.llms.openai import OpenAI\n'), ((5538, 5554), 'llama_index.core.query_pipeline.InputComponent', 'InputComponent', ([], {}), '()\n', (5552, 5554), False, 'from llama_index.core.query_pipeline import InputComponent\n'), ((6744, 6770), 'llama_index.core.PromptTemplate', 'PromptTemplate', (['prompt_str'], {}), '(prompt_str)\n', (6758, 6770), False, 'from llama_index.core import PromptTemplate\n'), ((6783, 6822), 'llama_index.core.query_pipeline.QueryPipeline', 'QueryPipeline', ([], {'chain': '[prompt_tmpl, llm]'}), '(chain=[prompt_tmpl, llm])\n', (6796, 6822), False, 'from llama_index.core.query_pipeline import QueryPipeline\n')]