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

!pip install -q llama-index

from google.colab import userdata
import os
import openai

openai.api_key = userdata.get('OPENAI_API_KEY')

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)

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

ドウデュースの主な勝ち鞍は、2021年の朝日杯フューチュリティステークス、2022年の東京優駿、2023年の有馬記念です。

import logging
import sys
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, force=True)     # force=Trueを追加
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

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

DEBUG:openai._base_client:Request options: {'method': 'post', 'url': '/embeddings', 'files': None, 'post_parser': <function Embeddings.create.<locals>.parser at 0x7ce9088b20e0>, 'json_data': {'input': ['ドウデュースの主な勝ち鞍は?'], 'model': <OpenAIEmbeddingModeModel.TEXT_EMBED_ADA_002: 'text-embedding-ada-002'>, 'encoding_format': 'base64'}}

DEBUG:llama_index.indices.utils:> Top 2 nodes:
> [Node d8122452-779f-44bb-8b38-a210a4d9580c] [Similarity score:             0.857178] 1番人気には推されたものの単勝のオッズは3.9倍で、これは1984年のグレード制導入以降、1990年アイネスフウジンの4.1倍に次ぐ皐月賞1番人気の低支持率オッズであった。レースでは道中後方から...
> [Node 6faee8f1-9f12-4784-9224-f046e4d4bc9f] [Similarity score:             0.85632] ドウデュース（欧字名:Do Deuce、2019年5月7日 - ）は、日本の競走馬。主な勝ち鞍は2021年の朝日杯フューチュリティステークス、2022年の東京優駿、2023年の有馬記念。
馬名の...

DEBUG:openai._base_client:Request options: {'method': 'post', 'url': '/chat/completions', 'files': None, 'json_data': {'messages': [{'role': <MessageRole.SYSTEM: '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': <MessageRole.USER: 'user'>, 'content': 'Context information is below.\n---------------------\nfile_path: data/ドウデュース.txt\n\n1番人気には推されたものの単勝のオッズは3.9倍で、これは1984年のグレード制導入以降、1990年アイネスフウジンの4.1倍に次ぐ皐月賞1番人気の低支持率オッズであった。レースでは道中後方からじっくり運び、最後の直線は外からメンバー最速となる上がり (競馬)3ハロン33秒8の鋭い脚を使って追い詰めたものの、3着に敗れた。\n次に、5月29日の東京優駿（日本ダービー）に出走。単勝4.2倍の3番人気でレースを迎えた。レースではスタートから後方に控え、前半1000メートルを58秒9という早いペースの中落ち着いてレースを進め、直線に入ってスパート。鞍上武豊の右ムチに応え加速し先団を捉え、大外から伸びてくるイクイノックスをクビ差抑え込み優勝。朝日杯以来のGI制覇となり、勝ちタイムは2分21秒9のダービーレコード(コースレコードは2018年ジャパンカップでアーモンドアイが記録した2分20秒6)。鞍上の武は、歴代最多を更新する2013年のキズナ以来となるダービー6勝目、かつ歴代最年長、史上初の50代でのダービージョッキーの名誉となった。また朝日杯フューチュリティステークスの勝ち馬が日本ダービーを制したのは、1994年のナリタブライアン（前身である朝日杯3歳ステークスを勝利）以来28年ぶりとなった。\n6月10日に国際競馬統括機関連盟が発表した「ロンジンワールドベストレースホースランキング」において、ドウデュースは日本ダービーを勝利した功績を評価され、シャフリヤールやエンブレムロードと並ぶレーティング120で第15位タイに位置づけられた。\n次走として凱旋門賞への出走を表明した。ドウデュースののんびりとした性格上、フランスへ赴いた場合に放牧と勘違いする可能性を踏まえ、日本ダービー直後の時点で友道は直行する予定と語っていたが、のちに松島はニエル賞を前哨戦として使ってから凱旋門賞に出走するプランもあり、その後はアメリカのブリーダーズカップ・ターフへの出走も検討していると述べた。最終的に友道は凱旋門賞に直行ではなくニエル賞を経由して凱旋門賞に出走すると発表した。いずれも後方からの競馬となったが休み明けや重馬場の影響もあり十分なパフォーマンスを発揮出来ず、それぞれ4着、19着に終わった。\n\nfile_path: data/ドウデュース.txt\n\nドウデュース（欧字名:Do Deuce、2019年5月7日 - ）は、日本の競走馬。主な勝ち鞍は2021年の朝日杯フューチュリティステークス、2022年の東京優駿、2023年の有馬記念。\n馬名の意味は「する＋テニス用語（勝利目前の意味）」。2021年のJRA賞最優秀2歳牡馬である。\n\n\n== 戦績 ==\n\n\n=== デビュー前 ===\n2019年5月7日、北海道安平町のノーザンファームで誕生。松島正昭が代表を務める株式会社キーファーズの所有馬となり、ノーザンファーム空港牧場で育成の後、栗東トレーニングセンターの友道康夫厩舎に入厩した。\n\n\n=== 2歳（2021年） ===\n9月5日に小倉競馬場で行われた2歳新馬戦（芝1800メートル）に武豊鞍上で出走。1番人気に推されると、レースは直線でガイアフォースとの追い比べをクビ差制してデビュー勝ちを果たした。\n次走はリステッド競走のアイビーステークスを選択。2番人気に推され、レースでは追い比べから抜け出すと、最後は追い込んできたグランシエロをクビ差凌いで優勝、デビュー2連勝とした。\n続いて朝日杯フューチュリティステークスに出走。重賞勝ち馬セリフォスやジオグリフをはじめとした自身と同じ無敗馬が多く顔を揃える中、3番人気に支持される。レースでは直線で外に出すと、先に抜け出していたセリフォスを半馬身差で差し切り優勝、無傷3連勝でGI初制覇を果たした。鞍上の武豊はこの競走22回目の挑戦で初制覇となり、日本の中央競馬 (JRA) の平地GI完全制覇までホープフルステークスを残すのみとした。また馬主である松島及びキーファーズにとっては初の単独所有馬によるGI勝利、並びに国内GI初制覇となった。\n\n\n=== 3歳（2022年） ===\n3歳初戦として、弥生賞ディープインパクト記念に出走。単勝オッズ2.2倍の1番人気に推された。道中は勝ち馬アスクビクターモアを見る形で追走。残り800メートル過ぎに後方からロジハービンが一気に進出したため、いったんポジションを下げる。そこから立て直し、ゴール前では勝ち馬を懸命に追い上げたがクビ差届かず2着に。デビューからの連勝は3でストップした。\n続いて、4月17日に行われた皐月賞に出走。\n---------------------\nGiven the context information and not prior knowledge, answer the query.\nQuery: ドウデュースの主な勝ち鞍は?\nAnswer: '}], 'model': 'gpt-3.5-turbo', 'stream': False, 'temperature': 0.1}}

ドウデュースの主な勝ち鞍は、2021年の朝日杯フューチュリティステークス、2022年の東京優駿、2023年の有馬記念です。

from llama_index import set_global_handler

set_global_handler("simple")

response = query_engine.query("ドウデュースの主な勝ち鞍は?")
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

1番人気には推されたものの単勝のオッズは3.9倍で、これは1984年のグレード制導入以降、1990年アイネスフウジンの4.1倍に次ぐ皐月賞1番人気の低支持率オッズであった。レースでは道中後方からじっくり運び、最後の直線は外からメンバー最速となる上がり (競馬)3ハロン33秒8の鋭い脚を使って追い詰めたものの、3着に敗れた。(snip)

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

ドウデュース（欧字名:Do Deuce、2019年5月7日 - ）は、日本の競走馬。主な勝ち鞍は2021年の朝日杯フューチュリティステークス、2022年の東京優駿、2023年の有馬記念。
馬名の意味は「する＋テニス用語（勝利目前の意味）」。2021年のJRA賞最優秀2歳牡馬である。(snip)

---------------------
Given the context information and not prior knowledge, answer the query.
Query: ドウデュースの主な勝ち鞍は?
Answer: 
**************************************************
** Response: **
assistant: ドウデュースの主な勝ち鞍は、2021年の朝日杯フューチュリティステークス、2022年の東京優駿、2023年の有馬記念です。
**************************************************


ドウデュースの主な勝ち鞍は、2021年の朝日杯フューチュリティステークス、2022年の東京優駿、2023年の有馬記念です。

!pip install wandb

from llama_index import set_global_handler
import llama_index

set_global_handler("wandb", run_args={"project": "llama-index-trace-test"})
wandb_callback = llama_index.global_handler

wandb: Appending key for api.wandb.ai to your netrc file: /root/.netrc
wandb: Streaming LlamaIndex events to W&B at https://wandb.ai/kun432/llama-index-trace-test/runs/XXXXXXX

wandb: `WandbCallbackHandler` is currently in beta.
wandb: Please report any issues to https://github.com/wandb/wandb/issues with the tag `llamaindex`.

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=1024, 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)

wandb: Logged trace tree to W&B.

wandb_callback.persist_index(index, index_name="simple_vector_store")

wandb: Adding directory to artifact (/content/wandb/run-20240104_175113-XXXXXXXX/files/storage)... Done. 0.0s

wandb: Logged trace tree to W&B.
ドウデュースの主な勝ち鞍は、2021年の朝日杯フューチュリティステークス、2022年の東京優駿、2023年の有馬記念です。

from llama_index.callbacks import CallbackManager
from llama_index.callbacks import LlamaDebugHandler, WandbCallbackHandler

llama_debug = LlamaDebugHandler(print_trace_on_end=True)
wandb_callback = WandbCallbackHandler(run_args={"project": "llama-index-trace-test"})

callback_manager = CallbackManager([llama_debug, wandb_callback])

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=1024, chunk_overlap=20)

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

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

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

**********
Trace: index_construction
    |_embedding ->  1.684219 seconds
**********
wandb: Logged trace tree to W&B.

query_engine = index.as_query_engine()

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

**********
Trace: query
    |_query ->  2.235186 seconds
      |_retrieve ->  0.201549 seconds
        |_embedding ->  0.198536 seconds
      |_synthesize ->  2.03339 seconds
        |_templating ->  3.1e-05 seconds
        |_llm ->  2.027773 seconds
**********
wandb: Logged trace tree to W&B.
ドウデュースの主な勝ち鞍は、2021年の朝日杯フューチュリティステークス、2022年の東京優駿、2023年の有馬記念です。

llama_debug = LlamaDebugHandler(print_trace_on_end=True)

llama_debug.print_trace_map()

**********
Trace: query
    |_query ->  2.235186 seconds
      |_retrieve ->  0.201549 seconds
        |_embedding ->  0.198536 seconds
      |_synthesize ->  2.03339 seconds
        |_templating ->  3.1e-05 seconds
        |_llm ->  2.027773 seconds
**********

llama_debug.get_event_time_info()

EventStats(total_secs=8.380683999999999, average_secs=1.1972405714285712, total_count=7)

from llama_index.callbacks import CBEventType

llama_debug.get_event_time_info(CBEventType.QUERY)

EventStats(total_secs=2.235186, average_secs=2.235186, total_count=1)

event_pairs = llama_debug.get_llm_inputs_outputs()
print(event_pairs[0][0])
print(event_pairs[0][1].payload.keys())
print(event_pairs[0][1].payload["response"])

CBEvent(event_type=<CBEventType.LLM: 'llm'>, payload={<EventPayload.MESSAGES: 'messages'>: "system: 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.\nuser: Context information is below.\n---------------------\nfile_path: data/ドウデュース.txt\n\n1番人気には推されたものの単勝のオッズは3.9倍で、(snip)\n---------------------\nGiven the context information and not prior knowledge, answer the query.\nQuery: ドウデュースの主な勝ち鞍は?\nAnswer: ", <EventPayload.ADDITIONAL_KWARGS: 'additional_kwargs'>: {}, <EventPayload.SERIALIZED: 'serialized'>: {'system_prompt': None, 'pydantic_program_mode': <PydanticProgramMode.DEFAULT: 'default'>, 'model': 'gpt-3.5-turbo', 'temperature': 0.1, 'max_tokens': None, 'additional_kwargs': {}, 'max_retries': 3, 'timeout': 60.0, 'default_headers': None, 'reuse_client': True, 'api_base': 'https://api.openai.com/v1', 'api_version': '', 'class_name': 'openai_llm'}}, time='01/04/2024, 18:20:49.627220', id_='3d56c414-e57e-4e8b-b0bd-6c1ddcc127dc')

dict_keys([<EventPayload.MESSAGES: 'messages'>, <EventPayload.RESPONSE: 'response'>])

assistant: ドウデュースの主な勝ち鞍は、2021年の朝日杯フューチュリティステークス、2022年の東京優駿、2023年の有馬記念です。

print(llama_debug.get_event_pairs(CBEventType.LLM)[0][0])
print(llama_debug.get_event_pairs(CBEventType.LLM)[0][1].payload.keys())
print(llama_debug.get_event_pairs(CBEventType.LLM)[0][1].payload["response"])