42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
"""Shared collection and kind validation for vector stores."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from tht.ports.vector import VectorStoreError
|
|
|
|
COLLECTION_KINDS = {
|
|
"schema_records": {"schema_table", "schema_column", "schema_relationship"},
|
|
"evidence": {"evidence"},
|
|
"memory": {"memory", "solved_question"},
|
|
}
|
|
ALLOWED_COLLECTIONS = frozenset(COLLECTION_KINDS)
|
|
ALLOWED_KINDS = frozenset().union(*COLLECTION_KINDS.values())
|
|
|
|
|
|
def validate_collection(collection: str) -> str:
|
|
if collection not in ALLOWED_COLLECTIONS:
|
|
raise VectorStoreError(f"Collection not allowed: {collection}")
|
|
return collection
|
|
|
|
|
|
def validate_collection_kinds(collection: str, kinds: list[str]) -> None:
|
|
invalid = set(kinds) - COLLECTION_KINDS[collection]
|
|
if invalid:
|
|
raise VectorStoreError(f"Kind not allowed for {collection}: {', '.join(sorted(invalid))}")
|
|
|
|
|
|
def validate_known_kinds(kinds: list[str]) -> None:
|
|
invalid = set(kinds) - ALLOWED_KINDS
|
|
if invalid:
|
|
raise VectorStoreError(f"Kind not allowed: {', '.join(sorted(invalid))}")
|
|
|
|
|
|
__all__ = [
|
|
"ALLOWED_COLLECTIONS",
|
|
"ALLOWED_KINDS",
|
|
"COLLECTION_KINDS",
|
|
"validate_collection",
|
|
"validate_collection_kinds",
|
|
"validate_known_kinds",
|
|
]
|