diff --git a/oss.py b/oss.py new file mode 100644 index 0000000..5c3b6d3 --- /dev/null +++ b/oss.py @@ -0,0 +1,89 @@ +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 \ No newline at end of file