refactor: extract model redaction logic and add model params

This commit is contained in:
Paul Gauthier (aider) 2024-10-31 10:05:10 -07:00
parent 2fd1681fab
commit 09a9fac91e

View file

@ -108,21 +108,31 @@ class Analytics:
"machine": platform.machine(), "machine": platform.machine(),
} }
#ai add weak_model and editor_model as args, redact them too def _redact_model_name(self, model):
def event(self, event_name, main_model=None, **kwargs): if not model:
return None
info = model_info_manager.get_model_from_cached_json_db(model.name)
if info:
return model.name
elif "/" in model.name:
return model.name.split("/")[0] + "/REDACTED"
return None
def event(self, event_name, main_model=None, weak_model=None, editor_model=None, **kwargs):
if not (self.mp or self.ph) and not self.logfile: if not (self.mp or self.ph) and not self.logfile:
return return
properties = {} properties = {}
if main_model: if main_model:
# ai: refactor this into a method! properties["main_model"] = self._redact_model_name(main_model)
# Redact the main model name unless it is in the public litellm db
info = model_info_manager.get_model_from_cached_json_db(main_model.name) if weak_model:
if info: properties["weak_model"] = self._redact_model_name(weak_model)
properties["main_model"] = main_model.name
elif "/" in main_model.name: if editor_model:
properties["main_model"] = main_model.name.split("/")[0] + "/REDACTED" properties["editor_model"] = self._redact_model_name(editor_model)
properties.update(kwargs) properties.update(kwargs)
properties.update(self.get_system_info()) # Add system info to all events properties.update(self.get_system_info()) # Add system info to all events