src/query_processing/query_processor.py

import spacy
import json
from langchain.prompts import ChatPromptTemplate
from src.utils.api_key_manager import with_api_manager
from src.helpers.helper import remove_markdown

class QueryProcessor:
    def __init__(self):
        self.nlp = spacy.load("en_core_web_sm")

    def extract_entities(self, query):
        doc = self.nlp(query)
        return [ent.text for ent in doc.ents]

    @with_api_manager()
    async def classify_query(self, query, *, llm):
        template = \
"""You are an expert at classifying queries. Your task is to classify the given user query (descriptions are provided below) into one of the following categories:        
1. Basic -> The query is a simple question that can be answered without external internet resources. It is the simplest category.
2. Advanced -> The query requires up to 5 external internet resource references for a quick fact-check before answering. It is a moderate category.
3. Pro -> The query needs complex reasoning and up to 20 external internet resource references to answer accurately. It is a complex category.
4. Super -> The query requires very complex reasoning, dynamic sub-query generation, and up to 50 external internet resource references to answer accurately. It is a very complex category.
5. Ultra -> The query requires extensive research, extremely complex reasoning, and up to 500 external internet resource references to answer accurately. It is the most complex category.

Rules:
1. Only classify the query into one of the categories provided above.
2. Classify the query based on its complexity.
3. If the query contains "category: <category_name>" in it, then the query should be classified into the category that is mentioned in the query.

Examples:
1. Basic
    - "What is the capital of France?"
    - "How many continents are there in the world?"
    - "What is the chemical formula for water?"
    - "Who wrote the play 'Hamlet'?"
    - "Who painted the Mona Lisa?"
    - "What is Newton's first law of motion?"
    - "What is the sum of the angles in a triangle?"

2. Advanced
    - "What is the current market price of gold?"
    - "What is the current population of Japan?"
    - "Who won the Nobel Peace Prize in 2023?"
    - "What is the average temperature in London in October?"
    - "How many countries are in the European Union?"
    - "What are the top five best-selling novels of 2024?"

3. Pro
    - "Compare the political structures of ancient Greece and ancient Rome"
    - "How do climate change and deforestation contribute to the decline in biodiversity?"
    - "What are the major differences in healthcare systems between the U.S., Canada, and the U.K.?"
    - "How have cryptocurrencies evolved since their inception, and what are the key events that shaped their development?"
    - "What are the economic and social impacts of artificial intelligence on the job market?"
    - "What are the pros and cons of genetically modified organisms in agriculture?"

4. Super
    - "What are the long-term psychological and societal effects of internet addiction on teenagers?"
    - "How do various world religions view the concept of life after death, and what cultural practices reflect these beliefs?"
    - "What are the implications of quantum computing on encryption technologies, and how might they evolve in the next decade?"
    - "How have trade policies between China and the U.S. affected global supply chains since 2000?"
    - "What are the top five electric SUVs available in the market as of October 2024 for under $60,000, and how do they compare in terms of range, features, and performance?"
    - "What are the major obstacles to achieving carbon neutrality in heavy industries like steel and cement? What are the potential solutions?"

5. Ultra
    - "Create a comprehensive study on the economic, cultural, and environmental impacts of megacities around the world, including case studies from at least five different continents."
    - "Conduct a detailed comparative analysis of various renewable energy policies across G20 countries, evaluating their effectiveness, challenges, and future projections."
    - "Research the evolution of space exploration programs in the last 60 years, highlighting the key milestones, technological advancements, and geopolitical factors that shaped them."
    - "Analyze the factors that contribute to the spread of misinformation on social media platforms, and propose a detailed multi-faceted strategy for combating this issue globally, including technological, legal, and educational solutions."
    - "Create a comprehensive report on the evolution and future of human-machine interaction, covering past technological milestones, current trends in AI, and the potential ethical implications of future advancements."
    - "Evaluate the long-term effects of colonialism on economic development in Africa, including case studies of at least five different countries, and analyze how these effects differ based on colonial power, time of independence, and resource distribution."

Your response should ONLY contain the category in plain text without any formatting.

Query: {query}"""
        prompt = ChatPromptTemplate.from_template(template)
        response = await llm.ainvoke(prompt.format_messages(query=query))
        return response.content.strip()

    @with_api_manager()
    async def decompose_query(self, query, intent=None, *, llm):
        if intent:
            template = \
"""You are an expert in information retrieval and query analysis.
Your task is to decompose the given user query into smaller sub-queries based on the provided intent.

Rules:
1. Analyse if the query can be decomposed into smaller sub-queries first, given that the sub-queries will be used to search the internet.
    - Only decompose a query into sub-queries if:
        1. It is complex enough to benefit from decomposition. 
        For example, "What is the capital of France?", "Who is the current CEO of Apple?", and "What is the weather in Tokyo?" 
        are simple queries and do not benefit from decomposition, whereas "What are the benefits of AI in healthcare?", 
        "What are the latest advancements in AI for healthcare diagnostics?", and "What are the key factors influencing the stock market?" 
        are complex queries that benefit from decomposition.
        2. It is a nested query (i.e. the query contains sub-queries within it). 
        For example, "What are the latest reviews for the iPhone 15, Samsung Galaxy S23, and Google Pixel 7?" is a nested query as it contains three sub-queries: 
        "latest reviews for iPhone 15", "latest reviews for Samsung Galaxy S23", and "latest reviews for Google Pixel 7".
2. The sub-queries should be precise and specific to the information the user is seeking.
3. The sub-queries should be independent and not overlap with each other.
4. The sub-queries should be exhaustive and cover all the information the user is seeking.
5. The sub-queries should follow a funneling approach and filter out irrelevant information with each successive sub-query.
6. Each sub-query should be a standalone question that can be answered independently.
7. Each sub-query should be logically connected to the original query and the intent provided.

Your response should be in JSON format with the following structure (do not include any markdown):
{{
    "originalquery": "<original query>",
    "subqueries": [
        {{"subquery": "<sub query>", "explanation": "<explanation>"}},
        {{"subquery": "<sub query>", "explanation": "<explanation>"}},
        ...
    ]
}}

**OR**

[IMPORTANT] If the query is simple or if it is not beneficial to decompose the query, 
for internet search purposes, return the original query in plain text without any formatting and/or markdown.

Intent: {intent}
Query: {query}"""
            prompt = ChatPromptTemplate.from_template(template)
            messages = prompt.format_messages(intent=intent, query=query)
        else:
            template = \
"""You are an expert in information retrieval and query analysis.
Your task is to decompose the given user query into smaller sub-queries.

Rules:
1. Analyse if the query can be decomposed into smaller sub-queries first, given that the sub-queries will be used to search the internet.
    - Only decompose a query into sub-queries if:
        1. It is complex enough to benefit from decomposition. 
        For example, "What is the capital of France?", "Who is the current CEO of Apple?", and "What is the weather in Tokyo?" 
        are simple queries and do not benefit from decomposition, whereas "What are the benefits of AI in healthcare?", 
        "What are the latest advancements in AI for healthcare diagnostics?", and "What are the key factors influencing the stock market?" 
        are complex queries that benefit from decomposition.
        2. It is a nested query (i.e. the query contains sub-queries within it). 
        For example, "What are the latest reviews for the iPhone 15, Samsung Galaxy S23, and Google Pixel 7?" is a nested query as it contains three sub-queries: 
        "latest reviews for iPhone 15", "latest reviews for Samsung Galaxy S23", and "latest reviews for Google Pixel 7".
2. The sub-queries should be precise and specific to the information the user is seeking.
3. The sub-queries should be independent and not overlap with each other.
4. The sub-queries should be exhaustive and cover all the information the user is seeking.
5. The sub-queries should follow a funneling approach and filter out irrelevant information with each successive sub-query.
6. Each sub-query should be a standalone question that can be answered independently.
7. Each sub-query should be logically connected to the original query.

Your response should be in JSON format with the following structure (do not include any markdown):
{{
    "originalquery": "<original query>",
    "subqueries": [
        {{"subquery": "<sub query>", "explanation": "<explanation>"}},
        {{"subquery": "<sub query>", "explanation": "<explanation>"}},
        ...
    ]
}}

**OR**

[IMPORTANT] If the query is simple or if it is not beneficial to decompose the query, 
for internet search purposes, return the original query in plain text without any formatting and/or markdown.

Query: {query}"""
            prompt = ChatPromptTemplate.from_template(template)
            messages = prompt.format_messages(query=query)

        try:
            response = await llm.ainvoke(messages)
            cleaned_response = remove_markdown(response.content.strip())

            # Attempt to parse the response as JSON
            try:
                data = json.loads(cleaned_response)
            
                # Validate the presence of 'subqueries' in the JSON
                if "subqueries" in data and isinstance(data["subqueries"], list):
                    sub_queries = [item["subquery"] for item in data["subqueries"] if "subquery" in item]
                    original_query = data["originalquery"]

                    if len(sub_queries) > 1:
                        return sub_queries, original_query
                    else:
                        raise ValueError("Response JSON does not contain 'sub_queries'.")
                else:
                    raise ValueError("Response JSON does not contain 'sub_queries'.")            
            except json.JSONDecodeError:
                # If JSON parsing fails, assume the response is the original query in plain text
                print(f"JSON parsing failed. Assuming the response is the original query in plain text.")
                return [query], None
            except ValueError as ve:
                # Handle other potential value errors during parsing
                print(f"ValueError during JSON parsing: {ve}")
                return [query], None
        except Exception as e:
            print(f"Error decomposing query: {e}")
            return [query], None
    
    @with_api_manager()
    async def decompose_query_with_dependencies(self, query, intent=None, context=None, *, llm):
        if intent:
            if context:
                template = \
"""You are an expert in information retrieval and query analysis.

Your task is to decompose the given user query into smaller sub-queries, 
based on the provided intent and previous queries, and assign each sub-query a logical role.

Logical Roles:
1. Pre-requisite -> The sub-query is a pre-requisite for other sub-queries and must be answered first.
2. Dependent -> The sub-query is dependent on answers from one or more pre-requisite sub-queries and must be modified accordingly.
3. Independent -> The sub-query can be answered independently of the other sub-queries.

Rules:
1. Analyse if the query can be decomposed into smaller sub-queries first, given that the sub-queries will be used to search the internet.
    - Only decompose a query into sub-queries if:
        1. It is complex enough to benefit from decomposition. 
        For example, "What is the capital of France?", "Who is the current CEO of Apple?", and "What is the weather in Tokyo?" 
        are simple queries and do not benefit from decomposition, whereas "What are the benefits of AI in healthcare?", 
        "What are the latest advancements in AI for healthcare diagnostics?", and "What are the key factors influencing the stock market?" 
        are complex queries that benefit from decomposition.
        2. It is a nested query (i.e. the query contains sub-queries within it). 
        For example, "What are the latest reviews for the iPhone 15, Samsung Galaxy S23, and Google Pixel 7?" is a nested query as it contains three sub-queries: 
        "latest reviews for iPhone 15", "latest reviews for Samsung Galaxy S23", and "latest reviews for Google Pixel 7".
2. The sub-queries should be precise and specific to the information the user is seeking.
3. The sub-queries should be as distinct as possible, 
minimizing unnecessary overlap while acknowledging that some dependencies may exist due to the logical roles assigned.
4. The sub-queries should be exhaustive and cover all the information the user is seeking.
5. The sub-queries should be organized in a logical manner, 
with the pre-requisite sub-queries addressing foundational information needed for dependent sub-queries.
6. The sub-queries should not repeat any of the previous queries.
7. The dependent sub-queries can overlap with previous sub-queries and/or sub-sub-queries, 
i.e. the dependent sub-query can be a subset of the previous sub-queries (excluding previous independent sub-queries), 
other than also being dependent on the current lot of pre-requisite sub-queries. 
This essentially means that the previous dependent sub-queries/sub-sub-queries can also play a pre-requisite role for the current dependent sub-query,
without "officially" changing their role.
8. Do not decompose sub-sub-queries further into sub-sub-sub-queries.

Your response should be in JSON format with the following structure (do not include any markdown):
{{
    "originalquery": "original query",
    "subqueries": [
        {{"subquery": "sub query", "role": "logical role", 
        "dependson": [[This sub-list is for the dependencies from the previous queries, containing indices of the sub-queries from the previous context that this sub-query depends on 
        (if role is 'dependent' AND this sub-query is dependent on previous queries, else empty list)], 
        [This sub-list is for the dependencies from the current lot of sub-queries containing indices of the sub-queries from the current query that this sub-query depends on 
        (if role is 'dependent', else empty list)]]}},
        {{"subquery": "sub query", "role": "logical role", 
        "dependson": [[This sub-list is for the dependencies from the previous queries, containing indices of the sub-queries from the previous context that this sub-query depends on 
        (if role is 'dependent' AND this sub-query is dependent on previous queries, else empty list)], 
        [This sub-list is for the dependencies from the current lot of sub-queries containing indices of the sub-queries from the current query that this sub-query depends on 
        (if role is 'dependent', else empty list)]]}},
        ...
    ]
}}

**OR**

[IMPORTANT] If the query is simple or if it is not beneficial to decompose the query, 
for internet search purposes, return the original query in plain text without any formatting and/or markdown.

---------------------------------------------------------------------------------------------------------------

The following is an example where the previous context is a list of three dictionaries/JSON objects, with each of the dictionaries/JSON objects in the same format mentioned above:
{{
    "originalquery": "What are the top three electric SUVs available in the market as of October 2024 for under $60,000?",
    "subqueries": [
        {{"subquery": "What are the top three electric SUVs available in the market as of October 2024?", "role": "independent", 
        "dependson": []}},
        {{"subquery": "How do they compare in terms of range, features, and performance?", "role": "dependent", 
        "dependson": [[0, None, None] <This means that this sub-query depends on the first sub-query (i.e. index 0) from the first lot of previous queries (i.e. first dictionary/JSON object), 
        none from the second lot of previous queries (i.e. second dictionary/JSON object), and none from the third lot of sub-queries (i.e. third dictionary/JSON object)>, 
        []<This means there are no sub-queries this sub-query depends on from the current lot of sub-queries (i.e. current query)>]}},
        {{"subquery": "What are the key features of the top three electric SUVs available in the market as of October 2024?", "role": "dependent", 
        "dependson": [[0, 1, None] <This means that this sub-query depends on the first and second sub-queries (i.e. indices 0 and 1) from the first and second lot of previous queries (i.e. first and second dictionary/JSON objects) respectively, 
        none from the third lot of previous queries (i.e. third dictionary/JSON object)>, [0] <This means that this sub-query depends on the first sub-query (i.e. index 0) from the current lot of sub-queries (i.e. current query)>
        >}},
        ...
    ]
}}

[IMPORTANT] Note: The length of the first "dependson" sub-list must be equal to the number of dictionaries/JSON objects in the list of previous queries, 
and the length of the second "dependson" sub-list must be equal to the number of sub-queries in the current query.

---------------------------------------------------------------------------------------------------------------

Intent: 
{intent}

Previous Queries:
{context}

Query: 
{query}"""
                prompt = ChatPromptTemplate.from_template(template)
                messages = prompt.format_messages(intent=intent, context=context, query=query)
            else:
                template = \
"""You are an expert in information retrieval and query analysis.

Your task is to decompose the given user query into smaller sub-queries, 
based on the provided intent, and assign each sub-query a logical role.

Logical Roles:
1. Pre-requisite -> The sub-query is a pre-requisite for other sub-queries and must be answered first.
2. Dependent -> The sub-query is dependent on answers from one or more pre-requisite sub-queries and must be modified accordingly.
3. Independent -> The sub-query can be answered independently of the other sub-queries.

Rules:
1. Analyse if the query can be decomposed into smaller sub-queries first, given that the sub-queries will be used to search the internet.
    - Only decompose a query into sub-queries if:
        1. It is complex enough to benefit from decomposition. 
        For example, "What is the capital of France?", "Who is the current CEO of Apple?", and "What is the weather in Tokyo?" 
        are simple queries and do not benefit from decomposition, whereas "What are the benefits of AI in healthcare?", 
        "What are the latest advancements in AI for healthcare diagnostics?", and "What are the key factors influencing the stock market?" 
        are complex queries that benefit from decomposition.
        2. It is a nested query (i.e. the query contains sub-queries within it). 
        For example, "What are the latest reviews for the iPhone 15, Samsung Galaxy S23, and Google Pixel 7?" is a nested query as it contains three sub-queries: 
        "latest reviews for iPhone 15", "latest reviews for Samsung Galaxy S23", and "latest reviews for Google Pixel 7".
2. The sub-queries should be precise and specific to the information the user is seeking.
3. The sub-queries should be as distinct as possible, 
minimizing unnecessary overlap while acknowledging that some dependencies may exist due to the logical roles assigned.
4. The sub-queries should be exhaustive and cover all the information the user is seeking.
5. The sub-queries should be organized in a logical manner, 
with the pre-requisite sub-queries addressing foundational information needed for dependent sub-queries.

Your response should be in JSON format with the following structure (do not include any markdown):
{{
    "originalquery": "original query",
    "subqueries": [
        {{"subquery": "sub query", "role": "logical role", 
        "dependson": [list of indices of sub_queries this sub query depends on (if role is 'dependent', else empty list)]}},
        {{"subquery": "sub query", "role": "logical role", 
        "dependson": [list of indices of sub_queries this sub query depends on (if role is 'dependent', else empty list)]}},
        ...
    ]
}}

**OR**

[IMPORTANT] If the query is simple or if it is not beneficial to decompose the query, 
for internet search purposes, return the original query in plain text without any formatting and/or markdown.

Intent: 
{intent}

Query: 
{query}"""
                prompt = ChatPromptTemplate.from_template(template)
                messages = prompt.format_messages(intent=intent, query=query)
        else:
            if context:
                template = \
"""You are an expert in information retrieval and query analysis.

Your task is to decompose the given user query into smaller sub-queries, 
based on the provided context, and assign each sub-query a logical role.

Logical Roles:
1. Pre-requisite -> The sub-query is a pre-requisite for other sub-queries and must be answered first.
2. Dependent -> The sub-query is dependent on answers from one or more pre-requisite sub-queries and must be modified accordingly.
3. Independent -> The sub-query can be answered independently of the other sub-queries.

Rules:
1. Analyse if the query can be decomposed into smaller sub-queries first, given that the sub-queries will be used to search the internet.
    - Only decompose a query into sub-queries if:
        1. It is complex enough to benefit from decomposition. 
        For example, "What is the capital of France?", "Who is the current CEO of Apple?", and "What is the weather in Tokyo?" 
        are simple queries and do not benefit from decomposition, whereas "What are the benefits of AI in healthcare?", 
        "What are the latest advancements in AI for healthcare diagnostics?", and "What are the key factors influencing the stock market?" 
        are complex queries that benefit from decomposition.
        2. It is a nested query (i.e. the query contains sub-queries within it). 
        For example, "What are the latest reviews for the iPhone 15, Samsung Galaxy S23, and Google Pixel 7?" is a nested query as it contains three sub-queries: 
        "latest reviews for iPhone 15", "latest reviews for Samsung Galaxy S23", and "latest reviews for Google Pixel 7".
2. The sub-queries should be precise and specific to the information the user is seeking.
3. The sub-queries should be as distinct as possible, 
minimizing unnecessary overlap while acknowledging that some dependencies may exist due to the logical roles assigned.
4. The sub-queries should be exhaustive and cover all the information the user is seeking.
5. The sub-queries should be organized in a logical manner, 
with the pre-requisite sub-queries addressing foundational information needed for dependent sub-queries.
6. The sub-queries should not repeat any of the previous queries.
7. The dependent sub-queries can overlap with previous sub-queries and/or sub-sub-queries, 
i.e. the dependent sub-query can be a subset of the previous sub-queries (excluding previous independent sub-queries), 
other than also being dependent on the current lot of pre-requisite sub-queries.
This essentially means that the previous dependent sub-queries/sub-sub-queries can also play a pre-requisite role for the current dependent sub-query,
without "officially" changing their role.
8. Do not decompose sub-sub-queries further into sub-sub-sub-queries.

Your response should be in JSON format with the following structure (do not include any markdown):
{{
    "originalquery": "original query",
    "subqueries": [
        {{"subquery": "sub query", "role": "logical role", 
        "dependson": [[This sub-list is for the dependencies from the previous queries, containing indices of the sub-queries from the previous context that this sub-query depends on 
        (if role is 'dependent' AND this sub-query is dependent on previous queries, else empty list)], 
        [This sub-list is for the dependencies from the current lot of sub-queries containing indices of the sub-queries from the current query that this sub-query depends on 
        (if role is 'dependent', else empty list)]]}},
        {{"subquery": "sub query", "role": "logical role", 
        "dependson": [[This sub-list is for the dependencies from the previous queries, containing indices of the sub-queries from the previous context that this sub-query depends on 
        (if role is 'dependent' AND this sub-query is dependent on previous queries, else empty list)], 
        [This sub-list is for the dependencies from the current lot of sub-queries containing indices of the sub-queries from the current query that this sub-query depends on 
        (if role is 'dependent', else empty list)]]}},
        ...
    ]
}}

**OR**

[IMPORTANT] If the query is simple or if it is not beneficial to decompose the query, 
for internet search purposes, return the original query in plain text without any formatting and/or markdown.

---------------------------------------------------------------------------------------------------------------

The following is an example where the previous context is a list of three dictionaries/JSON objects, with each of the dictionaries/JSON objects in the same format mentioned above:
{{
    "originalquery": "What are the top three electric SUVs available in the market as of October 2024 for under $60,000?",
    "subqueries": [
        {{"subquery": "What are the top three electric SUVs available in the market as of October 2024?", "role": "independent", 
        "dependson": []}},
        {{"subquery": "How do they compare in terms of range, features, and performance?", "role": "dependent", 
        "dependson": [[0, None, None] <This means that this sub-query depends on the first sub-query (i.e. index 0) from the first lot of previous queries (i.e. first dictionary/JSON object), 
        none from the second lot of previous queries (i.e. second dictionary/JSON object), and none from the third lot of sub-queries (i.e. third dictionary/JSON object)>, 
        []<This means there are no sub-queries this sub-query depends on from the current lot of sub-queries (i.e. current query)>]}},
        {{"subquery": "What are the key features of the top three electric SUVs available in the market as of October 2024?", "role": "dependent", 
        "depends_on": [[0, 1, None] <This means that this sub-query depends on the first and second sub-queries (i.e. indices 0 and 1) from the first and second lot of previous queries (i.e. first and second dictionary/JSON objects) respectively, 
        none from the third lot of previous queries (i.e. third dictionary/JSON object)>, [0] <This means that this sub-query depends on the first sub-query (i.e. index 0) from the current lot of sub-queries (i.e. current query)>
        >}},
        ...
    ]
}}

[IMPORTANT] Note: The length of the first "dependson" sub-list must be equal to the number of dictionaries/JSON objects in the list of previous queries, 
and the length of the second "dependson" sub-list must be equal to the number of sub-queries in the current query.

---------------------------------------------------------------------------------------------------------------

Previous Queries:
{context}

Query: 
{query}"""
                prompt = ChatPromptTemplate.from_template(template)
                messages = prompt.format_messages(context=context, query=query)
            else:
                template = \
"""You are an expert in information retrieval and query analysis.

Your task is to decompose the given user query into smaller sub-queries, and assign each sub-query a logical role.

Logical Roles:
1. Pre-requisite -> The sub-query is a pre-requisite for other sub-queries and must be answered first.
2. Dependent -> The sub-query is dependent on answers from one or more pre-requisite sub-queries and must be modified accordingly.
3. Independent -> The sub-query can be answered independently of the other sub-queries.

Rules:
1. Analyse if the query can be decomposed into smaller sub-queries first, given that the sub-queries will be used to search the internet.
    - Only decompose a query into sub-queries if:
        1. It is complex enough to benefit from decomposition. 
        For example, "What is the capital of France?", "Who is the current CEO of Apple?", and "What is the weather in Tokyo?" 
        are simple queries and do not benefit from decomposition, whereas "What are the benefits of AI in healthcare?", 
        "What are the latest advancements in AI for healthcare diagnostics?", and "What are the key factors influencing the stock market?" 
        are complex queries that benefit from decomposition.
        2. It is a nested query (i.e. the query contains sub-queries within it). 
        For example, "What are the latest reviews for the iPhone 15, Samsung Galaxy S23, and Google Pixel 7?" is a nested query as it contains three sub-queries: 
        "latest reviews for iPhone 15", "latest reviews for Samsung Galaxy S23", and "latest reviews for Google Pixel 7".
2. The sub-queries should be precise and specific to the information the user is seeking.
3. The sub-queries should be as distinct as possible, 
minimizing unnecessary overlap while acknowledging that some dependencies may exist due to the logical roles assigned.
4. The sub-queries should be exhaustive and cover all the information the user is seeking.
5. The sub-queries should be organized in a logical manner, 
with the pre-requisite sub-queries addressing foundational information needed for dependent sub-queries.

Your response should be in JSON format with the following structure (do not include any markdown):
{{
    "originalquery": "original query",
    "subqueries": [
        {{"subquery": "sub query", "role": "logical role", 
        "dependson": [list of indices of sub_queries this sub query depends on (if role is 'dependent', else empty list)]}},
        {{"subquery": "sub query", "role": "logical role", 
        "dependson": [list of indices of sub_queries this sub query depends on (if role is 'dependent', else empty list)]}},
        ...
    ]
}}

**OR**

[IMPORTANT] If the query is simple or if it is not beneficial to decompose the query, 
for internet search purposes, return the original query in plain text without any formatting and/or markdown.

Query: 
{query}"""
                prompt = ChatPromptTemplate.from_template(template)
                messages = prompt.format_messages(query=query)

        try:
            response = await llm.ainvoke(messages)
            cleaned_response = remove_markdown(response.content.strip())
            print(f"Cleaned Response: {cleaned_response}")

            # Attempt to parse the response as JSON 
            try:
                data = json.loads(cleaned_response)
            
                # Validate the presence of 'subqueries' in the JSON
                if "subqueries" in data and isinstance(data["subqueries"], list):
                    sub_queries = [item["subquery"] for item in data["subqueries"] if "subquery" in item]

                    if len(sub_queries) > 1:
                        roles = [item["role"] for item in data["subqueries"] if "role" in item]
                        dependencies = [item["dependson"] for item in data["subqueries"] if "dependson" in item]
                        return data, sub_queries, roles, dependencies
                    else:
                        # Assume the model made a mistake and was supposed to return the original query
                        print("No sub-queries found in the response.")
                        return None, [query], ['root'], []
                else:
                    # Assume the model made a mistake and was supposed to return the original query
                    print("Response JSON does not contain 'sub_queries'.")
                    return None, [query], ['root'], []            
            except json.JSONDecodeError:
                # If JSON parsing fails, assume the response is the original query in plain text
                print(f"JSON parsing failed. Assuming the response is the original query in plain text.")
                return None, [query], ['root'], []
            except ValueError as ve:
                # Handle other potential value errors during parsing
                print(f"ValueError during JSON parsing: {ve}")
                return None, [query], ['root'], []
        except Exception as e:
            print(f"Error decomposing query: {e}")
            return None, [query], ['root'], []
    
    @with_api_manager()
    async def modify_query(self, query, context, *, llm):
        combined_context = "\n\n".join(f"Context {i}:\n\n{content}" for i, content in enumerate(context, 1))
        template = \
"""[IMPORTANT]
This is a very important task. 
Please take a deep breath, read the rules and the context VERY carefully, 
and then think step-by-step to reason your way to the answer.

[PROMPT]
You are an expert in query refinement.
Your task is to modify the given query based on the context(s) provided.

Rules:
1. Assume that the user is unable to see the query and the context(s) provided so he has no idea what the original query is.
2. The modified query should maintain the intent of the original query.
3. The modified query should be precise and specific to the information the query is seeking.
4. The modified query should be concise but include essential information from the context.

Your response should ONLY contain the modified query in plain text without any formatting.

Query: {query}

Context:
{context}"""
        prompt = ChatPromptTemplate.from_template(template)
        messages = prompt.format_messages(query=query, context=combined_context)
        response = await llm.ainvoke(messages)
        return response.content.strip()
    
    @with_api_manager()
    async def get_query_intent(self, query, *, llm):
        template = \
"""You are the best psychologist in the world with many decades of experience. 
Your task is to think step-by-step and create 10 different versions of the given user query,
considering that these versions should be semantically clear and help describe the user's intention behind the query.

Your goal is to think step-by-step and look at the 10 versions of the query, along with the original query,
and then take a deep breath to conduct a thorough analysis of the user's intention.
Your analysis should be concise, detailed, and to the point. It should also contain the kind of answers and information the user is looking for.

[IMPORTANT] Your response should ONLY be the intent analysis of the user's query without any formatting. 
Do not include the reasoning process or the 10 versions of the query in your response.

Original query: {query}"""
  
        prompt = ChatPromptTemplate.from_template(template)
        response = await llm.ainvoke(prompt.format_messages(query=query))
        return response.content.strip()

if __name__ == "__main__":
    import asyncio

    query = "Given the current state of scientific knowledge, how can cold-fusion be achieved?" 
    contexts = []
    processor = QueryProcessor()

    intent = asyncio.run(processor.get_query_intent(query))
    data1, context1, _, _ = asyncio.run(processor.decompose_query_with_dependencies(query, intent=intent))
    contexts.append(data1)
    print(f"Intent:\n{intent}\n")
    print(f"Context 1:\n{contexts}\n")

    if len(context1) > 1:
        data2, context2, _, _ = asyncio.run(processor.decompose_query_with_dependencies(
            context1[4], 
            context=contexts
        )
    )
        contexts.append(data2)
        print(f"Context 2:\n{contexts}\n")

        if len(context2) > 1:
            _, sub_queries, roles, dependencies = asyncio.run(processor.decompose_query_with_dependencies(
                context2[1], 
                context=contexts
            ) 
        )
            print(f"Sub-queries:\n{sub_queries}\n")
            print(f"Roles:\n{roles}\n")
            print(f"Dependencies:\n{dependencies}\n")