from llama_index.core import Settings
from llama_index.core.base.base_selector import MultiSelection, SingleSelection
from llama_index.core.llms import MockLLM
from llama_index.core.query_engine import CustomQueryEngine, RouterQueryEngine
from llama_index.core.selectors import BaseSelector
from llama_index.core.tools import QueryEngineTool


Settings.llm = MockLLM(max_tokens=32)


class StaticQueryEngine(CustomQueryEngine):
    label: str
    answer: str

    def custom_query(self, query_str):
        return f"{self.label}: {self.answer}"