fix typing on opik.py

This commit is contained in:
Ishaan Jaff 2024-10-10 18:46:07 +05:30
parent aadbbe9841
commit 1a9d9e1cad

View file

@ -1,11 +1,13 @@
import configparser
import os import os
import time import time
from typing import Optional, Final, Dict, List from typing import Dict, Final, List, Optional
import configparser
from litellm.types.utils import ModelResponse from litellm.types.utils import ModelResponse
CONFIG_FILE_PATH_DEFAULT: Final[str] = "~/.opik.config" CONFIG_FILE_PATH_DEFAULT: Final[str] = "~/.opik.config"
def create_uuid7(): def create_uuid7():
ns = time.time_ns() ns = time.time_ns()
last = [0, 0, 0, 0] last = [0, 0, 0, 0]
@ -36,6 +38,7 @@ def create_uuid7():
rand = os.urandom(6) rand = os.urandom(6)
return f"{t1:>08x}-{t2:>04x}-{t3:>04x}-{t4:>04x}-{rand.hex()}" return f"{t1:>08x}-{t2:>04x}-{t3:>04x}-{t4:>04x}-{rand.hex()}"
def _read_opik_config_file() -> Dict[str, str]: def _read_opik_config_file() -> Dict[str, str]:
config_path = os.path.expanduser(CONFIG_FILE_PATH_DEFAULT) config_path = os.path.expanduser(CONFIG_FILE_PATH_DEFAULT)
@ -51,15 +54,15 @@ def _read_opik_config_file() -> Dict[str, str]:
return {} return {}
def _get_env_variable(key: str) -> str:
def _get_env_variable(key: str) -> Optional[str]:
env_prefix = "opik_" env_prefix = "opik_"
return os.getenv((env_prefix + key).upper(), None) return os.getenv((env_prefix + key).upper(), None)
def get_opik_config_variable( def get_opik_config_variable(
key: str, key: str, user_value: Optional[str] = None, default_value: Optional[str] = None
user_value: Optional[str] = None, ) -> Optional[str]:
default_value: Optional[str] = None
) -> str:
""" """
Get the configuration value of a variable, order priority is: Get the configuration value of a variable, order priority is:
1. user provided value 1. user provided value
@ -85,6 +88,7 @@ def get_opik_config_variable(
# Return default value if it is not None # Return default value if it is not None
return default_value return default_value
def create_usage_object(usage): def create_usage_object(usage):
usage_dict = {} usage_dict = {}
@ -96,10 +100,12 @@ def create_usage_object(usage):
usage_dict["total_tokens"] = usage.total_tokens usage_dict["total_tokens"] = usage.total_tokens
return usage_dict return usage_dict
def _remove_nulls(x): def _remove_nulls(x):
x_ = {k: v for k, v in x.items() if v is not None} x_ = {k: v for k, v in x.items() if v is not None}
return x_ return x_
def get_traces_and_spans_from_payload(payload: List): def get_traces_and_spans_from_payload(payload: List):
traces = [_remove_nulls(x) for x in payload if "type" not in x] traces = [_remove_nulls(x) for x in payload if "type" not in x]
spans = [_remove_nulls(x) for x in payload if "type" in x] spans = [_remove_nulls(x) for x in payload if "type" in x]