fix(prompt_registry): add __init__.py and registry for langfuse integ…#24648
fix(prompt_registry): add __init__.py and registry for langfuse integ…#24648thiago-carbonera wants to merge 2 commits intoBerriAI:mainfrom
Conversation
| The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile SummaryThis PR fixes the root cause of issue #23860 where Confidence Score: 3/5The discovery and isinstance-check fixes are sound, but the sync path silently drops user messages — a functional issue for the feature being fixed. The core issue (missing init.py + wrong inheritance) is correctly resolved and the test has been improved. However, the sync get_chat_completion_prompt override drops the caller's original messages (flagged in prior threads, not yet addressed), which would make the integration non-functional for synchronous callers. Portuguese comments are also present in production code. litellm/integrations/langfuse/langfuse_prompt_management.py — the get_chat_completion_prompt override and Portuguese inline comments need attention before merge.
|
| Filename | Overview |
|---|---|
| litellm/integrations/langfuse/init.py | New file that adds the missing init.py and exports prompt_initializer_registry, enabling dynamic discovery by get_prompt_initializer_from_integrations(). Implementation is clean and uses safe getattr access. |
| litellm/integrations/langfuse/langfuse_prompt_management.py | Key changes: class now inherits from CustomPromptManagement (fixing the isinstance check), should_run_prompt_management returns True when prompt_id is None (correctly enabling langfuse/ model prefix), and adds a sync get_chat_completion_prompt override. Contains Portuguese comments in production code. The sync override drops user messages — a concern flagged in prior review threads that remains unresolved. |
| tests/test_litellm/integrations/langfuse/test_langfuse_prompt_init.py | New test that validates Langfuse discovery and mocked initialization. Correctly patches langfuse_client_init to prevent network calls. Does not exercise the InMemoryPromptRegistry.initialize_prompt production path (previously flagged). |
| enterprise/litellm_enterprise/enterprise_callbacks/callback_controls.py | Pure reformatting — indentation corrected on the static method body. No logic changes. |
| enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py | Reformatting only — long lines wrapped for PEP 8 compliance, quote style normalised. No behavioral changes. |
| enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py | Reformatting only — trailing whitespace removed, long lines wrapped. No logic changes. |
Sequence Diagram
sequenceDiagram participant Proxy participant InMemoryPromptRegistry participant Scanner as get_prompt_initializer_from_integrations participant LangfuseInit as litellm/integrations/langfuse/__init__.py participant LangfusePromptManagement Proxy->>InMemoryPromptRegistry: initialize_prompt(PromptSpec) InMemoryPromptRegistry->>Scanner: scan integrations dir Scanner->>LangfuseInit: importlib.import_module(litellm.integrations.langfuse) LangfuseInit-->>Scanner: prompt_initializer_registry {langfuse: initialize_prompt} Scanner-->>InMemoryPromptRegistry: discovered_initializers InMemoryPromptRegistry->>LangfuseInit: initializer(litellm_params, prompt) LangfuseInit->>LangfusePromptManagement: LangfusePromptManagement(public_key, secret, host) LangfusePromptManagement-->>InMemoryPromptRegistry: instance (isinstance CustomPromptManagement check passes) Note over Proxy,LangfusePromptManagement: At request time Proxy->>LangfusePromptManagement: should_run_prompt_management(prompt_id=None) LangfusePromptManagement-->>Proxy: True Proxy->>LangfusePromptManagement: get_chat_completion_prompt(model, messages, ...) alt prompt_id is None AND model not langfuse/ prefix LangfusePromptManagement-->>Proxy: pass-through unchanged else model starts with langfuse/ OR prompt_id set LangfusePromptManagement->>LangfusePromptManagement: _compile_prompt_helper(prompt_id) LangfusePromptManagement-->>Proxy: (model, template_only, optional_params) - user messages dropped end Reviews (7): Last reviewed commit: "test: add mocked langfuse discovery test..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
tests/test_langfuse_prompt_init.py Outdated
| import pytest | ||
| from litellm.proxy.prompts.prompt_registry import get_prompt_initializer_from_integrations | ||
| from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec | ||
| |
There was a problem hiding this comment.
Test placed in wrong directory
The pre-submission checklist requires new tests to live under tests/test_litellm/. Existing Langfuse tests already live at tests/test_litellm/integrations/langfuse/ (e.g., test_langfuse_prompt_management.py). Please move this file there so it is picked up by make test-unit and co-located with the rest of the integration tests.
tests/test_langfuse_prompt_init.py Outdated
| ) | ||
| spec = PromptSpec(prompt_id="test-id", litellm_params=params) | ||
| | ||
| obj = init_func(params, spec) |
There was a problem hiding this comment.
Test instantiates real Langfuse client without mocking
init_func(params, spec) calls LangfusePromptManagement.__init__, which invokes langfuse_client_init(), which constructs a real Langfuse(...) SDK client. The Langfuse SDK spawns a background flush worker on construction and may attempt to reach https://cloud.langfuse.com — producing real outbound network calls in CI. The test will also fail on any environment where the langfuse package is not installed.
To keep this test hermetic, langfuse_client_init (defined in litellm/integrations/langfuse/langfuse_prompt_management.py) should be patched with unittest.mock.patch before calling init_func. This follows the same pattern used in tests/test_litellm/integrations/langfuse/test_langfuse_prompt_management.py.
Rule Used: What: prevent any tests from being added here that... (source)
tests/test_langfuse_prompt_init.py Outdated
| | ||
| obj = init_func(params, spec) | ||
| assert obj is not None, "Initiation function returned None!" | ||
| assert obj.integration_name == "langfuse" No newline at end of file |
There was a problem hiding this comment.
Missing newline at end of file
The file is missing a trailing newline, which many linters and the pre-commit hook flag.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| @greptileai please review this |
thiago-carbonera force-pushed the litellm_fix_langfuse_init branch from daa41f0 to 16fd5d4 Compare | | ||
| def initialize_prompt(litellm_params, prompt_spec): | ||
| """ | ||
| Initialization function that prompt_registry.py will call. | ||
| """ | ||
| return LangfusePromptManagement( | ||
| langfuse_public_key=getattr(litellm_params, "langfuse_public_key", None), | ||
| langfuse_secret=getattr(litellm_params, "langfuse_secret", None), | ||
| langfuse_host=getattr(litellm_params, "langfuse_host", None), | ||
| ) |
There was a problem hiding this comment.
initialize_prompt returns instance that fails isinstance check in production
LangfusePromptManagement inherits from LangFuseLogger, PromptManagementBase, and CustomLogger — but NOT from CustomPromptManagement. Every other integration in this codebase (ArizePhoenixPromptManager, BitBucketPromptManager, etc.) explicitly inherits from CustomPromptManagement.
The call path in InMemoryPromptRegistry.initialize_prompt() (prompt_registry.py lines 137–147) does:
custom_prompt_callback = initializer(litellm_params, prompt) # returns LangfusePromptManagement if not isinstance(custom_prompt_callback, CustomPromptManagement): # → False raise ValueError(f"CustomPromptManagement is required, got {type(custom_prompt_callback)}")After this fix users will still get a ValueError when actually trying to use the langfuse prompt integration — just a different error than before. The test added in this PR tests initialize_prompt directly and never goes through InMemoryPromptRegistry.initialize_prompt(), so it does not catch this.
The fix is to make LangfusePromptManagement inherit from CustomPromptManagement (as the other manager classes do), or to adjust the prompt_registry.py isinstance check to accept PromptManagementBase instead of CustomPromptManagement.
| params = PromptLiteLLMParams( | ||
| prompt_integration="langfuse", | ||
| langfuse_public_key="test-key", | ||
| langfuse_secret="test-secret" | ||
| ) | ||
| spec = PromptSpec(prompt_id="test-id", litellm_params=params) | ||
| | ||
| # Initialize the class (this would call the network, but is now mocked) | ||
| obj = init_func(params, spec) | ||
| | ||
| assert obj is not None, "Initiation function returned None!" |
There was a problem hiding this comment.
Test does not cover the full production code path
The test calls init_func(params, spec) which is initialize_prompt from litellm/integrations/langfuse/__init__.py directly. This bypasses InMemoryPromptRegistry.initialize_prompt() in prompt_registry.py, which is the actual production entry point and contains the isinstance(custom_prompt_callback, CustomPromptManagement) guard.
A test that exercises the real path would call InMemoryPromptRegistry.initialize_prompt(prompt) and would surface the isinstance check failure described in the companion comment. Consider adding:
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY # ... with patch(...): prompt_spec = PromptSpec(prompt_id="test-id", litellm_params=params) result = IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt_spec) assert result is not None thiago-carbonera force-pushed the litellm_fix_langfuse_init branch from 16fd5d4 to c701314 Compare thiago-carbonera force-pushed the litellm_fix_langfuse_init branch from c701314 to 632f471 Compare thiago-carbonera force-pushed the litellm_fix_langfuse_init branch from 632f471 to 156b456 Compare | prompt_client = self._compile_prompt_helper( | ||
| prompt_id=prompt_id or model.replace("langfuse/", ""), | ||
| prompt_variables=prompt_variables, | ||
| dynamic_callback_params=dynamic_callback_params, | ||
| prompt_label=prompt_label, | ||
| prompt_version=prompt_version, | ||
| prompt_spec=prompt_spec | ||
| ) | ||
| | ||
| template = getattr(prompt_client, "prompt_template", None) or prompt_client["prompt_template"] | ||
| optional_params = getattr(prompt_client, "prompt_template_optional_params", None) or prompt_client["prompt_template_optional_params"] | ||
| | ||
| return ( | ||
| model, | ||
| template, | ||
| optional_params, | ||
| ) |
There was a problem hiding this comment.
User messages silently dropped in new sync override
The new get_chat_completion_prompt calls _compile_prompt_helper directly and returns (model, template, optional_params) — where template is the raw Langfuse prompt template with no user messages appended. This bypasses the compile_prompt step in PromptManagementBase that concatenates prompt_template + client_messages into completed_messages.
Before this PR, both the sync and async paths went through PromptManagementBase.get_chat_completion_prompt → compile_prompt → _compile_prompt_helper, which correctly merges the template with the caller's messages. After this PR, any call that reaches this code path (model starts with "langfuse/" or an explicit prompt_id is provided) silently drops the original user messages.
To fix this, delegate to compile_prompt (and then post_compile_prompt_processing) the same way the base class does, e.g.:
prompt_template = self.compile_prompt( prompt_id=prompt_id or model.replace("langfuse/", ""), prompt_variables=prompt_variables, client_messages=messages, dynamic_callback_params=dynamic_callback_params, prompt_label=prompt_label, prompt_version=prompt_version, prompt_spec=prompt_spec, ) return self.post_compile_prompt_processing( prompt_template=prompt_template, messages=messages, non_default_params=non_default_params, model=model, ignore_prompt_manager_model=ignore_prompt_manager_model, ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params, )![corridor-security[bot]](https://avatars.githubusercontent.com/in/1177339?s=60&v=4){kind=small}
There was a problem hiding this comment.
Security Issues
- Server-Side Request Forgery (SSRF) via Langfuse prompt management host from request metadata
Recent changes expand when Langfuse prompt management is invoked (now even when no prompt_id is provided), and allow triggering it via amodelstarting withlangfuse/. The prompt management client uses per-request dynamic callback params (e.g.,langfuse_host) sourced from request metadata without validation. A malicious caller can setlangfuse_hostto an internal or non-HTTPS URL and force the server to make outbound HTTP requests to arbitrary endpoints (e.g., cloud metadata, internal services), resulting in SSRF. This is exploitable by any authenticated proxy user and violates the SSRF guardrail (URLs from user input must be validated: HTTPS-only, block private IPs, domain allowlist).
Recommendations
- Validate and restrict
langfuse_host(HTTPS-only, block private IP ranges, enforce an allowlist of domains, or only allow server-configured host — ignore per-request host values). - Revert fail-open expansion: return False when
prompt_idis None, so prompt management is not invoked implicitly. - If per-request host is absolutely needed, run it through a centralized URL validator and reject unsafe targets.
| @@ -212,7 +252,7 @@ def should_run_prompt_management( | |||
| dynamic_callback_params: StandardCallbackDynamicParams, | |||
There was a problem hiding this comment.
The change below causes prompt management to run even when no prompt_id is provided:
if prompt_id is None: return TrueCombined with the fact that Langfuse client initialization consumes per-request dynamic callback params (e.g., langfuse_host) from request metadata, an attacker can now more easily trigger outbound HTTP requests to attacker-controlled or internal URLs by setting model to langfuse/<anything> or similar patterns. Without validating the host, this enables SSRF (e.g., requests to http://169.254.169.254/ or internal services).
Impact: Authenticated proxy users can coerce the server into making HTTP(S) requests to arbitrary endpoints.
Remediation: Validate/allowlist the host and enforce HTTPS and private IP blocking, or revert this logic to avoid implicit prompt management execution when prompt_id is absent.
| dynamic_callback_params: StandardCallbackDynamicParams, | |
| if prompt_id is None: | |
| return False |
For more details, see the finding in Corridor.
Provide feedback: Reply with whether this is a valid vulnerability or false positive to help improve Corridor's accuracy.
| template = getattr(prompt_client, "prompt_template", None) or prompt_client["prompt_template"] | ||
| optional_params = getattr(prompt_client, "prompt_template_optional_params", None) or prompt_client["prompt_template_optional_params"] | ||
| | ||
| return ( |
There was a problem hiding this comment.
This method invokes Langfuse prompt compilation with dynamic_callback_params taken from request metadata, without validating the target host (e.g., langfuse_host). A caller can set model to start with "langfuse/" (even without prompt_id) which leads here:
prompt_client = self._compile_prompt_helper( prompt_id=prompt_id or model.replace("langfuse/", ""), prompt_variables=prompt_variables, dynamic_callback_params=dynamic_callback_params, prompt_label=prompt_label, prompt_version=prompt_version, prompt_spec=prompt_spec )If dynamic_callback_params includes a malicious langfuse_host, the server will make outbound requests to that host during prompt resolution. With no URL validation (HTTPS-only, private IP blocks, allowlist), this is SSRF.
Remediation: Reject or sanitize per-request langfuse_host values (use a server-configured host only or validate against an allowlist and block private IPs). Ensure any outbound URL is validated before use.
For more details, see the finding in Corridor.
Provide feedback: Reply with whether this is a valid vulnerability or false positive to help improve Corridor's accuracy.
thiago-carbonera force-pushed the litellm_fix_langfuse_init branch from 156b456 to 576046f Compare | Hey! My local checks for the Langfuse integration are passing (both pytest and ruff). The current Lint failure in anthropic and a2a handlers seems to be pre-existing on the main branch and unrelated to this PR. Ready for review! |
1 participant
Relevant issues
Fixes #23860
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewType
🐛 Bug Fix
Changes
The
InMemoryPromptRegistryuses a dynamic discovery mechanism to find prompt management integrations (like Langfuse, Hub, etc.). While the logic for Langfuse existed inlangfuse_prompt_management.py, the directory was missing the__init__.pyfile required for the scanner to recognize it as a valid module.Root cause
The function
get_prompt_initializer_from_integrations()inprompt_registry.pyscans subdirectories for__init__.pyfiles. Sincelitellm/integrations/langfuse/lacked this file, the integration was silently skipped during the registry initialization, leading to an "Unsupported prompt" error when a user attempted to useprompt_integration:"langfuse"in their config.Fix
Added
litellm/integrations/langfuse/__init__.py.Implemented an
initialize_prompt()function in the__init__.pythat correctly instantiatesLangfusePromptManagementusinggetattrfor safe attribute access fromPromptLiteLLMParams.Exported the
prompt_initializer_registrydictionary so the scanner can automatically discover and register the Langfuse integration.Testing
Added a new unit test
tests/test_langfuse_prompt_init.pywhich verifies:test_langfuse_discovery_and_init. Confirmslangfuseis present in the registry after discovery and ensuresinitialize_promptcreates a validLangfusePromptManagementinstance without raisingAttributeError.