forked from basicmachines-co/basic-memory
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathjobs.py
More file actions
308 lines (241 loc) · 9.94 KB
/
Copy pathjobs.py
File metadata and controls
308 lines (241 loc) · 9.94 KB
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
298
299
300
301
302
303
304
305
306
307
308
"""Portable job-queue contracts for Basic Memory runtimes."""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import timedelta
from enum import StrEnum
from typing import Protocol
from basic_memory.runtime.projects import ProjectRuntimeReference
from basic_memory.runtime.storage import (
ProjectExternalId,
ProjectId,
ProjectName,
ProjectPath,
RuntimeEntityId,
RuntimeFileChecksum,
RuntimeFilePath,
StorageEtag,
StorageEventSource,
normalize_storage_etag,
)
type JobEntrypoint = str
type RuntimeJobDedupeKey = str
type RuntimeJobId = str | int
class RuntimeStorageFileIndexMode(StrEnum):
"""Producer mode for one runtime file-index job."""
observed_object = "observed_object"
current_file = "current_file"
@dataclass(frozen=True, slots=True)
class RuntimeStorageFileIndexContext:
"""Project context required before enqueueing one runtime file-index job."""
mode: RuntimeStorageFileIndexMode
project_external_id: ProjectExternalId | None = None
project_name: ProjectName | None = None
def require_enqueue_context(self) -> None:
"""Raise when an observed object job lacks UI-facing project context."""
if self.mode != RuntimeStorageFileIndexMode.observed_object:
return
if not self.project_external_id:
raise ValueError("observed_object index jobs require project_external_id")
if not self.project_name:
raise ValueError("observed_object index jobs require project_name")
@dataclass(frozen=True, slots=True)
class RuntimeStorageFileIndexJobIdentity:
"""Stable queue identity for one runtime file-index job."""
project_id: ProjectId
file_path: RuntimeFilePath
mode: RuntimeStorageFileIndexMode
object_etag: StorageEtag | None = None
object_size: int | None = None
def dedupe_key(self) -> str:
"""Return the existing logical work key for file-index queue requests."""
base = f"index-file:{self.project_id}:{self.file_path}"
if self.mode == RuntimeStorageFileIndexMode.current_file:
return f"{base}:current"
if self.object_etag is None:
raise ValueError("observed_object index jobs require object metadata")
return f"{base}:observed:{normalize_storage_etag(self.object_etag)}:{self.object_size}"
def routing_headers(self, headers: Mapping[str, str] | None = None) -> dict[str, str]:
"""Return queue routing headers for the file-index job."""
routing_headers = dict(headers or {})
routing_headers["project_id"] = str(self.project_id)
return routing_headers
def job_request(
self,
*,
entrypoint: JobEntrypoint,
payload: bytes | None = None,
headers: Mapping[str, str] | None = None,
priority: int = 0,
execute_after: timedelta | None = None,
) -> RuntimeJobRequest:
"""Build the runtime queue request for this file-index identity."""
return runtime_job_request_from_source(
self,
entrypoint=entrypoint,
payload=payload,
priority=priority,
execute_after=execute_after,
headers=headers,
)
@dataclass(frozen=True, slots=True)
class RuntimeStorageObjectObservation:
"""Storage object metadata observed before enqueueing one file-index job."""
etag: StorageEtag
size: int | None = None
def to_file_index_job_identity(
self,
*,
project_id: ProjectId,
file_path: RuntimeFilePath,
) -> RuntimeStorageFileIndexJobIdentity:
"""Build the queue identity for this observed storage object."""
return RuntimeStorageFileIndexJobIdentity(
project_id=project_id,
file_path=file_path,
mode=RuntimeStorageFileIndexMode.observed_object,
object_etag=self.etag,
object_size=self.size,
)
@dataclass(frozen=True, slots=True)
class RuntimeObservedIndexFile:
"""Storage metadata observed before a project-index batch is queued."""
path: RuntimeFilePath
checksum: RuntimeFileChecksum | None = None
size: int | None = None
@dataclass(frozen=True, slots=True)
class RuntimeProjectIndexJobRequest:
"""Queue-neutral request shape for coordinating a project-wide index."""
project: ProjectRuntimeReference
force_full: bool = False
search: bool = True
embeddings: bool = True
def dedupe_key(self) -> str:
"""Return the logical project-index coordinator queue identity."""
return f"index-project:{self.project.project_id}"
def routing_headers(self, headers: Mapping[str, str] | None = None) -> dict[str, str]:
"""Return queue routing headers for the project-index coordinator."""
routing_headers = dict(headers or {})
routing_headers.update(
{
"project_id": str(self.project.project_id),
"project_path": self.project.project_path,
}
)
return routing_headers
def plan_project_index_job_request(
*,
project: ProjectRuntimeReference,
force_full: bool = False,
search: bool = True,
embeddings: bool = True,
) -> RuntimeProjectIndexJobRequest:
"""Flatten project-index coordinator inputs into a queue-neutral request."""
return RuntimeProjectIndexJobRequest(
project=project,
force_full=force_full,
search=search,
embeddings=embeddings,
)
@dataclass(frozen=True, slots=True)
class RuntimeProjectDeleteJobRequest:
"""Queue-neutral request shape for hard-deleting one inactive project."""
project_id: ProjectId
project_external_id: ProjectExternalId
project_name: ProjectName
project_path: ProjectPath
delete_notes: bool = True
def dedupe_key(self) -> str:
"""Return the logical project-delete queue identity."""
return f"delete-project:{self.project_id}"
def routing_headers(self, headers: Mapping[str, str] | None = None) -> dict[str, str]:
"""Return queue routing headers for the project-delete job."""
routing_headers = dict(headers or {})
routing_headers["project_id"] = str(self.project_id)
return routing_headers
@dataclass(frozen=True, slots=True)
class RuntimeIndexFileBatchJobRequest:
"""Queue-neutral request shape for indexing one project file batch."""
project: ProjectRuntimeReference
batch_index: int
batch_count: int
file_paths: tuple[RuntimeFilePath, ...] = ()
observed_files: tuple[RuntimeObservedIndexFile, ...] = ()
index_embeddings: bool = True
force_full: bool = False
def dedupe_key(self) -> str:
"""Return the logical file-batch index queue identity."""
return f"index-file-batch:{self.project.project_id}:{self.batch_index}"
def routing_headers(self, headers: Mapping[str, str] | None = None) -> dict[str, str]:
"""Return queue routing headers for the file-batch index job."""
routing_headers = dict(headers or {})
routing_headers.update(
{
"project_id": str(self.project.project_id),
"project_external_id": self.project.project_external_id,
"project_path": self.project.project_path,
}
)
return routing_headers
def target_paths(self) -> tuple[RuntimeFilePath, ...]:
"""Return target paths using observed metadata when it is available."""
if self.observed_files:
return tuple(observed.path for observed in self.observed_files)
return self.file_paths
@dataclass(frozen=True, slots=True)
class RuntimeJobRequest:
"""Concrete queue request built after payload validation."""
entrypoint: JobEntrypoint
payload: bytes | None = None
priority: int = 0
execute_after: timedelta | None = None
dedupe_key: RuntimeJobDedupeKey | None = None
headers: Mapping[str, str] | None = None
class RuntimeJobRequestSource(Protocol):
"""Minimal typed source that can become a concrete runtime queue request."""
def dedupe_key(self) -> RuntimeJobDedupeKey: ...
def routing_headers(self, headers: Mapping[str, str] | None = None) -> dict[str, str]: ...
def runtime_job_request_from_source(
source: RuntimeJobRequestSource,
*,
entrypoint: JobEntrypoint,
payload: bytes | None = None,
headers: Mapping[str, str] | None = None,
priority: int = 0,
execute_after: timedelta | None = None,
) -> RuntimeJobRequest:
"""Build a concrete queue request from a typed runtime request source."""
return RuntimeJobRequest(
entrypoint=entrypoint,
payload=payload,
priority=priority,
execute_after=execute_after,
dedupe_key=source.dedupe_key(),
headers=source.routing_headers(headers),
)
@dataclass(frozen=True, slots=True)
class RuntimeScheduledVectorSyncTask:
"""Typed scheduler handoff for syncing one entity's vector index."""
entity_id: RuntimeEntityId
project_id: ProjectId
@dataclass(frozen=True, slots=True)
class RuntimeScheduledProjectIndexTask:
"""Typed scheduler handoff for indexing one project."""
project_id: ProjectId
force_full: bool = False
class JobRuntime(Protocol):
"""Capability for enqueueing runtime jobs without depending on one queue."""
async def enqueue(self, request: RuntimeJobRequest) -> RuntimeJobId: ...
@dataclass(frozen=True, slots=True)
class RuntimeCapabilities:
"""Internal adapter bundle selected for one runtime surface."""
job_runtime: JobRuntime | None = None
storage_event_source: StorageEventSource | None = None
def require_job_runtime(self) -> JobRuntime:
if self.job_runtime is None:
raise RuntimeError("Job runtime is not configured")
return self.job_runtime
def require_storage_event_source(self) -> StorageEventSource:
if self.storage_event_source is None:
raise RuntimeError("Storage event source is not configured")
return self.storage_event_source