mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(prompt_management): don't hijack vector-store requests without a prompt_id
An always-on vector_store_ids request with a prompt manager (e.g. dotprompt) registered raised "prompt_id is required for Prompt Management Base class" instead of running vector-store retrieval. should_run_prompt_management_hooks returns True for vector stores, but the fallback logger selection returned the first CustomPromptManagement callback even when prompt_id is None, and the sync get_chat_completion_prompt raised on prompt_id is None. Skip prompt managers that don't run without a prompt_id during fallback logger selection, and make the sync path no-op (matching async) when prompt_id is None. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
23de7a15d9
commit
249f59ca27
3 changed files with 106 additions and 3 deletions
|
|
@ -165,7 +165,7 @@ class PromptManagementBase(ABC):
|
|||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
if prompt_id is None:
|
||||
raise ValueError("prompt_id is required for Prompt Management Base class")
|
||||
return model, messages, non_default_params
|
||||
if not self.should_run_prompt_management(
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
|
|
|
|||
|
|
@ -800,6 +800,23 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _prompt_management_logger_runs_without_prompt_id(
|
||||
logger: CustomLogger,
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
dynamic_callback_params: Optional[StandardCallbackDynamicParams],
|
||||
) -> bool:
|
||||
if not isinstance(logger, CustomPromptManagement):
|
||||
return False
|
||||
try:
|
||||
return logger.should_run_prompt_management(
|
||||
prompt_id=None,
|
||||
prompt_spec=prompt_spec,
|
||||
dynamic_callback_params=dynamic_callback_params or StandardCallbackDynamicParams(),
|
||||
)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def get_custom_logger_for_prompt_management(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -850,8 +867,13 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
callback_type=CustomPromptManagement
|
||||
)
|
||||
|
||||
if prompt_management_loggers:
|
||||
logger = prompt_management_loggers[0]
|
||||
for logger in prompt_management_loggers:
|
||||
if prompt_id is None and not self._prompt_management_logger_runs_without_prompt_id(
|
||||
logger=logger,
|
||||
prompt_spec=prompt_spec,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
):
|
||||
continue
|
||||
self.model_call_details["prompt_integration"] = logger.__class__.__name__
|
||||
return logger
|
||||
|
||||
|
|
|
|||
|
|
@ -4153,3 +4153,84 @@ def test_pre_call_does_not_pin_request_in_module_state(logging_obj):
|
|||
logging_obj.post_call(original_response='{"ok": true}', input=big_input, api_key="sk-test")
|
||||
|
||||
assert litellm.error_logs == {}
|
||||
|
||||
|
||||
def test_vector_store_hook_not_hijacked_by_prompt_manager(logging_obj, tmp_path, monkeypatch):
|
||||
"""
|
||||
Regression: a model with always-on `vector_store_ids` and a registered prompt manager
|
||||
(e.g. dotprompt) sent the request through the prompt manager, which raised
|
||||
"prompt_id is required for Prompt Management Base class" instead of running vector store retrieval.
|
||||
"""
|
||||
from litellm.integrations.dotprompt.dotprompt_manager import DotpromptManager
|
||||
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
|
||||
VectorStorePreCallHook,
|
||||
)
|
||||
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
|
||||
from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
|
||||
|
||||
(tmp_path / "stem.prompt").write_text("---\nmodel: gpt-4\n---\nyou are a stem tutor\n")
|
||||
dotprompt_manager = DotpromptManager(prompt_directory=str(tmp_path))
|
||||
litellm.logging_callback_manager.add_litellm_callback(dotprompt_manager)
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"vector_store_registry",
|
||||
VectorStoreRegistry(
|
||||
vector_stores=[
|
||||
LiteLLM_ManagedVectorStore(vector_store_id="vs_123", custom_llm_provider="openai")
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
non_default_params = {"vector_store_ids": ["vs_123"]}
|
||||
assert logging_obj.should_run_prompt_management_hooks(
|
||||
prompt_id=None, non_default_params=non_default_params
|
||||
)
|
||||
|
||||
selected_logger = logging_obj.get_custom_logger_for_prompt_management(
|
||||
model="claude-opus-4-6",
|
||||
non_default_params=non_default_params,
|
||||
prompt_id=None,
|
||||
dynamic_callback_params={},
|
||||
)
|
||||
assert isinstance(selected_logger, VectorStorePreCallHook)
|
||||
|
||||
messages = [{"role": "user", "content": "what is in my study notes?"}]
|
||||
model, returned_messages, returned_params = logging_obj.get_chat_completion_prompt(
|
||||
model="claude-opus-4-6",
|
||||
messages=messages,
|
||||
non_default_params=non_default_params,
|
||||
prompt_variables=None,
|
||||
prompt_id=None,
|
||||
)
|
||||
assert (model, returned_messages, returned_params) == (
|
||||
"claude-opus-4-6",
|
||||
messages,
|
||||
non_default_params,
|
||||
)
|
||||
|
||||
assert dotprompt_manager.get_chat_completion_prompt(
|
||||
model="claude-opus-4-6",
|
||||
messages=messages,
|
||||
non_default_params=non_default_params,
|
||||
prompt_id=None,
|
||||
prompt_variables=None,
|
||||
dynamic_callback_params={},
|
||||
) == ("claude-opus-4-6", messages, non_default_params)
|
||||
|
||||
assert isinstance(
|
||||
logging_obj.get_custom_logger_for_prompt_management(
|
||||
model="claude-opus-4-6",
|
||||
non_default_params=non_default_params,
|
||||
prompt_id="stem",
|
||||
dynamic_callback_params={},
|
||||
),
|
||||
DotpromptManager,
|
||||
)
|
||||
finally:
|
||||
litellm.logging_callback_manager.remove_callback_from_list_by_object(
|
||||
litellm.callbacks, dotprompt_manager
|
||||
)
|
||||
litellm.logging_callback_manager.remove_callback_from_list_by_object(
|
||||
litellm._async_success_callback, dotprompt_manager
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue