From 2cab6b5be99e08122aec45514da1d5832c16767f Mon Sep 17 00:00:00 2001 From: mountain Date: Thu, 25 Jun 2026 16:47:27 +0800 Subject: [PATCH] fix(llm): configurable per-call litellm params instead of mutating the global llm_completion / llm_acompletion (pageindex/utils.py and pageindex/index/utils.py) set `litellm.drop_params = True` on the litellm module. litellm is a process-wide singleton, so this leaked into every other library sharing it (e.g. a host app like OpenKB that exposes its own litellm config) and could not be turned off. Replace the hardcoded `temperature=0` + global `drop_params` with a single PageIndex-owned per-call kwargs mechanism (config._LLM_PARAMS): - defaults preserve behavior: {"temperature": 0, "drop_params": True}; - passed per call via **get_llm_params(), never writing litellm's globals, so nothing leaks into other litellm users in the same process; - externally configurable: pageindex.set_llm_params(drop_params=False, temperature=1, num_retries=5, ...) or the PAGEINDEX_DROP_PARAMS env shortcut; - model/messages are reserved (PageIndex supplies them) and rejected. --- pageindex/__init__.py | 3 ++- pageindex/config.py | 42 ++++++++++++++++++++++++++++++++++++++++ pageindex/index/utils.py | 10 ++++++---- pageindex/utils.py | 9 ++++++--- 4 files changed, 56 insertions(+), 8 deletions(-) diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 4f2418e..802befa 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -6,7 +6,7 @@ from .retrieve import get_document, get_document_structure, get_page_content # SDK exports from .client import PageIndexClient, LocalClient, CloudClient -from .config import IndexConfig +from .config import IndexConfig, set_llm_params from .collection import Collection from .parser.protocol import ContentNode, ParsedDocument, DocumentParser from .storage.protocol import StorageEngine @@ -26,6 +26,7 @@ __all__ = [ "LocalClient", "CloudClient", "IndexConfig", + "set_llm_params", "Collection", "ContentNode", "ParsedDocument", diff --git a/pageindex/config.py b/pageindex/config.py index fd3b12f..2accaf8 100644 --- a/pageindex/config.py +++ b/pageindex/config.py @@ -1,5 +1,8 @@ # pageindex/config.py from __future__ import annotations + +import os + from pydantic import BaseModel @@ -20,3 +23,42 @@ class IndexConfig(BaseModel): if_add_node_summary: bool = True if_add_doc_description: bool = True if_add_node_text: bool = False + + +def _env_drop_params_default() -> bool: + return os.getenv("PAGEINDEX_DROP_PARAMS", "true").strip().lower() not in ( + "0", "false", "no", "off", + ) + + +# Per-call kwargs PageIndex passes to every litellm completion. These are +# PageIndex-OWNED and applied PER CALL — never written to litellm's shared module +# globals, so they don't leak into other libraries sharing the litellm module. +# Defaults preserve historical behavior: temperature=0 keeps structure +# extraction deterministic; drop_params=True lets a provider that rejects a param +# (e.g. temperature on some local / reasoning models) succeed by dropping it. +# Override/extend via set_llm_params(); the common drop_params case also has the +# PAGEINDEX_DROP_PARAMS env shortcut. +_LLM_PARAMS: dict = {"temperature": 0, "drop_params": _env_drop_params_default()} + +# Structural kwargs PageIndex always supplies itself — not overridable here. +_RESERVED_LLM_PARAMS = ("model", "messages") + + +def get_llm_params() -> dict: + """Return a copy of the per-call kwargs PageIndex passes to litellm.""" + return dict(_LLM_PARAMS) + + +def set_llm_params(**kwargs) -> None: + """Override or extend the litellm completion kwargs PageIndex sends per call. + + e.g. ``set_llm_params(drop_params=False, temperature=1, num_retries=5)``. + Applied per call; never writes litellm's global state, so it can't leak into + other litellm users in the same process. ``model`` / ``messages`` are + reserved (PageIndex supplies them) and rejected. + """ + reserved = [k for k in kwargs if k in _RESERVED_LLM_PARAMS] + if reserved: + raise ValueError(f"cannot override reserved litellm kwargs: {reserved}") + _LLM_PARAMS.update(kwargs) diff --git a/pageindex/index/utils.py b/pageindex/index/utils.py index f416d6d..5e9700b 100644 --- a/pageindex/index/utils.py +++ b/pageindex/index/utils.py @@ -7,6 +7,8 @@ import re import asyncio import PyPDF2 +from ..config import get_llm_params + logger = logging.getLogger(__name__) @@ -23,11 +25,12 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}] for i in range(max_retries): try: - litellm.drop_params = True response = litellm.completion( model=model, messages=messages, - temperature=0, + # Per-call litellm kwargs (default temperature=0, drop_params=True); + # configure via config.set_llm_params(...) — never the litellm global. + **get_llm_params(), ) content = response.choices[0].message.content if return_finish_reason: @@ -52,11 +55,10 @@ async def llm_acompletion(model, prompt): messages = [{"role": "user", "content": prompt}] for i in range(max_retries): try: - litellm.drop_params = True response = await litellm.acompletion( model=model, messages=messages, - temperature=0, + **get_llm_params(), # per-call kwargs; never the litellm global ) return response.choices[0].message.content except Exception as e: diff --git a/pageindex/utils.py b/pageindex/utils.py index 8cfe184..cfa8782 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -18,11 +18,12 @@ from pathlib import Path from pprint import pprint from types import SimpleNamespace as config +from .config import get_llm_params + # Backward compatibility: support CHATGPT_API_KEY as alias for OPENAI_API_KEY if not os.getenv("OPENAI_API_KEY") and os.getenv("CHATGPT_API_KEY"): os.environ["OPENAI_API_KEY"] = os.getenv("CHATGPT_API_KEY") -litellm.drop_params = True async def call_llm(prompt, api_key, model="gpt-4.1", temperature=0): """Call an LLM to generate a response to a prompt. @@ -56,7 +57,9 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) response = litellm.completion( model=model, messages=messages, - temperature=0, + # Per-call litellm kwargs (default temperature=0, drop_params=True); + # configure via config.set_llm_params(...) — never the litellm global. + **get_llm_params(), ) content = response.choices[0].message.content if return_finish_reason: @@ -86,7 +89,7 @@ async def llm_acompletion(model, prompt): response = await litellm.acompletion( model=model, messages=messages, - temperature=0, + **get_llm_params(), # per-call kwargs; never the litellm global ) return response.choices[0].message.content except Exception as e: