from __future__ import annotations

from typing import Any, TYPE_CHECKING
from ketju.rag.llama_index.contracts import LlamaIndexQueryHandler
from llama_index.core.base.response.schema import RESPONSE_TYPE

if TYPE_CHECKING:
    from ketju.rag.llama_index.variants.simple import LlamaIndexRag

class SimpleLlamaIndexQueryHandler(LlamaIndexQueryHandler):

    def query(self, rag: LlamaIndexRag, question: str, **kwargs: Any) -> RESPONSE_TYPE:
        vector_index: Any = rag.vector_index
        engine = vector_index.as_query_engine()
        return engine.query(question)