pip install llama-index llama-index-agent-openai
pip install llama-index-vector-stores-chroma

from llama_index.core import SimpleDirectoryReader

# ディレクトリからドキュメントを読み込む
documents = SimpleDirectoryReader(
    input_dir="./data",
    recursive=True,
    filename_as_id=True
).load_data()

print(f"Loaded {len(documents)} documents")

from llama_index.core import VectorStoreIndex, Settings
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI

# モデルを設定
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
Settings.llm = OpenAI(model="gpt-4-turbo")

# ベクトルインデックスを作成
index = VectorStoreIndex.from_documents(documents)

# 後で使用するために永続化
index.storage_context.persist(persist_dir="./storage")

from llama_index.core.tools import QueryEngineTool

# クエリエンジンを作成
query_engine = index.as_query_engine(
    similarity_top_k=5,
    response_mode="compact"
)

# ツールとしてラップ
knowledge_tool = QueryEngineTool.from_defaults(
    query_engine=query_engine,
    name="knowledge_base",
    description="会社のポリシーや手順に関する情報について、社内ナレッジベースを検索します。"
)

from llama_index.agent.openai import OpenAIAgent

agent = OpenAIAgent.from_tools(
    tools=[knowledge_tool],
    verbose=True,
    system_prompt="""あなたは会社のナレッジベースにアクセスできる親切なアシスタントです。
    ポリシーに関する質問には、必ずナレッジベースを検索してから答えてください。"""
)

# エージェントにクエリを実行
response = agent.chat("What is our vacation policy?")
print(response)

from llama_index.core.tools import FunctionTool

# 追加のツールを作成
def get_current_date() -> str:
    """現在の日付を返します。"""
    from datetime import datetime
    return datetime.now().strftime("%Y-%m-%d")

date_tool = FunctionTool.from_defaults(fn=get_current_date)

# 複数のツールを持つエージェントを作成
agent = OpenAIAgent.from_tools(
    tools=[knowledge_tool, date_tool, calculator_tool],
    verbose=True
)

from llama_index.core.query_engine import SubQuestionQueryEngine
from llama_index.core.tools import QueryEngineTool

# 異なるインデックス用のツールを作成
tools = [
    QueryEngineTool.from_defaults(
        query_engine=hr_index.as_query_engine(),
        name="hr_policies",
        description="人事ポリシーと従業員ハンドブック"
    ),
    QueryEngineTool.from_defaults(
        query_engine=tech_index.as_query_engine(),
        name="technical_docs",
        description="技術文書と API"
    )
]

# サブクエスチョンエンジンを作成
query_engine = SubQuestionQueryEngine.from_defaults(
    query_engine_tools=tools
)

from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector

query_engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),
    query_engine_tools=tools
)

import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore

# Chroma をセットアップ
chroma_client = chromadb.PersistentClient(path="./chroma_db")
collection = chroma_client.get_or_create_collection("knowledge_base")

# Chroma でインデックスを作成
vector_store = ChromaVectorStore(chroma_collection=collection)
index = VectorStoreIndex.from_vector_store(vector_store)

from llama_index.core import Settings

# キャッシングを有効化
Settings.cache_enabled = True

# 本番環境ではより高速な埋め込みモデルを使用
Settings.embed_model = OpenAIEmbedding(
    model="text-embedding-3-small",
    embed_batch_size=100
)