apunkt/PageIndex
1
0
Fork 0
mirror of https://github.com/VectifyAI/PageIndex.git synced 2026-05-16 18:25:14 +02:00
Code Issues Projects Releases Packages Wiki Activity

use async with for client calls

Browse source
This commit is contained in:
Ray 2025-06-26 01:37:20 +08:00
parent 6d6e92d56e
commit d9568ff4ed
1 changed files with 10 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

19
pageindex/utils.py
View file

@ -87,25 +87,26 @@ def ChatGPT_API(model, prompt, api_key=CHATGPT_API_KEY, chat_history=None):
async def ChatGPT_API_async(model, prompt, api_key=CHATGPT_API_KEY): async def ChatGPT_API_async(model, prompt, api_key=CHATGPT_API_KEY):
max_retries = 10 max_retries = 10
client = openai.AsyncOpenAI(api_key=api_key) messages = [{"role": "user", "content": prompt}]
for i in range(max_retries): for i in range(max_retries):
try: try:
messages = [{"role": "user", "content": prompt}] async with openai.AsyncOpenAI(api_key=api_key) as client:
response = await client.chat.completions.create( response = await client.chat.completions.create(
model=model, model=model,
messages=messages, messages=messages,
temperature=0, temperature=0,
) )
return response.choices[0].message.content return response.choices[0].message.content
except Exception as e: except Exception as e:
print('************* Retrying *************') print('************* Retrying *************')
logging.error(f"Error: {e}") logging.error(f"Error: {e}")
if i < max_retries - 1: if i < max_retries - 1:
await asyncio.sleep(1) # Wait for 1秒 before retrying await asyncio.sleep(1) # Wait for 1s before retrying
else: else:
logging.error('Max retries reached for prompt: ' + prompt) logging.error('Max retries reached for prompt: ' + prompt)
return "Error" return "Error"
def get_json_content(response): def get_json_content(response):
start_idx = response.find("```json") start_idx = response.find("```json")
if start_idx != -1: if start_idx != -1: