31 lines
1.1 KiB
Python
31 lines
1.1 KiB
Python
|
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))
|