trending/PageIndex
1
0
Fork 0
mirror of https://github.com/VectifyAI/PageIndex.git synced 2026-10-02 07:44:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #372 from VectifyAI/openai-sdk-fast-path

Browse Source 
Use OpenAI SDK directly for OpenAI models
This commit is contained in:
Ray
2026-07-31 02:38:45 +08:00
committed by GitHub
parent 66d5b9ac9b 548d3201b7
commit d178bc557d
2 changed files with 55 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pageindex/config.yaml
+3 -1
View File
@@ -1,6 +1,8 @@
# Models without a provider prefix use the OpenAI SDK directly.
# For other providers, use "provider/model" format (e.g. "anthropic/claude-sonnet-4-6").
model: "gpt-4o-2024-11-20"
# model: "anthropic/claude-sonnet-4-6"
summary_model: "gpt-4.1-nano"
summary_model: "gpt-5.6-luna"
retrieve_model: "gpt-5.4" # defaults to `model` if not set
toc_check_page_num: 20
max_page_num_each_node: 10
pageindex/utils.py
+52 -15
View File
@@ -31,20 +31,45 @@ def count_tokens(text, model=None):
return litellm.token_counter(model=model, text=text)
def _is_openai_model(model):
"""Models without a provider prefix (no '/') use the openai SDK directly.
For other providers, use 'provider/model' format (e.g. 'anthropic/claude-sonnet-4-6')."""
if not model or model.startswith('litellm/'):
return False
return '/' not in model or model.startswith('openai/')
_openai_sync_client = None
_openai_async_client = None
def llm_completion(model, prompt, chat_history=None, return_finish_reason=False):
import litellm
use_openai_sdk = _is_openai_model(model)
if model:
model = model.removeprefix("litellm/")
if use_openai_sdk:
model = model.removeprefix("openai/")
max_retries = 10
messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}]
for i in range(max_retries):
try:
response = litellm.completion(
model=model,
messages=messages,
temperature=0,
drop_params=True,
)
if use_openai_sdk:
global _openai_sync_client
if _openai_sync_client is None:
import openai
_openai_sync_client = openai.OpenAI(max_retries=0)
response = _openai_sync_client.chat.completions.create(
model=model,
messages=messages,
)
else:
import litellm
response = litellm.completion(
model=model,
messages=messages,
temperature=0,
drop_params=True,
)
content = response.choices[0].message.content
if return_finish_reason:
finish_reason = "max_output_reached" if response.choices[0].finish_reason == "length" else "finished"
@@ -62,21 +87,33 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False)
return ""
async def llm_acompletion(model, prompt):
import litellm
use_openai_sdk = _is_openai_model(model)
if model:
model = model.removeprefix("litellm/")
if use_openai_sdk:
model = model.removeprefix("openai/")
max_retries = 10
messages = [{"role": "user", "content": prompt}]
for i in range(max_retries):
try:
response = await litellm.acompletion(
model=model,
messages=messages,
temperature=0,
drop_params=True,
)
if use_openai_sdk:
global _openai_async_client
if _openai_async_client is None:
import openai
_openai_async_client = openai.AsyncOpenAI(max_retries=0)
response = await _openai_async_client.chat.completions.create(
model=model,
messages=messages,
)
else:
import litellm
response = await litellm.acompletion(
model=model,
messages=messages,
temperature=0,
drop_params=True,
)
return response.choices[0].message.content
except Exception as e:
print('************* Retrying *************')