llamaindex_openai_question_generator_query_decomposition.py

1from llama_index.question_gen.openai import OpenAIQuestionGenerator
2from llama_index.core.tools import ToolMetadata
3from llama_index.core.indices.query.schema import QueryBundle
4
5# Initialize the OpenAI Question Generator
6# Note: Ensure OPENAI_API_KEY is set in your environment variables
7question_gen = OpenAIQuestionGenerator.from_defaults()
8
9# Define the tools (metadata) that the generator can use to answer sub-questions
10tools = [
11    ToolMetadata(
12        name="lyft_10k",
13        description="Provides information about Lyft financials for year 2021",
14    ),
15    ToolMetadata(
16        name="uber_10k",
17        description="Provides information about Uber financials for year 2021",
18    ),
19]
20
21# Define a complex query that needs decomposition
22query_str = "Compare and contrast the revenue growth and risk factors of Uber and Lyft in 2021"
23query_bundle = QueryBundle(query_str)
24
25# Generate sub-questions
26sub_questions = question_gen.generate(tools=tools, query=query_bundle)
27
28# Print the generated sub-questions
29for i, sub_question in enumerate(sub_questions):
30    print(f"Sub-question {i+1}: {sub_question.sub_question}")
31    print(f"Tool to use: {sub_question.tool_name}")
32    print("-" * 20)