Adding budget values

This commit is contained in:
jknapp 2025-03-16 21:59:44 -07:00
parent 391253ec2f
commit 0a4b52c933

View File

@ -139,12 +139,32 @@ class Pipeline:
processed_messages.append({"role": message["role"], "content": processed_content})
reasoning_config = {
"thinking": {
"type": "enabled",
"budget_tokens": 4096
# Set budget tokens for reasoning
reasoning_effort = body.get("reasoning_effort", "medium")
budget_tokens = REASONING_EFFORT_BUDGET_TOKEN_MAP.get(reasoning_effort)
# Allow users to input an integer value representing budget tokens
if (
not budget_tokens
and reasoning_effort not in REASONING_EFFORT_BUDGET_TOKEN_MAP.keys()
):
try:
budget_tokens = int(reasoning_effort)
except ValueError as e:
print("Failed to convert reasoning effort to int", e)
budget_tokens = 4096
# Do not use thinking if budget_tokens is set to None
if budget_tokens is not None:
reasoning_config = {
"thinking": {
"type": "enabled",
"budget_tokens": budget_tokens
}
}
}
else:
reasoning_config = {}
payload = {"modelId": model_id,
"messages": processed_messages,
"system": [{'text': system_message if system_message else 'you are an intelligent ai assistant'}],