Add cognee config to telemetry

This commit is contained in:
Rita Aleksziev 2024-12-06 12:55:25 +01:00
parent acf036818e
commit dbfa91b635

View file

@ -1,10 +1,12 @@
import json
import inspect import inspect
import json
import logging import logging
from cognee.modules.settings import get_current_settings from cognee.modules.settings import get_current_settings
from cognee.shared.utils import send_telemetry
from cognee.modules.users.models import User
from cognee.modules.users.methods import get_default_user from cognee.modules.users.methods import get_default_user
from cognee.modules.users.models import User
from cognee.shared.utils import send_telemetry
from ..tasks.Task import Task from ..tasks.Task import Task
logger = logging.getLogger("run_tasks(tasks: [Task], data)") logger = logging.getLogger("run_tasks(tasks: [Task], data)")
@ -159,14 +161,17 @@ async def run_tasks_base(tasks: list[Task], data = None, user: User = None):
}) })
raise error raise error
async def run_tasks_with_telemetry(tasks: list[Task], data, pipeline_name: str): async def run_tasks_with_telemetry(tasks: list[Task], data, pipeline_name: str, config: dict):
logger.debug("\nRunning pipeline with configuration:\n%s\n", json.dumps(config, indent = 1))
user = await get_default_user() user = await get_default_user()
try: try:
logger.info("Pipeline run started: `%s`", pipeline_name) logger.info("Pipeline run started: `%s`", pipeline_name)
send_telemetry("Pipeline Run Started", user.id, { send_telemetry("Pipeline Run Started", user.id, {
"pipeline_name": pipeline_name, "pipeline_name": pipeline_name,
}) } | config
)
async for result in run_tasks_base(tasks, data, user): async for result in run_tasks_base(tasks, data, user):
yield result yield result
@ -190,7 +195,6 @@ async def run_tasks_with_telemetry(tasks: list[Task], data, pipeline_name: str):
async def run_tasks(tasks: list[Task], data = None, pipeline_name: str = "default_pipeline"): async def run_tasks(tasks: list[Task], data = None, pipeline_name: str = "default_pipeline"):
config = get_current_settings() config = get_current_settings()
logger.debug("\nRunning pipeline with configuration:\n%s\n", json.dumps(config, indent = 1))
async for result in run_tasks_with_telemetry(tasks, data, pipeline_name): async for result in run_tasks_with_telemetry(tasks, data, pipeline_name, config):
yield result yield result