mirror of
https://github.com/khoj-ai/khoj.git
synced 2024-11-27 17:35:07 +01:00
Add basic telemetry server for khoj
This commit is contained in:
parent
134cce9d32
commit
d42f0f5055
5 changed files with 77 additions and 5 deletions
|
@ -229,13 +229,14 @@ def save_chat_session():
|
|||
@schedule.repeat(schedule.every(1).minutes)
|
||||
def upload_telemetry():
|
||||
if not state.config.app.should_log_telemetry or not state.telemetry:
|
||||
print("No telemetry to upload") if not state.telemetry else print("Telemetry logging disabled")
|
||||
message = "📡 No telemetry to upload" if not state.telemetry else "📡 Telemetry logging disabled"
|
||||
logger.debug(message)
|
||||
return
|
||||
|
||||
try:
|
||||
logger.debug(f"📡 Upload usage telemetry to {constants.telemetry_server}: {state.telemetry}")
|
||||
logger.debug(f"📡 Upload usage telemetry to {constants.telemetry_server}:\n{state.telemetry}")
|
||||
requests.post(constants.telemetry_server, json=state.telemetry)
|
||||
except Exception as e:
|
||||
logger.error(f"Error uploading telemetry: {e}")
|
||||
logger.error(f"📡 Error uploading telemetry: {e}")
|
||||
else:
|
||||
state.telemetry = None
|
||||
state.telemetry = []
|
||||
|
|
|
@ -168,7 +168,10 @@ def search(
|
|||
# Cache results
|
||||
state.query_cache[query_cache_key] = results
|
||||
|
||||
state.telemetry += [log_telemetry(telemetry_type="api", api="search", app_config=state.config.app)]
|
||||
# Only log telemetry if query is new and not a continuation of previous query
|
||||
if state.previous_query is None or state.previous_query not in user_query:
|
||||
state.telemetry += [log_telemetry(telemetry_type="api", api="search", app_config=state.config.app)]
|
||||
state.previous_query = user_query
|
||||
|
||||
return results
|
||||
|
||||
|
|
|
@ -26,6 +26,7 @@ query_cache = LRU()
|
|||
search_index_lock = threading.Lock()
|
||||
SearchType = utils_config.SearchType
|
||||
telemetry: List[Dict[str, str]] = []
|
||||
previous_query: str = None
|
||||
|
||||
if torch.cuda.is_available():
|
||||
# Use CUDA GPU
|
||||
|
|
2
src/telemetry/requirements.txt
Normal file
2
src/telemetry/requirements.txt
Normal file
|
@ -0,0 +1,2 @@
|
|||
uvicorn
|
||||
fastapi
|
65
src/telemetry/telemetry.py
Normal file
65
src/telemetry/telemetry.py
Normal file
|
@ -0,0 +1,65 @@
|
|||
# Standard Packages
|
||||
import argparse
|
||||
import logging
|
||||
from typing import Dict, List
|
||||
|
||||
# External Packages
|
||||
from fastapi import FastAPI
|
||||
from fastapi import HTTPException
|
||||
import sqlite3
|
||||
import uvicorn
|
||||
|
||||
|
||||
# Initialize Global App Variables
|
||||
app = FastAPI()
|
||||
sqlfile = "khoj.sqlite"
|
||||
logger = logging.getLogger()
|
||||
logger.setLevel(logging.DEBUG)
|
||||
|
||||
|
||||
@app.post("/v1/telemetry")
|
||||
def v1_telemetry(telemetry_data: List[Dict[str, str]]):
|
||||
# Throw exception if no telemetry data received in POST request body
|
||||
if len(telemetry_data) == 0:
|
||||
error_message = "Post body is empty. It should contain some telemetry data"
|
||||
logger.error(error_message)
|
||||
raise HTTPException(status_code=500, detail=error_message)
|
||||
|
||||
# Insert recieved telemetry data into SQLite db
|
||||
logger.info(f"Insert row into telemetry table: {telemetry_data}")
|
||||
with sqlite3.connect(sqlfile) as conn:
|
||||
cur = conn.cursor()
|
||||
|
||||
# Create a table if it doesn't exist
|
||||
cur.execute(
|
||||
"""CREATE TABLE IF NOT EXISTS usage (id INTEGER PRIMARY KEY, time TIMESTAMP, type TEXT, server_id TEXT, os TEXT, api TEXT, client TEXT)"""
|
||||
)
|
||||
|
||||
# Log telemetry data
|
||||
for item in telemetry_data:
|
||||
cur.execute(
|
||||
"INSERT INTO usage (time, type, server_id, os, api, client) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
(
|
||||
item["timestamp"],
|
||||
item["telemetry_type"],
|
||||
item["server_id"],
|
||||
item["os"],
|
||||
item.get("api"),
|
||||
item.get("client"),
|
||||
),
|
||||
)
|
||||
# Commit the changes
|
||||
conn.commit()
|
||||
|
||||
return {"status": "ok", "message": "Logged usage telemetry"}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Setup Argument Parser
|
||||
parser = argparse.ArgumentParser(description="Start Khoj Telemetry Server")
|
||||
parser.add_argument("--host", default="127.0.0.1", type=str, help="I.P of telemetry server")
|
||||
parser.add_argument("--port", "-p", default=80, type=int, help="Port of telemetry server")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Start Application Server
|
||||
uvicorn.run(app, host=args.host, port=args.port, log_level="debug")
|
Loading…
Reference in a new issue