Integrate litellm for multi-provider LLM support

[KylinMountain]

Summary

• Replace openai + tiktoken with litellm for unified multi-provider LLM support (OpenAI, Azure, Gemini, etc.)
• Rename ChatGPT_API → llm_completion, ChatGPT_API_async → llm_acompletion (using litellm.acompletion for native async).
• Add backward compatibility: CHATGPT_API_KEY auto-aliased to OPENAI_API_KEY
• Remove tiktoken dependency; use litellm.token_counter instead

Test plan

• Set OPENAI_API_KEY and run existing PDF indexing — should behave identically to before
• Set CHATGPT_API_KEY (legacy) — should still work via alias
• Try a non-OpenAI provider (e.g. gemini/gemini-pro) by setting the appropriate env var

🤖 Generated with Claude Code

[KylinMountain]

Code review

Found 1 issue:

1. Two functions retain hard-coded model="gpt-4o-2024-11-20" defaults while the rest of the file was updated to model=None (which falls back to the configured provider). This defeats the PR's multi-provider goal for these call sites: callers using a non-OpenAI provider will still have OpenAI forced as the model unless they explicitly override it.

PageIndex/pageindex/page_index.py

Lines 501 to 503 in 847782c

	### add verify completeness
	def generate_toc_continue(toc_content, part, model="gpt-4o-2024-11-20"):
	print('start generate_toc_continue')

PageIndex/pageindex/page_index.py

Lines 734 to 736 in 847782c

	################### fix incorrect toc #########################################################
	def single_toc_item_index_fixer(section_title, content, model="gpt-4o-2024-11-20"):
	toc_extractor_prompt = """

🤖 Generated with Claude Code

_{If this code review was useful, please react with 👍. Otherwise, react with 👎.}

[KylinMountain]

Code review

Found 1 issue:

1. Two functions retain hard-coded model="gpt-4o-2024-11-20" defaults while the rest of the file was updated to model=None (which falls back to the configured provider). This defeats the PR's multi-provider goal for these call sites: callers using a non-OpenAI provider will still have OpenAI forced as the model unless they explicitly override it.

PageIndex/pageindex/page_index.py

Lines 501 to 503 in 847782c

	### add verify completeness
	def generate_toc_continue(toc_content, part, model="gpt-4o-2024-11-20"):
	print('start generate_toc_continue')

PageIndex/pageindex/page_index.py

Lines 734 to 736 in 847782c

	################### fix incorrect toc #########################################################
	def single_toc_item_index_fixer(section_title, content, model="gpt-4o-2024-11-20"):
	toc_extractor_prompt = """

🤖 Generated with Claude Code

If this code review was useful, please react with 👍. Otherwise, react with 👎.

resolved

[KylinMountain]

you can try using another llm provider like anthropic like this with env ANTHROPIC_API_KEY.

python run_pageindex.py --pdf_path tests/pdfs/q1-fy25-earnings.pdf --model anthropic/claude-opus-4-6

[rejojer]

Code review

Found 5 issues:

1. get_page_tokens retains hardcoded OpenAI model default despite body being updated for multi-provider support. The PR replaced tiktoken with litellm.token_counter inside this function and fixed the same hardcoded-default pattern in generate_toc_continue and single_toc_item_index_fixer (changing them to model=None), but missed get_page_tokens. The sole caller at page_index_main calls get_page_tokens(doc) without passing model, so it will always use "gpt-4o-2024-11-20" for tokenization regardless of config. (Confidence: 75/100)

PageIndex/pageindex/utils.py

Lines 381 to 391 in 6bda6e7

	def get_page_tokens(pdf_path, model="gpt-4o-2024-11-20", pdf_parser="PyPDF2"):
	if pdf_parser == "PyPDF2":
	pdf_reader = PyPDF2.PdfReader(pdf_path)
	page_list = []
	for page_num in range(len(pdf_reader.pages)):
	page = pdf_reader.pages[page_num]
	page_text = page.extract_text()
	token_length = litellm.token_counter(model=model, text=page_text)
	page_list.append((page_text, token_length))
	return page_list

2. openai==1.101.0 remains as a pinned direct dependency in requirements.txt despite no longer being directly imported. The PR removes import openai from utils.py and replaces all OpenAI calls with litellm, but keeps the pinned openai version. This could cause version conflicts with litellm's own openai dependency. (Confidence: 60/100)

PageIndex/requirements.txt

Lines 1 to 2 in 6bda6e7

	openai==1.101.0
	litellm==1.82.0

3. llm_acompletion returns bare "Error" string on max retry exhaustion, inconsistent with llm_completion which returns "". PR fix: make ChatGPT_API_with_finish_reason return consistent tuple #63 specifically fixed this pattern in the sync version (ChatGPT_API_with_finish_reason) because callers pass the return value to extract_json(). The async variant carries over the old behavior. (Confidence: 40/100 -- pre-existing, not newly introduced)

PageIndex/pageindex/utils.py

Lines 58 to 77 in 6bda6e7

	async def llm_acompletion(model, prompt):
	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,
	)
	return response.choices[0].message.content
	except Exception as e:
	print('************* Retrying *************')
	logging.error(f"Error: {e}")
	if i < max_retries - 1:
	await asyncio.sleep(1)
	else:
	logging.error('Max retries reached for prompt: ' + prompt)
	return "Error"

4. check_token_limit hardcodes model='gpt-4o', which is inconsistent with the PR's multi-provider goal. Not modified by the PR, but count_tokens (which it calls) was changed to use litellm.token_counter. (Confidence: 30/100 -- pre-existing, on unmodified lines)

PageIndex/pageindex/utils.py

Lines 501 to 510 in 6bda6e7

	def check_token_limit(structure, limit=110000):
	list = structure_to_list(structure)
	for node in list:
	num_tokens = count_tokens(node['text'], model='gpt-4o')
	if num_tokens > limit:
	print(f"Node ID: {node['node_id']} has {num_tokens} tokens")
	print("Start Index:", node['start_index'])
	print("End Index:", node['end_index'])
	print("Title:", node['title'])
	print("\n")

5. Unbounded while loop in toc_transformer has no max iteration guard. If llm_completion persistently returns ("", "error"), finish_reason will never be "finished" and the loop runs forever. The same pattern in extract_toc_content was already fixed by commit ac9ceaf. (Confidence: 30/100 -- pre-existing, PR only changed the function call inside)

PageIndex/pageindex/page_index.py

Lines 302 to 326 in 6bda6e7

	last_complete = get_json_content(last_complete)
	while not (if_complete == "yes" and finish_reason == "finished"):
	position = last_complete.rfind('}')
	if position != -1:
	last_complete = last_complete[:position+2]
	prompt = f"""
	Your task is to continue the table of contents json structure, directly output the remaining part of the json structure.
	The response should be in the following JSON format:
	The raw table of contents json structure is:
	{toc_content}
	The incomplete transformed table of contents json structure is:
	{last_complete}
	Please continue the json structure, directly output the remaining part of the json structure."""
	new_complete, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True)
	if new_complete.startswith('```json'):
	new_complete = get_json_content(new_complete)
	last_complete = last_complete+new_complete
	if_complete = check_if_toc_transformation_is_complete(toc_content, last_complete, model)

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}

[KylinMountain]

Code review

Found 5 issues:

1. get_page_tokens retains hardcoded OpenAI model default despite body being updated for multi-provider support. The PR replaced tiktoken with litellm.token_counter inside this function and fixed the same hardcoded-default pattern in generate_toc_continue and single_toc_item_index_fixer (changing them to model=None), but missed get_page_tokens. The sole caller at page_index_main calls get_page_tokens(doc) without passing model, so it will always use "gpt-4o-2024-11-20" for tokenization regardless of config. (Confidence: 75/100)

PageIndex/pageindex/utils.py

Lines 381 to 391 in 6bda6e7

	def get_page_tokens(pdf_path, model="gpt-4o-2024-11-20", pdf_parser="PyPDF2"):
	if pdf_parser == "PyPDF2":
	pdf_reader = PyPDF2.PdfReader(pdf_path)
	page_list = []
	for page_num in range(len(pdf_reader.pages)):
	page = pdf_reader.pages[page_num]
	page_text = page.extract_text()
	token_length = litellm.token_counter(model=model, text=page_text)
	page_list.append((page_text, token_length))
	return page_list

2. openai==1.101.0 remains as a pinned direct dependency in requirements.txt despite no longer being directly imported. The PR removes import openai from utils.py and replaces all OpenAI calls with litellm, but keeps the pinned openai version. This could cause version conflicts with litellm's own openai dependency. (Confidence: 60/100)

PageIndex/requirements.txt

Lines 1 to 2 in 6bda6e7

	openai==1.101.0
	litellm==1.82.0

3. llm_acompletion returns bare "Error" string on max retry exhaustion, inconsistent with llm_completion which returns "". PR fix: make ChatGPT_API_with_finish_reason return consistent tuple #63 specifically fixed this pattern in the sync version (ChatGPT_API_with_finish_reason) because callers pass the return value to extract_json(). The async variant carries over the old behavior. (Confidence: 40/100 -- pre-existing, not newly introduced)

PageIndex/pageindex/utils.py

Lines 58 to 77 in 6bda6e7

	async def llm_acompletion(model, prompt):
	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,
	)
	return response.choices[0].message.content
	except Exception as e:
	print('************* Retrying *************')
	logging.error(f"Error: {e}")
	if i < max_retries - 1:
	await asyncio.sleep(1)
	else:
	logging.error('Max retries reached for prompt: ' + prompt)
	return "Error"

4. check_token_limit hardcodes model='gpt-4o', which is inconsistent with the PR's multi-provider goal. Not modified by the PR, but count_tokens (which it calls) was changed to use litellm.token_counter. (Confidence: 30/100 -- pre-existing, on unmodified lines)

PageIndex/pageindex/utils.py

Lines 501 to 510 in 6bda6e7

	def check_token_limit(structure, limit=110000):
	list = structure_to_list(structure)
	for node in list:
	num_tokens = count_tokens(node['text'], model='gpt-4o')
	if num_tokens > limit:
	print(f"Node ID: {node['node_id']} has {num_tokens} tokens")
	print("Start Index:", node['start_index'])
	print("End Index:", node['end_index'])
	print("Title:", node['title'])
	print("\n")

5. Unbounded while loop in toc_transformer has no max iteration guard. If llm_completion persistently returns ("", "error"), finish_reason will never be "finished" and the loop runs forever. The same pattern in extract_toc_content was already fixed by commit ac9ceaf. (Confidence: 30/100 -- pre-existing, PR only changed the function call inside)

PageIndex/pageindex/page_index.py

Lines 302 to 326 in 6bda6e7

	last_complete = get_json_content(last_complete)
	while not (if_complete == "yes" and finish_reason == "finished"):
	position = last_complete.rfind('}')
	if position != -1:
	last_complete = last_complete[:position+2]
	prompt = f"""
	Your task is to continue the table of contents json structure, directly output the remaining part of the json structure.
	The response should be in the following JSON format:
	The raw table of contents json structure is:
	{toc_content}
	The incomplete transformed table of contents json structure is:
	{last_complete}
	Please continue the json structure, directly output the remaining part of the json structure."""
	new_complete, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True)
	if new_complete.startswith('```json'):
	new_complete = get_json_content(new_complete)
	last_complete = last_complete+new_complete
	if_complete = check_if_toc_transformation_is_complete(toc_content, last_complete, model)

• If this code review was useful, please react with 👍. Otherwise, react with 👎.

resolved

[BukeLy]

LGTM