from llama_index.question_gen.openai import OpenAIQuestionGenerator
from llama_index.core.tools import RetrieverTool, ToolMetadata
from llama_index.core import VectorStoreIndex, SummaryIndex, SimpleDirectoryReader, QueryBundle

documents = SimpleDirectoryReader("pliki").load_data()
vector_index = VectorStoreIndex.from_documents([documents[0]])
summary_index = SummaryIndex.from_documents([documents[1]])

vector_tool_metadata = ToolMetadata(
    name="Narzędzie wektorowe", 
    description="Użyj tej opcji do odpowiadania na pytania o starożytny Rzym."
)
summary_tool_metadata = ToolMetadata(
    name="Narzędzie podsumowujące", 
    description="Użyj ten opcji do odpowiadania na pytania o psy."
)

vector_tool = RetrieverTool(
    retriever=vector_index.as_retriever(), 
    metadata=vector_tool_metadata
)
summary_tool = RetrieverTool(
    retriever=summary_index.as_retriever(), 
    metadata=summary_tool_metadata
)

question_generator = OpenAIQuestionGenerator.from_defaults()
query_bundle = QueryBundle(query_str="Opowiedz mi o psach i starożytnym Rzymie")
sub_questions = question_generator.generate(
    tools=[vector_tool.metadata, summary_tool.metadata], 
    query=query_bundle
)

for sub_question in sub_questions:
    print(f"{sub_question.tool_name}: {sub_question.sub_question}")
