|
"""Base schema for data structures.""" |
|
from abc import abstractmethod |
|
from dataclasses import dataclass |
|
from typing import Any, Dict, List, Optional |
|
|
|
from dataclasses_json import DataClassJsonMixin |
|
|
|
|
|
@dataclass |
|
class BaseDocument(DataClassJsonMixin): |
|
"""Base document. |
|
|
|
Generic abstract interfaces that captures both index structs |
|
as well as documents. |
|
|
|
""" |
|
|
|
|
|
text: Optional[str] = None |
|
doc_id: Optional[str] = None |
|
embedding: Optional[List[float]] = None |
|
|
|
|
|
extra_info: Optional[Dict[str, Any]] = None |
|
|
|
@classmethod |
|
@abstractmethod |
|
def get_type(cls) -> str: |
|
"""Get Document type.""" |
|
|
|
def get_text(self) -> str: |
|
"""Get text.""" |
|
if self.text is None: |
|
raise ValueError("text field not set.") |
|
return self.text |
|
|
|
def get_doc_id(self) -> str: |
|
"""Get doc_id.""" |
|
if self.doc_id is None: |
|
raise ValueError("doc_id not set.") |
|
return self.doc_id |
|
|
|
@property |
|
def is_doc_id_none(self) -> bool: |
|
"""Check if doc_id is None.""" |
|
return self.doc_id is None |
|
|
|
def get_embedding(self) -> List[float]: |
|
"""Get embedding. |
|
|
|
Errors if embedding is None. |
|
|
|
""" |
|
if self.embedding is None: |
|
raise ValueError("embedding not set.") |
|
return self.embedding |
|
|
|
@property |
|
def extra_info_str(self) -> Optional[str]: |
|
"""Extra info string.""" |
|
if self.extra_info is None: |
|
return None |
|
|
|
return "\n".join([f"{k}: {str(v)}" for k, v in self.extra_info.items()]) |
|
|
