From e11998c89d59cb8c1ec3f4f83d682cc066c02860 Mon Sep 17 00:00:00 2001 From: LeonOstrez Date: Thu, 10 Oct 2024 11:41:04 +0200 Subject: [PATCH] fix telemetry --- core/agents/architect.py | 1 + core/agents/spec_writer.py | 1 + core/cli/main.py | 38 ++++++++++++++++++++++++++++--- core/state/state_manager.py | 15 ++++++++++-- core/telemetry/__init__.py | 6 ++--- tests/telemetry/test_telemetry.py | 4 ++-- 6 files changed, 55 insertions(+), 10 deletions(-) diff --git a/core/agents/architect.py b/core/agents/architect.py index c1691d4c..8deb96f2 100644 --- a/core/agents/architect.py +++ b/core/agents/architect.py @@ -186,6 +186,7 @@ class Architect(BaseAgent): spec.templates = {t.name: t.options_dict for t in templates.values()} spec.system_dependencies = [d.model_dump() for d in arch.system_dependencies] spec.package_dependencies = [d.model_dump() for d in arch.package_dependencies] + telemetry.set("architecture", arch) async def check_compatibility(self, arch: Architecture) -> bool: warn_system_deps = [dep.name for dep in arch.system_dependencies if dep.name.lower() in WARN_SYSTEM_DEPS] diff --git a/core/agents/spec_writer.py b/core/agents/spec_writer.py index d1cdb98f..998c1d77 100644 --- a/core/agents/spec_writer.py +++ b/core/agents/spec_writer.py @@ -123,6 +123,7 @@ class SpecWriter(BaseAgent): llm = self.get_llm(SPEC_WRITER_AGENT_NAME) convo = AgentConvo(self).template("prompt_complexity", prompt=prompt) llm_response: str = await llm(convo, temperature=0, parser=StringParser()) + log.info(f"Complexity check response: {llm_response}") return llm_response.lower() async def prepare_example_project(self, example_name: str): diff --git a/core/cli/main.py b/core/cli/main.py index 39c74f1e..e69d5d42 100644 --- a/core/cli/main.py +++ b/core/cli/main.py @@ -1,4 +1,6 @@ import asyncio +import atexit +import signal import sys from argparse import Namespace from asyncio import run @@ -17,6 +19,21 @@ from core.ui.base import UIBase, UIClosedError, UserInput, pythagora_source log = get_logger(__name__) +telemetry_sent = False + + +async def cleanup(ui: UIBase): + global telemetry_sent + if not telemetry_sent: + await telemetry.send() + telemetry_sent = True + await ui.stop() + + +def sync_cleanup(ui: UIBase): + asyncio.run(cleanup(ui)) + + async def run_project(sm: StateManager, ui: UIBase) -> bool: """ Work on the project. @@ -194,6 +211,7 @@ async def async_main( :param args: Command-line arguments. :return: True if the application ran successfully, False otherwise. """ + global telemetry_sent if args.list: await list_projects(db) @@ -223,9 +241,23 @@ async def async_main( return False telemetry.start() - success = await run_pythagora_session(sm, ui, args) - await telemetry.send() - await ui.stop() + + # Set up signal handlers + def signal_handler(sig, frame): + if not telemetry_sent: + sync_cleanup(ui) + sys.exit(0) + + for sig in (signal.SIGINT, signal.SIGTERM): + signal.signal(sig, signal_handler) + + # Register the cleanup function + atexit.register(sync_cleanup, ui) + + try: + success = await run_pythagora_session(sm, ui, args) + finally: + await cleanup(ui) return success diff --git a/core/state/state_manager.py b/core/state/state_manager.py index 5f1abf78..927f2d4f 100644 --- a/core/state/state_manager.py +++ b/core/state/state_manager.py @@ -9,7 +9,7 @@ from tenacity import retry, stop_after_attempt, wait_fixed from core.config import FileSystemType, get_config from core.db.models import Branch, ExecLog, File, FileContent, LLMRequest, Project, ProjectState, UserInput -from core.db.models.specification import Specification +from core.db.models.specification import Complexity, Specification from core.db.session import SessionManager from core.disk.ignore import IgnoreMatcher from core.disk.vfs import LocalDiskVFS, MemoryVFS, VirtualFileSystem @@ -136,7 +136,7 @@ class StateManager: The returned ProjectState will have branch and branch.project relationships preloaded. All other relationships must be - excplicitly loaded using ProjectState.awaitable_attrs or + explicitly loaded using ProjectState.awaitable_attrs or AsyncSession.refresh. :param project_id: Project ID (keyword-only, optional). @@ -211,6 +211,17 @@ class StateManager: self.current_state.tasks, ) + telemetry.set( + "architecture", + { + "system_dependencies": self.current_state.specification.system_dependencies, + "package_dependencies": self.current_state.specification.package_dependencies, + }, + ) + telemetry.set("example_project", self.current_state.specification.example_project) + telemetry.set("is_complex_app", self.current_state.specification.complexity != Complexity.SIMPLE) + telemetry.set("templates", self.current_state.specification.templates) + return self.current_state @retry(stop=stop_after_attempt(3), wait=wait_fixed(1)) diff --git a/core/telemetry/__init__.py b/core/telemetry/__init__.py index 16e62f60..4f5195af 100644 --- a/core/telemetry/__init__.py +++ b/core/telemetry/__init__.py @@ -86,8 +86,8 @@ class Telemetry: "updated_prompt": None, # App complexity "is_complex_app": None, - # Optional template used for the project - "template": None, + # Optional templates used for the project + "templates": None, # Optional, example project selected by the user "example_project": None, # Optional user contact email @@ -319,7 +319,7 @@ class Telemetry: "median_time": sorted(self.slow_requests)[n_slow // 2] if n_slow > 0 else None, } - async def send(self, event: str = "pilot-telemetry"): + async def send(self, event: str = "pythagora-core-telemetry"): """ Send telemetry data to the phone-home endpoint. diff --git a/tests/telemetry/test_telemetry.py b/tests/telemetry/test_telemetry.py index 19888c0a..c8cb267d 100644 --- a/tests/telemetry/test_telemetry.py +++ b/tests/telemetry/test_telemetry.py @@ -147,7 +147,7 @@ async def test_send_enabled_and_successful(mock_settings, mock_getenv, mock_http expected = { "pathId": "test-id", - "event": "pilot-telemetry", + "event": "pythagora-core-telemetry", "data": telemetry.data, } mock_httpx_post.assert_awaited_once_with("test-endpoint", json=expected) @@ -167,7 +167,7 @@ async def test_send_enabled_but_post_fails(mock_settings, mock_getenv, mock_http expected = { "pathId": "test-id", - "event": "pilot-telemetry", + "event": "pythagora-core-telemetry", "data": telemetry.data, } mock_httpx_post.assert_awaited_once_with(telemetry.endpoint, json=expected)