From 76930363eea2aa8367a62a1afbb6ed3b15f8070a Mon Sep 17 00:00:00 2001 From: Stan Girard Date: Mon, 10 Jun 2024 15:28:14 +0200 Subject: [PATCH] feat: telemetry improved (#2651) # Description Please include a summary of the changes and the related issue. Please also include relevant motivation and context. ## Checklist before requesting a review Please delete options that are not relevant. - [ ] My code follows the style guidelines of this project - [ ] I have performed a self-review of my code - [ ] I have commented hard-to-understand areas - [ ] I have ideally added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [ ] Any dependent changes have been merged ## Screenshots (if appropriate): --- .../modules/brain/controller/brain_routes.py | 13 ++++++++----- backend/modules/chat/controller/chat_routes.py | 5 ++--- backend/packages/utils/telemetry.py | 18 ++++++++++++++---- 3 files changed, 24 insertions(+), 12 deletions(-) diff --git a/backend/modules/brain/controller/brain_routes.py b/backend/modules/brain/controller/brain_routes.py index df70c94a9..3d5a27659 100644 --- a/backend/modules/brain/controller/brain_routes.py +++ b/backend/modules/brain/controller/brain_routes.py @@ -1,7 +1,7 @@ from typing import Dict from uuid import UUID -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException, Request from logger import get_logger from middlewares.auth.auth_bearer import AuthBearer, get_current_user from modules.brain.dto.inputs import ( @@ -14,7 +14,9 @@ from modules.brain.entity.integration_brain import IntegrationDescriptionEntity from modules.brain.service.brain_authorization_service import has_brain_authorization from modules.brain.service.brain_service import BrainService from modules.brain.service.brain_user_service import BrainUserService -from modules.brain.service.get_question_context_from_brain import get_question_context_from_brain +from modules.brain.service.get_question_context_from_brain import ( + get_question_context_from_brain, +) from modules.brain.service.integration_brain_service import ( IntegrationBrainDescriptionService, ) @@ -58,7 +60,6 @@ async def retrieve_public_brains() -> list[PublicBrain]: return brain_service.get_public_brains() - @brain_router.get( "/brains/{brain_id}/", dependencies=[ @@ -84,7 +85,9 @@ async def retrieve_brain_by_id( @brain_router.post("/brains/", dependencies=[Depends(AuthBearer())], tags=["Brain"]) async def create_new_brain( - brain: CreateBrainProperties, current_user: UserIdentity = Depends(get_current_user) + brain: CreateBrainProperties, + request: Request, + current_user: UserIdentity = Depends(get_current_user), ): """Create a new brain for the user.""" user_brains = brain_user_service.get_user_brains(current_user.id) @@ -99,7 +102,7 @@ async def create_new_brain( status_code=429, detail=f"Maximum number of brains reached ({user_settings.get('max_brains', 5)}).", ) - maybe_send_telemetry("create_brain", {"brain_name": brain.name}) + maybe_send_telemetry("create_brain", {"brain_name": brain.name}, request) new_brain = brain_service.create_brain( brain=brain, user_id=current_user.id, diff --git a/backend/modules/chat/controller/chat_routes.py b/backend/modules/chat/controller/chat_routes.py index b226ed02c..5eaa95195 100644 --- a/backend/modules/chat/controller/chat_routes.py +++ b/backend/modules/chat/controller/chat_routes.py @@ -77,9 +77,6 @@ def get_answer_generator( brain, metadata_brain = brain_service.find_brain_from_question( brain_id, chat_question.question, current_user, chat_id, history, vector_store ) - - maybe_send_telemetry("question_asked", {"model_name": brain.model}) - gpt_answer_generator = chat_instance.get_answer_generator( brain=brain, chat_id=str(chat_id), @@ -219,6 +216,7 @@ async def create_question_handler( chat_answer = gpt_answer_generator.generate_answer( chat_id, chat_question, save_answer=True ) + maybe_send_telemetry("question_asked", {"streaming": False}, request) return chat_answer except HTTPException as e: @@ -253,6 +251,7 @@ async def create_stream_question_handler( gpt_answer_generator = get_answer_generator( chat_id, chat_question, brain_id, current_user ) + maybe_send_telemetry("question_asked", {"streaming": True}, request) try: return StreamingResponse( diff --git a/backend/packages/utils/telemetry.py b/backend/packages/utils/telemetry.py index afb4954c0..0cd71d265 100644 --- a/backend/packages/utils/telemetry.py +++ b/backend/packages/utils/telemetry.py @@ -4,6 +4,10 @@ import os import threading import httpx +from fastapi import Request +from logger import get_logger + +logger = get_logger(__name__) # Assume these are your Supabase Function endpoint and any necessary headers TELEMETRY_URL = "https://ovbvcnwemowuuuaebizd.supabase.co/functions/v1/telemetry" @@ -22,10 +26,14 @@ def generate_machine_key(): return unique_key -def send_telemetry(event_name: str, event_data: dict): +def send_telemetry(event_name: str, event_data: dict, request: Request = None): # Generate a unique machine key machine_key = generate_machine_key() - + domain = None + if request: + domain = request.url.hostname + logger.info(f"Domain: {domain}") + event_data = {**event_data, "domain": domain} # Prepare the payload payload = json.dumps( { @@ -40,13 +48,15 @@ def send_telemetry(event_name: str, event_data: dict): _ = client.post(TELEMETRY_URL, headers=HEADERS, data=payload) -def maybe_send_telemetry(event_name: str, event_data: dict): +def maybe_send_telemetry(event_name: str, event_data: dict, request: Request = None): enable_telemetry = os.getenv("TELEMETRY_ENABLED", "false") if enable_telemetry.lower() != "true": return - threading.Thread(target=send_telemetry, args=(event_name, event_data)).start() + threading.Thread( + target=send_telemetry, args=(event_name, event_data, request) + ).start() async def main():