schemas/entity.py

from typing import List, Optional

from pydantic import BaseModel


class EntitySearchRequest(BaseModel):
    """Схема запроса для поиска сущностей."""
    query: str
    dataset_id: int


class EntitySearchResponse(BaseModel):
    """Схема ответа с результатами поиска сущностей."""
    scores: List[float]
    entity_ids: List[str]


class EntitySearchWithTextRequest(BaseModel):
    """Схема запроса для поиска сущностей с текстами."""
    query: str
    dataset_id: int


class ChunkInfo(BaseModel):
    """Информация о чанке."""
    id: str
    text: str
    score: float
    type: str
    in_search_text: str


class EntitySearchWithTextResponse(BaseModel):
    """Схема ответа с результатами поиска сущностей и их текстами."""
    chunks: List[ChunkInfo]


class EntityTextRequest(BaseModel):
    """Схема запроса для сборки текста из сущностей."""
    entities: List[str]
    chunk_scores: Optional[dict[str, float]] = None
    include_tables: bool = True
    max_documents: Optional[int] = None


class EntityTextResponse(BaseModel):
    """Схема ответа со сборкой текста из сущностей."""
    text: str


class EntityNeighborsRequest(BaseModel):
    """Схема запроса для получения соседних чанков."""
    entities: List[str]
    max_distance: int = 1


class EntityNeighborsResponse(BaseModel):
    """Схема ответа с соседними чанками."""
    entities: List[str]