Create-knowledge-find-decode/query_planning/ports.py

"""Ports owned by the query planning bounded context."""
from __future__ import annotations

from typing import Any, Protocol
from uuid import UUID

from query_planning.domain import GenerationRequest, GeneratorKind, GeneratorOutput, KnowledgeNeedDraft


class QueryGenerator(Protocol):
    kind: GeneratorKind

    def generate(self, request: GenerationRequest) -> GeneratorOutput:
        ...


class LegacyQueryBatchSink(Protocol):
    def create_query_batch(self, **kwargs: Any) -> Any:
        ...

    def add_query(self, **kwargs: Any) -> Any:
        ...


class QueryPlanningStore(Protocol):
    def create_plan(self, **kwargs: Any) -> UUID:
        ...

    def add_knowledge_need(self, *, plan_id: UUID, need: KnowledgeNeedDraft) -> UUID:
        ...

    def link_plan_batch(self, *, plan_id: UUID, batch_id: UUID) -> None:
        ...

    def link_query_need(self, *, query_id: UUID, knowledge_need_id: UUID) -> None:
        ...