Spaces:
Running
Running
File size: 10,497 Bytes
e70c280 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 | """
File: /data.py
Created Date: Sunday October 19th 2025
Author: Christian Nonis <alch.infoemail@gmail.com>
-----
Last Modified: Wednesday May 20th 2026 8:41:53 pm
Modified By: Christian Nonis <alch.infoemail@gmail.com>
-----
"""
from abc import ABC, abstractmethod
from typing import List, Literal, Optional, Tuple
from pydantic import BaseModel
from src.constants.data import Brain, KGChanges, Observation, StructuredData, TextChunk
class SearchResult(BaseModel):
"""
Search result model.
"""
text_chunks: List[TextChunk]
observations: List[Observation]
class DataClient(ABC):
"""
Abstract base class for data clients.
"""
@abstractmethod
def save_text_chunk(self, text_chunk: TextChunk, brain_id: str) -> TextChunk:
"""
Save data to the data client.
"""
raise NotImplementedError("save method not implemented")
@abstractmethod
def save_observations(
self, observations: List[Observation], brain_id: str
) -> List[Observation]:
"""
Save a list of observations to the data client.
"""
raise NotImplementedError("save_observations method not implemented")
@abstractmethod
def search(self, text: str, brain_id: str) -> SearchResult:
"""
Search data by text and return a list of text chunks and observations.
"""
raise NotImplementedError("search method not implemented")
def search_bm25(
self, text: str, brain_id: str, limit: int = 10
) -> List[Tuple[TextChunk, float]]:
raise NotImplementedError(
"search_bm25 requires SEARCH_ENABLED=true and DATA_DB=postgresql"
)
@abstractmethod
def get_text_chunks_by_ids(
self, ids: List[str], with_observations: bool, brain_id: str
) -> Tuple[List[TextChunk], List[Observation]]:
"""
Get data by their IDs.
"""
raise NotImplementedError("get_by_ids method not implemented")
@abstractmethod
def get_text_chunks(
self,
brain_id: str,
limit: int = 10,
skip: int = 0,
query_text: str = None,
metadata_eq: dict[str, str] | None = None,
order: Literal["asc", "desc"] = "desc",
) -> Tuple[List[TextChunk], int]:
"""
Get text chunks with pagination, optional text search, and metadata equality filters.
"""
raise NotImplementedError("get_text_chunks method not implemented")
@abstractmethod
def save_structured_data(
self, structured_data: StructuredData, brain_id: str
) -> StructuredData:
"""
Save a structured data to the data client.
"""
raise NotImplementedError("save_structured_data method not implemented")
@abstractmethod
def create_brain(self, name_key: str) -> Brain:
"""
Create a new brain in the data client.
"""
raise NotImplementedError("create_brain method not implemented")
@abstractmethod
def get_brain(self, name_key: str) -> Brain:
"""
Get a brain from the data client.
"""
raise NotImplementedError("get_brain method not implemented")
@abstractmethod
def get_brain_by_pat(self, pat: str) -> Optional[Brain]:
"""
Get a brain from the data client by its PAT.
"""
raise NotImplementedError("get_brain_by_pat method not implemented")
@abstractmethod
def get_brains_list(self) -> List[Brain]:
"""
Get the list of brains from the data client.
"""
raise NotImplementedError("get_brains_list method not implemented")
@abstractmethod
def save_kg_changes(self, kg_changes: KGChanges, brain_id: str) -> KGChanges:
"""
Save a KG changes to the data client.
"""
raise NotImplementedError("save_kg_changes method not implemented")
@abstractmethod
def get_structured_data_by_id(self, id: str, brain_id: str) -> StructuredData:
"""
Get structured data by ID.
"""
raise NotImplementedError("get_structured_data_by_id method not implemented")
@abstractmethod
def get_structured_data_list(
self,
brain_id: str,
limit: int = 10,
skip: int = 0,
types: list[str] = None,
query_text: str = None,
) -> Tuple[list[StructuredData], int]:
"""
Retrieve structured data entries for a specific brain with optional filtering and pagination.
Parameters:
brain_id (str): Identifier of the brain to query.
limit (int): Maximum number of items to return.
skip (int): Number of items to skip (offset) for pagination.
types (list[str] | None): If provided, restrict results to these structured data types.
query_text (str | None): If provided, filter entries that match the text query.
Returns:
Tuple[list[StructuredData], int]: Tuple containing the list of matching StructuredData objects and the total count of matching objects.
"""
raise NotImplementedError("get_structured_data_list method not implemented")
@abstractmethod
def get_structured_data_types(self, brain_id: str) -> list[str]:
"""
Get all unique types from structured data.
"""
raise NotImplementedError("get_structured_data_types method not implemented")
@abstractmethod
def get_observation_by_id(self, id: str, brain_id: str) -> Observation:
"""
Get observation by ID.
"""
raise NotImplementedError("get_observation_by_id method not implemented")
@abstractmethod
def get_observations_list(
self,
brain_id: str,
limit: int = 10,
skip: int = 0,
resource_id: str = None,
labels: list[str] = None,
query_text: str = None,
) -> list[Observation]:
"""
Retrieve a filtered, paginated list of observations for a specific brain.
Parameters:
brain_id (str): Identifier of the brain to query.
limit (int): Maximum number of observations to return.
skip (int): Number of observations to skip (offset) for pagination.
resource_id (str | None): If provided, restrict results to observations for this resource.
labels (list[str] | None): If provided, include only observations that have at least one of these labels.
query_text (str | None): If provided, include only observations whose text matches this query.
Returns:
observations (list[Observation]): A list of matching Observation objects ordered according to the backend's default sort.
"""
raise NotImplementedError("get_observations_list method not implemented")
@abstractmethod
def get_observation_labels(self, brain_id: str) -> list[str]:
"""
Retrieve all unique observation labels for the specified brain.
Parameters:
brain_id (str): Identifier of the brain whose observation labels should be retrieved.
Returns:
list[str]: Unique label strings present in the brain's observations.
"""
raise NotImplementedError("get_observation_labels method not implemented")
@abstractmethod
def get_changelog_by_id(self, id: str, brain_id: str) -> KGChanges:
"""
Retrieve a changelog entry by its identifier within a brain.
Parameters:
id (str): Identifier of the changelog to retrieve.
brain_id (str): Identifier of the brain that contains the changelog.
Returns:
KGChanges: The changelog entry matching the provided `id` in the specified `brain_id`.
"""
raise NotImplementedError("get_changelog_by_id method not implemented")
@abstractmethod
def get_changelogs_list(
self,
brain_id: str,
limit: int = 10,
skip: int = 0,
types: list[str] = None,
query_text: str = None,
) -> list[KGChanges]:
"""
Retrieve a paginated list of knowledge-graph changelogs for a brain.
Parameters:
brain_id (str): Identifier of the brain whose changelogs to retrieve.
limit (int): Maximum number of changelogs to return.
skip (int): Number of changelogs to skip (offset) for pagination.
types (list[str] | None): Optional list of changelog types to filter results by.
query_text (str | None): Optional text to filter changelogs by relevance or content.
Returns:
list[KGChanges]: A list of KGChanges matching the provided filters and pagination.
"""
raise NotImplementedError("get_changelogs_list method not implemented")
@abstractmethod
def get_changelog_types(self, brain_id: str) -> list[str]:
"""
Return the unique changelog types present for the specified brain.
Parameters:
brain_id (str): Identifier of the brain to query.
Returns:
list[str]: List of unique changelog type names for the brain (order not guaranteed).
"""
raise NotImplementedError("get_changelog_types method not implemented")
@abstractmethod
def update_structured_data(
self, structured_data: StructuredData, brain_id: str
) -> StructuredData:
"""
Update an existing StructuredData record for a specific brain.
Parameters:
structured_data (StructuredData): StructuredData object containing the updated content and identifier of the record to update.
brain_id (str): Identifier of the brain that owns the structured data.
Returns:
StructuredData: The stored StructuredData after the update.
"""
raise NotImplementedError("update_structured_data method not implemented")
@abstractmethod
def get_last_text_chunks(self, brain_id: str, limit: int = 10) -> list[TextChunk]:
"""
Get the last text chunks from the data client.
"""
raise NotImplementedError("get_last_text_chunks method not implemented")
@abstractmethod
def get_last_structured_data(
self, brain_id: str, limit: int = 10
) -> list[StructuredData]:
"""
Get the last structured data from the data client.
"""
raise NotImplementedError("get_last_structured_data method not implemented")
|