harness_zbt/backend/src/zbt/domains/knowledge/repository.py

from typing import Protocol

from zbt.domains.knowledge.models import KnowledgeDocument, KnowledgeSearchTest


class KnowledgeRepository(Protocol):
    def save_document(self, document: KnowledgeDocument) -> None: ...

    def get_document(self, document_id: str) -> KnowledgeDocument | None: ...

    def list_documents(self) -> list[KnowledgeDocument]: ...

    def save_search_test(self, search_test: KnowledgeSearchTest) -> None: ...

    def list_search_tests(self, document_id: str) -> list[KnowledgeSearchTest]: ...


class InMemoryKnowledgeRepository:
    def __init__(self) -> None:
        self._documents: dict[str, KnowledgeDocument] = {}
        self._search_tests: dict[str, KnowledgeSearchTest] = {}

    def save_document(self, document: KnowledgeDocument) -> None:
        self._documents[document.id] = document

    def get_document(self, document_id: str) -> KnowledgeDocument | None:
        return self._documents.get(document_id)

    def list_documents(self) -> list[KnowledgeDocument]:
        return list(self._documents.values())

    def save_search_test(self, search_test: KnowledgeSearchTest) -> None:
        self._search_tests[search_test.id] = search_test

    def list_search_tests(self, document_id: str) -> list[KnowledgeSearchTest]:
        return [
            search_test
            for search_test in self._search_tests.values()
            if search_test.document_id == document_id
        ]