72 lines
1.6 KiB
Python
72 lines
1.6 KiB
Python
from pydantic import BaseModel, Field
|
|
|
|
|
|
class DateRange(BaseModel):
|
|
from_: str = Field(alias="from")
|
|
to: str
|
|
|
|
|
|
class Entities(BaseModel):
|
|
people: list[str] | None = None
|
|
emails: list[str] | None = None
|
|
documents: list[str] | None = None
|
|
names: list[str] | None = None
|
|
links: list[str] | None = None
|
|
|
|
|
|
class Question(BaseModel):
|
|
text: str
|
|
asker: str = ""
|
|
asked_on: str = ""
|
|
variants: list[str] | None = None
|
|
hyde: list[str] | None = None
|
|
keywords: list[str] | None = None
|
|
entities: Entities | None = None
|
|
date_mentions: list[str] | None = None
|
|
date_range: DateRange | None = None
|
|
search_text: str = ""
|
|
|
|
|
|
class SearchAPIRequest(BaseModel):
|
|
question: Question
|
|
|
|
|
|
class SearchAPIItem(BaseModel):
|
|
message_ids: list[str]
|
|
|
|
|
|
class SearchAPIResponse(BaseModel):
|
|
results: list[SearchAPIItem]
|
|
|
|
|
|
class DenseEmbeddingItem(BaseModel):
|
|
index: int
|
|
embedding: list[float]
|
|
|
|
|
|
class DenseEmbeddingResponse(BaseModel):
|
|
data: list[DenseEmbeddingItem]
|
|
|
|
|
|
class SparseVector(BaseModel):
|
|
indices: list[int] = Field(default_factory=list)
|
|
values: list[float] = Field(default_factory=list)
|
|
|
|
|
|
class SparseEmbeddingResponse(BaseModel):
|
|
vectors: list[SparseVector]
|
|
|
|
|
|
class ChunkMetadata(BaseModel):
|
|
chat_name: str
|
|
chat_type: str
|
|
chat_id: str
|
|
chat_sn: str
|
|
thread_sn: str | None = None
|
|
message_ids: list[str]
|
|
start: str
|
|
end: str
|
|
participants: list[str] = Field(default_factory=list)
|
|
mentions: list[str] = Field(default_factory=list)
|
|
contains_forward: bool = False
|
|
contains_quote: bool = False
|