gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #265 from topoteretes/feat/COG-418-log-config-to-telemetry

Browse source 
Add cognee config to telemetry
This commit is contained in:
Vasilije 2024-12-07 09:45:49 +01:00 • committed by GitHub
commit ce96431055
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 25 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
cognee/modules/pipelines/operations/run_tasks.py
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)")
@ -160,21 +162,28 @@ 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):
user = await get_default_user()
config = get_current_settings()
logger.debug("\nRunning pipeline with configuration:\n%s\n", json.dumps(config, indent = 1))
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",
"pipeline_name": pipeline_name, user.id,
}) additional_properties = {"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
logger.info("Pipeline run completed: `%s`", pipeline_name) logger.info("Pipeline run completed: `%s`", pipeline_name)
send_telemetry("Pipeline Run Completed", user.id, { send_telemetry("Pipeline Run Completed",
"pipeline_name": pipeline_name, user.id,
}) additional_properties = {"pipeline_name": pipeline_name, }
)
except Exception as error: except Exception as error:
logger.error( logger.error(
"Pipeline run errored: `%s`\n%s\n", "Pipeline run errored: `%s`\n%s\n",
@ -182,15 +191,14 @@ async def run_tasks_with_telemetry(tasks: list[Task], data, pipeline_name: str):
str(error), str(error),
exc_info = True, exc_info = True,
) )
send_telemetry("Pipeline Run Errored", user.id, { send_telemetry("Pipeline Run Errored",
"pipeline_name": pipeline_name, user.id,
}) additional_properties = {"pipeline_name": pipeline_name, } | config
)
raise error raise error
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()
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):
yield result yield result