keble-task 2.22.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,100 @@
1
+ """Package-owned pydantic-ai MUTATION tool registrar for generic task actions.
2
+
3
+ This module is BEHAVIOR ONLY: the tool-registration config
4
+ (``TaskAgentMutationToolsConfig``) lives in ``keble_task.schemas.for_agent``; the
5
+ single-action tool payload (``TaskAction``) and its result
6
+ (``CreateRelatedTaskActionedResult``) live in the external action module
7
+ ``keble_task.actions``. One tool call = one action. Nothing here defines a
8
+ ``BaseModel``.
9
+
10
+ Side effect if changes:
11
+ - Tool NAME (``mutate_task_workspace``) is persisted in room diagnostics and
12
+ mapped to frontend labels — never rename casually.
13
+ - ``register_mutation_tools`` is consumed by keble.backend chat composition.
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ from typing import Any, TypeVar
19
+
20
+ import keble_exceptions
21
+ from keble_helpers import AgentToolApprovalMode, AgentToolRegistrationConfig
22
+ from pydantic_ai import Agent, ModelRetry, RunContext
23
+
24
+ from keble_task.actions import CreateRelatedTaskActionedResult, TaskAction
25
+ from keble_task.main import TaskClient
26
+ from keble_task.schemas.for_agent import TaskAgentMutationToolsConfig
27
+
28
+ from ..deps import TaskAgentDeps
29
+
30
+ DepsT = TypeVar("DepsT", bound=TaskAgentDeps)
31
+
32
+
33
+ def register_mutation_tools(
34
+ agent: Agent[DepsT, Any],
35
+ *,
36
+ task_client: TaskClient,
37
+ tools_config: TaskAgentMutationToolsConfig | dict[str, Any] | None = None,
38
+ ) -> None:
39
+ """Register task-owned mutation tools on a pydantic-ai agent.
40
+
41
+ Step by step this registrar:
42
+ 1. builds tool metadata through schema-owned `build(...)` methods;
43
+ 2. registers one `mutate_task_workspace` tool owned by this package;
44
+ 3. reads Mongo/Redis from `AgentDbDeps`;
45
+ 4. delegates ONE action per call to `TaskClient.aapply_action(...)`.
46
+ """
47
+
48
+ config = TaskAgentMutationToolsConfig.build(tools_config)
49
+ # Keep the resolved tool config under a distinct name so it does not shadow the
50
+ # decorated `mutate_task_workspace` tool function below. Sharing one name makes the
51
+ # type checker unify the symbol and flag the registration config assignment
52
+ # against the tool-function type; the rename keeps both strongly typed.
53
+ mutate_task_workspace_config = AgentToolRegistrationConfig.build(
54
+ config.mutate_task_workspace
55
+ )
56
+
57
+ @agent.tool(
58
+ name=mutate_task_workspace_config.name,
59
+ description=mutate_task_workspace_config.description,
60
+ requires_approval=(
61
+ mutate_task_workspace_config.approval
62
+ is AgentToolApprovalMode.REQUIRE_APPROVAL
63
+ ),
64
+ )
65
+ async def mutate_task_workspace(
66
+ ctx: RunContext[DepsT],
67
+ *,
68
+ payload: TaskAction,
69
+ ) -> CreateRelatedTaskActionedResult:
70
+ """Apply ONE typed task workspace action through `TaskClient`.
71
+
72
+ One tool call = one action = one approval card. To create several
73
+ related tasks the agent emits several `mutate_task_workspace` calls; the
74
+ deferred-batch runtime holds them and the frontend resolves one card per
75
+ click (in order). Do NOT re-introduce a `actions: list[...]` payload.
76
+ """
77
+
78
+ try:
79
+ return await task_client.aapply_action(
80
+ amongo=ctx.deps.amongo,
81
+ current_task=ctx.deps.task.current_task,
82
+ action=payload,
83
+ ux_context=ctx.deps.task.ux_context,
84
+ event_emitter=ctx.deps.task.event_emitter,
85
+ extended_aredis=ctx.deps.extended_aredis,
86
+ )
87
+ except (
88
+ keble_exceptions.ClientSideInvalidParams,
89
+ keble_exceptions.ClientSideMissingParams,
90
+ keble_exceptions.ServerSideInvalidParams,
91
+ ) as exc:
92
+ # The model supplied invalid action input (a hallucinated parent_task_id /
93
+ # from_task_id, a cross-room reference, or a missing required field). Surface it
94
+ # as a retryable signal so the model corrects its own arguments instead of
95
+ # aborting the run. Wiring/infra faults (e.g. ServerSideMissingParams) are NOT
96
+ # caught here and propagate as real failures.
97
+ raise ModelRetry(str(exc)) from exc
98
+
99
+
100
+ __all__ = ["register_mutation_tools"]
@@ -0,0 +1,160 @@
1
+ """Package-owned pydantic-ai READ/QUERY tool registrar for generic task browsing.
2
+
3
+ These tools let a chat agent discover and inspect tasks without mutating them.
4
+ They are thin, owner-scoped wrappers over `TaskClient` read methods:
5
+
6
+ 1. ``list_tasks`` lists the chat owner's tasks (optionally one task type);
7
+ 2. ``get_task`` loads one owner-visible task by id.
8
+
9
+ All tools are owner-scoped to the chat's ``current_task.owner`` so the agent can
10
+ only ever read its own owner's tasks. They are non-mutating and never require
11
+ approval. Background-session inspection is intentionally NOT duplicated here —
12
+ that is owned by the host-neutral general tools (``check_background_sessions``)
13
+ in ``keble_agentic_chat`` so there is one canonical session-poll path.
14
+
15
+ This module is BEHAVIOR ONLY: every tool I/O contract lives in
16
+ ``keble_task.schemas.for_agent`` (the ``TaskSummaryForAgent`` projection + the
17
+ ``TaskAgentQueryToolsConfig``); nothing here defines a ``BaseModel``.
18
+
19
+ Side effect if changes:
20
+ - Tool NAMES (``list_tasks`` / ``get_task``) are persisted in room diagnostics
21
+ and mapped to frontend labels — never rename casually.
22
+ - Return projection comes from ``schemas/for_agent.py``; field renames change the
23
+ agent-visible JSON.
24
+ """
25
+
26
+ from __future__ import annotations
27
+
28
+ from typing import Any, TypeVar
29
+
30
+ from keble_helpers import (
31
+ AgentToolRegistrationConfig,
32
+ parse_object_id_or_model_retry,
33
+ require_object_or_model_retry,
34
+ )
35
+ from pydantic_ai import Agent, RunContext
36
+
37
+ from keble_task.main import TaskClient
38
+ from keble_task.schemas.for_agent import (
39
+ TaskAgentQueryToolsConfig,
40
+ TaskSummaryForAgent,
41
+ )
42
+
43
+ from ..deps import TaskAgentDeps
44
+
45
+ DepsT = TypeVar("DepsT", bound=TaskAgentDeps)
46
+
47
+
48
+ def register_query_tools(
49
+ agent: Agent[DepsT, Any],
50
+ *,
51
+ task_client: TaskClient,
52
+ tools_config: TaskAgentQueryToolsConfig | dict[str, Any] | None = None,
53
+ ) -> None:
54
+ """Register task-owned READ/QUERY tools on a pydantic-ai agent.
55
+
56
+ Step by step this registrar:
57
+ 1. resolves per-tool name/description overrides through schema ``build(...)``;
58
+ 2. registers owner-scoped, non-mutating ``list_tasks`` / ``get_task`` tools;
59
+ 3. reads Mongo/Redis and the chat owner from ``TaskAgentDeps``;
60
+ 4. surfaces a bad model-supplied id as ``ModelRetry`` so the model self-corrects.
61
+ """
62
+
63
+ config = TaskAgentQueryToolsConfig.build(tools_config)
64
+ list_cfg = AgentToolRegistrationConfig.build(config.list_tasks)
65
+ get_cfg = AgentToolRegistrationConfig.build(config.get_task)
66
+
67
+ @agent.tool(
68
+ name=list_cfg.name or "list_tasks",
69
+ description=(
70
+ list_cfg.description
71
+ or "List your TOP-LEVEL tasks (most recent first). Optional filters: "
72
+ "task_type (e.g. AMZ_PRODUCT_REPORT or POSITIONING), stage (e.g. "
73
+ "SUCCESS, FAILURE, PROCESSING), title_contains (case-insensitive "
74
+ "substring), and parent_task_id (lists that task's DIRECT children "
75
+ "instead of top-level tasks). Returns "
76
+ "compact summaries with id, type, title, stage, and failure facts. "
77
+ "IMPORTANT: background sessions started inside a chat are CHILD tasks "
78
+ "and will NOT appear here — read them with check_background_sessions, "
79
+ "or open any session_id directly with get_task (a session_id IS a "
80
+ "task id)."
81
+ ),
82
+ requires_approval=False,
83
+ )
84
+ async def list_tasks(
85
+ ctx: RunContext[DepsT],
86
+ *,
87
+ task_type: str | None = None,
88
+ stage: str | None = None,
89
+ title_contains: str | None = None,
90
+ parent_task_id: str | None = None,
91
+ limit: int = 20,
92
+ ) -> list[TaskSummaryForAgent]:
93
+ """List the chat owner's tasks with optional indexed-shape filters.
94
+
95
+ Step by step (2.10.0 enrichment):
96
+ 1. `task_type` / `stage` are exact-value filters; `title_contains`
97
+ is a case-insensitive substring match (owner-bounded);
98
+ 2. `parent_task_id` lists the direct children of one task —
99
+ malformed ids are ModelRetry so the model self-corrects;
100
+ 3. results stay bounded compact summaries (limit <= 50).
101
+ """
102
+
103
+ owner = ctx.deps.task.current_task.owner
104
+ bounded_limit = max(1, min(limit, 50))
105
+ parent_object_id = (
106
+ parse_object_id_or_model_retry(
107
+ parent_task_id, field_name="parent_task_id"
108
+ )
109
+ if parent_task_id is not None
110
+ else None
111
+ )
112
+ tasks = await task_client.aowner_get_multi(
113
+ amongo=ctx.deps.amongo,
114
+ extended_aredis=ctx.deps.extended_aredis,
115
+ owner=owner,
116
+ skip=0,
117
+ limit=bounded_limit,
118
+ task_types=[task_type] if task_type is not None else None,
119
+ sharing_scopes=None,
120
+ stages=[stage] if stage is not None else None,
121
+ title_contains=title_contains,
122
+ parent_task=parent_object_id,
123
+ )
124
+ return [TaskSummaryForAgent.from_task(task) for task in tasks]
125
+
126
+ @agent.tool(
127
+ name=get_cfg.name or "get_task",
128
+ description=(
129
+ get_cfg.description
130
+ or "Load one of your tasks by its id and return a compact summary "
131
+ "(type, title, stage, parent/root ids, plus error and exception_type "
132
+ "when it failed). Works for ANY of your tasks, including chat-started "
133
+ "background sessions — a session_id from the start/check session "
134
+ "tools IS a valid task_id here."
135
+ ),
136
+ requires_approval=False,
137
+ )
138
+ async def get_task(
139
+ ctx: RunContext[DepsT],
140
+ *,
141
+ task_id: str,
142
+ ) -> TaskSummaryForAgent:
143
+ """Load one owner-visible task by id; raise ObjectNotFound when missing."""
144
+
145
+ object_id = parse_object_id_or_model_retry(task_id, field_name="task_id")
146
+ task = require_object_or_model_retry(
147
+ await task_client.aowner_get(
148
+ amongo=ctx.deps.amongo,
149
+ owner=ctx.deps.task.current_task.owner,
150
+ task_id=object_id,
151
+ task_type=None,
152
+ sharing_scope=None,
153
+ ),
154
+ field_name="task_id",
155
+ raw=task_id,
156
+ )
157
+ return TaskSummaryForAgent.from_task(task)
158
+
159
+
160
+ __all__ = ["register_query_tools"]
keble_task/crud.py ADDED
@@ -0,0 +1,347 @@
1
+ import keble_exceptions
2
+ from keble_db import MongoCRUDBase, QueryBase
3
+ from keble_helpers import ObjectId
4
+ from motor.motor_asyncio import AsyncIOMotorClient
5
+ from pymongo import ASCENDING, DESCENDING, IndexModel
6
+ from pymongo import MongoClient
7
+ from pymongo.results import InsertOneResult
8
+
9
+ from . import schemas
10
+
11
+
12
+ class CRUDTask(MongoCRUDBase[schemas.TaskMongoObject]):
13
+ async def aensure_public_id_index(self, amongo: AsyncIOMotorClient) -> None:
14
+ await amongo[self.database][self.collection].create_index(
15
+ [("public_id", 1)],
16
+ unique=True,
17
+ name="public_id_unique",
18
+ partialFilterExpression={"public_id": {"$type": "string"}},
19
+ )
20
+
21
+ async def aensure_task_indexes(self, amongo: AsyncIOMotorClient) -> None:
22
+ """Create indexes for root-tree, root-list, and stage-scan task reads.
23
+
24
+ Step by step:
25
+ 1. `root_task + created` backs the `$or` root-tree branch (`aget_tasks_in_root_tree`)
26
+ and descendant/lineage scans that filter by root and sort newest-first;
27
+ 2. `parent_task + created` backs root-task listing (`parent_task=None`) and
28
+ immediate-child lookups sorted by creation time;
29
+ 3. `stage + created` backs the retry/timeout sweeps (`aretry_all_undone`,
30
+ `aget_tasks_by_stage_since`) that match `stage $in [...]` over a `created` window;
31
+ 4. `owner + parent_task + created` backs owner-scoped root lists
32
+ (`aowner_get_multi` -> `_aget_root_tasks`, which fix `parent_task=None`), so `owner`
33
+ is an indexed equality instead of a residual post-filter;
34
+ 5. `sharing_scope + parent_task + created` backs public root lists
35
+ (`apublic_get_multi`, fixed `sharing_scope=PUBLIC`, `parent_task=None`).
36
+ 6. `sharing_scope + task_type + stage + updated` backs the dynamic-page
37
+ sitemap read (`apublic_list_indexable`, fixed `sharing_scope=PUBLIC`,
38
+ `task_type $in`, `stage $in`, sorted `updated` DESC).
39
+
40
+ Field order is scope/equality first then the `created`/`updated` sort/range
41
+ field so each index serves its query shape (`task_type`/`root_task` remain
42
+ residual filters where not led). Creation is idempotent via `create_indexes`.
43
+ """
44
+
45
+ collection = amongo[self.database][self.collection]
46
+ await collection.create_indexes(
47
+ [
48
+ IndexModel(
49
+ [("root_task", ASCENDING), ("created", DESCENDING)],
50
+ name="root_task_created_desc_idx",
51
+ ),
52
+ IndexModel(
53
+ [("parent_task", ASCENDING), ("created", DESCENDING)],
54
+ name="parent_task_created_desc_idx",
55
+ ),
56
+ IndexModel(
57
+ [("stage", ASCENDING), ("created", DESCENDING)],
58
+ name="stage_created_desc_idx",
59
+ ),
60
+ IndexModel(
61
+ [
62
+ ("owner", ASCENDING),
63
+ ("parent_task", ASCENDING),
64
+ ("created", DESCENDING),
65
+ ],
66
+ name="owner_parent_task_created_desc_idx",
67
+ ),
68
+ IndexModel(
69
+ [
70
+ ("sharing_scope", ASCENDING),
71
+ ("parent_task", ASCENDING),
72
+ ("created", DESCENDING),
73
+ ],
74
+ name="sharing_scope_parent_task_created_desc_idx",
75
+ ),
76
+ IndexModel(
77
+ [
78
+ ("sharing_scope", ASCENDING),
79
+ ("task_type", ASCENDING),
80
+ ("stage", ASCENDING),
81
+ ("updated", DESCENDING),
82
+ ],
83
+ name="sharing_scope_task_type_stage_updated_desc_idx",
84
+ ),
85
+ ]
86
+ )
87
+
88
+ async def aget_tasks_in_root_tree(
89
+ self,
90
+ amongo: AsyncIOMotorClient,
91
+ *,
92
+ root_id: ObjectId,
93
+ base_filter: dict,
94
+ ) -> list[schemas.TaskMongoObject]:
95
+ filter_ = dict(base_filter)
96
+ filter_["$or"] = [{"_id": root_id}, {"root_task": root_id}]
97
+ return await self.aget_multi(
98
+ amongo,
99
+ query=QueryBase(
100
+ filters=filter_,
101
+ order_by=[("created", DESCENDING)],
102
+ ),
103
+ )
104
+
105
+ def increment_attempts(
106
+ self, mongo: MongoClient, *, _id: ObjectId, increment: int = 1
107
+ ) -> schemas.TaskMongoObject:
108
+ existed_obj = self.first_by_id(mongo, _id=_id)
109
+ if existed_obj is None:
110
+ raise keble_exceptions.ServerSideInvalidParams(
111
+ admin_note={"task_id": _id},
112
+ alert_admin=True,
113
+ but_got="db_obj is None",
114
+ expected="db_obj is not None",
115
+ invalid_params="task_id",
116
+ )
117
+ self.update(
118
+ mongo, _id=_id, obj_in={"attempts": existed_obj.attempts + increment}
119
+ )
120
+ db_obj = self.first_by_id(mongo, _id=_id)
121
+ assert db_obj is not None
122
+ return db_obj
123
+
124
+ async def aincrement_attempts(
125
+ self, amongo: AsyncIOMotorClient, *, _id: ObjectId, increment: int = 1
126
+ ) -> schemas.TaskMongoObject:
127
+ """Async increment attempts for a task"""
128
+ existed_obj = await self.afirst_by_id(amongo, _id=_id)
129
+ if existed_obj is None:
130
+ raise keble_exceptions.ServerSideInvalidParams(
131
+ admin_note={"task_id": _id},
132
+ alert_admin=True,
133
+ but_got="db_obj is None",
134
+ expected="db_obj is not None",
135
+ invalid_params="task_id",
136
+ )
137
+ await self.aupdate(
138
+ amongo, _id=_id, obj_in={"attempts": existed_obj.attempts + increment}
139
+ )
140
+ db_obj = await self.afirst_by_id(amongo, _id=_id)
141
+ assert db_obj is not None
142
+ return db_obj
143
+
144
+
145
+ class CRUDTaskRelation(MongoCRUDBase[schemas.TaskRelationMongoObject]):
146
+ """Mongo CRUD helper for task-relation rows.
147
+
148
+ Step by step:
149
+ 1. keep relation storage separate from the task tree collection;
150
+ 2. expose index creation for workspace and edge lookup patterns;
151
+ 3. let `TaskClient` own relation validation before insertion.
152
+ """
153
+
154
+ async def aensure_relation_indexes(self, amongo: AsyncIOMotorClient) -> None:
155
+ """Create indexes required by root-scoped task-relation queries.
156
+
157
+ Step by step:
158
+ 1. index each endpoint for diagnostics and filtered lookups;
159
+ 2. index root/to/type for child-derived lineage queries;
160
+ 3. make the full edge unique so exact duplicate relation rows cannot drift.
161
+ """
162
+
163
+ collection = amongo[self.database][self.collection]
164
+ await collection.create_index([("root_task", 1)], name="root_task_1")
165
+ await collection.create_index([("from_task_id", 1)], name="from_task_id_1")
166
+ await collection.create_index([("to_task_id", 1)], name="to_task_id_1")
167
+ await collection.create_index(
168
+ [("root_task", 1), ("to_task_id", 1), ("relation_type", 1)],
169
+ name="root_task_1_to_task_id_1_relation_type_1",
170
+ )
171
+ await collection.create_index(
172
+ [
173
+ ("root_task", 1),
174
+ ("from_task_id", 1),
175
+ ("to_task_id", 1),
176
+ ("relation_type", 1),
177
+ ],
178
+ name="root_from_to_relation_type_unique",
179
+ unique=True,
180
+ )
181
+
182
+ async def acreate_relation(
183
+ self,
184
+ amongo: AsyncIOMotorClient,
185
+ *,
186
+ obj_in: schemas.TaskRelationCreate,
187
+ ) -> InsertOneResult:
188
+ """Insert one validated relation create payload into relation storage.
189
+
190
+ Step by step:
191
+ 1. accept the create schema before Mongo has assigned a durable `_id`;
192
+ 2. serialize the payload exactly like the generic CRUD insert path;
193
+ 3. return Mongo's insert result so callers can reload the persisted object.
194
+ """
195
+
196
+ return await amongo[self.database][self.collection].insert_one(
197
+ obj_in.model_dump()
198
+ )
199
+
200
+
201
+ class CRUDTaskCost(MongoCRUDBase[schemas.TaskCostMongoObject]):
202
+ """Mongo CRUD helper for task-cost rows.
203
+
204
+ Step by step:
205
+ 1. keep cost storage separate from task state and relation edges;
206
+ 2. expose indexed list/count query shapes for admin reports;
207
+ 3. centralize index creation so APIs and workers do not create indexes inline.
208
+ """
209
+
210
+ async def aensure_cost_indexes(self, amongo: AsyncIOMotorClient) -> None:
211
+ """Create indexes required by task-cost list and aggregate reads.
212
+
213
+ Step by step:
214
+ 1. index direct task and root-task time windows;
215
+ 2. index owner and owner/task-type reporting windows;
216
+ 3. index tags and occurrence-time scans used by admin aggregation.
217
+ """
218
+
219
+ collection = amongo[self.database][self.collection]
220
+ await collection.create_indexes(
221
+ [
222
+ IndexModel(
223
+ [("task_id", ASCENDING), ("occurred_at", DESCENDING)],
224
+ name="task_id_occurred_at_desc_idx",
225
+ ),
226
+ IndexModel(
227
+ [("root_task", ASCENDING), ("occurred_at", DESCENDING)],
228
+ name="root_task_occurred_at_desc_idx",
229
+ ),
230
+ IndexModel(
231
+ [("owner", ASCENDING), ("occurred_at", DESCENDING)],
232
+ name="owner_occurred_at_desc_idx",
233
+ ),
234
+ IndexModel(
235
+ [
236
+ ("owner", ASCENDING),
237
+ ("task_type", ASCENDING),
238
+ ("occurred_at", DESCENDING),
239
+ ],
240
+ name="owner_task_type_occurred_at_desc_idx",
241
+ ),
242
+ IndexModel(
243
+ [("tags", ASCENDING), ("occurred_at", DESCENDING)],
244
+ name="tags_occurred_at_desc_idx",
245
+ ),
246
+ IndexModel(
247
+ [("source", ASCENDING), ("occurred_at", DESCENDING)],
248
+ name="source_occurred_at_desc_idx",
249
+ ),
250
+ IndexModel(
251
+ [("occurred_at", DESCENDING)],
252
+ name="occurred_at_desc_idx",
253
+ ),
254
+ ]
255
+ )
256
+
257
+ def build_cost_filters(
258
+ self,
259
+ *,
260
+ query: schemas.TaskCostFilterBase,
261
+ ) -> dict[str, object]:
262
+ """Build the indexed Mongo filter from the canonical cost filter.
263
+
264
+ Step by step:
265
+ 1. accept the shared filter base used by list and aggregate reads;
266
+ 2. delegate field mapping to the schema that owns the filter contract;
267
+ 3. keep CRUD methods focused on storage access and stable sorting.
268
+ """
269
+
270
+ return query.to_mongo_filters()
271
+
272
+ async def acreate_cost(
273
+ self,
274
+ amongo: AsyncIOMotorClient,
275
+ *,
276
+ obj_in: schemas.TaskCostCreate,
277
+ ) -> InsertOneResult:
278
+ """Insert one validated task-cost row.
279
+
280
+ Step by step:
281
+ 1. accept the denormalized create schema produced by `TaskClient`;
282
+ 2. serialize Pydantic dataclasses and nested money values to Mongo-safe dicts;
283
+ 3. return Mongo's insert result so callers can reload the typed row.
284
+ """
285
+
286
+ return await amongo[self.database][self.collection].insert_one(
287
+ obj_in.model_dump()
288
+ )
289
+
290
+ async def alist_costs(
291
+ self,
292
+ amongo: AsyncIOMotorClient,
293
+ *,
294
+ query: schemas.TaskCostListRequest,
295
+ ) -> list[schemas.TaskCostMongoObject]:
296
+ """List task-cost rows with indexed filters and stable newest-first order.
297
+
298
+ Step by step:
299
+ 1. build the shared indexed filter from the request schema;
300
+ 2. sort by `occurred_at` descending and `_id` descending for stable pages;
301
+ 3. parse Mongo documents into `TaskCostMongoObject`.
302
+ """
303
+
304
+ return await self.aget_multi(
305
+ amongo,
306
+ query=QueryBase(
307
+ filters=self.build_cost_filters(query=query),
308
+ order_by=[("occurred_at", DESCENDING), ("_id", DESCENDING)],
309
+ offset=query.skip,
310
+ limit=query.limit,
311
+ ),
312
+ )
313
+
314
+ async def acount_costs(
315
+ self,
316
+ amongo: AsyncIOMotorClient,
317
+ *,
318
+ query: schemas.TaskCostFilterBase,
319
+ ) -> int:
320
+ """Count task-cost rows for a shared indexed filter."""
321
+
322
+ return await self.acount(
323
+ amongo,
324
+ query=QueryBase(filters=self.build_cost_filters(query=query)),
325
+ )
326
+
327
+ async def alist_costs_for_aggregate(
328
+ self,
329
+ amongo: AsyncIOMotorClient,
330
+ *,
331
+ query: schemas.TaskCostAggregateRequest,
332
+ ) -> list[schemas.TaskCostMongoObject]:
333
+ """Load filtered task-cost rows for package-owned aggregation.
334
+
335
+ Step by step:
336
+ 1. use the same indexed filter as list/count reads;
337
+ 2. sort ascending by occurrence so response buckets are deterministic;
338
+ 3. leave currency conversion and sub-cent math to schema methods.
339
+ """
340
+
341
+ return await self.aget_multi(
342
+ amongo,
343
+ query=QueryBase(
344
+ filters=self.build_cost_filters(query=query),
345
+ order_by=[("occurred_at", ASCENDING), ("_id", ASCENDING)],
346
+ ),
347
+ )
@@ -0,0 +1,62 @@
1
+ import json
2
+ import traceback
3
+ from enum import Enum
4
+ from typing import Optional
5
+
6
+ from keble_exceptions import KebleException
7
+
8
+
9
+ class TaskExceptionType(str, Enum):
10
+ UNKNOWN = "UNKNOWN"
11
+ FAILED_TO_START = "FAILED_TO_START"
12
+ NO_SUFFICIENT_DATA = "NO_SUFFICIENT_DATA"
13
+ TIMEOUT = "TIMEOUT"
14
+
15
+
16
+ class TaskException(KebleException):
17
+ def __init__(
18
+ self,
19
+ *,
20
+ exception_type: TaskExceptionType = TaskExceptionType.UNKNOWN,
21
+ # More information regarding the exception
22
+ error: Optional[str] = None,
23
+ ):
24
+ self.exception_type = exception_type
25
+ self.error = error
26
+ super().__init__(
27
+ admin_note=json.dumps(
28
+ {
29
+ "exception_type": str(exception_type),
30
+ "traceback": str(traceback.format_exc()),
31
+ "error": error,
32
+ }
33
+ ),
34
+ alert_admin=True,
35
+ function_identifier="TaskException.__init__",
36
+ )
37
+
38
+ def __str__(self):
39
+ return self.exception_type.value
40
+
41
+
42
+ class TaskNoSufficientDataException(TaskException):
43
+ def __init__(self, data_type: str):
44
+ super().__init__(
45
+ exception_type=TaskExceptionType.NO_SUFFICIENT_DATA,
46
+ error=f"No enough: {data_type}",
47
+ )
48
+
49
+
50
+ class TaskFailedToStartException(TaskException):
51
+ def __init__(self, error: str):
52
+ super().__init__(exception_type=TaskExceptionType.FAILED_TO_START, error=error)
53
+
54
+
55
+ class TaskTimeoutException(TaskException):
56
+ def __init__(self, error: str):
57
+ super().__init__(exception_type=TaskExceptionType.TIMEOUT, error=error)
58
+
59
+
60
+ class TaskUnknownException(TaskException):
61
+ def __init__(self, error: str):
62
+ super().__init__(exception_type=TaskExceptionType.UNKNOWN, error=error)