from praisonaiagents.knowledge.query_engine import (
    QueryMode,
    QueryResult,
    QueryEngineProtocol,
    get_query_engine_registry,
    decompose_question,
    SimpleQueryEngine,
    SubQuestionEngine
)

# Decompose complex questions
sub_questions = decompose_question(
    "What is Python and how do I install it?"
)
# Returns: ["What is Python?", "How do I install Python?"]

# Use sub-question engine
engine = SubQuestionEngine()
result = engine.query(
    "Compare Python and Java for web development",
    retriever=my_retriever
)

from praisonaiagents.knowledge.query_engine import QueryMode

class QueryMode(Enum):
    DEFAULT = "default"           # Direct retrieval + synthesis
    SUB_QUESTION = "sub_question" # Decompose into sub-questions
    SUMMARIZE = "summarize"       # Summarize retrieved content
    SQL = "sql"                   # SQL query generation
    ROUTER = "router"             # Route to appropriate engine

@dataclass
class QueryResult:
    answer: str
    sources: List[Dict[str, Any]] = field(default_factory=list)
    sub_questions: Optional[List[str]] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

class QueryEngineProtocol(Protocol):
    name: str
    mode: QueryMode
    
    def query(
        self,
        question: str,
        retriever: Any,
        top_k: int = 10,
        **kwargs
    ) -> QueryResult:
        """Answer a question using retrieved knowledge."""
        ...

from praisonaiagents.knowledge.query_engine import SimpleQueryEngine

engine = SimpleQueryEngine()

result = engine.query(
    question="What is machine learning?",
    retriever=my_retriever,
    top_k=5
)

print(result.answer)
print(f"Sources: {len(result.sources)}")

from praisonaiagents.knowledge.query_engine import SubQuestionEngine

engine = SubQuestionEngine()

result = engine.query(
    question="Compare the performance and ease of use of Python vs Java",
    retriever=my_retriever
)

print(f"Sub-questions: {result.sub_questions}")
print(f"Answer: {result.answer}")

from praisonaiagents.knowledge.query_engine import decompose_question

# Simple decomposition (keyword-based)
questions = decompose_question(
    "What is Python and how do I install it on Windows?"
)
# Returns: ["What is Python?", "How do I install Python on Windows?"]

# With LLM decomposition
questions = decompose_question(
    "Compare the pros and cons of microservices vs monolith",
    use_llm=True
)

from praisonaiagents.knowledge.query_engine import get_query_engine_registry

registry = get_query_engine_registry()

# List available engines
engines = registry.list_engines()  # ['default', 'sub_question', ...]

# Get engine by name
engine = registry.get("sub_question")

# Register custom engine
registry.register("custom", MyEngine)

from praisonaiagents import Agent, Knowledge

# Configure query mode
agent = Agent(
    instructions="You are a helpful assistant",
        knowledge={
        "sources": ["docs/"],
        "query_mode": "sub_question"  # or "default", "summarize"
    }
        "query_mode": "sub_question",  # or "default", "summarize"
    }
)

response = agent.chat("What are the benefits and drawbacks of this approach?")

from praisonaiagents.knowledge.query_engine import (
    QueryMode,
    QueryResult,
    get_query_engine_registry
)
from typing import Any

class MyQueryEngine:
    name = "my_engine"
    mode = QueryMode.DEFAULT
    
    def query(
        self,
        question: str,
        retriever: Any,
        top_k: int = 10,
        **kwargs
    ) -> QueryResult:
        # Retrieve relevant documents
        docs = retriever.retrieve(question, top_k=top_k)
        
        # Custom answer synthesis
        answer = self._synthesize(question, docs)
        
        return QueryResult(
            answer=answer,
            sources=[{"text": d.text} for d in docs]
        )
    
    def _synthesize(self, question: str, docs: list) -> str:
        # Your synthesis logic
        ...

# Register
registry = get_query_engine_registry()
registry.register("my_engine", MyQueryEngine)