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)
|
@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 = []
|
||||||
|
|
|
@ -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
|
||||||
|
|
||||||
|
|
|
@ -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
|
||||||
|
|
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