Add basic telemetry server for khoj

This commit is contained in:
Debanjum Singh Solanky 2023-05-15 22:55:49 +08:00
parent 134cce9d32
commit d42f0f5055
5 changed files with 77 additions and 5 deletions

View file

@ -229,13 +229,14 @@ def save_chat_session():
@schedule.repeat(schedule.every(1).minutes) @schedule.repeat(schedule.every(1).minutes)
def upload_telemetry(): def upload_telemetry():
if not state.config.app.should_log_telemetry or not state.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 return
try: 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) requests.post(constants.telemetry_server, json=state.telemetry)
except Exception as e: except Exception as e:
logger.error(f"Error uploading telemetry: {e}") logger.error(f"📡 Error uploading telemetry: {e}")
else: else:
state.telemetry = None state.telemetry = []

View file

@ -168,7 +168,10 @@ def search(
# Cache results # Cache results
state.query_cache[query_cache_key] = 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 return results

View file

@ -26,6 +26,7 @@ query_cache = LRU()
search_index_lock = threading.Lock() search_index_lock = threading.Lock()
SearchType = utils_config.SearchType SearchType = utils_config.SearchType
telemetry: List[Dict[str, str]] = [] telemetry: List[Dict[str, str]] = []
previous_query: str = None
if torch.cuda.is_available(): if torch.cuda.is_available():
# Use CUDA GPU # Use CUDA GPU

View file

@ -0,0 +1,2 @@
uvicorn
fastapi

View 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")