from llama_index.core.selectors import PydanticMultiSelector
from llama_index.core.retrievers import RouterRetriever
from llama_index.core.tools import RetrieverTool
from llama_index.core import VectorStoreIndex, SummaryIndex, SimpleDirectoryReader
documents = SimpleDirectoryReader("pliki").load_data()

vector_index = VectorStoreIndex.from_documents([documents[0]])
summary_index = SummaryIndex.from_documents([documents[1]])

vector_retriever = vector_index.as_retriever()
summary_retriever = summary_index.as_retriever()

vector_tool = RetrieverTool.from_defaults(
    retriever=vector_retriever,
    description="Użyj tej opcji do odpowiadania na pytania o starożytny Rzym."
)
summary_tool = RetrieverTool.from_defaults(
    retriever=summary_retriever,
    description="Użyj ten opcji do odpowiadania na pytania o psy."
)

retriever = RouterRetriever(
    selector=PydanticMultiSelector.from_defaults(),
    retriever_tools=[
        vector_tool,
        summary_tool
    ]
)
response = retriever.retrieve(
    "Co możesz mi powiedzieć o staorożytnym Rzymie?"
)
for r in response:
    print(r.text)

'''
#inne możliwe zapytania
retriever.retrieve("Co możesz mi powiedzieć o starożytnym Rzymie?") 
retriever.retrieve("Powiedz wszystko co wiesz o psach.") 
retriever.retrieve("Opowiedz mi o psach w starożytnym Rzymie.")
'''