LlamaIndexモジュールガイドを試してみる: Querying

query_engine = index.as_query_engine()
response = query_engine.query("イクイノックスの主な勝ち鞍は？")
print(response)

from llama_index import get_response_synthesizer
from llama_index.retrievers import VectorIndexRetriever
from llama_index.query_engine import RetrieverQueryEngine
from llama_index.postprocessor import SimilarityPostprocessor

retriever = VectorIndexRetriever(
    index=index,
    similarity_top_k=10,
)

response_synthesizer = get_response_synthesizer(response_mode="refine")

query_engine = RetrieverQueryEngine(
    retriever=retriever,
    response_synthesizer=response_synthesizer,
    node_postprocessors=[SimilarityPostprocessor(similarity_cutoff=0.8)],
)

response = query_engine.query("イクイノックスの主なGI勝ち鞍は？")
print(response)

イクイノックスの主なGI勝ち鞍は、天皇賞（秋）、有馬記念、そしてドバイシーマクラシックです。

retriever = index.as_retriever(
    similarity_top_k=10,
)
nodes = retriever.retrieve("ドウデュースの主な勝ち鞍は?")

for n in nodes:
    print("Score: {}\nText: {}".format(n.get_score(), n.get_content()[:50].replace("\n","") + "..."))

Score: 0.8710878825048197
Text: 鞍上の武は、歴代最多を更新する2013年のキズナ以来となるダービー6勝目、かつ歴代最年長、史上初の5...
Score: 0.8586540345620574
Text: まずまずのスタートを決めると、道中は馬群を見て最後方グループで待機し、向正面から外へ出て徐々に捲りを...
Score: 0.8566146252006821
Text: 鞍上には武豊が復帰した。鞍上横山和生と本レースがラストランとなっており、レースを牽引して尚直線粘る...
Score: 0.8541785521180076
Text: さらに、ダービー馬による有馬記念制覇はオルフェーヴル以来10年ぶり9頭目で、三冠馬以外ではハクチカラ...
Score: 0.8489308913433957
Text: その後夏は治療と休養にあて、秋初戦として10月29日に東京競馬場で開催される天皇賞（秋）に出走。当日...
Score: 0.847351191459949
Text: また馬主である松島及びキーファーズにとっては初の単独所有馬によるGI勝利、並びに国内GI初制覇となっ...
Score: 0.8472923877327061
Text: ドウデュース（欧字名:Do Deuce、2019年5月7日 - ）は、日本の競走馬。主な勝ち鞍は20...
Score: 0.8425653266200219
Text: レースでは道中後方からじっくり運び、最後の直線は外からメンバー最速となる上がり (競馬)3ハロン33...
Score: 0.8407498004580377
Text: 1番人気に推されると、レースは直線でガイアフォースとの追い比べをクビ差制してデビュー勝ちを果たした。...
Score: 0.8374629379221759
Text: == 血統表 ==母ダストアンドダイヤモンズはアメリカで重賞2勝を挙げ、2012年のGI・ブリーダ...

from llama_index.retrievers import VectorIndexRetriever

retriever = VectorIndexRetriever(
    index=index,
    similarity_top_k=10,
)

nodes = retriever.retrieve("ドウデュースの主な勝ち鞍は?")

from llama_index.response_synthesizers import (
    ResponseMode,
    get_response_synthesizer,
)

response_synthesizer = get_response_synthesizer(
    response_mode=ResponseMode.COMPACT
)

response = response_synthesizer.synthesize(
    "ドウデュースの主な勝ち鞍は?",
    nodes=nodes
)

print(response)

朝日杯フューチュリティステークス、東京優駿、有馬記念

from llama_index.response_synthesizers import (
    ResponseMode,
    get_response_synthesizer,
)

response_synthesizer = get_response_synthesizer(
    response_mode=ResponseMode.COMPACT
)

query_engine = index.as_query_engine(response_synthesizer=response_synthesizer)
response = query_engine.query("ドウデュースの主な勝ち鞍は？")

from llama_index.response_synthesizers import (
    ResponseMode,
    get_response_synthesizer,
)

response_synthesizer = get_response_synthesizer(
    response_mode=ResponseMode.REFINE,     # refineモードで有効
    structured_answer_filtering=True,     # ここ
    verbose=True     # Trueにするとstructured_answer_filteringの結果が標準出力に出力される
)

response = response_synthesizer.synthesize(
    "ドウデュースの主な勝ち鞍を教えて。",
    nodes=nodes
)

print(response)

** Messages: **
system: You are an expert Q&A system that is trusted around the world.
Always answer the query using the provided context information, and not prior knowledge.
Some rules to follow:
1. Never directly reference the given context in your answer.
2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.
user: Context information is below.
---------------------
file_path: data/ドウデュース.txt

鞍上の武は、歴代最多を更新する2013年のキズナ以来となるダービー6勝目、かつ歴代最年長、史上初の50代でのダービージョッキーの名誉となった。また朝日杯フューチュリティステークスの勝ち馬が日本ダービーを制したのは、1994年のナリタブライアン（前身である朝日杯3歳ステークスを勝利）以来28年ぶりとなった。
6月10日に国際競馬統括機関連盟が発表した「ロンジンワールドベストレースホースランキング」において、ドウデュースは日本ダービーを勝利した功績を評価され、シャフリヤールやエンブレムロードと並ぶレーティング120で第15位タイに位置づけられた。
次走として凱旋門賞への出走を表明した。
---------------------
Given the context information and not prior knowledge, answer the query.
Query: ドウデュースの主な勝ち鞍を教えて。
Answer: 
**************************************************
** Response: **
assistant: None
**************************************************

Function call: StructuredRefineResponse with args: {
  "answer": "ドウデュースの主な勝ち鞍は、2013年のキズナ以来となるダービー6勝目、かつ歴代最年長、史上初の50代でのダービージョッキーの名誉となった。",
  "query_satisfied": true
}
> Refine context: file_path: data/ドウデュース.txt

** Messages: **
user: You are an expert Q&A system that strictly operates in two modes when refining existing answers:
1. **Rewrite** an original answer using the new context.
2. **Repeat** the original answer if the new context isn't useful.
Never reference the original answer or context directly in your answer.
When in doubt, just repeat the original answer.New Context: file_path: data/ドウデュース.txt

まずまずのスタートを決めると、道中は馬群を見て最後方グループで待機し、向正面から外へ出て徐々に捲りをかけ、直線に入ると早々と馬群から抜け出して後続を突き放し、3馬身半差の快勝でGIを含む重賞3勝目を挙げた。日本ダービーを勝利した馬が京都記念を勝利するのは、1948年春の京都記念のマツミドリ以来75年ぶりとなった。
次走は3月25日にドバイのメイダン競馬場で行われるドバイターフとし、同月15日（現地時間同月14日）にメイダン競馬場に到着した。しかし出馬投票後の同月24日、調教後に左前肢跛行を発症しドバイターフへの出走を取り消した。友道は「調教後に左腕節に違和感を認め、競馬に向けて進めておりましたが、将来のある馬なのでここでは無理をせず、取り消すことを決断いたしました」と語った。
Query: ドウデュースの主な勝ち鞍を教えて。
Original Answer: ドウデュースの主な勝ち鞍は、2013年のキズナ以来となるダービー6勝目、かつ歴代最年長、史上初の50代でのダービージョッキーの名誉となった。
New Answer: 

** Messages: **
user: You are an expert Q&A system that strictly operates in two modes when refining existing answers:
1. **Rewrite** an original answer using the new context.
2. **Repeat** the original answer if the new context isn't useful.
Never reference the original answer or context directly in your answer.
When in doubt, just repeat the original answer.New Context: file_path: data/ドウデュース.txt

鞍上には武豊が復帰した。
鞍上横山和生と本レースがラストランとなっており、レースを牽引して尚直線粘るタイトルホルダーと、二番手から先に仕掛けていた鞍上クリストフ・ルメールのスターズオンアースを抜き去り、先頭でゴール板を駆け抜けて見事復活の勝利を挙げた。武豊は2017年のキタサンブラック以来、6年ぶりの制覇となった。また、54歳9カ月10日での有馬記念勝利は最年長勝利記録であり、同時に自身が持つJRA・GI最年長勝利記録（54歳19日）を更新。ドリームジャーニーやオルフェーヴル、ブラストワンピースに騎乗し歴代最多の有馬記念4勝を挙げている池添謙一に並ぶ勝利数となった。馬番5番での優勝は、1970年のスピードシンボリ、1972年のイシノヒカルに次いで、51年ぶり3度目となった。
Query: ドウデュースの主な勝ち鞍を教えて。
Original Answer: ドウデュースの主な勝ち鞍は、GIを含む重賞3勝目です。
New Answer: 
**************************************************
** Response: **
assistant: None
**************************************************


Function call: StructuredRefineResponse with args: {
  "answer": "ドウデュースの主な勝ち鞍は、GIを含む重賞3勝目です。",
  "query_satisfied": false
}
> Refine context: file_path: data/ドウデュース.txt

from llama_index.response_synthesizers import TreeSummarize

tree_summarizer = TreeSummarize(verbose=True)
response = await tree_summarizer.aget_response("イクイノックスの主な勝ち鞍を教えて", [documents[0].text])

print(response)

3 text chunks after repacking

** Messages: **
system: You are an expert Q&A system that is trusted around the world.
Always answer the query using the provided context information, and not prior knowledge.
Some rules to follow:
1. Never directly reference the given context in your answer.
2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.
user: Context information from multiple sources is below.
---------------------
4月14日、ロンジンワールドベストレースホースランキング（2023年1月1日から4月9日までの世界の主要レースを対象）が発表され、イクイノックスは2着馬に3馬身1／2差をつけコースレコードで逃げ切ったドバイシーマクラシックの内容が高く評価され、レーティングは129ポンドで第1位となった。(snip)
---------------------
Given the information from multiple sources and not prior knowledge, answer the query.
Query: イクイノックスの主な勝ち鞍を教えて
Answer: 
**************************************************
** Response: **
assistant: イクイノックスの主な勝ち鞍は、ドバイシーマクラシック、宝塚記念、天皇賞（秋）、ジャパンカップです。
**************************************************

1 text chunks after repacking

** Messages: **
system: You are an expert Q&A system that is trusted around the world.
Always answer the query using the provided context information, and not prior knowledge.
Some rules to follow:
1. Never directly reference the given context in your answer.
2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.
user: Context information from multiple sources is below.
---------------------
イクイノックスの主な勝ち鞍は、2022年・2023年の天皇賞（秋）連覇、2022年の有馬記念、2023年のドバイシーマクラシック、宝塚記念、ジャパンカップです。

イクイノックスの主な勝ち鞍は、ドバイシーマクラシック、宝塚記念、天皇賞（秋）、ジャパンカップです。

イクイノックスの主な勝ち鞍は、ジャパンカップ、天皇賞（春）、宝塚記念、有馬記念、安田記念、マイルチャンピオンシップなどです。
---------------------
Given the information from multiple sources and not prior knowledge, answer the query.
Query: イクイノックスの主な勝ち鞍を教えて
Answer: 
**************************************************
** Response: **
assistant: イクイノックスの主な勝ち鞍は、2022年・2023年の天皇賞（秋）連覇、2022年の有馬記念、ドバイシーマクラシック、宝塚記念、ジャパンカップ、天皇賞（春）、有馬記念、安田記念、マイルチャンピオンシップなどです。
**************************************************

from llama_index.retrievers import VectorIndexRetriever

retriever = VectorIndexRetriever(
    index=index,
    similarity_top_k=10,
)

nodes = retriever.retrieve("ドウデュースの主な勝ち鞍は?")

for n in nodes:
    print("Score: {}\nText: {}".format(n.get_score(), n.get_content()[:50].replace("\n","") + "..."))

Score: 0.8711197622628892
Text: 鞍上の武は、歴代最多を更新する2013年のキズナ以来となるダービー6勝目、かつ歴代最年長、史上初の5...
Score: 0.8586644710637568
Text: まずまずのスタートを決めると、道中は馬群を見て最後方グループで待機し、向正面から外へ出て徐々に捲りを...
Score: 0.856530004414721
Text: 鞍上には武豊が復帰した。鞍上横山和生と本レースがラストランとなっており、レースを牽引して尚直線粘る...
Score: 0.8543041231544019
Text: さらに、ダービー馬による有馬記念制覇はオルフェーヴル以来10年ぶり9頭目で、三冠馬以外ではハクチカラ...
Score: 0.8489308913433957
Text: その後夏は治療と休養にあて、秋初戦として10月29日に東京競馬場で開催される天皇賞（秋）に出走。当日...
Score: 0.8472844128737264
Text: また馬主である松島及びキーファーズにとっては初の単独所有馬によるGI勝利、並びに国内GI初制覇となっ...
Score: 0.8472790958966984
Text: ドウデュース（欧字名:Do Deuce、2019年5月7日 - ）は、日本の競走馬。主な勝ち鞍は20...
Score: 0.8425653266200219
Text: レースでは道中後方からじっくり運び、最後の直線は外からメンバー最速となる上がり (競馬)3ハロン33...
Score: 0.8385146201009842
Text: 1番人気に推されると、レースは直線でガイアフォースとの追い比べをクビ差制してデビュー勝ちを果たした。...
Score: 0.8376079308869422
Text: == 血統表 ==母ダストアンドダイヤモンズはアメリカで重賞2勝を挙げ、2012年のGI・ブリーダ...

from llama_index.postprocessor import SimilarityPostprocessor

processor = SimilarityPostprocessor(similarity_cutoff=0.85)
filtered_nodes = processor.postprocess_nodes(nodes)

for n in filtered_nodes:
    print("Score: {}\nText: {}".format(n.get_score(), n.get_content()[:50].replace("\n","") + "..."))

Score: 0.8711197622628892
Text: 鞍上の武は、歴代最多を更新する2013年のキズナ以来となるダービー6勝目、かつ歴代最年長、史上初の5...
Score: 0.8586644710637568
Text: まずまずのスタートを決めると、道中は馬群を見て最後方グループで待機し、向正面から外へ出て徐々に捲りを...
Score: 0.856530004414721
Text: 鞍上には武豊が復帰した。鞍上横山和生と本レースがラストランとなっており、レースを牽引して尚直線粘る...
Score: 0.8543041231544019
Text: さらに、ダービー馬による有馬記念制覇はオルフェーヴル以来10年ぶり9頭目で、三冠馬以外ではハクチカラ...

!pip install cohere

from llama_index.postprocessor import CohereRerank

reranker = CohereRerank(api_key=os.environ["COHERE_API_KEY"], top_n=5)
reranked_nodes = reranker.postprocess_nodes(filtered_nodes, query_str="ドウデュースの主な勝ち鞍は?")

for n in reranked_nodes:
    print("Score: {}\nText: {}".format(n.get_score(), n.get_content()[:50].replace("\n","") + "..."))

Score: 0.97257143
Text: 鞍上の武は、歴代最多を更新する2013年のキズナ以来となるダービー6勝目、かつ歴代最年長、史上初の5...
Score: 0.03926875
Text: さらに、ダービー馬による有馬記念制覇はオルフェーヴル以来10年ぶり9頭目で、三冠馬以外ではハクチカラ...
Score: 0.03442355
Text: まずまずのスタートを決めると、道中は馬群を見て最後方グループで待機し、向正面から外へ出て徐々に捲りを...
Score: 0.027220903
Text: 鞍上には武豊が復帰した。鞍上横山和生と本レースがラストランとなっており、レースを牽引して尚直線粘る...

query_engine = index.as_query_engine(
    similarity_top_k=10,
    node_postprocessors=[
        SimilarityPostprocessor(similarity_cutoff=0.8),
        CohereRerank(api_key=os.environ["COHERE_API_KEY"], top_n=5),
    ]
)

response = query_engine.query("ドウデュースの主な勝ち鞍を教えて。")

from llama_index import VectorStoreIndex, SimpleDirectoryReader

documents = SimpleDirectoryReader("data").load_data()
index = VectorStoreIndex.from_documents(documents)

query_engine = index.as_query_engine()
response = query_engine.query("ドウデュースの主な勝ち鞍を教えて。")
print(response)

query_engine = index.as_query_engine(streaming=True)
streaming_response = query_engine.query("ドウデュースの主な勝ち鞍は?")
streaming_response.print_response_stream()

ドウデュースの主な勝ち鞍は、日本ダービー、朝日杯フューチュリティステークス、京都記念です。

query_engine = index.as_query_engine(streaming=True)
streaming_response = query_engine.query("ドウデュースの主な勝ち鞍は?")

for text in streaming_response.response_gen:
    print(text)
    pass

ド
ウ
デ
ュ
ース
の
主
な
勝
ち
鞍
は
、
日
本
ダ
ービ
ー
、
朝
日
杯
フ
ュ
ーチ
ュ
リ
テ
ィ
ス
テ
ー
ク
ス
、
京
都
記
念
です
。

chat_engine = index.as_chat_engine()
response = chat_engine.chat("ドウデュースの主な勝ち鞍は?")
print(response)

from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext
from llama_index.text_splitter import SentenceSplitter
from llama_index.embeddings import OpenAIEmbedding
from llama_index.llms import OpenAI
from llama_index import set_global_handler

set_global_handler("simple")

llm = OpenAI(model="gpt-3.5-turbo", temperarture=0)
embed_model = OpenAIEmbedding()
node_parser = SentenceSplitter(chunk_size=400, chunk_overlap=20)

service_context = ServiceContext.from_defaults(
  llm=llm,
  embed_model=embed_model,
  node_parser=node_parser,
)

documents = SimpleDirectoryReader("data").load_data()

index = VectorStoreIndex.from_documents(documents, service_context=service_context)

chat_engine = index.as_chat_engine()

response = chat_engine.chat("ドウデュースの主な勝ち鞍について教えて。")
print(f"### 回答出力 ###\n\n{response}")

** Messages: **
user: ドウデュースの主な勝ち鞍について教えて。
**************************************************
** Response: **
assistant: None
**************************************************


** Messages: **
system: You are an expert Q&A system that is trusted around the world.
Always answer the query using the provided context information, and not prior knowledge.
Some rules to follow:
1. Never directly reference the given context in your answer.
2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.
user: Context information is below.
---------------------
file_path: data/ドウデュース.txt

鞍上の武は、歴代最多を更新する2013年のキズナ以来となるダービー6勝目、かつ歴代最年長、史上初の50代でのダービージョッキーの名誉となった。また朝日杯フューチュリティステークスの勝ち馬が日本ダービーを制したのは、1994年のナリタブライアン（前身である朝日杯3歳ステークスを勝利）以来28年ぶりとなった。
6月10日に国際競馬統括機関連盟が発表した「ロンジンワールドベストレースホースランキング」において、ドウデュースは日本ダービーを勝利した功績を評価され、シャフリヤールやエンブレムロードと並ぶレーティング120で第15位タイに位置づけられた。
次走として凱旋門賞への出走を表明した。

file_path: data/ドウデュース.txt

まずまずのスタートを決めると、道中は馬群を見て最後方グループで待機し、向正面から外へ出て徐々に捲りをかけ、直線に入ると早々と馬群から抜け出して後続を突き放し、3馬身半差の快勝でGIを含む重賞3勝目を挙げた。日本ダービーを勝利した馬が京都記念を勝利するのは、1948年春の京都記念のマツミドリ以来75年ぶりとなった。
次走は3月25日にドバイのメイダン競馬場で行われるドバイターフとし、同月15日（現地時間同月14日）にメイダン競馬場に到着した。しかし出馬投票後の同月24日、調教後に左前肢跛行を発症しドバイターフへの出走を取り消した。友道は「調教後に左腕節に違和感を認め、競馬に向けて進めておりましたが、将来のある馬なのでここでは無理をせず、取り消すことを決断いたしました」と語った。
---------------------
Given the context information and not prior knowledge, answer the query.
Query: ドウデュースの主な勝ち鞍は何ですか？
Answer: 
**************************************************
** Response: **
assistant: ドウデュースの主な勝ち鞍は、日本ダービー、朝日杯フューチュリティステークス、京都記念です。
**************************************************


** Messages: **
user: ドウデュースの主な勝ち鞍について教えて。
assistant: None
tool: ドウデュースの主な勝ち鞍は、日本ダービー、朝日杯フューチュリティステークス、京都記念です。
**************************************************
** Response: **
assistant: ドウデュースの主な勝ち鞍は、日本ダービー、朝日杯フューチュリティステークス、京都記念です。
**************************************************


### 回答出力 ###

ドウデュースの主な勝ち鞍は、日本ダービー、朝日杯フューチュリティステークス、京都記念です。

response = chat_engine.chat("へぇ、すごいですね。")
print(f"### 回答出力 ###\n\n{response}")

** Messages: **
user: ドウデュースの主な勝ち鞍について教えて。
assistant: None
tool: ドウデュースの主な勝ち鞍は、日本ダービー、朝日杯フューチュリティステークス、京都記念です。
assistant: ドウデュースの主な勝ち鞍は、日本ダービー、朝日杯フューチュリティステークス、京都記念です。
user: へぇ、すごいですね。
**************************************************
** Response: **
assistant: はい、ドウデュースはこれらのレースで優れたパフォーマンスを見せました。彼の能力と才能は素晴らしいものであり、多くの人々に感銘を与えています。
**************************************************


### 回答出力 ###

はい、ドウデュースはこれらのレースで優れたパフォーマンスを見せました。彼の能力と才能は素晴らしいものであり、多くの人々に感銘を与えています。

chat_engine.reset()

chat_engine = index.as_chat_engine()

response = chat_engine.chat("レース内容について詳しく教えて")
print(f"### 回答出力 ###\n\n{response}")

** Messages: **
user: レース内容について詳しく教えて
**************************************************
** Response: **
assistant: もちろんです！どのレースの内容について詳しく知りたいですか？例えば、F1レースやオリンピックの陸上競技、自転車レースなど、具体的な競技やイベントを教えてください。
**************************************************


### 回答出力 ###

もちろんです！どのレースの内容について詳しく知りたいですか？例えば、F1レースやオリンピックの陸上競技、自転車レースなど、具体的な競技やイベントを教えてください。

chat_engine = index.as_chat_engine()

chat_engine.chat_repl()

from llama_index import set_global_handler

set_global_handler("simple")

chat_modes = ["best", "condense_question", "context", "condense_plus_context", "simple", "react", "openai"]
messages = ["ドウデュースの主な勝ち鞍について教えて。", "へえ、すごいですね", "そのうちGI勝利はいくつですか？"]

for chat_mode in chat_modes:
    print("### MODE: {}\n".format(chat_mode.upper()))
    
    chat_engine = index.as_chat_engine(chat_mode=chat_mode, verbose=True)
    chat_engine.reset()

    for i, message in enumerate(messages, start=1):
        print("#### {} 回目\n".format(i))

        response = chat_engine.chat(message)
        print("\n⇒ 回答: {}\n".format(response.response.replace("\n","")))

from llama_index.prompts import PromptTemplate, ChatPromptTemplate
from llama_index.llms import ChatMessage, MessageRole
from llama_index.chat_engine.condense_question import CondenseQuestionChatEngine
from llama_index.query_engine import RetrieverQueryEngine

custom_text_qa_system_prompt = ChatMessage(
    content="""\
あなたは世界中で信頼されているQAシステムです。
事前知識ではなく、常に提供されたコンテキスト情報を使用してクエリに回答してください
従うべきいくつかのルール:
1. 回答内で指定されたコンテキストを直接参照しないでください。
2. 「コンテキストに基づいて、...」や「コンテキスト情報は...」、またはそれに類するような記述は避けてください。""",
    role=MessageRole.SYSTEM,
)

custom_text_qa_prompt_template_messages = [
    custom_text_qa_system_prompt,
    ChatMessage(
        content="""コンテキスト情報は以下のとおりです。
---------------------
{context_str}
---------------------
事前知識ではなくコンテキスト情報を考慮して、クエリに答えます。
Query: {query_str}
Answer: """,
        role=MessageRole.USER,
    ),
]

custom_chat_text_qa_prompt = ChatPromptTemplate(message_templates=custom_text_qa_prompt_template_messages)

# Indexは事前に作成済みとする
query_engine = index.as_query_engine(
    text_qa_template=custom_chat_text_qa_prompt,
)

custom_condense_question_prompt = PromptTemplate(
    """\
以下のHumanとAssistantの会話履歴（Chat History）と、それに続く、人からのメッセージ（Follow Up Message）に基づいて,
メッセージを、会話に関連するすべての文脈をとらえた独立した質問（Standalone question）に書き換えてください。

<Chat History>
{chat_history}

<Follow Up Message>
{question}

<Standalone question>
"""
)

custom_chat_history = [
    ChatMessage(
        role=MessageRole.USER,
        content="こんにちは、今日は競馬についていろいろ語りましょう。",
    ),
    ChatMessage(role=MessageRole.ASSISTANT, content="わかりました。いいですね。では始めましょう。"),
]


chat_engine = CondenseQuestionChatEngine.from_defaults(
    query_engine=query_engine,
    condense_question_prompt=custom_condense_question_prompt,
    chat_history=custom_chat_history,
    verbose=False,
)

chat_engine.chat_repl()

streaming_response = chat_engine.stream_chat("ドウデュースの主な勝ち鞍は？")

for token in streaming_response.response_gen:
    print(token, end="\n")

ド
ウ
デ
ュ
ース
の
主
な
勝
ち
鞍
は
、
日
本
ダ
ービ
ー
、
京
都
記
念
、
朝
日
杯
フ
ュ
ーチ
ュ
リ
テ
ィ
ス
テ
ー
ク
ス
です
。

from llama_index.storage.chat_store import SimpleChatStore
from llama_index.memory import ChatMemoryBuffer

chat_store = SimpleChatStore()

chat_memory = ChatMemoryBuffer.from_defaults(
    token_limit=3000,
    chat_store=chat_store,
    chat_store_key="user1",
)

chat_engine = index.as_chat_engine(memory=chat_memory)
chat_engine.chat_repl()

===== Entering Chat REPL =====
Type "exit" to exit.

Human: ドウデュースの主な勝ち鞍を教えて。
Assistant: ドウデュースの主な勝ち鞍は、日本ダービー、朝日杯フューチュリティステークス、京都記念です。

Human: 朝日杯フューチュリティステークスで騎乗していた騎手は誰？
Assistant: ドウデュースの朝日杯フューチュリティステークスでの騎手は戸崎圭太騎手でした。

Human: 朝日杯フューチュリティステークスは何番人気だった？
Assistant: ドウデュースの朝日杯フューチュリティステークスの人気は3番でした。

Human: exit

chat_store.persist(persist_path="chat_store.json")

{
    "store": {
        "chat_history": [
            {
                "role": "user",
                "content": "ドウデュースの主な勝ち鞍を教えて。",
                "additional_kwargs": {}
            },
            {
                "role": "assistant",
                "content": null,
                "additional_kwargs": {
                    "tool_calls": [
                        {
                            "id": "call_mCPlIwTE1jGVQCWK1mjqUEN6",
                            "function": {
                                "arguments": {
                                    "input": "ドウデュースの主な勝ち鞍は何ですか？"
                                },
                                "name": "query_engine_tool"
                            },
                            "type": "function"
                        }
                    ]
                }
            },
            {
                "role": "tool",
                "content": "ドウデュースの主な勝ち鞍は、日本ダービー、朝日杯フューチュリティステークス、京都記念です。",
                "additional_kwargs": {
                    "name": "query_engine_tool",
                    "tool_call_id": "call_mCPlIwTE1jGVQCWK1mjqUEN6"
                }
            },
            {
                "role": "assistant",
                "content": "ドウデュースの主な勝ち鞍は、日本ダービー、朝日杯フューチュリティステークス、京都記念です。",
                "additional_kwargs": {}
            },
            {
                "role": "user",
                "content": "朝日杯フューチュリティステークスで騎乗していた騎手は誰？",
                "additional_kwargs": {}
            },
            {
                "role": "assistant",
                "content": null,
                "additional_kwargs": {
                    "tool_calls": [
                        {
                            "id": "call_VDasPIQ0K0kMs8nhwYVf3v6l",
                            "function": {
                                "arguments": {
                                    "input": "ドウデュースの朝日杯フューチュリティステークスでの騎手は誰でしたか？"
                                },
                                "name": "query_engine_tool"
                            },
                            "type": "function"
                        }
                    ]
                }
            },
            {
                "role": "tool",
                "content": "戸崎圭太",
                "additional_kwargs": {
                    "name": "query_engine_tool",
                    "tool_call_id": "call_VDasPIQ0K0kMs8nhwYVf3v6l"
                }
            },
            {
                "role": "assistant",
                "content": "ドウデュースの朝日杯フューチュリティステークスでの騎手は戸崎圭太騎手でした。",
                "additional_kwargs": {}
            },
            {
                "role": "user",
                "content": "朝日杯フューチュリティステークスは何番人気だった？",
                "additional_kwargs": {}
            },
            {
                "role": "assistant",
                "content": null,
                "additional_kwargs": {
                    "tool_calls": [
                        {
                            "id": "call_te4UM7uSnxiplZ4guU3JIOEO",
                            "function": {
                                "arguments": {
                                    "input": "ドウデュースの朝日杯フューチュリティステークスの人気は何番でしたか？"
                                },
                                "name": "query_engine_tool"
                            },
                            "type": "function"
                        }
                    ]
                }
            },
            {
                "role": "tool",
                "content": "ドウデュースの朝日杯フューチュリティステークスの人気は3番でした。",
                "additional_kwargs": {
                    "name": "query_engine_tool",
                    "tool_call_id": "call_te4UM7uSnxiplZ4guU3JIOEO"
                }
            },
            {
                "role": "assistant",
                "content": "ドウデュースの朝日杯フューチュリティステークスの人気は3番でした。",
                "additional_kwargs": {}
            }
        ]
    },
    "class_name": "SimpleChatStore"
}

chat_engine.reset()

loaded_chat_store = SimpleChatStore.from_persist_path(
    persist_path="chat_store.json"
)

chat_memory = ChatMemoryBuffer.from_defaults(
    token_limit=3000,
    chat_store=loaded_chat_store,
    chat_store_key="user1",
)

chat_engine = index.as_chat_engine(memory=chat_memory)

for i in chat_engine.memory.get_all():
    print(i)

user: ドウデュースの主な勝ち鞍を教えて。
assistant: None
tool: ドウデュースの主な勝ち鞍は、日本ダービー、朝日杯フューチュリティステークス、京都記念です。
assistant: ドウデュースの主な勝ち鞍は、日本ダービー、朝日杯フューチュリティステークス、京都記念です。
user: 朝日杯フューチュリティステークスで騎乗していた騎手は誰？
assistant: None
tool: 戸崎圭太
assistant: ドウデュースの朝日杯フューチュリティステークスでの騎手は戸崎圭太騎手でした。
user: 朝日杯フューチュリティステークスは何番人気だった？
assistant: None
tool: ドウデュースの朝日杯フューチュリティステークスの人気は3番でした。
assistant: ドウデュースの朝日杯フューチュリティステークスの人気は3番でした。

response = chat_engine.chat("へぇ、すごいですね")
print(response)

はい、ドウデュースは朝日杯フューチュリティステークスでの活躍が素晴らしかったです。人気のある競走で好成績を収めることは、競走馬にとって大きな成果です。ドウデュースの才能と騎手の技術が光りましたね。

from llama_index.tools import FunctionTool
from llama_index.llms import OpenAI
from llama_index.agent import ReActAgent

# ツールで使用する関数の定義
def multiply(a: int, b: int) -> int:
    """Multiple two integers and returns the result integer"""
    return a * b

# 関数をツールとして登録
multiply_tool = FunctionTool.from_defaults(fn=multiply)

# LLM初期化
llm = OpenAI(model="gpt-3.5-turbo")

# ReAct エージェントを初期化、上記のツールを渡す
agent = ReActAgent.from_tools([multiply_tool], llm=llm, verbose=True)

response = agent.query("2123 かける 215123 は？")
print("回答: ", response)

Thought: I need to use the multiply tool to calculate the product of 2123 and 215123.
Action: multiply
Action Input: {'a': 2123, 'b': 215123}
Observation: 456706129
Thought: I have the result of multiplying 2123 by 215123.
Response: The product of 2123 and 215123 is 456706129.
回答:  The product of 2123 and 215123 is 456706129.

response = agent.query("2123 たす 215123 は？")
print("回答: ", response)

Thought: The user wants to add 2123 and 215123.
Action: multiply
Action Input: {'a': 2123, 'b': 215123}
Observation: 456706129
Thought: I have the result of adding 2123 and 215123.
Response: The sum of 2123 and 215123 is 217246.
回答:  The sum of 2123 and 215123 is 217246.

from llama_index.tools import FunctionTool
from llama_index.llms import OpenAI
from llama_index.agent import ReActAgent

def multiply(a: int, b: int) -> int:
    """Multiple two integers and returns the result integer"""
    return a * b

multiply_tool = FunctionTool.from_defaults(fn=multiply)

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

agent = ReActAgent.from_tools([multiply_tool], llm=llm, verbose=True)

response = agent.query("2123 ひく 215123 は？")
print("回答: ", response)

multiply_tool.call(2,4)

ToolOutput(content='8', tool_name='multiply', raw_input={'args': (2, 4), 'kwargs': {}}, raw_output=8)

multiply_tool.metadata

ToolMetadata(description='multiply(a: int, b: int) -> int\nMultiple two integers and returns the result integer.', name='multiply', fn_schema=<class 'pydantic.main.multiply'>)

multiply_tool.metadata.to_openai_tool()

{
  'type': 'function',
  'function': {
    'name': 'multiply',
    'description': 'multiply(a: int, b: int) -> int\nMultiple two integers and returns the result integer.',
    'parameters': {
      'title': 'multiply',
      'type': 'object',
      'properties': {
        'a': {'title': 'A', 'type': 'integer'},
        'b': {'title': 'B', 'type': 'integer'}
      },
      'required': ['a', 'b']
    }
  }
}

from llama_index import VectorStoreIndex, ServiceContext
from llama_index.readers.file.flat_reader import FlatReader

from llama_index.text_splitter import SentenceSplitter
from llama_index.embeddings import OpenAIEmbedding
from llama_index.llms import OpenAI
from llama_index import set_global_handler

llm = OpenAI(model="gpt-3.5-turbo", temperarture=0)
embed_model = OpenAIEmbedding()
node_parser = SentenceSplitter(chunk_size=400, chunk_overlap=20)

service_context = ServiceContext.from_defaults(
  llm=llm,
  embed_model=embed_model,
  node_parser=node_parser,
)

reader = FlatReader()

docs_d = reader.load_data(Path("data/ドウデュース.txt"))
docs_e = reader.load_data(Path("data/イクイノックス.txt"))

index_d = VectorStoreIndex.from_documents(docs_d, service_context=service_context)
index_e = VectorStoreIndex.from_documents(docs_e, service_context=service_context)

query_engine_d = index_d.as_query_engine(similarity_top_k=5)
query_engine_e = index_e.as_query_engine(similarity_top_k=5)

from llama_index.agent import ReActAgent
from llama_index.tools import QueryEngineTool, ToolMetadata
from llama_index.llms import OpenAI

query_engine_tools = [
    QueryEngineTool(
        query_engine=query_engine_d,
        metadata=ToolMetadata(
            name="query_engine_doduece",
            description="Provides information in Japanese about a race horse named 'Do Duece(ドウデュース)'"
            "Use a detailed plain text question as input to the tool.",
        ),
    ),
    QueryEngineTool(
        query_engine=query_engine_e,
        metadata=ToolMetadata(
            name="query_engine_equinox",
            description="Provides information in Japanese about a race horse named 'Equinox(イクイノックス)'"
            "Use a detailed plain text question as input to the tool.",
        ),
    ),
]

llm = OpenAI(model="gpt-3.5-turbo", temperature=0)

agent = ReActAgent.from_tools(query_engine_tools, llm=llm, verbose=True)

response = agent.query("ドウデュースの主な勝ち鞍は？")

Thought: I need to use a tool to help me answer the question.
Action: query_engine_doduece
Action Input: {'input': 'ドウデュースの主な勝ち鞍は？'}
Observation: ドウデュースの主な勝ち鞍は、日本ダービー、有馬記念、京都記念などです。
Thought: I can answer without using any more tools.
Response: ドウデュースの主な勝ち鞍は、日本ダービー、有馬記念、京都記念などです。

response = agent.query("イクイノックスの主な勝ち鞍は？")

Thought: I need to use a tool to help me answer the question.
Action: query_engine_equinox
Action Input: {'input': 'イクイノックスの主な勝ち鞍は？'}
Observation: イクイノックスの主な勝ち鞍は、天皇賞（秋）と有馬記念です。
Thought: I can answer without using any more tools.
Response: イクイノックスの主な勝ち鞍は、天皇賞（秋）と有馬記念です。

from llama_index.agent import ReActAgent
from llama_index.tools import FunctionTool, QueryEngineTool, ToolMetadata
from llama_index.llms import OpenAI


def multiply(a: int, b: int) -> int:
    """Multiple two integers and returns the result integer"""
    return a * b


multiply_tool = FunctionTool.from_defaults(fn=multiply)

llm = OpenAI(model="gpt-3.5-turbo", temperature=0)

multiply_agent = ReActAgent.from_tools([multiply_tool], llm=llm, verbose=True)

tools = [
    QueryEngineTool(
        query_engine=query_engine_d,
        metadata=ToolMetadata(
            name="query_engine_doduece",
            description="Provides information in Japanese about a race horse named 'Do Duece(ドウデュース)'"
            "Use a detailed plain text question as input to the tool.",
        ),
    ),
    QueryEngineTool(
        query_engine=query_engine_e,
        metadata=ToolMetadata(
            name="query_engine_equinox",
            description="Provides information in Japanese about a race horse named 'Equinox(イクイノックス)'"
            "Use a detailed plain text question as input to the tool.",
        ),
    ),
    QueryEngineTool(
        query_engine=multiply_agent,
        metadata=ToolMetadata(
            name="multiply_agent",
            description="Multiple two integers and returns the result integer."
        ),
    ),]


agent = ReActAgent.from_tools(tools, llm=llm, verbose=True)

response = agent.query("ドウデュースの主な勝ち鞍は？")

Thought: I need to use a tool to help me answer the question.
Action: query_engine_doduece
Action Input: {'input': 'ドウデュースの主な勝ち鞍は？'}
Observation: ドウデュースの主な勝ち鞍は、日本ダービー、有馬記念、京都記念などです。
Thought: The query_engine_doduece tool provided the following information about Do Duece's major victories: Japanese Derby, Arima Kinen, Kyoto Kinen, and more.
Response: ドウデュースの主な勝ち鞍は、日本ダービー、有馬記念、京都記念などです。

response = agent.query("イクイノックスの主な勝ち鞍は？")

Thought: I need to use a tool to help me answer the question.
Action: query_engine_equinox
Action Input: {'input': 'イクイノックスの主な勝ち鞍は？'}
Observation: イクイノックスの主な勝ち鞍は、天皇賞（秋）と有馬記念です。
Thought: I can answer without using any more tools.
Response: イクイノックスの主な勝ち鞍は、天皇賞（秋）と有馬記念です。

response = agent.query("2156 かける 256 は？")

Thought: I can use the multiply_agent tool to multiply 2156 by 256.
Action: multiply_agent
Action Input: {'input': '2156, 256'}
Thought: I need to multiply the numbers 2156 and 256.
Action: multiply
Action Input: {'a': 2156, 'b': 256}
Observation: 551936
Thought: I have multiplied 2156 and 256 and the result is 551936.
Response: The product of 2156 and 256 is 551936.
Observation: The product of 2156 and 256 is 551936.
Thought: I can answer without using any more tools.
Response: The product of 2156 and 256 is 551936.

from llama_index.tools.tool_spec.base import BaseToolSpec
from typing import Optional
import urllib

# Wolfram Alpha Short Answers APIのエンドポイントを定義
QUERY_URL_TMPL = "http://api.wolframalpha.com/v1/result?appid={app_id}&i={query}"

# BaseToolSpecを継承して、Wolfram Alpha用のToolSpecを定義
class WolframAlphaToolSpec(BaseToolSpec):

    # LLMに公開する関数名を定義
    spec_functions = ["wolfram_alpha_query"]

    # コンストラクタ。Wolfram AlphaのAPIキーで初期化する。
    def __init__(self, app_id: Optional[str] = None) -> None:
        """Initialize with parameters."""
        self.token = app_id
  
    # エージェントから呼ばれる関数の定義。多分docstringは英語の方がいいとは思う。
    def wolfram_alpha_query(self, query: str) -> str:
        """
        数学的・科学的な問題について、Wolfram Alphaにクエリを送信する。
        
        入力例:
            "(7 * 12 ^ 10) / 321"
            "1ポンドのいちごのカロリー"
  
        引数:
            query (str): Wolfram Alphaにに渡されるクエリ文字列
  
        """
        response = requests.get(QUERY_URL_TMPL.format(app_id=self.token, query=urllib.parse.quote_plus(query)))
        return response.text

# Tool SpecからToolインスタンス初期化
wolfram_spec = WolframAlphaToolSpec(app_id="XXXXXXXXXX")

# Tool SpecをToolのリストに変換。今回は1つだけ。
tools = wolfram_spec.to_tool_list()

# ToolをOpenAI Function Callingの定義で出力して確認
for i in tools:
    print(i.metadata.to_openai_tool())

{
  'type': 'function',
  'function': {
    'name': 'wolfram_alpha_query',
    'description': 'wolfram_alpha_query(query: str) -> str\n\n        数学的・科学的な問題について、Wolfram Alphaにクエリを送信する。\n\n        入力例:\n            "(7 * 12 ^ 10) / 321"\n            "1ポンドのいちごのカロリー"\n\n        引数:\n            query (str): Wolfram Alphaにに渡されるクエリ文字列\n\n        ',
    'parameters': {
      'title': 'wolfram_alpha_query',
      'type': 'object',
      'properties': {
        'query': {
          'title': 'Query',
          'type': 'string'
        }
      },
      'required': ['query']
    }
  }
}

from llama_index.agent import ReActAgent
from llama_index.llms import OpenAI

agent = ReActAgent.from_tools(tools, llm=llm, verbose=True)

print(agent.query('(7 * 12 ^ 10) / 321 は？'))

Thought: The user wants to know the result of the expression (7 * 12 ^ 10) / 321.
Action: wolfram_alpha_query
Action Input: {'query': '(7 * 12 ^ 10) / 321'}
Observation: 144473849856/107
Thought: I have the result of the expression.
Response: The result of (7 * 12 ^ 10) / 321 is approximately 1,350,000,000.
The result of (7 * 12 ^ 10) / 321 is approximately 1,350,000,000.

print(agent.query("1ポンドのいちごのカロリー"))

Thought: The user is asking for the calorie content of 1 pound of strawberries. I can use the Wolfram Alpha tool to find this information.
Action: wolfram_alpha_query
Action Input: {'query': '1 pound of strawberries calorie'}
Observation: about 145 dietary Calories
Thought: I have obtained the answer to the user's question.
Response: The calorie content of 1 pound of strawberries is approximately 145 dietary Calories.
The calorie content of 1 pound of strawberries is approximately 145 dietary Calories.

!pip install wikipedia

from llama_hub.wikipedia.base import WikipediaReader
from llama_index.tools.ondemand_loader_tool import OnDemandLoaderTool
from llama_index.llms import OpenAI
from llama_index.agent import OpenAIAgent

reader = WikipediaReader()

tool = OnDemandLoaderTool.from_defaults(
    reader,
    name="Wikipedia_Search",
    description="A tool for loading data and querying articles from Wikipedia",
)

llm = OpenAI(model="gpt-3.5-turbo", temperature=0)

agent = OpenAIAgent.from_tools([tool], llm=llm, verbose=True)

print(agent.query("日本の競走馬、ドウデュースの主な勝ち鞍を教えて"))

Added user message to memory: 日本の競走馬、ドウデュースの主な勝ち鞍を教えて
=== Calling Function ===
Calling function: Wikipedia_Search with args: {
  "pages": ["ドウデュース"]
}
Got output: Error: Missing query_str in kwargs with parameter name: query_str
========================

=== Calling Function ===
Calling function: Wikipedia_Search with args: {
  "pages": ["ドウデュース"],
  "query_str": "競走成績"
}
Got output: Do Deuce had a successful racing career. As a two-year-old, he won three races, including the Ivy Stakes and the Asahi Hai Futurity Stakes. In his three-year-old season, he finished second in the Grade 2 Yayoi Sho and third in the Satsuki Sho before winning the Tokyo Yushun. As a four-year-old, he won the Kyoto Kinen and the Arima Kinen. He also competed in races in France, finishing fourth in the Prix Niel and 19th in the Prix de l'Arc de Triomphe. Overall, Do Deuce had a strong racing record.
========================

ドウデュースは、競走馬として成功したキャリアを持っています。2歳の時には、アイビーステークスや朝日杯フューチュリティステークスなど3つのレースに勝利しました。3歳のシーズンでは、Grade 2の弥生賞で2着、皐月賞で3着となり、東京優駿を制しました。4歳の時には、京都記念や有馬記念などを制覇しました。また、フランスのレースにも参戦し、プリ・ニエルで4着、凱旋門賞で19着となりました。総じて、ドウデュースは強力な競走成績を持っています。

from llama_hub.tools.wikipedia.base import WikipediaToolSpec
from llama_index.tools.tool_spec.load_and_search.base import LoadAndSearchToolSpec
from llama_index.agent import OpenAIAgent

# Wikipedia用tool specからtool specインスタンスを初期化
wiki_spec = WikipediaToolSpec()

# tool specインスタンスからtoolを取り出し。loadとsearchのうちsearchだけを取り出しているように見える。
tool = wiki_spec.to_tool_list()[1]

# LoadAndSearchToolSpecでtoolをリスト化して、エージェントを作成
agent = OpenAIAgent.from_tools(
    LoadAndSearchToolSpec.from_defaults(tool).to_tool_list(), verbose=True
)

agent.query("ドウデュースの主な勝ち鞍を教えて")

Added user message to memory: ドウデュースの主な勝ち鞍を教えて
=== Calling Function ===
Calling function: search_data with args: {
  "query": "ドウデュース"
}
Got output: Error: 'str' object has no attribute 'get_doc_id'
========================

=== Calling Function ===
Calling function: search_data with args: {
  "query": "ドウデュース",
  "lang": "ja"
}
Got output: Error: 'str' object has no attribute 'get_doc_id'
========================

(snip)

=== Calling Function ===
Calling function: search_data with args: {
  "query": "ドウデュース",
  "lang": "ja"
}
Got output: Error: 'str' object has no attribute 'get_doc_id'
========================

Response(response='None', source_nodes=[], metadata=None)

tools = wiki_spec.to_tool_list()

for i in tools:
    print(i.metadata.to_openai_tool())

{
  'type': 'function',
  'function': {
    'name': 'load_data',
    'description': "load_data(page: str, lang: str = 'en', **load_kwargs: Dict[str, Any]) -> str\n\n        Retrieve a Wikipedia page. Useful for learning about a particular concept that isn't private information.\n\n        Args:\n            page (str): Title of the page to read.\n            lang (str): Language of Wikipedia to read. (default: English)\n        ",
    'parameters': {
      'title': 'load_data',
      'type': 'object',
      'properties': {
        'page': {'title': 'Page','type': 'string'},
        'lang': {'title': 'Lang', 'default': 'en', 'type': 'string'},
        'load_kwargs': {'title': 'Load Kwargs', 'type': 'object'}
      },
      'required': ['page', 'load_kwargs']
    }
  }
}
{
  'type': 'function',
  'function': {
    'name': 'search_data',
    'description': "search_data(query: str, lang: str = 'en') -> str\n\n        Search Wikipedia for a page related to the given query.\n        Use this tool when `load_data` returns no results.\n\n        Args:\n            query (str): the string to search for\n        ",
    'parameters': {
      'title': 'search_data',
      'type': 'object',
      'properties': {
        'query': {'title': 'Query', 'type': 'string'},
        'lang': {'title': 'Lang', 'default': 'en', 'type': 'string'}
      },
      'required': ['query']
    }
  }
}

from llama_hub.tools.wikipedia.base import WikipediaToolSpec
from llama_index.agent import OpenAIAgent

wiki_spec = WikipediaToolSpec()
tools = wiki_spec.to_tool_list()

agent = OpenAIAgent.from_tools(tools, verbose=True)

agent.query("日本の競走馬、ドウデュースの主な勝ち鞍を教えて")

Added user message to memory: 日本の競走馬、ドウデュースの主な勝ち鞍を教えて
=== Calling Function ===
Calling function: search_data with args: {
  "query": "ドウデュース"
}
Got output: Do Deuce (Japanese: ドウデュース, foaled 7 May 2019) is a Japanese Thoroughbred racehorse. He was the best two-year-old colt in Japan in 2021 when he was undefeated in three races including the Ivy Stakes and the Asahi Hai Futurity Stakes. (snip)

(snip)

== Pedigree ==
Do Deuce in inbred 4 × 4 to Lyphard, meaning that this stallion appears twice in the fourth generation of his pedigree.


== References ==
========================
Response(response='ドウデュースは、日本の競走馬であり、2019年に生まれました。彼は2021年に日本の最優秀2歳牡馬となり、アイビーステークスや朝日杯フューチュリティステークスなど3つのレースで無敗の成績を収めました。彼の2年目のキャンペーンのハイライトは、東京優駿でエクイノックスを破って優勝したことです。2023年の4歳シーズンには、2月の京都記念と年末の有馬記念を制しました。\n\n参考文献: Wikipedia', source_nodes=[], metadata=None)

from pydantic import BaseModel
from typing import List

class Song(BaseModel):
    """Data model for a song."""

    title: str
    length_seconds: int


class Album(BaseModel):
    """Data model for an album."""

    name: str
    artist: str
    songs: List[Song]

from llama_index.program import LLMTextCompletionProgram

prompt_template_str = """\
アーティストと曲のリストを含む、アルバムの例を生成してください。\
映画 {movie_name} をヒントとして使ってください。
"""

program = LLMTextCompletionProgram.from_defaults(
    output_cls=Album,
    prompt_template_str=prompt_template_str,
    verbose=True,
)

output = program(movie_name="侍戦隊シンケンジャー 銀幕版 天下分け目の戦")
output

Album(name='侍戦隊シンケンジャー 銀幕版 天下分け目の戦', artist='Various Artists', songs=[Song(title='侍戦隊シンケンジャー 銀幕版 天下分け目の戦', length_seconds=240), Song(title='勇者たちの誓い', length_seconds=210), Song(title='戦国の風', length_seconds=180), Song(title='魂の剣', length_seconds=195), Song(title='炎の心', length_seconds=220)])

from llama_index.output_parsers import PydanticOutputParser
from llama_index.program import LLMTextCompletionProgram

program = LLMTextCompletionProgram.from_defaults(
    output_parser=PydanticOutputParser(output_cls=Album),
    prompt_template_str=prompt_template_str,
    verbose=True,
)

output = program(movie_name="侍戦隊シンケンジャー 銀幕版 天下分け目の戦")
output

user: アーティストと曲のリストを含む、アルバムの例を生成してください。映画 侍戦隊シンケンジャー 銀幕版 天下分け目の戦 をヒントとして使ってください。



Here's a JSON schema to follow:
{{"title": "Album", "description": "Data model for an album.", "type": "object", "properties": {{"name": {{"title": "Name", "type": "string"}}, "artist": {{"title": "Artist", "type": "string"}}, "songs": {{"title": "Songs", "type": "array", "items": {{"$ref": "#/definitions/Song"}}}}}}, "required": ["name", "artist", "songs"], "definitions": {{"Song": {{"title": "Song", "description": "Data model for a song.", "type": "object", "properties": {{"title": {{"title": "Title", "type": "string"}}, "length_seconds": {{"title": "Length Seconds", "type": "integer"}}}}, "required": ["title", "length_seconds"]}}}}}}

Output a valid JSON object but do not repeat the schema.

assistant: {
  "name": "侍戦隊シンケンジャー 銀幕版 天下分け目の戦",
  "artist": "Various Artists",
  "songs": [
    {
      "title": "侍戦隊シンケンジャー 銀幕版 天下分け目の戦",
      "length_seconds": 240
    },
    {
      "title": "勇者たちの誓い",
      "length_seconds": 210
    },
    {
      "title": "戦国の風",
      "length_seconds": 180
    },
    {
      "title": "魂の剣",
      "length_seconds": 195
    },
    {
      "title": "炎の心",
      "length_seconds": 220
    }
  ]
}

from llama_index.types import BaseOutputParser
from pydantic import BaseModel
from typing import List

class Song(BaseModel):
    """Data model for a song."""

    title: str
    length_seconds: int


class Album(BaseModel):
    """Data model for an album."""

    name: str
    artist: str
    songs: List[Song]


class CustomAlbumOutputParser(BaseOutputParser):
    """Custom Album output parser.

    Assume first line is name and artist.

    Assume each subsequent line is the song.

    """

    def __init__(self, verbose: bool = False):
        self.verbose = verbose

    def parse(self, output: str) -> Album:
        """Parse output."""
        if self.verbose:
            print(f"> Raw output: {output}")
        lines = output.split("\n")
        name, artist = lines[0].split(",")
        songs = []
        for i in range(1, len(lines)):
            title, length_seconds = lines[i].split(",")
            songs.append(Song(title=title, length_seconds=length_seconds))

        return Album(name=name, artist=artist, songs=songs)

prompt_template_str = """\
アーティストと曲のリストを含む、アルバムの例を生成してください。\
映画 {movie_name} をヒントとして使ってください。\

以下のフォーマットで回答を返す。
最初の行: <album_name>,<album_artist> 
2行目以降: <song_title>,<song_length_seconds>
"""

program = LLMTextCompletionProgram.from_defaults(
    output_parser=CustomAlbumOutputParser(verbose=True),
    output_cls=Album,
    prompt_template_str=prompt_template_str,
    verbose=True,
)

output = program(movie_name="侍戦隊シンケンジャー 銀幕版 天下分け目の戦")
output

> Raw output: 侍戦隊シンケンジャー 銀幕版 天下分け目の戦,侍戦隊シンケンジャー
天下分け目の戦,240
勇者の証,210
炎神戦隊ゴーオンジャー BGM,180
炎神戦隊ゴーオンジャー,200
炎神戦隊ゴーオンジャー エンディングテーマ,190
炎神戦隊ゴーオンジャー インストゥルメンタル,180
炎神戦隊ゴーオンジャー インストゥルメンタル2,170
炎神戦隊ゴーオンジャー インストゥルメンタル3,160
炎神戦隊ゴーオンジャー インストゥルメンタル4,150
炎神戦隊ゴーオンジャー インストゥルメンタル5,140

Album(name='侍戦隊シンケンジャー 銀幕版 天下分け目の戦', artist='侍戦隊シンケンジャー', songs=[Song(title='天下分け目の戦', length_seconds=240), Song(title='勇者の証', length_seconds=210), Song(title='炎神戦隊ゴーオンジャー BGM', length_seconds=180), Song(title='炎神戦隊ゴーオンジャー', length_seconds=200), Song(title='炎神戦隊ゴーオンジャー エンディングテーマ', length_seconds=190), Song(title='炎神戦隊ゴーオンジャー インストゥルメンタル', length_seconds=180), Song(title='炎神戦隊ゴーオンジャー インストゥルメンタル2', length_seconds=170), Song(title='炎神戦隊ゴーオンジャー インストゥルメンタル3', length_seconds=160), Song(title='炎神戦隊ゴーオンジャー インストゥルメンタル4', length_seconds=150), Song(title='炎神戦隊ゴーオンジャー インストゥルメンタル5', length_seconds=140)])

from pydantic import BaseModel
from typing import List
from llama_index.program import OpenAIPydanticProgram


class Song(BaseModel):
    title: str
    length_seconds: int


class Album(BaseModel):
    name: str
    artist: str
    songs: List[Song]


prompt_template_str = """\
アーティストと曲のリストを含む、アルバムの例を生成してください。\
映画 {movie_name} をヒントとして使ってください。
"""
program = OpenAIPydanticProgram.from_defaults(
    output_cls=Album,
    prompt_template_str=prompt_template_str,
    verbose=True
)

output = program(
    movie_name="仮面ライダーW FOREVER AtoZ/運命のガイアメモリ",
    description="アルバムのデータモデル"
)

Function call: Album with args: {
  "name": "映画 仮面ライダーW FOREVER AtoZ/運命のガイアメモリ",
  "artist": "Various Artists",
  "songs": [
    {
      "title": "W-B-X 〜W-Boiled Extreme〜",
      "length_seconds": 251
    },
    {
      "title": "Finger on the Trigger",
      "length_seconds": 252
    },
    {
      "title": "W-B-X 〜W-Boiled Extreme〜 (TVサイズ)",
      "length_seconds": 92
    },
    {
      "title": "Finger on the Trigger (TVサイズ)",
      "length_seconds": 92
    },
    {
      "title": "W-B-X 〜W-Boiled Extreme〜 (Instrumental)",
      "length_seconds": 251
    },
    {
      "title": "Finger on the Trigger (Instrumental)",
      "length_seconds": 252
    }
  ]
}

output

Album(name='映画 仮面ライダーW FOREVER AtoZ/運命のガイアメモリ', artist='Various Artists', songs=[Song(title='W-B-X 〜W-Boiled Extreme〜', length_seconds=251), Song(title='Finger on the Trigger', length_seconds=252), Song(title='W-B-X 〜W-Boiled Extreme〜 (TVサイズ)', length_seconds=92), Song(title='Finger on the Trigger (TVサイズ)', length_seconds=92), Song(title='W-B-X 〜W-Boiled Extreme〜 (Instrumental)', length_seconds=251), Song(title='Finger on the Trigger (Instrumental)', length_seconds=252)])

from pydantic import BaseModel
from typing import List
from llama_index.program import OpenAIPydanticProgram

class Song(BaseModel):
    """
    曲のデータモデル
    """
    title: str
    length_seconds: int


class Album(BaseModel):
    """
    アルバムのデータモデル
    """
    name: str
    artist: str
    songs: List[Song]

prompt_template_str = """\
アーティストと曲のリストを含む、アルバムの例を生成してください。\
映画 {movie_name} をヒントとして使ってください。
"""
program = OpenAIPydanticProgram.from_defaults(
    output_cls=Album, prompt_template_str=prompt_template_str, verbose=True
)

output = program(
    movie_name="仮面ライダーW FOREVER AtoZ/運命のガイアメモリ"
)

ValueError: Must provide description for your Pydantic Model. Either provide a docstring or add `description=<your_description>` to the method. Required to convert Pydantic Model to OpenAI Function.

{
    'messages': [
        {
            'role': 'user',
            'content': 'アーティストと曲のリストを含む、アルバムの例を生成してください。映画 仮面ライダーW FOREVER AtoZ/運命のガイアメモリ をヒントとして使ってください。\n'
        }
    ],
    'model': 'gpt-3.5-turbo-0613',
    'stream': False,
    'temperature': 0.1,
    'tool_choice': {
        'type': 'function',
        'function': {
            'name': 'Album'
        }
    },
    'tools': [
        {
            'type': 'function',
            'function': {
                'name': 'Album',
                'description': 'アルバムのデータモデル',
                'parameters': {
                    'title': 'Album',
                    'description': 'アルバムのデータモデル',
                    'type': 'object',
                    'properties': {
                        'name': {
                            'title': 'Name',
                            'type': 'string'
                        },
                        'artist': {
                            'title': 'Artist',
                            'type': 'string'
                        },
                        'songs': {
                            'title': 'Songs',
                            'type': 'array',
                            'items': {
                                '$ref': '#/definitions/Song'
                            }
                        }
                    },
                    'required': ['name', 'artist', 'songs'],
                    'definitions': {
                        'Song': {
                            'title': 'Song',
                            'description': '曲のデータモデル',
                            'type': 'object',
                            'properties': {
                                'title': {
                                    'title': 'Title',
                                    'type': 'string'
                                },
                                'length_seconds': {
                                    'title': 'Length Seconds',
                                    'type': 'integer'
                                }
                            },
                            'required': ['title', 'length_seconds']
                        }
                    }
                }
            }
        }
    ]
}

from pydantic import BaseModel
from typing import List
from llama_index.program import OpenAIPydanticProgram
from llama_index.llms import OpenAI

class Song(BaseModel):
    """
    曲のデータモデル
    """
    title: str
    length_seconds: int


class Album(BaseModel):
    """
    アルバムのデータモデル
    """
    name: str
    artist: str
    songs: List[Song]

prompt_template_str = """\
アーティストと曲のリストを含む、3つのアルバムの例を生成してください。\
それぞれのテーマは、仮面ライダーW、侍戦隊シンケンジャー、 ウルトラマンゼロ、です。
"""

program = OpenAIPydanticProgram.from_defaults(
    output_cls=Album,
    llm=OpenAI(model="gpt-3.5-turbo-1106"),      # 1106以降のモデルを指定
    prompt_template_str=prompt_template_str,
    allow_multiple=True,                            # allow_multiple=Trueで並列処理を有効化
    verbose=True
)

output = program()
output

Function call: Album with args: {"name": "仮面ライダーW", "artist": "仮面ライダーW", "songs": [{"title": "W-B-X", "length_seconds": 240}, {"title": "Cyclone Effect", "length_seconds": 210}, {"title": "Finger on the Trigger", "length_seconds": 195}]}
Function call: Album with args: {"name": "侍戦隊シンケンジャー", "artist": "侍戦隊シンケンジャー", "songs": [{"title": "Samurai Sentai Shinkenger", "length_seconds": 220}, {"title": "Uramasa", "length_seconds": 190}, {"title": "Shirokuji", "length_seconds": 205}]}
Function call: Album with args: {"name": "ウルトラマンゼロ", "artist": "ウルトラマンゼロ", "songs": [{"title": "Go! Fight! Ultraman Zero!!", "length_seconds": 230}, {"title": "Ultra Zero Fight", "length_seconds": 215}, {"title": "Believe in Zero", "length_seconds": 200}]}

output

[
  Album(name='仮面ライダーW', artist='仮面ライダーW', songs=[Song(title='W-B-X', length_seconds=240), Song(title='Cyclone Effect', length_seconds=210), Song(title='Finger on the Trigger', length_seconds=195)]),
  Album(name='侍戦隊シンケンジャー', artist='侍戦隊シンケンジャー', songs=[Song(title='Samurai Sentai Shinkenger', length_seconds=220), Song(title='Uramasa', length_seconds=190), Song(title='Shirokuji', length_seconds=205)]),
  Album(name='ウルトラマンゼロ', artist='ウルトラマンゼロ', songs=[Song(title='Go! Fight! Ultraman Zero!!', length_seconds=230), Song(title='Ultra Zero Fight', length_seconds=215), Song(title='Believe in Zero', length_seconds=200)])
]

output = program.stream_list()

for obj in output:
    print(obj.json(indent=2))

from llama_index.program import OpenAIPydanticProgram, DataFrame

prompt_template_str = """\
以下のクエリを構造化データに展開してください。：

{input_str}.

列名の集合と行の集合の両方を抽出してください。
"""

program = OpenAIPydanticProgram.from_defaults(
    output_cls=DataFrame,
    prompt_template_str=prompt_template_str,
    verbose=True,
)

input_str = """\
僕の名前はジョンで、25歳です。ニューヨーク バスケットボールが好きです。\
彼の名前は マイクで30歳。サンフランシスコに住んでいて、野球をするのが好きです。\
サラは20歳でロサンゼルス在住。テニスが好きです。\
彼女の名前はメアリーで35歳。シカゴに住んでいる。\
"""

response_obj = program(input_str=input_str)

Function call: DataFrame with args: {
  "columns": [
    {"column_name": "名前", "column_desc": "名前"},
    {"column_name": "年齢", "column_desc": "年齢"},
    {"column_name": "居住地", "column_desc": "居住地"},
    {"column_name": "趣味", "column_desc": "趣味"}
  ],
  "rows": [
    {"row_values": ["ジョン", 25, "ニューヨーク", "バスケットボール"]},
    {"row_values": ["マイク", 30, "サンフランシスコ", "野球"]},
    {"row_values": ["サラ", 20, "ロサンゼルス", "テニス"]},
    {"row_values": ["メアリー", 35, "シカゴ", ""]}
  ]
}

DataFrame(description=None, columns=[DataFrameColumn(column_name='名前', column_desc='名前'), DataFrameColumn(column_name='年齢', column_desc='年齢'), DataFrameColumn(column_name='居住地', column_desc='居住地'), DataFrameColumn(column_name='趣味', column_desc='趣味')], rows=[DataFrameRow(row_values=['ジョン', 25, 'ニューヨーク', 'バスケットボール']), DataFrameRow(row_values=['マイク', 30, 'サンフランシスコ', '野球']), DataFrameRow(row_values=['サラ', 20, 'ロサンゼルス', 'テニス']), DataFrameRow(row_values=['メアリー', 35, 'シカゴ', ''])])

response_obj.to_df()

{
    'messages': [
        {
            'role': 'user', 'content': '以下のクエリを構造化データに展開してください。：\n\n僕の名前はジョンで、25歳です。ニューヨーク バスケットボールが好きです。彼の名前は マイクで30歳。サンフランシスコに住んでいて、野球をするのが好きです。サラは20歳でロサンゼルス在住。テニスが好きです。彼女の名前はメアリーで35歳。シカゴに住んでいる。.\n\n列名の集合と行の集合の両方を抽出してください。\n'
        }
    ],
    'model': 'gpt-3.5-turbo-0613',
    'stream': False,
    'temperature': 0.1,
    'tool_choice': {
        'type': 'function',
        'function': {
            'name': 'DataFrame'
        }
    },
    'tools': [
        {
            'type': 'function',
            'function': {
                'name': 'DataFrame',
                'description': 'Data-frame class.\n\nConsists of a `rows` field which is a list of dictionaries,\nas well as a `columns` field which is a list of column names.',
                'parameters': {
                    'title': 'DataFrame',
                    'description': 'Data-frame class.\n\nConsists of a `rows` field which is a list of dictionaries,\nas well as a `columns` field which is a list of column names.',
                    'type': 'object',
                    'properties': {
                        'description': {
                            'title': 'Description',
                            'type': 'string'
                        },
                        'columns': {
                            'title': 'Columns',
                            'description': 'List of column names.',
                            'type': 'array',
                            'items': {'$ref': '#/definitions/DataFrameColumn'}
                        },
                        'rows': {
                            'title': 'Rows',
                            'description': 'List of DataFrameRow objects. Each DataFrameRow contains         valuesin order of the data frame column.',
                            'type': 'array',
                            'items': {'$ref': '#/definitions/DataFrameRow'}
                        }
                    },
                    'required': ['columns', 'rows'],
                    'definitions': {
                        'DataFrameColumn': {
                            'title': 'DataFrameColumn',
                            'description': 'Column in a DataFrame.',
                            'type': 'object',
                            'properties': {
                                'column_name': {
                                    'title': 'Column Name',
                                    'description': 'Column name.',
                                    'type': 'string'
                                },
                                'column_desc': {
                                    'title': 'Column Desc',
                                    'description': 'Column description.',
                                    'type': 'string'
                                }
                            },
                            'required': ['column_name', 'column_desc']
                        },
                        'DataFrameRow': {
                            'title': 'DataFrameRow',
                            'description': 'Row in a DataFrame.',
                            'type': 'object',
                            'properties': {
                                'row_values': {
                                    'title': 'Row Values',
                                    'description': 'List of row values, where each value corresponds to a row key.',
                                    'type': 'array',
                                    'items': {}
                                }
                            },
                            'required': ['row_values']
                        }
                    }
                }
            }
        }
    ]
}

from llama_index.program import OpenAIPydanticProgram, DataFrameRowsOnly

prompt_template_str = """\
以下のクエリを構造化データに展開してください:

{input_str}

カラム名: ['名前', '年齢', '都市', '好きなスポーツ']。
関数スキーマにない追加のパラメータを指定しないでください。
"""

program = OpenAIPydanticProgram.from_defaults(
    output_cls=DataFrameRowsOnly,
    prompt_template_str=prompt_template_str,
    verbose=True,
)

input_str = """\
僕の名前はジョンで、25歳です。ニューヨーク バスケットボールが好きです。\
彼の名前は マイクで30歳。サンフランシスコに住んでいて、野球をするのが好きです。\
サラは20歳でロサンゼルス在住。テニスが好きです。\
彼女の名前はメアリーで35歳。シカゴに住んでいる。\
"""

response_obj = program(input_str=input_str)
response_obj

Function call: DataFrameRowsOnly with args: {
  "rows": [
    {
      "row_values": ["ジョン", 25, "ニューヨーク", "バスケットボール"]
    },
    {
      "row_values": ["マイク", 30, "サンフランシスコ", "野球"]
    },
    {
      "row_values": ["サラ", 20, "ロサンゼルス", "テニス"]
    },
    {
      "row_values": ["メアリー", 35, "シカゴ", null]
    }
  ]
}

DataFrameRowsOnly(rows=[DataFrameRow(row_values=['ジョン', 25, 'ニューヨーク', 'バスケットボール']), DataFrameRow(row_values=['マイク', 30, 'サンフランシスコ', '野球']), DataFrameRow(row_values=['サラ', 20, 'ロサンゼルス', 'テニス']), DataFrameRow(row_values=['メアリー', 35, 'シカゴ', None])])

response_obj.to_df()

from llama_index.program import OpenAIPydanticProgram, DFRowsProgram
import pandas as pd

df = pd.DataFrame(
    {
        "Name": pd.Series(dtype="str"),
        "Age": pd.Series(dtype="int"),
        "City": pd.Series(dtype="str"),
        "Favorite Sport": pd.Series(dtype="str"),
    }
)

df_rows_program = DFRowsProgram.from_defaults(
    pydantic_program_cls=OpenAIPydanticProgram,
    df=df,
    verbose=True
)

input_str = """\
僕の名前はジョンで、25歳です。ニューヨーク バスケットボールが好きです。\
彼の名前は マイクで30歳。サンフランシスコに住んでいて、野球をするのが好きです。\
サラは20歳でロサンゼルス在住。テニスが好きです。\
彼女の名前はメアリーで35歳。シカゴに住んでいる。\
"""

result_obj = df_rows_program(
    input_str=input_str,
)
result_obj

Function call: DataFrameRowsOnly with args: {
  "rows": [
    {
      "row_values": ["ジョン", 25, "ニューヨーク", "バスケットボール"]
    },
    {
      "row_values": ["マイク", 30, "サンフランシスコ", "野球"]
    },
    {
      "row_values": ["サラ", 20, "ロサンゼルス", "テニス"]
    },
    {
      "row_values": ["メアリー", 35, "シカゴ", null]
    }
  ]
}

DataFrameRowsOnly(rows=[DataFrameRow(row_values=['ジョン', 25, 'ニューヨーク', 'バスケットボール']), DataFrameRow(row_values=['マイク', 30, 'サンフランシスコ', '野球']), DataFrameRow(row_values=['サラ', 20, 'ロサンゼルス', 'テニス']), DataFrameRow(row_values=['メアリー', 35, 'シカゴ', None])])

result_obj.to_df()

result_obj.to_df(existing_df=df)

from llama_index.program import OpenAIPydanticProgram, DFFullProgram
import pandas as pd


df_full_program = DFFullProgram.from_defaults(
    pydantic_program_cls=OpenAIPydanticProgram,
)

input_str = """\
僕の名前はジョンで、25歳です。ニューヨーク バスケットボールが好きです。\
彼の名前は マイクで30歳。サンフランシスコに住んでいて、野球をするのが好きです。\
サラは20歳でロサンゼルス在住。テニスが好きです。\
彼女の名前はメアリーで35歳。シカゴに住んでいる。\
"""

result_obj = df_full_program(
    input_str=input_str,
)
result_obj

DataFrame(description=None, columns=[DataFrameColumn(column_name='Name', column_desc='Name of the person'), DataFrameColumn(column_name='Age', column_desc='Age of the person'), DataFrameColumn(column_name='Location', column_desc='City where the person lives'), DataFrameColumn(column_name='Hobby', column_desc='Hobby of the person')], rows=[DataFrameRow(row_values=['ジョン', 25, 'ニューヨーク', 'バスケットボール']), DataFrameRow(row_values=['マイク', 30, 'サンフランシスコ', '野球']), DataFrameRow(row_values=['サラ', 20, 'ロサンゼルス', 'テニス']), DataFrameRow(row_values=['メアリー', 35, 'シカゴ', None])])

result_obj.to_df()

from pathlib import Path
import requests

wiki_titles = ["ドウデュース", "イクイノックス"]
for title in wiki_titles:
    response = requests.get(
        "https://ja.wikipedia.org/w/api.php",
        params={
            "action": "query",
            "format": "json",
            "titles": title,
            "prop": "extracts",
            # 'exintro': True,
            "explaintext": True,
        },
    ).json()
    page = next(iter(response["query"]["pages"].values()))
    wiki_text = page["extract"]

    data_path = Path("data")
    if not data_path.exists():
        Path.mkdir(data_path)

    with open(data_path / f"{title}.txt", "w") as fp:
        fp.write(wiki_text)

from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext
from llama_index.text_splitter import SentenceSplitter
from llama_index.embeddings import OpenAIEmbedding
from llama_index.llms import OpenAI

llm = OpenAI(model="gpt-4", temperarture=0.3)
embed_model = OpenAIEmbedding()
node_parser = SentenceSplitter(chunk_size=512, chunk_overlap=20)

service_context = ServiceContext.from_defaults(
  llm=llm,
  embed_model=embed_model,
  node_parser=node_parser,
)

documents = SimpleDirectoryReader("data").load_data()

index = VectorStoreIndex.from_documents(documents, service_context=service_context)

from typing import List
from pydantic import BaseModel


class RaceHorseBiography(BaseModel):
    """Data model for a biography of a race horse."""

    name: str
    major_wins: List[str]
    primary_jockey: List[str]
    best_known_for: List[str]
    extra_info: str

query_engine = index.as_query_engine(
    output_cls=RaceHorseBiography,
    response_mode="refine"
)

response = query_engine.query("ドウデュースについて教えて。")

print(response.name)
print(response.major_wins)
print(response.primary_jockey)
print(response.best_known_for)
print(response.extra_info)

ドウデュース
['日本ダービー', '京都記念']
['武豊']
['日本ダービーと京都記念の勝利', 'ダービーレコードの更新', 'ロンジンワールドベストレースホースランキングで第15位']
ドウデュースは、日本ダービーと京都記念を勝利したことで知られています。特に、日本ダービーでは2分21秒9のダービーレコードを更新し、鞍上の武豊は史上初の50代でのダービージョッキーの名誉を得ました。また、朝日杯フューチュリティステークスの勝ち馬が日本ダービーを制したのは、1994年のナリタブライアン以来28年ぶりとなりました。さらに、国際競馬統括機関連盟が発表した「ロンジンワールドベストレースホースランキング」では、日本ダービーの勝利が評価され、シャフリヤールやエンブレムロードと並ぶレーティング120で第15位に位置づけられました。しかし、ドバイターフへの出走を予定していたものの、調教後に左前脚を負傷し、出走を取り消しました。その後は治療と休養を経て、天皇賞（秋）に出走しました。

{
    'messages': [
        {
            'role': 'system',
            'content': "You are an expert Q&A system that is trusted around the world.\nAlways answer the query using the provided context information, and not prior knowledge.\nSome rules to follow:\n1. Never directly reference the given context in your answer.\n2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines."
        },
        {
            'role': 'user',
            'content': 'Context information is below.\n---------------------\nfile_path: data/ドウデュース.txt\n\nまずまずのスタートを決めると、道中は馬群を見て最後方グループで待機し、向正面から外へ出て徐々に捲りをかけ、直線に入ると早々と馬群から抜け出して後続を突き放し、3馬身半差の快勝でGIを含む重賞3勝目を挙げた。日本ダービーを勝利した馬が京都記念を勝利するのは、1948年春の京都記念のマツミドリ以来75年ぶりとなった。\n次走は3月25日にドバイのメイダン競馬場で行われるドバイターフとし、同月15日（現地時間同月14日）にメイダン競馬場に到着した。しかし出馬投票後の同月24日、調教後に左前肢跛行を発症しドバイターフへの出走を取り消した。友道は「調教後に左腕節に違和感を認め、競馬に向けて進めておりましたが、将来のある馬なのでここでは無理をせず、取り消すことを決断いたしました」と語った。\nその後夏は治療と休養にあて、秋初戦として10月29日に東京競馬場で開催される天皇賞（秋）に出走。\n---------------------\nGiven the context information and not prior knowledge, answer the query.\nQuery: ドウデュースについて教えて。\nAnswer: '
        }
    ],
    'model': 'gpt-4',
    'stream': False,
    'temperature': 0.1,
    'tool_choice': {
        'type': 'function',
        'function': {
            'name': 'RaceHorseBiography'
        }
    },
    'tools': [
        {
            'type': 'function',
            'function': {
                'name': 'RaceHorseBiography',
                'description': 'Data model for a biography of a race horse.',
                'parameters': {
                    'title': 'RaceHorseBiography',
                    'description': 'Data model for a biography of a race horse.',
                    'type': 'object',
                    'properties': {
                        'name': {
                            'title': 'Name',
                            'type': 'string'
                        },
                        'major_wins': {
                            'title': 'Major Wins',
                            'type': 'array',
                            'items': {
                                'type': 'string'
                            }
                        },
                        'primary_jockey': {
                            'title': 'Primary Jockey',
                            'type': 'array',
                            'items': {
                                'type': 'string'
                            }
                        },
                        'best_known_for': {
                            'title': 'Best Known For',
                            'type': 'array',
                            'items': {
                                'type': 'string'
                            }
                        },
                        'extra_info': {
                            'title': 'Extra Info',
                            'type': 'string'
                        }
                    },
                    'required': ['name', 'major_wins', 'primary_jockey', 'best_known_for', 'extra_info']
                }
            }
        }
    ]
}

{
    'messages': [
        {
            'role': 'user',
            'content': 'You are an expert Q&A system that strictly operates in two modes when refining existing answers:\n1. **Rewrite** an original answer using the new context.\n2. **Repeat** the original answer if the new context isn\'t useful.\nNever reference the original answer or context directly in your answer.\nWhen in doubt, just repeat the original answer.New Context: file_path: data/ドウデュース.txt\n\n2倍の3番人気でレースを迎えた。レースではスタートから後方に控え、前半1000メートルを58秒9という早いペースの中落ち着いてレースを進め、直線に入ってスパート。鞍上武豊の右ムチに応え加速し先団を捉え、大外から伸びてくるイクイノックスをクビ差抑え込み優勝。朝日杯以来のGI制覇となり、勝ちタイムは2分21秒9のダービーレコード(コースレコードは2018年ジャパンカップでアーモンドアイが記録した2分20秒6)。鞍上の武は、歴代最多を更新する2013年のキズナ以来となるダービー6勝目、かつ歴代最年長、史上初の50代でのダービージョッキーの名誉となった。また朝日杯フューチュリティステークスの勝ち馬が日本ダービーを制したのは、1994年のナリタブライアン（前身である朝日杯3歳ステークスを勝利）以来28年ぶりとなった。\n6月10日に国際競馬統括機関連盟が発表した「ロンジンワールドベストレースホースランキング」において、ドウデュースは日本ダービーを勝利した功績を評価され、シャフリヤールやエンブレムロードと並ぶレーティング120で第15位タイに位置づけられた。\nQuery: ドウデュースについて教えて。\nOriginal Answer: {"name": "ドウデュース", "major_wins": ["日本ダービー", "京都記念"], "primary_jockey": [], "best_known_for": ["日本ダービーと京都記念の勝利"], "extra_info": "ドウデュースは、日本ダービーと京都記念を勝利したことで知られています。特に、日本ダービーを勝利した馬が京都記念を勝利するのは、1948年春の京都記念のマツミドリ以来75年ぶりとなったことが注目されています。しかし、ドバイターフへの出走を予定していたものの、調教後に左前肢跛行を発症し、出走を取り消しました。その後は治療と休養を経て、天皇賞（秋）に出走しました。"}\nNew Answer: '
        }
    ],
    'model': 'gpt-4',
    'stream': False,
    'temperature': 0.1,
    'tool_choice': {
        'type': 'function',
        'function': {
            'name': 'RaceHorseBiography'
        }
    },
    'tools': [
        {
            'type': 'function',
            'function': {
                'name': 'RaceHorseBiography',
                'description': 'Data model for a biography of a race horse.',
                'parameters': {
                    'title': 'RaceHorseBiography',
                    'description': 'Data model for a biography of a race horse.',
                    'type': 'object',
                    'properties': {
                        'name': {
                            'title': 'Name',
                            'type': 'string'
                        },
                        'major_wins': {
                            'title': 'Major Wins',
                            'type': 'array',
                            'items': {
                                'type': 'string'
                            }
                        },
                        'primary_jockey': {
                            'title': 'Primary Jockey',
                            'type': 'array',
                            'items': {
                                'type': 'string'
                            }
                        },
                        'best_known_for': {
                            'title': 'Best Known For',
                            'type': 'array',
                            'items': {
                                'type': 'string'
                            }
                        },
                        'extra_info': {
                            'title': 'Extra Info',
                            'type': 'string'
                        }
                    },
                    'required': ['name', 'major_wins', 'primary_jockey', 'best_known_for', 'extra_info']
                }
            }
        }
    ]
}

!pip install -q llama-index arize-phoenix

from google.colab import userdata
import os

os.environ["OPENAI_API_KEY"] = userdata.get('OPENAI_API_KEY')

import phoenix as px
import llama_index

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

🌍 To view the Phoenix app in your browser, visit https://XXXXXXXX-colab.googleusercontent.com/

from llama_index.query_pipeline.query import QueryPipeline
from llama_index.llms import OpenAI
from llama_index.prompts import PromptTemplate

prompt_str = "{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="仮面ライダーW")

> Running module 712ac40a-1a8f-453f-85e3-43fe4f81ecaa with input: 
movie_name: 仮面ライダーW

> Running module 34461f34-3853-44d9-a83e-72e91fd66ce1 with input: 
messages: 仮面ライダーW に関係する映画名を生成してください。

print(str(output))

assistant: 1. 仮面ライダーW FOREVER AtoZ/運命のガイアメモリ
2. 仮面ライダーW RETURNS 仮面ライダーアクセル/帰還のゲイツ
3. 仮面ライダーW RETURNS 仮面ライダーエターナル/永遠のメモリー
4. 仮面ライダーW RETURNS 仮面ライダーサイクロン/再生のガイアメモリ
5. 仮面ライダーW RETURNS 仮面ライダーアクセル/帰還のゲイツ
6. 仮面ライダーW RETURNS 仮面ライダーエターナル/永遠のメモリー
7. 仮面ライダーW RETURNS 仮面ライダーサイクロン/再生のガイアメモリ
8. 仮面ライダーW RETURNS 仮面ライダーアクセル/帰還のゲイツ
9. 仮面ライダーW RETURNS 仮面ライダーエターナル/永遠のメモリー
10. 仮面ライダーW RETURNS 仮面ライダーサイクロン/再生のガイアメモリ

from typing import List
from pydantic import BaseModel, Field
from llama_index.output_parsers import PydanticOutputParser


class Movie(BaseModel):
    """1つの映画を表すオブジェクト"""
    name: str = Field(..., description="映画のタイトル")
    year: int = Field(..., description="映画の上映年")


class Movies(BaseModel):
    """映画のリストを表すオブジェクト"""
    movies: List[Movie] = Field(..., description="映画のリスト")

llm = OpenAI(model="gpt-3.5-turbo")
output_parser = PydanticOutputParser(Movies)

json_prompt_str = """\
{movie_name} に関係する映画名を生成してください。出力は以下のJSON形式で:
"""

json_prompt_str = output_parser.format(json_prompt_str)

{movie_name} に関係する映画名を生成してください。出力は以下のJSON形式で:



Here's a JSON schema to follow:
{{"title": "Movies", "description": "映画のリストを表すオブジェクト", "type": "object", "properties": {{"movies": {{"title": "Movies", "description": "映画のリスト", "type": "array", "items": {{"$ref": "#/definitions/Movie"}}}}}}, "required": ["movies"], "definitions": {{"Movie": {{"title": "Movie", "description": "1つの映画を表すオブジェクト", "type": "object", "properties": {{"name": {{"title": "Name", "description": "映画のタイトル", "type": "string"}}, "year": {{"title": "Year", "description": "映画の上映年", "type": "integer"}}}}, "required": ["name", "year"]}}}}}}

Output a valid JSON object but do not repeat the schema.

json_prompt_tmpl = PromptTemplate(json_prompt_str)

p = QueryPipeline(chain=[json_prompt_tmpl, llm, output_parser], verbose=True)
output = p.run(movie_name="仮面ライダーW")

> Running module 66a47d95-4c63-43dd-b428-1ba2e362ee4e with input: 
movie_name: 仮面ライダーW

> Running module ccd40fcb-f4c2-4d64-bd4c-f7ca232fee54 with input: 
messages: 仮面ライダーW に関係する映画名を生成してください。出力は以下のJSON形式で:



Here's a JSON schema to follow:
{"title": "Movies", "description": "\u6620\u753b\u306e\u30ea\u30b9\u30c8\u3092\u8868\u3059\u30aa\u30d6\u30b8\u30a7\u30af\u30...

> Running module 9ed5de10-d63b-42f4-98b2-a841f12251ff with input: 
input: assistant: {
  "movies": [
    {
      "name": "仮面ライダーW FOREVER AtoZ/運命のガイアメモリ",
      "year": 2010
    },
    {
      "name": "仮面ライダーW RETURNS 仮面ライダーアクセル",
      "year": 2011
    },
    {
      "name...

output

Movies(movies=[Movie(name='仮面ライダーW FOREVER AtoZ/運命のガイアメモリ', year=2010), Movie(name='仮面ライダーW RETURNS 仮面ライダーアクセル', year=2011), Movie(name='仮面ライダーW RETURNS 仮面ライダーエターナル', year=2011), Movie(name='仮面ライダーW RETURNS 仮面ライダーアクセル/仮面ライダーエターナル', year=2011), Movie(name='仮面ライダーW RETURNS 仮面ライダーサンバ', year=2011)])

prompt_str = "{movie_name} に関係する映画名を生成してください。"
prompt_tmpl = PromptTemplate(prompt_str)

# let's add some subsequent prompts for fun
prompt_str2 = """\
以下に文章があります:

{text}

各映画の概要を書き直してください。
"""
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="仮面ライダーW")

> Running module 52a13d6f-9059-497f-8468-a503103eae29 with input: 
movie_name: 仮面ライダーW

> Running module 73c824b1-7c1d-47ef-846a-207315478c21 with input: 
messages: 仮面ライダーW に関係する映画名を生成してください。

> Running module 26411af1-2924-4898-85a7-917838b0e540 with input: 
text: <generator object llm_chat_callback.<locals>.wrap.<locals>.wrapped_llm_chat.<locals>.wrapped_gen at 0x7c4f29e49070>

> Running module f6ab8d12-ca18-4e8f-b988-92d1b1c206d9 with input: 
messages: 以下に文章があります:

1. 仮面ライダーW: フォーエバー AtoZ/運命のガイアメモリ
2. 仮面ライダーW: リターンズ アクセル/エターナル
3. 仮面ライダーW: ゴーフォーサイクロン
4. 仮面ライダーW: サイクロンジョーカー
5. 仮面ライダーW: エクストリームサイクロン
6. 仮面ライダーW: サイクロンアクセル
7. 仮面ライダーW: ジョーカーエターナル
8. 仮面ライダ...

for o in output:
    print(o.delta, end="\n")

1
.
 
仮
面
ラ
イ
ダ
ー
W
:
 フ
ォ
ーエ
バ
ー
 A
to
Z
/
運
命
の
ガ
イ
ア
メ
モ
リ
 -
 
仮
面
ラ
イ
ダ
ー
W
が
、
新
た
な
敵
と
(snip)

from pathlib import Path
import requests

wiki_titles = ["ドウデュース"]
for title in wiki_titles:
    response = requests.get(
        "https://ja.wikipedia.org/w/api.php",
        params={
            "action": "query",
            "format": "json",
            "titles": title,
            "prop": "extracts",
            # 'exintro': True,
            "explaintext": True,
        },
    ).json()
    page = next(iter(response["query"]["pages"].values()))
    wiki_text = page["extract"]

    data_path = Path("data")
    if not data_path.exists():
        Path.mkdir(data_path)

    with open(data_path / f"{title}.txt", "w") as fp:
        fp.write(wiki_text)

from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext
from llama_index.text_splitter import SentenceSplitter
from llama_index.embeddings import OpenAIEmbedding
from llama_index.llms import OpenAI

llm = OpenAI(model="gpt-3.5-turbo", temperarture=0.3)
embed_model = OpenAIEmbedding()
node_parser = SentenceSplitter(chunk_size=512, chunk_overlap=20)

service_context = ServiceContext.from_defaults(
  llm=llm,
  embed_model=embed_model,
  node_parser=node_parser,
)

documents = SimpleDirectoryReader("data").load_data()

index = VectorStoreIndex.from_documents(documents, service_context=service_context)

from llama_index.query_pipeline.query import QueryPipeline
from llama_index.llms import OpenAI
from llama_index.prompts import PromptTemplate

prompt_str1 = "競走馬 ドウデュース について、次のトピックに関する簡潔な質問を作成してください: 「{topic}」"
prompt_tmpl1 = PromptTemplate(prompt_str1)

# use HyDE to hallucinate answer.
prompt_str2 = '''\
以下の質問に答える文章を書いてください。
できるだけ多くの重要な詳細を含めるようにしてください。


{query_str}


文章: """
'''

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="主な重賞勝ち鞍")

> Running module 124bf7de-83b8-43a7-8c05-d10e60330e6e with input: 
topic: 主な重賞勝ち鞍

> Running module c7de191c-3e23-4b96-b8b0-e228a222d381 with input: 
messages: 競走馬 ドウデュース について、次のトピックに関する簡潔な質問を作成してください: 「主な重賞勝ち鞍」

> Running module 404174ae-8d21-4816-94bb-af6372b831f3 with input: 
query_str: assistant: ドウデュースはどの重賞を勝ちましたか？

> Running module 5e928f17-b91e-4fe7-bc8c-3c59a719e6d8 with input: 
messages: 以下の質問に答える文章を書いてください。
できるだけ多くの重要な詳細を含めるようにしてください。


assistant: ドウデュースはどの重賞を勝ちましたか？


文章: """


> Running module a9e05c2d-99db-4a6d-a909-8a4d77dccff1 with input: 
input: assistant: ドウデュースは、2021年に行われた日本の競馬の重賞レースであるジャパンカップを勝ちました。ジャパンカップは、東京競馬場で行われる芝2400メートルのレースで、日本国内外からトップクラスの競走馬が集まる国際的なレースです。ドウデュースは、このレースで見事な走りを見せ、優勝しました。この勝利により、ドウデュースはその年の最も優れた競走馬として称えられ、競馬界での名声を高めました...

5

len(nodes)

5

!pip install cohere

from google.colab import userdata
import os

os.environ["COHERE_API_KEY"] = userdata.get('COHERE_API_KEY')

!pip install pyvis

from llama_index.prompts import PromptTemplate
from llama_index.llms import OpenAI
from llama_index import ServiceContext
from llama_index.postprocessor import CohereRerank
from llama_index.response_synthesizers import TreeSummarize

# モジュール定義

## プロンプトテンプレート
prompt_str = "競走馬 ドウデュース について、次のトピックに関する簡潔な質問を作成してください: 「{topic}」"
prompt_tmpl = PromptTemplate(prompt_str)

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

## 検索
retriever = index.as_retriever(similarity_top_k=3)

## リランク
reranker = CohereRerank()

## 要約
summarizer = TreeSummarize(
    service_context=ServiceContext.from_defaults(llm=llm)
)

from llama_index.query_pipeline.query import QueryPipeline

# クエリパイプラインの定義
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")

for m in ["llm", "prompt_tmpl", "retriever", "summarizer", "reranker"]:
    print(f"##### {m} #####\n")
    obj = locals().get(m)
    print("INPUT:", obj.as_query_component().input_keys)
    print("OUTPUT:", obj.as_query_component().input_keys)
    print()

##### llm #####

INPUT: required_keys={'messages'} optional_keys=set()
OUTPUT: required_keys={'messages'} optional_keys=set()

##### prompt_tmpl #####

INPUT: required_keys={'topic'} optional_keys=set()
OUTPUT: required_keys={'topic'} optional_keys=set()

##### retriever #####

INPUT: required_keys={'input'} optional_keys=set()
OUTPUT: required_keys={'input'} optional_keys=set()

##### summarizer #####

INPUT: required_keys={'nodes', 'query_str'} optional_keys=set()
OUTPUT: required_keys={'nodes', 'query_str'} optional_keys=set()

##### reranker #####

INPUT: required_keys={'nodes'} optional_keys={'query_str'}
OUTPUT: required_keys={'nodes'} optional_keys={'query_str'}

from pyvis.network import Network
from IPython.display import HTML

net = Network(notebook=True, cdn_resources="in_line", directed=True)
net.from_nx(p.dag)
net.save_graph("rag_dag.html")
HTML(filename="rag_dag.html")

response = p.run(topic="主な重賞勝ち鞍")

> Running module prompt_tmpl with input: 
topic: 主な重賞勝ち鞍

> Running module llm with input: 
messages: 競走馬 ドウデュース について、次のトピックに関する簡潔な質問を作成してください: 「主な重賞勝ち鞍」

> Running module retriever with input: 
input: assistant: ドウデュースはどの重賞を勝ちましたか？

> Running module reranker with input: 
query_str: assistant: ドウデュースはどの重賞を勝ちましたか？
nodes: [NodeWithScore(node=TextNode(id_='eac6bdf4-b0db-4fc1-a2f7-bf3ac757147f', embedding=None, metadata={'file_path': 'data/ドウデュース.txt', 'file_name': 'ドウデュース.txt', 'file_type': 'text/plain', 'file_size': 95...

> Running module summarizer with input: 
query_str: assistant: ドウデュースはどの重賞を勝ちましたか？
nodes: [NodeWithScore(node=TextNode(id_='eac6bdf4-b0db-4fc1-a2f7-bf3ac757147f', embedding=None, metadata={'file_path': 'data/ドウデュース.txt', 'file_name': 'ドウデュース.txt', 'file_type': 'text/plain', 'file_size': 95...

print(str(response))

ドウデュースは日本ダービーを勝ちました。

response = await p.arun(topic="主な重賞勝ち鞍")
print(str(response))

from llama_index.response_synthesizers import TreeSummarize
from llama_index import ServiceContext
from llama_index.query_pipeline import InputComponent

retriever = index.as_retriever(similarity_top_k=5)
summarizer = TreeSummarize(
    service_context=ServiceContext.from_defaults(
        llm=OpenAI(model="gpt-3.5-turbo")
    )
)

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")

response = p.run(input="ドウデュースの主な重賞勝ち鞍を教えて。")

> Running module input with input: 
input: ドウデュースの主な重賞勝ち鞍を教えて。

> Running module retriever with input: 
input: ドウデュースの主な重賞勝ち鞍を教えて。

> Running module summarizer with input: 
query_str: ドウデュースの主な重賞勝ち鞍を教えて。
nodes: [NodeWithScore(node=TextNode(id_='eac6bdf4-b0db-4fc1-a2f7-bf3ac757147f', embedding=None, metadata={'file_path': 'data/ドウデュース.txt', 'file_name': 'ドウデュース.txt', 'file_type': 'text/plain', 'file_size': 95...

print(str(response))

ドウデュースの主な重賞勝ち鞍は、日本ダービー、有馬記念、京都記念です。

from llama_index.query_pipeline import (
    CustomQueryComponent,
    InputKeys,
    OutputKeys,
)
from llama_index.query_pipeline.query import QueryPipeline
from llama_index.prompts import PromptTemplate
from typing import Dict, Any
from llama_index.llms.llm import BaseLLM
from pydantic import Field


class RelatedMovieComponent(CustomQueryComponent):
    """関係する映画名を出力するコンポーネント"""

    llm: BaseLLM = Field(..., description="OpenAI LLM")

    def _validate_component_inputs(
        self, input: Dict[str, Any]
    ) -> Dict[str, Any]:
        """run_component実行時に入力をバリデーションする"""
        # 注記: ここはオプションだがバリデーションの例として実装
        return input

    @property
    def _input_keys(self) -> set:
        """入力キーの辞書"""
        # 注記: 必要な入力をリストアップ。オプション入力の場合は
        # `optional_input_keys_dict`をオーバーライドする必要がある。
        return {"movie"}

    @property
    def _output_keys(self) -> set:
        """出力キーの辞書"""
        return {"output"}

    def _run_component(self, **kwargs) -> Dict[str, Any]:
        """コンポーネントを実行する"""
        # 簡単のためQueryPipeline自体を使用する。
        prompt_str = "{movie_name} に関係する映画名を生成してください。"
        prompt_tmpl = PromptTemplate(prompt_str)
        p = QueryPipeline(chain=[prompt_tmpl, llm])
        return {"output": p.run(movie_name=kwargs["movie"])}

from llama_index.llms import OpenAI
from llama_index.query_pipeline.query import QueryPipeline
from llama_index.prompts import PromptTemplate

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

prompt_str = """\
以下に文章があります:

{text}

これらにシェークスピアのコメントを追加して書き直してください。コメントは日本語で。
"""
prompt_tmpl = PromptTemplate(prompt_str)

p = QueryPipeline(chain=[component, prompt_tmpl, llm], verbose=True)

output = p.run(movie="仮面ライダーW")

> Running module 37c59e50-218d-4fdb-a06b-74317b7aa58c with input: 
movie: 仮面ライダーW

> Running module e4f98fec-18b0-4b36-b774-35b3ec1a268d with input: 
text: assistant: 1. 仮面ライダーW: ダブルフォーエバー
2. 仮面ライダーW: エターナルサーガ
3. 仮面ライダーW: フォーゼ×スカル
4. 仮面ライダーW: ジョーカーの逆襲
5. 仮面ライダーW: サイクロンの復讐
6. 仮面ライダーW: ヒートアップ！サイクロンジョーカー
7. 仮面ライダーW: フォーゼとの共闘
8. 仮面ライダーW: サイクロンの秘密
9. 仮面ライダーW:...

> Running module dc9e8314-efc7-4db8-9f5d-f22800fef2ea with input: 
messages: 以下に文章があります:

assistant: 1. 仮面ライダーW: ダブルフォーエバー
2. 仮面ライダーW: エターナルサーガ
3. 仮面ライダーW: フォーゼ×スカル
4. 仮面ライダーW: ジョーカーの逆襲
5. 仮面ライダーW: サイクロンの復讐
6. 仮面ライダーW: ヒートアップ！サイクロンジョーカー
7. 仮面ライダーW: フォーゼとの共闘
8. 仮面ライダーW: サイクロンの秘...

print(str(output))

assistant: 1. 仮面ライダーW: ダブルフォーエバー - "愛と勇気が永遠に続くように。"
2. 仮面ライダーW: エターナルサーガ - "運命の糸が絡み合い、悲劇が幕を開ける。"
3. 仮面ライダーW: フォーゼ×スカル - "星々の輝きが骸と共に舞い踊る。"
4. 仮面ライダーW: ジョーカーの逆襲 - "闇の中に潜む笑いが、復讐の始まりを告げる。"
5. 仮面ライダーW: サイクロンの復讐 - "怒りの嵐が巻き起こり、復讐の刃が振り下ろされる。"
6. 仮面ライダーW: ヒートアップ！サイクロンジョーカー - "炎と風が融合し、新たな力が解き放たれる。"
7. 仮面ライダーW: フォーゼとの共闘 - "異なる運命が交差し、共に戦う絆が生まれる。"
8. 仮面ライダーW: サイクロンの秘密 - "風が運ぶ秘密が、闇に包まれた真実を明かす。"
9. 仮面ライダーW: ジョーカーの誕生 - "闇の中で生まれた笑いが、新たな英雄を生み出す。"
10. 仮面ライダーW: ダブルの絆 - "二つの心が交わり、絆が深まる。"