local-llm-server/llm_server/database/log_to_db.py

import pickle
from typing import Union

from redis import Redis


def log_to_db(ip: str, token: str, prompt: str, response: Union[str, None], gen_time: Union[int, float, None], parameters: dict, headers: dict, backend_response_code: int, request_url: str, backend_url: str, response_tokens: int = None, is_error: bool = False):
    assert isinstance(prompt, str)
    assert isinstance(backend_url, str)

    r = Redis(host='localhost', port=6379, db=3)
    data = {
        'function': 'log_prompt',
        'args': [],
        'kwargs': {
            'ip': ip,
            'token': token,
            'prompt': prompt,
            'response': response,
            'gen_time': gen_time,
            'parameters': parameters,
            'headers': dict(headers) if headers else headers,
            'backend_response_code': backend_response_code,
            'request_url': request_url,
            'backend_url': backend_url,
            'response_tokens': response_tokens,
            'is_error': is_error
        }
    }
    r.publish('database-logger', pickle.dumps(data))