feat: Consolidate extra parameters in sendchat.py

This commit is contained in:
Paul Gauthier (aider) 2024-09-27 13:02:44 -07:00
parent eb0331baed
commit 74f615bbb4

View file

@ -52,9 +52,7 @@ def send_completion(
functions, functions,
stream, stream,
temperature=0, temperature=0,
extra_headers=None, extra_params=None,
extra_body=None,
max_tokens=None,
): ):
from aider.llm import litellm from aider.llm import litellm
@ -70,12 +68,9 @@ def send_completion(
function = functions[0] function = functions[0]
kwargs["tools"] = [dict(type="function", function=function)] kwargs["tools"] = [dict(type="function", function=function)]
kwargs["tool_choice"] = {"type": "function", "function": {"name": function["name"]}} kwargs["tool_choice"] = {"type": "function", "function": {"name": function["name"]}}
if extra_headers is not None:
kwargs["extra_headers"] = extra_headers if extra_params is not None:
if extra_body is not None: kwargs.update(extra_params)
kwargs["extra_body"] = extra_body
if max_tokens is not None:
kwargs["max_tokens"] = max_tokens
key = json.dumps(kwargs, sort_keys=True).encode() key = json.dumps(kwargs, sort_keys=True).encode()
@ -85,8 +80,6 @@ def send_completion(
if not stream and CACHE is not None and key in CACHE: if not stream and CACHE is not None and key in CACHE:
return hash_object, CACHE[key] return hash_object, CACHE[key]
# del kwargs['stream']
res = litellm.completion(**kwargs) res = litellm.completion(**kwargs)
if not stream and CACHE is not None: if not stream and CACHE is not None:
@ -96,7 +89,7 @@ def send_completion(
@lazy_litellm_retry_decorator @lazy_litellm_retry_decorator
def simple_send_with_retries(model_name, messages, extra_headers=None, extra_body=None): def simple_send_with_retries(model_name, messages, extra_params=None):
try: try:
kwargs = { kwargs = {
"model_name": model_name, "model_name": model_name,
@ -104,10 +97,8 @@ def simple_send_with_retries(model_name, messages, extra_headers=None, extra_bod
"functions": None, "functions": None,
"stream": False, "stream": False,
} }
if extra_headers is not None: if extra_params is not None:
kwargs["extra_headers"] = extra_headers kwargs["extra_params"] = extra_params
if extra_body is not None:
kwargs["extra_body"] = extra_body
_hash, response = send_completion(**kwargs) _hash, response = send_completion(**kwargs)
return response.choices[0].message.content return response.choices[0].message.content