LlamaIndexのQuery Pipelineを試す

!pip install llama-index-llms-openai llama-index-callbacks-arize-phoenix python-dotenv pygraphviz

$ sudo apt-get install -y graphviz graphviz-dev

import phoenix as px
import llama_index.core

px.launch_app()
llama_index.core.set_global_handler("arize_phoenix")

from dotenv import load_dotenv

load_dotenv(verbose=True)

from llama_index.core.query_pipeline import QueryPipeline, Link, InputComponent
from llama_index.core.query_engine.pandas import PandasInstructionParser
from llama_index.llms.openai import OpenAI
from llama_index.core import PromptTemplate

!wget 'https://raw.githubusercontent.com/jerryjliu/llama_index/main/docs/docs/examples/data/csv/titanic_train.csv' -O 'titanic_train.csv'

import pandas as pd

df = pd.read_csv("./titanic_train.csv")

instruction_str = """\
1. Pandasを使ってクエリを実行可能なPythonコードに変換してください。
2. コードの最終行は、`eval()`関数で呼び出せるPython式にしてください。
3. コードはクエリの解を表してください。
4. 式だけを出力して下さい。
5. 式を引用しないでください。 \
"""

pandas_prompt_str = """\
Pythonでpandasのデータフレームを操作しています。
データフレームの名前は'df'です。
`print(df.head())`の結果はこれです:
{df_str}

以下の指示に従ってください:
{instruction_str}
質問: {query_str}

式: \
"""

response_synthesis_prompt_str = """\
入力された質問に対して、クエリー結果から回答を生成して下さい。
質問: {query_str}

Pandasへの命令(オプション): 
{pandas_instructions}

Pandasの出力t: {pandas_output}
回答: \
"""

pandas_prompt = PromptTemplate(pandas_prompt_str).partial_format(
    instruction_str=instruction_str, df_str=df.head(5)
)
pandas_output_parser = PandasInstructionParser(df)
response_synthesis_prompt = PromptTemplate(response_synthesis_prompt_str)

pandas_output_parser = PandasInstructionParser(df)
llm = OpenAI(model="gpt-3.5-turbo")

qp = QueryPipeline(
    modules={
        "input": InputComponent(),
        "pandas_prompt": pandas_prompt,
        "llm1": llm,
        "pandas_output_parser": pandas_output_parser,
        "response_synthesis_prompt": response_synthesis_prompt,
        "llm2": llm,
    },
    verbose=True,
)
qp.add_chain(["input", "pandas_prompt", "llm1", "pandas_output_parser"])
qp.add_links(
    [
        Link("input", "response_synthesis_prompt", dest_key="query_str"),
        Link(
            "llm1", "response_synthesis_prompt", dest_key="pandas_instructions"
        ),
        Link(
            "pandas_output_parser",
            "response_synthesis_prompt",
            dest_key="pandas_output",
        ),
    ]
)
qp.add_link("response_synthesis_prompt", "llm2")

from networkx.drawing.nx_agraph import to_agraph
from IPython.display import Image

agraph = to_agraph(qp.dag)
agraph.layout(prog="dot")
agraph.draw('pandas_dag.png')
display(Image('pandas_dag.png'))

response = qp.run(query_str="生存率と年齢の相関関係は？")
print(response.message.content)

!pip install llama-index-llms-openai llama-index-embeddings-openai llama-index-callbacks-arize-phoenix python-dotenv pygraphviz

import phoenix as px
import llama_index.core

px.launch_app()
llama_index.core.set_global_handler("arize_phoenix")

from dotenv import load_dotenv

load_dotenv(verbose=True)

!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt' -O pg_essay.txt

from llama_index.core import SimpleDirectoryReader

reader = SimpleDirectoryReader(input_files=["pg_essay.txt"])
documents = reader.load_data()

from llama_index.llms.openai import OpenAI
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core import VectorStoreIndex

llm = OpenAI(model="gpt-3.5-turbo")

chunk_sizes = [128, 256, 512, 1024]

query_engines = {}
for chunk_size in chunk_sizes:
    splitter = SentenceSplitter(chunk_size=chunk_size, chunk_overlap=0)
    nodes = splitter.get_nodes_from_documents(documents)
    vector_index = VectorStoreIndex(nodes, show_progress=True)
    query_engines[str(chunk_size)] = vector_index.as_query_engine(llm=llm)

for i in query_engines:
    print(f"===== {i} =====")
    response = query_engines[i].retrieve("著者はYC時代に何をしていたのか？日本語で簡潔に回答して。")
    print(response[0])

===== 128 =====
Node ID: ab9b9821-fd06-4774-91b5-9378bd6a9f53
Text: When I was dealing with some urgent problem during YC, there was
about a 60% chance it had to do with HN, and a 40% chance it had do
with everything else combined. [17]  As well as HN, I wrote all of
YC's internal software in Arc. But while I continued to work a good
deal in Arc, I gradually stopped working on Arc, partly because I
didn't have t...
Score:  0.780

===== 256 =====
Node ID: 516079d9-d722-4cff-8291-dff7720bbd18
Text: Whereas I was like someone who was in pain while running a
marathon not from the exertion of running, but because I had a blister
from an ill-fitting shoe. When I was dealing with some urgent problem
during YC, there was about a 60% chance it had to do with HN, and a
40% chance it had do with everything else combined. [17]  As well as
HN, I wrot...
Score:  0.781

===== 512 =====
Node ID: eab4e6a8-c14a-404c-a3d3-92fad0fc87f6
Text: In the summer of 2006, Robert and I started working on a new
version of Arc. This one was reasonably fast, because it was compiled
into Scheme. To test this new Arc, I wrote Hacker News in it. It was
originally meant to be a news aggregator for startup founders and was
called Startup News, but after a few months I got tired of reading
about noth...
Score:  0.774

===== 1024 =====
Node ID: 31b333aa-1105-4640-ae8b-18f74351cc71
Text: [17]  As well as HN, I wrote all of YC's internal software in
Arc. But while I continued to work a good deal in Arc, I gradually
stopped working on Arc, partly because I didn't have time to, and
partly because it was a lot less attractive to mess around with the
language now that we had all this infrastructure depending on it. So
now my three pr...
Score:  0.790

from llama_index.core.query_pipeline import (
    QueryPipeline,
    InputComponent,
    ArgPackComponent,
)
from llama_index.core.response_synthesizers import TreeSummarize
from llama_index.core.schema import NodeWithScore, TextNode

p = QueryPipeline(verbose=True)
module_dict = {
    **query_engines,
    "input": InputComponent(),
    "summarizer": TreeSummarize(),
    "join": ArgPackComponent(
        convert_fn=lambda x: NodeWithScore(node=TextNode(text=str(x)))
    ),
}
p.add_modules(module_dict)
for chunk_size in chunk_sizes:
    p.add_link("input", str(chunk_size))
    p.add_link(str(chunk_size), "join", dest_key=str(chunk_size))
p.add_link("join", "summarizer", dest_key="nodes")
p.add_link("input", "summarizer", dest_key="query_str")

from networkx.drawing.nx_agraph import to_agraph
from IPython.display import Image

agraph = to_agraph(p.dag)
agraph.layout(prog="dot")
agraph.draw('rag_dag.png')
display(Image('rag_dag.png'))

import time

start_time = time.time()
response = p.run(input="著者はYC時代に何をしていたのか？日本語で回答して。")
print(str(response))
end_time = time.time()
print(f"Time taken: {end_time - start_time}")

import time

start_time = time.time()
response = await p.arun(input="著者はYC時代に何をしていたのか？日本語で回答して。")     # ここ
print(str(response))
end_time = time.time()
print(f"Time taken: {end_time - start_time}")

!pip install llama-index-llms-openai llama-index-embeddings-openai llama-index-callbacks-arize-phoenix python-dotenv pygraphviz

import phoenix as px
import llama_index.core

px.launch_app()
llama_index.core.set_global_handler("arize_phoenix")

from dotenv import load_dotenv

load_dotenv(verbose=True)

!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt' -O pg_essay.txt

from llama_index.core import SimpleDirectoryReader

reader = SimpleDirectoryReader(input_files=["pg_essay.txt"])
documents = reader.load_data()

from llama_index.llms.openai import OpenAI
from llama_index.core import PromptTemplate
from llama_index.core import VectorStoreIndex, SummaryIndex
from llama_index.core.response_synthesizers import TreeSummarize

# HyDEのプロンプトテンプレート定義
hyde_str = '''\
以下の質問に答える文章を書いてください。
できるだけ多くの重要な詳細を含めるようにしてください。

質問: {query_str}
文章: """
'''
hyde_prompt = PromptTemplate(hyde_str)

# LLMの定義
llm = OpenAI(model="gpt-3.5-turbo")

# シンセサイザーの定義
summarizer = TreeSummarize(llm=llm)

# VectorStoreIndexを使ったクエリエンジンの定義
vector_index = VectorStoreIndex.from_documents(documents)
vector_query_engine = vector_index.as_query_engine(similarity_top_k=2)

# SummaryIndexを使用した、プロンプトテンプレートおよびクエリエンジンの定義
summary_index = SummaryIndex.from_documents(documents)
summary_qrewrite_str = """\
質問：{query_str}

あなたは、与えられたコンテキストで質問に答えようとするエージェントに質問を送る責任があります。
文脈は関係あるかもしれませんし、ないかもしれません。事実を強調するために質問を書き直してください。
コンテキストの一部だけが（あるいは全く）関連しているかもしれないという事実を強調するように質問を書き換えてください。
"""
summary_qrewrite_prompt = PromptTemplate(summary_qrewrite_str)
summary_query_engine = summary_index.as_query_engine()

from llama_index.core.query_pipeline import RouterComponent
from llama_index.core.selectors import LLMSingleSelector

vector_chain = QueryPipeline(chain=[vector_query_engine])

summary_chain = QueryPipeline(
    chain=[summary_qrewrite_prompt, llm, summary_query_engine], verbose=True
)

choices = [
    "このツールは、文書に関する特定の質問に答えます（文書全体の要約的な質問ではありません）",
    "このツールは、文書に関する要約的な質問に答えます（特定の質問ではありません）",
]

selector = LLMSingleSelector.from_defaults()

router_c = RouterComponent(
    selector=selector,
    choices=choices,
    components=[vector_chain, summary_chain],
    verbose=True,
)

qp = QueryPipeline(chain=[router_c], verbose=True)

response = qp.run("著者はYC時代に何をしていたのか？日本語で回答して。")
print(str(response))

Selecting component 0: The first choice is most relevant as it mentions answering specific questions related to the document, which would include providing information about what the author was doing during the YC era..

response = qp.run("この文書の趣旨は何？日本語で回答して。")
print(str(response))

Selecting component 1: Choice 2 is more relevant as it mentions answering summary questions related to the document, which aligns with the question asking for the main idea or summary of the document..

!pip install llama-index-llms-openai llama-index-embeddings-openai llama-index-postprocessor-cohere-rerank llama-index-readers-web llama-index-callbacks-arize-phoenix python-dotenv pygraphviz

import phoenix as px
import llama_index.core

px.launch_app()
llama_index.core.set_global_handler("arize_phoenix")

from dotenv import load_dotenv

load_dotenv(verbose=True)

from llama_index.readers.web import BeautifulSoupWebReader

reader = BeautifulSoupWebReader()

documents = reader.load_data(
    ["https://docs.anthropic.com/claude/docs/tool-use"]
)

lines = documents[0].text.split("\n")

# 空行が2行以上続くセクションを削除する
fixed_lines = [lines[0]]
for idx in range(1, len(lines)):
    if lines[idx].strip() == "" and lines[idx - 1].strip() == "":
        continue
    fixed_lines.append(lines[idx])

documents[0].text = "\n".join(fixed_lines)

from llama_index.core import VectorStoreIndex
from llama_index.embeddings.openai import OpenAIEmbedding

index = VectorStoreIndex.from_documents(
    documents,
    embed_model=OpenAIEmbedding(
        model="text-embedding-3-large", embed_batch_size=256
    ),
    show_progress=True
)

from llama_index.core.query_pipeline import (
    QueryPipeline,
    InputComponent,
    ArgPackComponent,
)
from llama_index.core.prompts import PromptTemplate
from llama_index.llms.openai import OpenAI
from llama_index.postprocessor.cohere_rerank import CohereRerank

# まず、ユーザークエリを取り込むための入力コンポーネントを作成する。
input_component = InputComponent()

# 次に、LLMを使ってユーザークエリを書き換える。
rewrite = '''\
現在の会話を使って、セマンティック検索エンジンへのクエリを書いてください。


{chat_history_str}


最新のメッセージ: {query_str}
クエリ: """\
'''
rewrite_template = PromptTemplate(rewrite)

# LLMの定義
llm = OpenAI(
    model="gpt-3.5-turbo",
    temperature=0.2,
)

# 2回取得するので、取得したノードを1つのリストにまとめる必要がある。
argpack_component = ArgPackComponent()

# それを使って検索を行う
retriever = index.as_retriever(similarity_top_k=6)

# 後処理としてCohereでリランクする
reranker = CohereRerank(model="rerank-multilingual-v2.0", top_n=3)

# 最後に、ノードとチャット履歴を使用してレスポンスを作成する必要がある
from typing import Any, Dict, List, Optional
from llama_index.core.bridge.pydantic import Field
from llama_index.core.llms import ChatMessage
from llama_index.core.query_pipeline import CustomQueryComponent
from llama_index.core.schema import NodeWithScore

DEFAULT_CONTEXT_PROMPT = """\
以下は以下は、関連しそうなコンテキストです:"
----
{node_context}
----
上記のコンテキストを使って、以下の質問に対する回答を書いてください:
{query_str}\
"""


class ResponseWithChatHistory(CustomQueryComponent):
    llm: OpenAI = Field(..., description="OpenAI LLM")
    system_prompt: Optional[str] = Field(
        default=None, description="System prompt to use for the LLM"
    )
    context_prompt: str = Field(
        default=DEFAULT_CONTEXT_PROMPT,
        description="Context prompt to use for the LLM",
    )

    def _validate_component_inputs(
        self, input: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Validate component inputs during run_component."""
        # NOTE: this is OPTIONAL but we show you where to do validation as an example
        return input

    @property
    def _input_keys(self) -> set:
        """Input keys dict."""
        # NOTE: These are required inputs. If you have optional inputs please override
        # `optional_input_keys_dict`
        return {"chat_history", "nodes", "query_str"}

    @property
    def _output_keys(self) -> set:
        return {"response"}

    def _prepare_context(
        self,
        chat_history: List[ChatMessage],
        nodes: List[NodeWithScore],
        query_str: str,
    ) -> List[ChatMessage]:
        node_context = ""
        for idx, node in enumerate(nodes):
            node_text = node.get_content(metadata_mode="llm")
            node_context += f"Context Chunk {idx}:\n{node_text}\n\n"

        formatted_context = self.context_prompt.format(
            node_context=node_context, query_str=query_str
        )
        user_message = ChatMessage(role="user", content=formatted_context)

        chat_history.append(user_message)

        if self.system_prompt is not None:
            chat_history = [
                ChatMessage(role="system", content=self.system_prompt)
            ] + chat_history

        return chat_history

    def _run_component(self, **kwargs) -> Dict[str, Any]:
        """Run the component."""
        chat_history = kwargs["chat_history"]
        nodes = kwargs["nodes"]
        query_str = kwargs["query_str"]

        prepared_context = self._prepare_context(
            chat_history, nodes, query_str
        )

        response = llm.chat(prepared_context)

        return {"response": response}

    async def _arun_component(self, **kwargs: Any) -> Dict[str, Any]:
        """Run the component asynchronously."""
        # NOTE: Optional, but async LLM calls are easy to implement
        chat_history = kwargs["chat_history"]
        nodes = kwargs["nodes"]
        query_str = kwargs["query_str"]

        prepared_context = self._prepare_context(
            chat_history, nodes, query_str
        )

        response = await llm.achat(prepared_context)

        return {"response": response}

response_component = ResponseWithChatHistory(
    llm=llm,
    system_prompt="あなたはQ&Aシステムです。ユーザーからのメッセージに答えるために、以前のチャット履歴や関連するコンテキストが提供されます。"
)

pipeline = QueryPipeline(
    modules={
        "input": input_component,
        "rewrite_template": rewrite_template,
        "llm": llm,
        "rewrite_retriever": retriever,
        "query_retriever": retriever,
        "join": argpack_component,
        "reranker": reranker,
        "response_component": response_component,
    },
    verbose=False,
)

# HyDE的なクエリーと本物のクエリーの両方のretrieverを実行
pipeline.add_link(
    "input", "rewrite_template", src_key="query_str", dest_key="query_str"
)
pipeline.add_link(
    "input",
    "rewrite_template",
    src_key="chat_history_str",
    dest_key="chat_history_str",
)
pipeline.add_link("rewrite_template", "llm")
pipeline.add_link("llm", "rewrite_retriever")
pipeline.add_link("input", "query_retriever", src_key="query_str")

# argpackコンポーネントへの各入力には、destキーが必要 - なんでもよい
# そして、argpackコンポーネントはすべての入力を1つのリストにまとめる。
pipeline.add_link("rewrite_retriever", "join", dest_key="rewrite_nodes")
pipeline.add_link("query_retriever", "join", dest_key="query_nodes")

# rerankerは、パックされたノードとクエリー文字列を必要とする
pipeline.add_link("join", "reranker", dest_key="nodes")
pipeline.add_link(
    "input", "reranker", src_key="query_str", dest_key="query_str"
)

# synthesizerは、リランク付けされたノードとクエリ文字列を必要とする
pipeline.add_link("reranker", "response_component", dest_key="nodes")
pipeline.add_link(
    "input", "response_component", src_key="query_str", dest_key="query_str"
)
pipeline.add_link(
    "input",
    "response_component",
    src_key="chat_history",
    dest_key="chat_history",
)

from networkx.drawing.nx_agraph import to_agraph
from IPython.display import Image

agraph = to_agraph(pipeline.dag)
agraph.layout(prog="dot")
agraph.draw('chat_memory.png')
display(Image('chat_memory.png'))

from llama_index.core.memory import ChatMemoryBuffer

# メモリを初期化
pipeline_memory = ChatMemoryBuffer.from_defaults(token_limit=8000)

user_inputs = [
    "こんにちは!",
    "Claude-3でのtool-useの使い方は？",
    "どのモデルがそれをサポートしていますか？",
    "ありがとう、それを知りたかったんだ！",
]

for msg in user_inputs:
    # メモリを取得
    chat_history = pipeline_memory.get()

    # 会話履歴を入力文字列に変換
    chat_history_str = "\n".join([str(x) for x in chat_history])

    # パイプライン実行
    response = pipeline.run(
        query_str=msg,
        chat_history=chat_history,
        chat_history_str=chat_history_str,
    )

    # メモリを更新
    user_msg = ChatMessage(role="user", content=msg)
    pipeline_memory.put(user_msg)
    print(str(user_msg))

    pipeline_memory.put(response.message)
    print(str(response.message))
    print()

for i in pipeline_memory.get_all():
    print(i.json(ensure_ascii=False))

{"role": "user", "content": "こんにちは!", "additional_kwargs": {}}
{"role": "assistant", "content": "こんにちは！どのようにお手伝いできますか？", "additional_kwargs": {}}
{"role": "user", "content": "Claude-3でのtool-useの使い方は？", "additional_kwargs": {}}
{"role": "assistant", "content": "Claude-3でのtool-useの使い方は、以下の手順に従います：\n\n1. Claudeにツールとユーザープロンプトを提供する（APIリクエスト）:\n   - Claudeがアクセスできるツールのセットを定義します。これには、ツールの名前、説明、および入力スキーマが含まれます。\n   - ユーザープロンプトを提供します。これは、例えば「サンフランシスコの天気はどうですか？」など、1つ以上のツールの使用を必要とするかもしれません。\n\n2. Claudeがツールを使用する（APIレスポンス）:\n   - Claudeはユーザープロンプトを評価し、利用可能なツールがユーザーのクエリやタスクに役立つかどうかを判断します。役立つ場合、どのツールを使用し、どの入力を使用するかも決定します。\n   - 適切にフォーマットされたツール使用リクエストを構築します。APIレスポンスには、Claudeが外部ツールを使用したいということを示すstop_reasonが含まれます。\n\n3. ツールの入力を抽出し、コードを実行し、結果を返す（APIリクエスト）:\n   - クライアント側で、Claudeのツール使用リクエストからツール名と入力を抽出します。\n   - クライアント側で実際のツールコードを実行します。\n   - ツール結果を返します。これは、新しいユーザーメッセージを含むツール結果コンテンツブロックを使用して会話を継続します。\n\n4. Claudeがツール結果を使用して応答を形成する（APIレスポンス）:\n   - ツール結果を受け取った後、Claudeはその情報を使用して元のユーザープロンプトに対する最終的な応答を形成します。\n\nこれらの手順に従うことで、Claude-3でのtool-useを効果的に活用することができます。", "additional_kwargs": {}}
{"role": "user", "content": "どのモデルがそれをサポートしていますか？", "additional_kwargs": {}}
{"role": "assistant", "content": "上記のコンテキストによると、Claude-3 Opusモデルが\"chain of thought tool use\"（思考の連鎖を示すツール使用）をサポートしています。この機能により、Claudeは問題を分解し、どのツールを使用するかを決定するためのステップバイステップの推論を表示します。SonnetモデルやHaikuモデルもこの機能を使用することができますが、デフォルトではOpusモデルほど頻繁に使用されません。", "additional_kwargs": {}}
{"role": "user", "content": "ありがとう、それを知りたかったんだ！", "additional_kwargs": {}}
{"role": "assistant", "content": "このコンテキストに基づいて、Claude-3のツール使用に関する情報を提供しました。Claude-3では、Opusモデルを使用して複雑なツール使用をナビゲートすることが推奨されています。Opusは最も多くの同時ツールを処理でき、他のモデルと比較して欠落している引数をキャッチする能力があります。また、Haikuモデルは、単純なツールを扱う場合に適しています。これらの情報を活用して、Claudeのツール使用を効果的に設計し、Claudeの機能を拡張することができます。", "additional_kwargs": {}}

!pip install llama-index-llms-openai llama-index-callbacks-arize-phoenix python-dotenv pygraphviz

import phoenix as px
import llama_index.core

px.launch_app()
llama_index.core.set_global_handler("arize_phoenix")

from dotenv import load_dotenv

load_dotenv(verbose=True)

!curl "https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip" -O ./chinook.zip
!unzip ./chinook.zip

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_engine import NLSQLTableQueryEngine
from llama_index.core.tools import QueryEngineTool
from llama_index.llms.openai import OpenAI

llm = OpenAI(model="gpt-3.5-turbo")

sql_query_engine = NLSQLTableQueryEngine(
    sql_database=sql_database,
    tables=["albums", "tracks", "artists"],
    llm=llm,
    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"
    ),
)

response = sql_query_engine.query("いくつのデータが存在しますか？")
print(response)

from llama_index.core.agent.react import ReActAgent

sample_agent = ReActAgent.from_tools([sql_tool], llm=llm, verbose=True)

response = sample_agent.query("いくつのデータが存在しますか？")
print(response)

Thought: The current language of the user is: Japanese. I need to use a tool to help me answer the question.
Action: sql_tool
Action Input: {'input': 'SELECT COUNT(*) FROM table_name'}
Observation: There are a total of 347 albums in the database.
Thought: I can answer without using any more tools. I'll use the user's language to answer
Answer: データベースには合計347枚のアルバムが存在します。

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:
        出力キーと値の辞書。
        このコンポーネントと他のコンポーネント間のリンクを定義する際に`src_key`を指定する場合は、
        src_keyと指定された 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


## ReActプロンプト関数の定義
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)
    # TODO: get output

    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
    # Now that we're done with this step, put into memory
    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)

# Reactプロンプトへの入力を、パースされたレスポンスにリンクする
qp.add_chain(["agent_input", "react_prompt", "llm", "react_output_parser"])

# ReActの出力からツール呼び出しを判断するためのリンクを追加
qp.add_link(
    "react_output_parser",
    "run_tool",
    condition_fn=lambda x: not x["done"],
    input_fn=lambda x: x["reasoning_step"],
)
# ReActの出力から最終回答を判断するためのリンクを追加
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 networkx.drawing.nx_agraph import to_agraph
from IPython.display import Image

agraph = to_agraph(qp.dag)
agraph.layout(prog="dot")
agraph.draw('agent.png')
display(Image('agent.png'))

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
)

agent.reset()
response = agent.chat(
    "アーティストAC/DCの楽曲を教えて下さい？3曲まで。"
)

agent.reset()

task = agent.create_task(
    "アーティストAC/DCの楽曲を教えて下さい？3曲まで。"
)

step_output = agent.run_step(task.task_id)
step_output.is_last

> Running step 34b5dca4-7445-40f7-9d64-e615ef074530. Step input: アーティストAC/DCの楽曲を教えて下さい？3曲まで。
> Running module agent_input with input: 
state: {'sources': [], 'memory': ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all'), chat_store=SimpleChatSto...
task: task_id='00668449-5f01-477a-a1de-61524cd3a324' input='アーティストAC/DCの楽曲を教えて下さい？3曲まで。' memory=ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding '...

> Running module react_prompt with input: 
input: アーティストAC/DCの楽曲を教えて下さい？3曲まで。

> Running module llm with input: 
messages: [ChatMessage(role=<MessageRole.SYSTEM: 'system'>, content='You are designed to help with a variety of tasks, from answering questions to providing summaries to other types of analyses.\n\n## Tools\n\n...

> Running module react_output_parser with input: 
chat_response: assistant: Thought: The current language of the user is: Japanese. I need to use a tool to help me answer the question.
Action: sql_tool
Action Input: {"input": "SELECT song_name FROM artists_songs WH...

> Running module run_tool with input: 
reasoning_step: thought='The current language of the user is: Japanese. I need to use a tool to help me answer the question.' action='sql_tool' action_input={'input': "SELECT song_name FROM artists_songs WHERE artist...

> Running module process_agent_response with input: 
response_dict: {'response_str': 'Observation: {\'output\': ToolOutput(content=\'Some of AC/DC\\\'s songs are "For Those About To Rock (We Salute You)", "Put The Finger On You", and "Let\\\'s Get It Up".\', tool_name...

False

response = agent.finalize_response(task.task_id)
print(str(response))

from llama_index.llms.openai import OpenAI

llm = OpenAI(model="gpt-3.5-turbo-0125")

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 = """\
あなたは、ユーザ入力から適切な自然言語クエリを生成しようとしています。

このクエリは、クエリをSQL文に変換する下流のtext-to-SQLエージェントによって解釈されます。
エージェントがエラーをトリガーした場合、現在の会話履歴に反映されます（下記参照）。

会話履歴がNoneの場合、ユーザ入力を使用します。
Noneでない場合は、前のSQLクエリの問題を回避した新しいSQLクエリを生成します。

入力: {input}
会話履歴 (失敗した試み): 
{convo_history}

新しい入力: """
retry_prompt = PromptTemplate(retry_prompt_str)

from llama_index.core import Response
from typing import Tuple

validate_prompt_str = """\
与えられたユーザークエリから、推論されたSQLクエリとクエリ実行によるレスポンスが正しいかどうか、そしてクエリに答えれているかかどうかを検証してください。

結果は YES または NO で回答してください。

クエリ: {input}
推論されたSQLクエリ: {sql_query}
SQLの結果: {sql_response}

結果: """
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 networkx.drawing.nx_agraph import to_agraph
from IPython.display import Image

agraph = to_agraph(qp.dag)
agraph.layout(prog="dot")
agraph.draw('agent_retry.png')
display(Image('agent_retry.png'))

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(
    "「Restless and Wild」を書いたアーティストは何枚のアルバムをリリースした？(答えはゼロではないはず）"
)
print(str(response))

> Running module input with input: 
state: {'sources': [], 'memory': ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all'), chat_store=SimpleChatSto...
task: task_id='f1d6032a-4963-46ca-97ca-31aa8678c9a3' input='「Restless and Wild」を書いたアーティストは何枚のアルバムをリリースした？(答えはゼロではないはず）' memory=ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method...

> Running module retry_prompt with input: 
input: 「Restless and Wild」を書いたアーティストは何枚のアルバムをリリースした？(答えはゼロではないはず）
convo_history: User: 「Restless and Wild」を書いたアーティストは何枚のアルバムをリリースした？(答えはゼロではないはず）

> Running module llm with input: 
messages: あなたは、ユーザ入力から適切な自然言語クエリを生成しようとしています。

このクエリは、クエリをSQL文に変換する下流のtext-to-SQLエージェントによって解釈されます。
エージェントがエラーをトリガーした場合、現在の会話履歴に反映されます（下記参照）。

会話履歴がNoneの場合、ユーザ入力を使用します。
Noneでない場合は、前のSQLクエリの問題を回避した新しいSQLクエリを生成します...

> Running module sql_query_engine with input: 
input: assistant: 「Restless and Wild」を書いたアーティストのアルバム数は何枚ですか？

> Running module output_component with input: 
output: I'm sorry, but there was an error in running the SQL query to find out how many albums the artist who wrote "Restless and Wild" has. Unfortunately, I am unable to provide you with the specific number ...

> Inferred SQL Query: SQL Query to run: 
SELECT artists.Name, COUNT(albums.AlbumId) AS AlbumCount
FROM artists
JOIN albums ON artists.ArtistId = albums.ArtbumId
WHERE albums.Title = "Restless and Wild"
GROUP BY artists.Name
> SQL Response: I'm sorry, but there was an error in running the SQL query to find out how many albums the artist who wrote "Restless and Wild" has. Unfortunately, I am unable to provide you with the specific number of albums at this time.
> Running module input with input: 
state: {'sources': [], 'memory': ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all'), chat_store=SimpleChatSto...
task: task_id='f1d6032a-4963-46ca-97ca-31aa8678c9a3' input='「Restless and Wild」を書いたアーティストは何枚のアルバムをリリースした？(答えはゼロではないはず）' memory=ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method...

> Running module retry_prompt with input: 
input: 「Restless and Wild」を書いたアーティストは何枚のアルバムをリリースした？(答えはゼロではないはず）
convo_history: User: 「Restless and Wild」を書いたアーティストは何枚のアルバムをリリースした？(答えはゼロではないはず）
Assistant (inferred SQL query): SQL Query to run: 
SELECT artists.Name, COUNT(albums.AlbumId) AS AlbumCount
FROM artists
JOIN albums ON...

> Running module llm with input: 
messages: あなたは、ユーザ入力から適切な自然言語クエリを生成しようとしています。

このクエリは、クエリをSQL文に変換する下流のtext-to-SQLエージェントによって解釈されます。
エージェントがエラーをトリガーした場合、現在の会話履歴に反映されます（下記参照）。

会話履歴がNoneの場合、ユーザ入力を使用します。
Noneでない場合は、前のSQLクエリの問題を回避した新しいSQLクエリを生成します...

> Running module sql_query_engine with input: 
input: assistant: 「Restless and Wild」を書いたアーティストは何枚のアルバムをリリースした？(答えはゼロではないはず）

> Running module output_component with input: 
output: 「Restless and Wild」を書いたアーティストはAcceptというバンドで、1枚のアルバムをリリースしています。

> Inferred SQL Query: SELECT artists.Name, COUNT(albums.AlbumId) AS AlbumCount
FROM artists
JOIN albums ON artists.ArtistId = albums.ArtistId
WHERE albums.Title = "Restless and Wild"
GROUP BY artists.Name
> SQL Response: 「Restless and Wild」を書いたアーティストはAcceptというバンドで、1枚のアルバムをリリースしています。
「Restless and Wild」を書いたアーティストはAcceptというバンドで、1枚のアルバムをリリースしています。