| | """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()]) |
| |
|
