Skip to content

Commit

Permalink
refactor: Return previous ingestion Tasks to add function
Browse files Browse the repository at this point in the history
Returned previous ignestion tasks to add function

Refactor #COG-337
  • Loading branch information
dexters1 committed Nov 8, 2024
1 parent 733caa7 commit 8c7fdc4
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions cognee/api/v1/add/add_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,9 @@
from cognee.modules.users.models import User
from cognee.modules.users.methods import get_default_user
from cognee.modules.pipelines import run_tasks, Task
from cognee.tasks.ingestion import save_data_to_storage, ingest_data
from cognee.infrastructure.databases.relational import create_db_and_tables as create_relational_db_and_tables
from cognee.infrastructure.databases.vector.pgvector import create_db_and_tables as create_pgvector_db_and_tables
from cognee.tasks.ingestion.ingest_data_with_metadata import ingest_data_with_metadata


async def add(data: Union[BinaryIO, list[BinaryIO], str, list[str]], dataset_name: str = "main_dataset", user: User = None):
await create_relational_db_and_tables()
Expand All @@ -15,10 +14,11 @@ async def add(data: Union[BinaryIO, list[BinaryIO], str, list[str]], dataset_nam
user = await get_default_user()

tasks = [
Task(ingest_data_with_metadata, dataset_name, user)
Task(save_data_to_storage, dataset_name),
Task(ingest_data, dataset_name, user)
]

pipeline = run_tasks(tasks, data, "add_pipeline")

async for result in pipeline:
print(result)
print(result)

0 comments on commit 8c7fdc4

Please sign in to comment.