code: from llama_index.core import SQLDatabase
  from sqlalchemy import (
      create_engine,
      MetaData,
      Table,
      Column,
      String,
      Integer,
      select,
      column,
  )
  
  engine = create_engine("sqlite:///chinook.db")
  sql_database = SQLDatabase(engine)
  
  
  from llama_index.core.query_pipeline import QueryPipeline
  
  
  
  get_ipython().run_line_magic('pip', 'install llama-index-llms-openai')
  
  
  get_ipython().system('curl "https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip" -O ./chinook.zip')
  get_ipython().system('unzip ./chinook.zip')
  
  
  
  from llama_index.core.settings import Settings
  from llama_index.core.callbacks import CallbackManager
  
  callback_manager = CallbackManager()
  Settings.callback_manager = callback_manager
  
  
  import phoenix as px
  import llama_index.core
  
  px.launch_app()
  llama_index.core.set_global_handler("arize_phoenix")
  
  
  
  from llama_index.core.query_engine import NLSQLTableQueryEngine
  from llama_index.core.tools import QueryEngineTool
  
  sql_query_engine = NLSQLTableQueryEngine(
      sql_database=sql_database,
      tables=["albums", "tracks", "artists"],
      verbose=True,
  )
  sql_tool = QueryEngineTool.from_defaults(
      query_engine=sql_query_engine,
      name="sql_tool",
      description=(
          "Useful for translating a natural language query into a SQL query"
      ),
  )
  
  
  
  from llama_index.core.query_pipeline import QueryPipeline as QP
  
  qp = QP(verbose=True)
  
  
  
  from llama_index.core.agent.react.types import (
      ActionReasoningStep,
      ObservationReasoningStep,
      ResponseReasoningStep,
  )
  from llama_index.core.agent import Task, AgentChatResponse
  from llama_index.core.query_pipeline import (
      AgentInputComponent,
      AgentFnComponent,
      CustomAgentComponent,
      QueryComponent,
      ToolRunnerComponent,
  )
  from llama_index.core.llms import MessageRole
  from typing import Dict, Any, Optional, Tuple, List, cast
  
  
  def agent_input_fn(task: Task, state: Dict[str, Any]) -> Dict[str, Any]:
      """Agent input function.
  
      Returns:
          A Dictionary of output keys and values. If you are specifying
          src_key when defining links between this component and other
          components, make sure the src_key matches the specified output_key.
  
      """
      if "current_reasoning" not in state:
          state["current_reasoning"] = []
      reasoning_step = ObservationReasoningStep(observation=task.input)
      state["current_reasoning"].append(reasoning_step)
      return {"input": task.input}
  
  
  agent_input_component = AgentInputComponent(fn=agent_input_fn)
  
  
  
  from llama_index.core.agent import ReActChatFormatter
  from llama_index.core.query_pipeline import InputComponent, Link
  from llama_index.core.llms import ChatMessage
  from llama_index.core.tools import BaseTool
  
  
  def react_prompt_fn(
      task: Task, state: Dict[str, Any], input: str, tools: List[BaseTool]
  ) -> List[ChatMessage]:
      chat_formatter = ReActChatFormatter()
      return chat_formatter.format(
          tools,
          chat_history=task.memory.get() + state["memory"].get_all(),
          current_reasoning=state["current_reasoning"],
      )
  
  
  react_prompt_component = AgentFnComponent(
      fn=react_prompt_fn, partial_dict={"tools": [sql_tool]}
  )
  
  
  
  from typing import Set, Optional
  from llama_index.core.agent.react.output_parser import ReActOutputParser
  from llama_index.core.llms import ChatResponse
  from llama_index.core.agent.types import Task
  
  
  def parse_react_output_fn(
      task: Task, state: Dict[str, Any], chat_response: ChatResponse
  ):
      """Parse ReAct output into a reasoning step."""
      output_parser = ReActOutputParser()
      reasoning_step = output_parser.parse(chat_response.message.content)
      return {"done": reasoning_step.is_done, "reasoning_step": reasoning_step}
  
  
  parse_react_output = AgentFnComponent(fn=parse_react_output_fn)
  
  
  def run_tool_fn(
      task: Task, state: Dict[str, Any], reasoning_step: ActionReasoningStep
  ):
      """Run tool and process tool output."""
      tool_runner_component = ToolRunnerComponent(
          [sql_tool], callback_manager=task.callback_manager
      )
      tool_output = tool_runner_component.run_component(
          tool_name=reasoning_step.action,
          tool_input=reasoning_step.action_input,
      )
      observation_step = ObservationReasoningStep(observation=str(tool_output))
      state["current_reasoning"].append(observation_step)
  
      return {"response_str": observation_step.get_content(), "is_done": False}
  
  
  run_tool = AgentFnComponent(fn=run_tool_fn)
  
  
  def process_response_fn(
      task: Task, state: Dict[str, Any], response_step: ResponseReasoningStep
  ):
      """Process response."""
      state["current_reasoning"].append(response_step)
      response_str = response_step.response
      state["memory"].put(ChatMessage(content=task.input, role=MessageRole.USER))
      state["memory"].put(
          ChatMessage(content=response_str, role=MessageRole.ASSISTANT)
      )
  
      return {"response_str": response_str, "is_done": True}
  
  
  process_response = AgentFnComponent(fn=process_response_fn)
  
  
  def process_agent_response_fn(
      task: Task, state: Dict[str, Any], response_dict: dict
  ):
      """Process agent response."""
      return (
          AgentChatResponse(response_dict["response_str"]),
          response_dict["is_done"],
      )
  
  
  process_agent_response = AgentFnComponent(fn=process_agent_response_fn)
  
  
  
  from llama_index.core.query_pipeline import QueryPipeline as QP
  from llama_index.llms.openai import OpenAI
  
  qp.add_modules(
      {
          "agent_input": agent_input_component,
          "react_prompt": react_prompt_component,
          "llm": OpenAI(model="gpt-4-1106-preview"),
          "react_output_parser": parse_react_output,
          "run_tool": run_tool,
          "process_response": process_response,
          "process_agent_response": process_agent_response,
      }
  )
  
  
  qp.add_chain(["agent_input", "react_prompt", "llm", "react_output_parser"])
  
  qp.add_link(
      "react_output_parser",
      "run_tool",
      condition_fn=lambda x: not x["done"],
      input_fn=lambda x: x["reasoning_step"],
  )
  qp.add_link(
      "react_output_parser",
      "process_response",
      condition_fn=lambda x: x["done"],
      input_fn=lambda x: x["reasoning_step"],
  )
  
  qp.add_link("process_response", "process_agent_response")
  qp.add_link("run_tool", "process_agent_response")
  
  
  
  from pyvis.network import Network
  
  net = Network(notebook=True, cdn_resources="in_line", directed=True)
  net.from_nx(qp.clean_dag)
  net.show("agent_dag.html")
  
  
  
  from llama_index.core.agent import QueryPipelineAgentWorker, AgentRunner
  from llama_index.core.callbacks import CallbackManager
  
  agent_worker = QueryPipelineAgentWorker(qp)
  agent = AgentRunner(
      agent_worker, callback_manager=CallbackManager([]), verbose=True
  )
  
  
  
  task = agent.create_task(
      "What are some tracks from the artist AC/DC? Limit it to 3"
  )
  
  
  step_output = agent.run_step(task.task_id)
  
  
  step_output = agent.run_step(task.task_id)
  
  
  step_output.is_last
  
  
  response = agent.finalize_response(task.task_id)
  
  
  print(str(response))
  
  
  agent.reset()
  response = agent.chat(
      "What are some tracks from the artist AC/DC? Limit it to 3"
  )
  
  
  print(str(response))
  
  
  
  
  from llama_index.llms.openai import OpenAI
  
  llm = OpenAI(model="gpt-4-1106-preview")
  
  
  from llama_index.core.agent import Task, AgentChatResponse
  from typing import Dict, Any
  from llama_index.core.query_pipeline import (
      AgentInputComponent,
      AgentFnComponent,
  )
  
  
  def agent_input_fn(task: Task, state: Dict[str, Any]) -> Dict:
      """Agent input function."""
      if "convo_history" not in state:
          state["convo_history"] = []
          state["count"] = 0
      state["convo_history"].append(f"User: {task.input}")
      convo_history_str = "\n".join(state["convo_history"]) or "None"
      return {"input": task.input, "convo_history": convo_history_str}
  
  
  agent_input_component = AgentInputComponent(fn=agent_input_fn)
  
  
  from llama_index.core import PromptTemplate
  
  retry_prompt_str = """\
  You are trying to generate a proper natural language query given a user input.
  
  This query will then be interpreted by a downstream text-to-SQL agent which
  will convert the query to a SQL statement. If the agent triggers an error,
  then that will be reflected in the current conversation history (see below).
  
  If the conversation history is None, use the user input. If its not None,
  generate a new SQL query that avoids the problems of the previous SQL query.
  
  Input: {input}
  Convo history (failed attempts): 
  {convo_history}
  
  New input: """
  retry_prompt = PromptTemplate(retry_prompt_str)
  
  
  from llama_index.core import Response
  from typing import Tuple
  
  validate_prompt_str = """\
  Given the user query, validate whether the inferred SQL query and response from executing the query is correct and answers the query.
  
  Answer with YES or NO.
  
  Query: {input}
  Inferred SQL query: {sql_query}
  SQL Response: {sql_response}
  
  Result: """
  validate_prompt = PromptTemplate(validate_prompt_str)
  
  MAX_ITER = 3
  
  
  def agent_output_fn(
      task: Task, state: Dict[str, Any], output: Response
  ) -> Tuple[AgentChatResponse, bool]:
      """Agent output component."""
      print(f"> Inferred SQL Query: {output.metadata['sql_query']}")
      print(f"> SQL Response: {str(output)}")
      state["convo_history"].append(
          f"Assistant (inferred SQL query): {output.metadata['sql_query']}"
      )
      state["convo_history"].append(f"Assistant (response): {str(output)}")
  
      validate_prompt_partial = validate_prompt.as_query_component(
          partial={
              "sql_query": output.metadata["sql_query"],
              "sql_response": str(output),
          }
      )
      qp = QP(chain=[validate_prompt_partial, llm])
      validate_output = qp.run(input=task.input)
  
      state["count"] += 1
      is_done = False
      if state["count"] >= MAX_ITER:
          is_done = True
      if "YES" in validate_output.message.content:
          is_done = True
  
      return AgentChatResponse(response=str(output)), is_done
  
  
  agent_output_component = AgentFnComponent(fn=agent_output_fn)
  
  
  from llama_index.core.query_pipeline import (
      QueryPipeline as QP,
      Link,
      InputComponent,
  )
  
  qp = QP(
      modules={
          "input": agent_input_component,
          "retry_prompt": retry_prompt,
          "llm": llm,
          "sql_query_engine": sql_query_engine,
          "output_component": agent_output_component,
      },
      verbose=True,
  )
  qp.add_link("input", "retry_prompt", src_key="input", dest_key="input")
  qp.add_link(
      "input", "retry_prompt", src_key="convo_history", dest_key="convo_history"
  )
  qp.add_chain(["retry_prompt", "llm", "sql_query_engine", "output_component"])
  
  
  
  from pyvis.network import Network
  
  net = Network(notebook=True, cdn_resources="in_line", directed=True)
  net.from_nx(qp.dag)
  net.show("agent_dag.html")
  
  
  
  from llama_index.core.agent import QueryPipelineAgentWorker, AgentRunner
  from llama_index.core.callbacks import CallbackManager
  
  agent_worker = QueryPipelineAgentWorker(qp)
  agent = AgentRunner(
      agent_worker, callback_manager=CallbackManager(), verbose=False
  )
  
  
  response = agent.chat(
      "How many albums did the artist who wrote 'Restless and Wild' release? (answer should be non-zero)?"
  )
  print(str(response))
apis: llama_index.core.query_engine.NLSQLTableQueryEngine llama_index.core.callbacks.CallbackManager llama_index.core.agent.QueryPipelineAgentWorker llama_index.core.agent.react.output_parser.ReActOutputParser llama_index.core.query_pipeline.ToolRunnerComponent llama_index.core.PromptTemplate llama_index.core.query_pipeline.QueryPipeline llama_index.core.agent.AgentChatResponse llama_index.core.agent.react.types.ObservationReasoningStep llama_index.core.agent.ReActChatFormatter llama_index.core.query_pipeline.AgentFnComponent llama_index.llms.openai.OpenAI llama_index.core.llms.ChatMessage llama_index.core.query_pipeline.AgentInputComponent llama_index.core.tools.QueryEngineTool.from_defaults llama_index.core.SQLDatabase
extract_api: [((183, 220), 'sqlalchemy.create_engine', 'create_engine', (['"""sqlite:///chinook.db"""'], {}), "('sqlite:///chinook.db')\n", (196, 220), False, 'from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, select, column\n'), ((236, 255), 'llama_index.core.SQLDatabase', 'SQLDatabase', (['engine'], {}), '(engine)\n', (247, 255), False, 'from llama_index.core import SQLDatabase\n'), ((679, 696), 'llama_index.core.callbacks.CallbackManager', 'CallbackManager', ([], {}), '()\n', (694, 696), False, 'from llama_index.core.callbacks import CallbackManager\n'), ((790, 805), 'phoenix.launch_app', 'px.launch_app', ([], {}), '()\n', (803, 805), True, 'import phoenix as px\n'), ((997, 1103), 'llama_index.core.query_engine.NLSQLTableQueryEngine', 'NLSQLTableQueryEngine', ([], {'sql_database': 'sql_database', 'tables': "['albums', 'tracks', 'artists']", 'verbose': '(True)'}), "(sql_database=sql_database, tables=['albums', 'tracks',\n    'artists'], verbose=True)\n", (1018, 1103), False, 'from llama_index.core.query_engine import NLSQLTableQueryEngine\n'), ((1126, 1293), 'llama_index.core.tools.QueryEngineTool.from_defaults', 'QueryEngineTool.from_defaults', ([], {'query_engine': 'sql_query_engine', 'name': '"""sql_tool"""', 'description': '"""Useful for translating a natural language query into a SQL query"""'}), "(query_engine=sql_query_engine, name=\n    'sql_tool', description=\n    'Useful for translating a natural language query into a SQL query')\n", (1155, 1293), False, 'from llama_index.core.tools import QueryEngineTool\n'), ((1388, 1404), 'llama_index.core.query_pipeline.QueryPipeline', 'QP', ([], {'verbose': '(True)'}), '(verbose=True)\n', (1390, 1404), True, 'from llama_index.core.query_pipeline import QueryPipeline as QP, Link, InputComponent\n'), ((2476, 2514), 'llama_index.core.query_pipeline.AgentInputComponent', 'AgentInputComponent', ([], {'fn': 'agent_input_fn'}), '(fn=agent_input_fn)\n', (2495, 2514), False, 'from llama_index.core.query_pipeline import AgentInputComponent, AgentFnComponent\n'), ((3093, 3165), 'llama_index.core.query_pipeline.AgentFnComponent', 'AgentFnComponent', ([], {'fn': 'react_prompt_fn', 'partial_dict': "{'tools': [sql_tool]}"}), "(fn=react_prompt_fn, partial_dict={'tools': [sql_tool]})\n", (3109, 3165), False, 'from llama_index.core.query_pipeline import AgentInputComponent, AgentFnComponent\n'), ((3738, 3780), 'llama_index.core.query_pipeline.AgentFnComponent', 'AgentFnComponent', ([], {'fn': 'parse_react_output_fn'}), '(fn=parse_react_output_fn)\n', (3754, 3780), False, 'from llama_index.core.query_pipeline import AgentInputComponent, AgentFnComponent\n'), ((4412, 4444), 'llama_index.core.query_pipeline.AgentFnComponent', 'AgentFnComponent', ([], {'fn': 'run_tool_fn'}), '(fn=run_tool_fn)\n', (4428, 4444), False, 'from llama_index.core.query_pipeline import AgentInputComponent, AgentFnComponent\n'), ((4936, 4976), 'llama_index.core.query_pipeline.AgentFnComponent', 'AgentFnComponent', ([], {'fn': 'process_response_fn'}), '(fn=process_response_fn)\n', (4952, 4976), False, 'from llama_index.core.query_pipeline import AgentInputComponent, AgentFnComponent\n'), ((5244, 5290), 'llama_index.core.query_pipeline.AgentFnComponent', 'AgentFnComponent', ([], {'fn': 'process_agent_response_fn'}), '(fn=process_agent_response_fn)\n', (5260, 5290), False, 'from llama_index.core.query_pipeline import AgentInputComponent, AgentFnComponent\n'), ((6286, 6348), 'pyvis.network.Network', 'Network', ([], {'notebook': '(True)', 'cdn_resources': '"""in_line"""', 'directed': '(True)'}), "(notebook=True, cdn_resources='in_line', directed=True)\n", (6293, 6348), False, 'from pyvis.network import Network\n'), ((6549, 6577), 'llama_index.core.agent.QueryPipelineAgentWorker', 'QueryPipelineAgentWorker', (['qp'], {}), '(qp)\n', (6573, 6577), False, 'from llama_index.core.agent import QueryPipelineAgentWorker, AgentRunner\n'), ((7133, 7167), 'llama_index.llms.openai.OpenAI', 'OpenAI', ([], {'model': '"""gpt-4-1106-preview"""'}), "(model='gpt-4-1106-preview')\n", (7139, 7167), False, 'from llama_index.llms.openai import OpenAI\n'), ((7770, 7808), 'llama_index.core.query_pipeline.AgentInputComponent', 'AgentInputComponent', ([], {'fn': 'agent_input_fn'}), '(fn=agent_input_fn)\n', (7789, 7808), False, 'from llama_index.core.query_pipeline import AgentInputComponent, AgentFnComponent\n'), ((8437, 8469), 'llama_index.core.PromptTemplate', 'PromptTemplate', (['retry_prompt_str'], {}), '(retry_prompt_str)\n', (8451, 8469), False, 'from llama_index.core import PromptTemplate\n'), ((8829, 8864), 'llama_index.core.PromptTemplate', 'PromptTemplate', (['validate_prompt_str'], {}), '(validate_prompt_str)\n', (8843, 8864), False, 'from llama_index.core import PromptTemplate\n'), ((9886, 9922), 'llama_index.core.query_pipeline.AgentFnComponent', 'AgentFnComponent', ([], {'fn': 'agent_output_fn'}), '(fn=agent_output_fn)\n', (9902, 9922), False, 'from llama_index.core.query_pipeline import AgentInputComponent, AgentFnComponent\n'), ((10034, 10224), 'llama_index.core.query_pipeline.QueryPipeline', 'QP', ([], {'modules': "{'input': agent_input_component, 'retry_prompt': retry_prompt, 'llm': llm,\n    'sql_query_engine': sql_query_engine, 'output_component':\n    agent_output_component}", 'verbose': '(True)'}), "(modules={'input': agent_input_component, 'retry_prompt': retry_prompt,\n    'llm': llm, 'sql_query_engine': sql_query_engine, 'output_component':\n    agent_output_component}, verbose=True)\n", (10036, 10224), True, 'from llama_index.core.query_pipeline import QueryPipeline as QP, Link, InputComponent\n'), ((10563, 10625), 'pyvis.network.Network', 'Network', ([], {'notebook': '(True)', 'cdn_resources': '"""in_line"""', 'directed': '(True)'}), "(notebook=True, cdn_resources='in_line', directed=True)\n", (10570, 10625), False, 'from pyvis.network import Network\n'), ((10820, 10848), 'llama_index.core.agent.QueryPipelineAgentWorker', 'QueryPipelineAgentWorker', (['qp'], {}), '(qp)\n', (10844, 10848), False, 'from llama_index.core.agent import QueryPipelineAgentWorker, AgentRunner\n'), ((2314, 2362), 'llama_index.core.agent.react.types.ObservationReasoningStep', 'ObservationReasoningStep', ([], {'observation': 'task.input'}), '(observation=task.input)\n', (2338, 2362), False, 'from llama_index.core.agent.react.types import ActionReasoningStep, ObservationReasoningStep, ResponseReasoningStep\n'), ((2868, 2888), 'llama_index.core.agent.ReActChatFormatter', 'ReActChatFormatter', ([], {}), '()\n', (2886, 2888), False, 'from llama_index.core.agent import ReActChatFormatter\n'), ((3545, 3564), 'llama_index.core.agent.react.output_parser.ReActOutputParser', 'ReActOutputParser', ([], {}), '()\n', (3562, 3564), False, 'from llama_index.core.agent.react.output_parser import ReActOutputParser\n'), ((3950, 4021), 'llama_index.core.query_pipeline.ToolRunnerComponent', 'ToolRunnerComponent', (['[sql_tool]'], {'callback_manager': 'task.callback_manager'}), '([sql_tool], callback_manager=task.callback_manager)\n', (3969, 4021), False, 'from llama_index.core.query_pipeline import AgentInputComponent, AgentFnComponent, CustomAgentComponent, QueryComponent, ToolRunnerComponent\n'), ((9535, 9575), 'llama_index.core.query_pipeline.QueryPipeline', 'QP', ([], {'chain': '[validate_prompt_partial, llm]'}), '(chain=[validate_prompt_partial, llm])\n', (9537, 9575), True, 'from llama_index.core.query_pipeline import QueryPipeline as QP, Link, InputComponent\n'), ((4698, 4752), 'llama_index.core.llms.ChatMessage', 'ChatMessage', ([], {'content': 'task.input', 'role': 'MessageRole.USER'}), '(content=task.input, role=MessageRole.USER)\n', (4709, 4752), False, 'from llama_index.core.llms import ChatMessage\n'), ((4787, 4848), 'llama_index.core.llms.ChatMessage', 'ChatMessage', ([], {'content': 'response_str', 'role': 'MessageRole.ASSISTANT'}), '(content=response_str, role=MessageRole.ASSISTANT)\n', (4798, 4848), False, 'from llama_index.core.llms import ChatMessage\n'), ((5127, 5175), 'llama_index.core.agent.AgentChatResponse', 'AgentChatResponse', (["response_dict['response_str']"], {}), "(response_dict['response_str'])\n", (5144, 5175), False, 'from llama_index.core.agent import Task, AgentChatResponse\n'), ((5533, 5567), 'llama_index.llms.openai.OpenAI', 'OpenAI', ([], {'model': '"""gpt-4-1106-preview"""'}), "(model='gpt-4-1106-preview')\n", (5539, 5567), False, 'from llama_index.llms.openai import OpenAI\n'), ((6634, 6653), 'llama_index.core.callbacks.CallbackManager', 'CallbackManager', (['[]'], {}), '([])\n', (6649, 6653), False, 'from llama_index.core.callbacks import CallbackManager\n'), ((10905, 10922), 'llama_index.core.callbacks.CallbackManager', 'CallbackManager', ([], {}), '()\n', (10920, 10922), False, 'from llama_index.core.callbacks import CallbackManager\n')]