Adding budget values
This commit is contained in:
parent
391253ec2f
commit
0a4b52c933
@ -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'}],
|
||||
|
Loading…
x
Reference in New Issue
Block a user