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.
- keble_task/__init__.py +143 -0
- keble_task/actions.py +304 -0
- keble_task/agent/__init__.py +27 -0
- keble_task/agent/chat_provider.py +117 -0
- keble_task/agent/deps.py +38 -0
- keble_task/agent/tools/__init__.py +14 -0
- keble_task/agent/tools/mutation.py +100 -0
- keble_task/agent/tools/query.py +160 -0
- keble_task/crud.py +347 -0
- keble_task/exceptions.py +62 -0
- keble_task/main.py +2708 -0
- keble_task/schemas/__init__.py +1295 -0
- keble_task/schemas/for_agent.py +177 -0
- keble_task/task_tree.py +106 -0
- keble_task/utils.py +58 -0
- keble_task-2.22.0.dist-info/METADATA +1083 -0
- keble_task-2.22.0.dist-info/RECORD +18 -0
- keble_task-2.22.0.dist-info/WHEEL +4 -0
|
@@ -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
|
+
)
|
keble_task/exceptions.py
ADDED
|
@@ -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)
|