diff --git a/py/autoevals/llm.py b/py/autoevals/llm.py index 870c2c2..9e97284 100644 --- a/py/autoevals/llm.py +++ b/py/autoevals/llm.py @@ -79,6 +79,7 @@ def __init__( max_tokens=None, temperature=None, engine=None, + api_key=None, ): self.name = name self.model = model @@ -90,6 +91,8 @@ def __init__( self.extra_args = {"temperature": temperature or 0} if max_tokens: self.extra_args["max_tokens"] = max(max_tokens, 5) + if api_key: + self.extra_args["api_key"] = api_key self.render_args = {} if render_args: @@ -191,6 +194,7 @@ def __init__( max_tokens=512, temperature=0, engine=None, + api_key=None, ): choice_strings = list(choice_scores.keys()) @@ -211,6 +215,7 @@ def __init__( max_tokens=max_tokens, temperature=temperature, engine=engine, + api_key=api_key, render_args={"__choices": choice_strings}, ) @@ -226,7 +231,7 @@ def from_spec_file(cls, name: str, path: str, **kwargs): class SpecFileClassifier(LLMClassifier): - def __new__(cls, model=None, engine=None, use_cot=None, max_tokens=None, temperature=None): + def __new__(cls, model=None, engine=None, use_cot=None, max_tokens=None, temperature=None, api_key=None): kwargs = {} if model is not None: kwargs["model"] = model @@ -238,6 +243,8 @@ def __new__(cls, model=None, engine=None, use_cot=None, max_tokens=None, tempera kwargs["max_tokens"] = max_tokens if temperature is not None: kwargs["temperature"] = temperature + if api_key is not None: + kwargs["api_key"] = api_key # convert FooBar to foo_bar template_name = re.sub(r"(?