d1eab97102
<!-- .github/pull_request_template.md --> ## Description <!-- Provide a clear description of the changes in this PR --> - Extracted run_tasks_base function into a new file run_tasks_base.py. - Extracted four executors that execute core logic based on the task type. - Extracted a task handler/wrapper that safely executes the core logic with logging and telemetry. - Fixed the inconsistency with the batches of size 1. ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin.
50 lines
1.3 KiB
Python
50 lines
1.3 KiB
Python
import asyncio
|
|
|
|
import cognee
|
|
from cognee.modules.pipelines.operations.run_tasks import run_tasks_base
|
|
from cognee.modules.pipelines.tasks.task import Task
|
|
from cognee.modules.users.methods import get_default_user
|
|
from cognee.infrastructure.databases.relational import create_db_and_tables
|
|
|
|
|
|
async def run_and_check_tasks():
|
|
await cognee.prune.prune_data()
|
|
await cognee.prune.prune_system(metadata=True)
|
|
|
|
def number_generator(num):
|
|
for i in range(num):
|
|
yield i + 1
|
|
|
|
async def add_one(nums):
|
|
for num in nums:
|
|
yield num + 1
|
|
|
|
async def multiply_by_two(nums):
|
|
yield nums[0] * 2
|
|
|
|
async def add_one_single(nums):
|
|
yield nums[0] + 1
|
|
|
|
await create_db_and_tables()
|
|
user = await get_default_user()
|
|
|
|
pipeline = run_tasks_base(
|
|
[
|
|
Task(number_generator),
|
|
Task(add_one, task_config={"batch_size": 5}),
|
|
Task(multiply_by_two, task_config={"batch_size": 1}),
|
|
Task(add_one_single),
|
|
],
|
|
data=10,
|
|
user=user,
|
|
)
|
|
|
|
results = [5, 7, 9, 11, 13, 15, 17, 19, 21, 23]
|
|
index = 0
|
|
async for result in pipeline:
|
|
assert result[0] == results[index], f"at {index = }: {result = } != {results[index] = }"
|
|
index += 1
|
|
|
|
|
|
def test_run_tasks():
|
|
asyncio.run(run_and_check_tasks())
|