89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
async def single_async_item(
|
|
api_url,
|
|
api_key,
|
|
item,
|
|
reasoning_effort,
|
|
temperature,
|
|
top_p,
|
|
semaphore_number,
|
|
model_name,
|
|
priority,
|
|
output_schema=None,
|
|
max_token=4096,
|
|
print_logs=False,
|
|
return_reason=False,
|
|
stop=None,
|
|
return_used_token=False,
|
|
timeout=300,
|
|
):
|
|
try:
|
|
async with AsyncOpenAI(
|
|
base_url=api_url, api_key=api_key
|
|
) as client:
|
|
semaphore = asyncio.Semaphore(semaphore_number)
|
|
async with semaphore:
|
|
messages = [{"role": "user", "content": item["user_prompt"]}]
|
|
if item.get("system_prompt"):
|
|
messages.insert(
|
|
0, {"role": "system", "content": item["system_prompt"]}
|
|
)
|
|
if item.get("assistant_prompt"):
|
|
messages.append(
|
|
{"role": "assistant", "content": item["assistant_prompt"]}
|
|
)
|
|
|
|
coro = client.chat.completions.parse(
|
|
model=model_name,
|
|
messages=messages,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
max_tokens=max_token,
|
|
stop=stop,
|
|
response_format=output_schema,
|
|
reasoning_effort=reasoning_effort,
|
|
extra_body={"priority": priority},
|
|
)
|
|
response = await asyncio.wait_for(coro, timeout=timeout)
|
|
|
|
if print_logs:
|
|
print(f"parse response ---- {response}")
|
|
|
|
parsed_obj = response.choices[0].message.parsed
|
|
# print(f'parsed_obj {parsed_obj}')
|
|
if parsed_obj is None:
|
|
return {
|
|
"error": "Failed to parse response",
|
|
"raw": str(response),
|
|
}
|
|
|
|
parsed_obj = output_schema.model_validate(parsed_obj)
|
|
# Validate just in case (optional, چون .parse already does it)
|
|
if return_reason:
|
|
reasoning_content = response.choices[
|
|
0
|
|
].message.reasoning_content
|
|
if return_used_token:
|
|
_total_token = response.usage.total_tokens
|
|
item["llm_output"] = (
|
|
parsed_obj.model_dump(),
|
|
str(reasoning_content),
|
|
int(_total_token),
|
|
)
|
|
return item
|
|
|
|
item["llm_output"] = (
|
|
parsed_obj.model_dump(),
|
|
str(reasoning_content)
|
|
)
|
|
return item
|
|
|
|
item["llm_output"] = parsed_obj.model_dump()
|
|
return item
|
|
|
|
except asyncio.TimeoutError:
|
|
print(f"⏳ Timeout on item {item}")
|
|
return None
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ Error __process_item {item}: {traceback.print_exc()}")
|
|
return None |