from collections.abc import Sequence
from typing import Protocol

from techdocs_graphrag.models import Chunk, GraphExtraction, GraphFact, RetrievedChunk


class LanguageModel(Protocol):
    def complete(self, system_prompt: str, user_prompt: str, *, json_mode: bool = False) -> str: ...

    def embed(self, texts: Sequence[str]) -> list[list[float]]: ...


class KnowledgeGraph(Protocol):
    def ensure_schema(self) -> None: ...

    def await_indexes(self) -> None: ...

    def upsert_chunks(self, chunks: Sequence[Chunk], embeddings: Sequence[list[float]]) -> None: ...

    def upsert_extraction(self, chunk: Chunk, extraction: GraphExtraction) -> None: ...

    def vector_search(self, embedding: list[float], top_k: int) -> list[RetrievedChunk]: ...

    def search_entities(self, text: str, limit: int) -> list[str]: ...

    def expand(self, seed_keys: Sequence[str], max_hops: int, limit: int) -> list[GraphFact]: ...