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):
This commit is contained in:
Stan Girard 2024-06-10 15:28:14 +02:00 committed by GitHub
parent aac7c15151
commit 76930363ee
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 24 additions and 12 deletions

View File

@ -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,

View File

@ -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(

View File

@ -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():