UN-3215 [FIX] Add LLMCompat bridge class to fix retriever LLM compatibility with llama-index#1788
Open
pk-zipstack wants to merge 31 commits intomainfrom
Open
UN-3215 [FIX] Add LLMCompat bridge class to fix retriever LLM compatibility with llama-index#1788pk-zipstack wants to merge 31 commits intomainfrom
pk-zipstack wants to merge 31 commits intomainfrom
Conversation
pk-zipstack
requested review from
gaya3-zipstack,
hari-kuriakose and
harini-venkataraman
Contributor
|
Important Review skippedReview was skipped due to path filters ⛔ Files ignored due to path filters (1)
CodeRabbit blocks several paths by default. You can override this behavior by explicitly including those paths in the path filters. For example, including ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
WalkthroughThe changes introduce a compatibility layer to adapt the SDK1 LLM interface for use with llama-index retriever components. This includes new emulated llama-index types (MessageRole, ChatMessage, ChatResponse, CompletionResponse, LLMMetadata) and a wrapper class LLMCompat in SDK1, along with a RetrieverLLM adapter in the prompt service that bridges the two interfaces. Changes
Sequence DiagramsequenceDiagram
participant Client
participant Retriever as BaseRetriever
participant Converter as _get_llm()
participant LLMCompat
participant RetrieverLLM
participant LlamaIndex as llama-index<br/>Components
Client->>Retriever: Initialize with LLM
Retriever->>Converter: _get_llm(llm)
Converter->>LLMCompat: Wrap LLM instance
LLMCompat-->>Converter: Return LLMCompat
Converter->>RetrieverLLM: Initialize with LLMCompat
RetrieverLLM-->>Converter: Return RetrieverLLM instance
Converter-->>Retriever: Return RetrieverLLM
Retriever->>LlamaIndex: Pass RetrieverLLM to retriever ops
LlamaIndex->>RetrieverLLM: Call chat/complete methods
RetrieverLLM->>LLMCompat: Delegate to internal LLMCompat
LLMCompat->>LLMCompat: Route to chat/complete/predict
LLMCompat-->>RetrieverLLM: Return ChatResponse/CompletionResponse
RetrieverLLM-->>LlamaIndex: Return response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment Tip You can validate your CodeRabbit configuration file in your editor.If your editor has YAML language server, you can enable auto-completion and validation by adding |
for more information, see https://pre-commit.ci
…endency (#1793) * [FEAT] Rewrite LLMCompat to emulate llama-index interface without dependency - Add emulated llama-index types (MessageRole, ChatMessage, ChatResponse, CompletionResponse, LLMMetadata) as dataclasses - Rewrite LLMCompat to match llama-index LLM interface without inheritance - Implement chat(), complete(), achat(), acomplete() for retriever compatibility - Follow EmbeddingCompat initialization pattern (takes adapter params directly) - Remove llama-index imports from llm.py - Raise NotImplementedError for streaming methods not needed by retrievers This allows SubQuestionQueryEngine, QueryFusionRetriever, and other llama-index components to use LLMCompat without requiring llama-index as a dependency. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Added retrieverLLM class to handle issues with retreivers * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: pk-zipstack <praveen@zipstack.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (6)
unstract/sdk1/src/unstract/sdk1/llm.py (4)
36-43: Preferenum.StrEnumover(str, Enum)dual inheritance.Python 3.11+ provides
StrEnumwhich is the idiomatic replacement. Since the codebase already uses|union types (Python 3.10+),StrEnumshould be available.♻️ Suggested change
-from enum import Enum +from enum import StrEnum -class MessageRole(str, Enum): +class MessageRole(StrEnum): """Emulates llama_index.core.base.llms.types.MessageRole."""🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@unstract/sdk1/src/unstract/sdk1/llm.py` around lines 36 - 43, Replace the MessageRole class to inherit from enum.StrEnum instead of (str, Enum): import StrEnum from enum and change the class definition for MessageRole to subclass StrEnum; keep the same member names and values (SYSTEM, USER, ASSISTANT, FUNCTION, TOOL) so behavior and comparisons remain identical.
659-676:predictandapredictassume aPromptTemplate-like interface onprompt.These methods call
prompt.format_messages(llm=self, ...)andprompt.format(llm=self, ...), expecting a llama-indexPromptTemplateobject. The type is annotated asAny, which hides this contract. Since llama-index'sLLM.predictreceives aBasePromptTemplate, this should work — but if ever called with a plain string, it'll raiseAttributeError.This is acceptable for the bridge use case, but documenting the expected type (even as a comment) would help maintainability.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@unstract/sdk1/src/unstract/sdk1/llm.py` around lines 659 - 676, predict and apredict assume a PromptTemplate-like object (they call prompt.format_messages(...) and prompt.format(...)) but are annotated as Any; update the contract by changing the type annotation from Any to the appropriate llama-index type (e.g. BasePromptTemplate or PromptTemplate) or at minimum add a clear inline comment/docstring stating the expected type, and add a simple runtime check in predict and apredict that raises a descriptive TypeError if the prompt lacks format/format_messages (e.g. if not hasattr(prompt, "format") and not hasattr(prompt, "format_messages")), referencing the predict/apredict methods and self.metadata.is_chat_model, so callers get a clear error instead of an AttributeError.
574-628:LLMCompatcreates a secondLLMinstance — consider accepting an existing one.
RetrieverLLM.__init__(inretriever_llm.py, Line 33-43) reads private attributes (llm._adapter_id,llm._adapter_metadata, etc.) from the SDK1LLMto constructLLMCompat, which then creates anotherLLMinternally (Line 613). This means every retriever bridge creates a duplicateLLMwith duplicate adapter validation, platform config fetching, and callback setup.Consider allowing
LLMCompatto accept an existingLLMinstance directly rather than always constructing a new one, to avoid this duplication and the reliance on private attributes.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@unstract/sdk1/src/unstract/sdk1/llm.py` around lines 574 - 628, LLMCompat currently always constructs a new LLM (in LLMCompat.__init__) causing duplicate LLMs; change the constructor to accept an optional existing LLM instance parameter (e.g., llm_instance: LLM | None = None) and if provided set self._llm_instance = llm_instance and skip creating a new LLM, while still setting self._tool, self._adapter_instance_id, self.model_name = self._llm_instance.get_model_name(), and self.callback_manager = self._llm_instance.callback_manager (or None) to preserve behavior; retain the original parameter list and only construct a new LLM when llm_instance is None so RetrieverLLM can pass the SDK1 LLM directly and avoid duplicated adapter validation and config fetching.
804-815: Add comment documenting the duck-typing contract for ChatMessage compatibility.
_to_litellm_messagesrelies on duck typing to accept both the emulatedChatMessage(defined in this file) and realllama_index.core.base.llms.types.ChatMessageobjects, sinceachat()andacomplete()may receive either. Both provide.roleand.contentattributes, with thegetattr(m.role, "value", ...)pattern handling both enum and string variants.While the
.contentattribute is maintained in llama-index as a backward-compatibility property (returningstr | Nonefrom the underlying block-based structure), a brief comment here clarifying the intentional duck-typing contract would improve maintainability and signal to future maintainers why both types are accepted.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@unstract/sdk1/src/unstract/sdk1/llm.py` around lines 804 - 815, Add a brief doc-comment above the _to_litellm_messages function explaining the duck-typing contract: that the function accepts both the emulated ChatMessage defined in this module and llama_index.core.base.llms.types.ChatMessage objects because both expose .role and .content, and that getattr(m.role, "value", str(m.role)) handles enum or string roles while m.content may be str | None; mention this is intentional to support inputs from achat() and acomplete().prompt-service/src/unstract/prompt_service/core/retrievers/retriever_llm.py (1)
15-15: Remove unusednoqadirectives flagged by Ruff.Static analysis (Ruff RUF100) reports 10 unused
noqadirectives across this file (N811,ANN401). These should be removed to keep the codebase clean.Also applies to: 33-33, 57-57, 72-72, 80-80, 88-88, 97-97, 112-112, 120-120, 128-128
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@prompt-service/src/unstract/prompt_service/core/retrievers/retriever_llm.py` at line 15, Remove the redundant noqa directives reported by Ruff (RUF100) in retriever_llm.py: delete the unused "# noqa: N811" on the import "from llama_index.core.llms.llm import LLM as LlamaIndexBaseLLM" and likewise remove the other unused "# noqa" tags on the affected lines (the import/annotation lines referencing LlamaIndexBaseLLM and any annotated functions/classes flagged with ANN401). Ensure you only remove the "# noqa" comments (not the imports or annotations themselves), run Ruff/pytest to verify no new linter errors, and keep the symbols like LlamaIndexBaseLLM unchanged so references in the file remain valid.llm-predict-error-investigation.md (1)
1-221: Debug artifact should not be committed tomain.This file reads as a transient investigation note, not stable documentation. Committing it to
mainintroduces noise with no durable value: it already contains stale content (see below) and will continue diverging from the codebase over time.If a record of the design decision is needed, consider:
- Converting it to an Architecture Decision Record (ADR) in a
docs/adr/folder, keeping only the decision, rationale, and final chosen approach.- Moving the content into the PR description (where it already lives as context) and dropping the file entirely.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@llm-predict-error-investigation.md` around lines 1 - 221, This investigation file is a debug artifact and should not be committed to main; remove llm-predict-error-investigation.md from the branch (or revert the commit) and either convert its essential outcome into a concise ADR under docs/adr/ (summarizing the chosen approach such as the SDK1LLMBridge, use of LiteLLM, or SimpleKeywordTableIndex and the BaseRetriever/KeywordTableIndex impact) or move the full investigative content into the PR description/temporary notes, keeping only the final decision and rationale in the repo.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@llm-predict-error-investigation.md`:
- Around line 150-184: Update the document's Option 1 example to match the
actual two-layer implementation: show LLMCompat (from
unstract/sdk1/src/unstract/sdk1/llm.py) that emulates llama-index types and
RetrieverLLM (from prompt-service/.../core/retrievers/retriever_llm.py) which
inherits from llama_index.core.llms.llm.LLM and delegates to LLMCompat; replace
the single-class SDK1LLMBridge/CustomLLM example with a short description and
class names matching LLMCompat -> RetrieverLLM delegation and note that
RetrieverLLM inherits from llama_index.core.llms.llm.LLM (not CustomLLM).
- Line 5: Add explicit language specifiers to the two fenced code blocks to
satisfy MD040: change the traceback block that starts with "ERROR:
unstract.prompt_service.core.retrievers.keyword_table:" to use ```text and
change the block containing "BaseLLM -> LLM (has predict()) ->
CustomLLM/FunctionCallingLLM -> Provider" to use ```python so markdownlint
recognizes their languages.
In `@prompt-service/src/unstract/prompt_service/core/retrievers/retriever_llm.py`:
- Around line 33-43: Declare _compat as a Pydantic PrivateAttr at the class
level for v2 compatibility: import PrivateAttr from pydantic and add a class
attribute like "_compat: Any = PrivateAttr()" (or similar typed PrivateAttr) on
the retriever class, then keep the existing assignment to self._compat in
__init__ where LLMCompat(...) is created (referencing __init__, LLMCompat, and
the _compat attribute).
In `@unstract/sdk1/src/unstract/sdk1/llm.py`:
- Around line 698-728: LLMCompat.chat() and LLMCompat.complete() call
litellm.completion() directly, skipping the SDK’s usage recording and structured
error wrapping used elsewhere; update these methods to (1) call into the
internal LLM instance’s usage recorder (e.g., invoke
self._llm_instance._record_usage(...) or the same _record_usage flow used by
LLM.complete) after getting the response so usage/audit data is emitted, (2)
wrap litellm.completion() in a try/except and re-raise errors as LLMError to
match callers' expectations, and (3) if metrics capture is required, apply the
same `@capture_metrics` behavior or add a TODO comment documenting intentional
omission; locate changes in the LLMCompat.chat and LLMCompat.complete methods
around the litellm.completion(...) calls and the response handling to implement
these fixes.
---
Nitpick comments:
In `@llm-predict-error-investigation.md`:
- Around line 1-221: This investigation file is a debug artifact and should not
be committed to main; remove llm-predict-error-investigation.md from the branch
(or revert the commit) and either convert its essential outcome into a concise
ADR under docs/adr/ (summarizing the chosen approach such as the SDK1LLMBridge,
use of LiteLLM, or SimpleKeywordTableIndex and the
BaseRetriever/KeywordTableIndex impact) or move the full investigative content
into the PR description/temporary notes, keeping only the final decision and
rationale in the repo.
In `@prompt-service/src/unstract/prompt_service/core/retrievers/retriever_llm.py`:
- Line 15: Remove the redundant noqa directives reported by Ruff (RUF100) in
retriever_llm.py: delete the unused "# noqa: N811" on the import "from
llama_index.core.llms.llm import LLM as LlamaIndexBaseLLM" and likewise remove
the other unused "# noqa" tags on the affected lines (the import/annotation
lines referencing LlamaIndexBaseLLM and any annotated functions/classes flagged
with ANN401). Ensure you only remove the "# noqa" comments (not the imports or
annotations themselves), run Ruff/pytest to verify no new linter errors, and
keep the symbols like LlamaIndexBaseLLM unchanged so references in the file
remain valid.
In `@unstract/sdk1/src/unstract/sdk1/llm.py`:
- Around line 36-43: Replace the MessageRole class to inherit from enum.StrEnum
instead of (str, Enum): import StrEnum from enum and change the class definition
for MessageRole to subclass StrEnum; keep the same member names and values
(SYSTEM, USER, ASSISTANT, FUNCTION, TOOL) so behavior and comparisons remain
identical.
- Around line 659-676: predict and apredict assume a PromptTemplate-like object
(they call prompt.format_messages(...) and prompt.format(...)) but are annotated
as Any; update the contract by changing the type annotation from Any to the
appropriate llama-index type (e.g. BasePromptTemplate or PromptTemplate) or at
minimum add a clear inline comment/docstring stating the expected type, and add
a simple runtime check in predict and apredict that raises a descriptive
TypeError if the prompt lacks format/format_messages (e.g. if not
hasattr(prompt, "format") and not hasattr(prompt, "format_messages")),
referencing the predict/apredict methods and self.metadata.is_chat_model, so
callers get a clear error instead of an AttributeError.
- Around line 574-628: LLMCompat currently always constructs a new LLM (in
LLMCompat.__init__) causing duplicate LLMs; change the constructor to accept an
optional existing LLM instance parameter (e.g., llm_instance: LLM | None = None)
and if provided set self._llm_instance = llm_instance and skip creating a new
LLM, while still setting self._tool, self._adapter_instance_id, self.model_name
= self._llm_instance.get_model_name(), and self.callback_manager =
self._llm_instance.callback_manager (or None) to preserve behavior; retain the
original parameter list and only construct a new LLM when llm_instance is None
so RetrieverLLM can pass the SDK1 LLM directly and avoid duplicated adapter
validation and config fetching.
- Around line 804-815: Add a brief doc-comment above the _to_litellm_messages
function explaining the duck-typing contract: that the function accepts both the
emulated ChatMessage defined in this module and
llama_index.core.base.llms.types.ChatMessage objects because both expose .role
and .content, and that getattr(m.role, "value", str(m.role)) handles enum or
string roles while m.content may be str | None; mention this is intentional to
support inputs from achat() and acomplete().
prompt-service/src/unstract/prompt_service/core/retrievers/retriever_llm.py
Outdated
Show resolved
Hide resolved
pk-zipstack
changed the title
harini-venkataraman
approved these changes
Mar 3, 2026
Contributor
harini-venkataraman
left a comment
harini-venkataraman
left a comment
There was a problem hiding this comment.
Added some comments. Please address. Otherwise LGTM
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
hari-kuriakose
requested changes
Mar 3, 2026
Contributor
hari-kuriakose
left a comment
hari-kuriakose
left a comment
There was a problem hiding this comment.
@pk-zipstack LGTM overall.
Let's address the minor comments though, from both humans and AI.
prompt-service/src/unstract/prompt_service/core/retrievers/base_retriever.py
Outdated
Show resolved
Hide resolved
prompt-service/src/unstract/prompt_service/core/retrievers/fusion.py
Outdated
Show resolved
Hide resolved
prompt-service/src/unstract/prompt_service/core/retrievers/base_retriever.py
Outdated
Show resolved
Hide resolved
…e_retriever.py Co-authored-by: Hari John Kuriakose <hari@zipstack.com> Signed-off-by: Praveen Kumar <praveen@zipstack.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@prompt-service/src/unstract/prompt_service/core/retrievers/base_retriever.py`:
- Around line 28-41: The constructor is using an undefined _get_llm and never
initializes the cache causing NameError/AttributeError and potential
double-wrapping of LLMs; fix by (1) replacing the call to _get_llm with the
correct factory (e.g., get_llm) when assigning self.llm (or import/define
_get_llm if that was intended), (2) initializing self._llama_index_llm = None in
__init__, and (3) adjust the llama_index_llm property to avoid double-bridging
by checking types: if self._llama_index_llm is None then if isinstance(self.llm,
RetrieverLLM) set self._llama_index_llm = get_llama_index_llm(self.llm) (or use
the existing RetrieverLLM directly) otherwise convert once with
get_llama_index_llm(self.llm), then return self._llama_index_llm.
ℹ️ Review info
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to Reviews > Disable Cache setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (1)
prompt-service/src/unstract/prompt_service/core/retrievers/base_retriever.py
prompt-service/src/unstract/prompt_service/core/retrievers/base_retriever.py
Outdated
Show resolved
Hide resolved
- Remove llm_helper.py — move conversion to _get_llm staticmethod in BaseRetriever - Convert SDK1 LLM to RetrieverLLM eagerly in constructor - Remove llama_index_llm lazy property - Revert all retrievers to use self.llm directly Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
for more information, see https://pre-commit.ci
|
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
prompt-service/src/unstract/prompt_service/core/retrievers/subquestion.py (1)
43-50: Fail fast whenself.llmis missing before building subquestion components.Add an explicit guard before constructing
question_genand query engine so error handling is deterministic and immediate.Suggested patch
query_bundle = QueryBundle(query_str=self.prompt) + if self.llm is None: + raise RetrievalError( + "SubquestionRetriever requires an initialized LLM instance." + ) + question_gen = LLMQuestionGenerator.from_defaults( llm=self.llm, ) query_engine = SubQuestionQueryEngine.from_defaults( query_engine_tools=query_engine_tools, question_gen=question_gen, use_async=True, llm=self.llm, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@prompt-service/src/unstract/prompt_service/core/retrievers/subquestion.py` around lines 43 - 50, Before calling LLMQuestionGenerator.from_defaults and SubQuestionQueryEngine.from_defaults, add an explicit guard that checks self.llm and raises a clear exception (e.g., ValueError) if it is None or falsy so construction fails fast; update the code path where question_gen and query_engine are created (LLMQuestionGenerator.from_defaults, SubQuestionQueryEngine.from_defaults) to perform this check first and return/raise immediately with a descriptive message referencing self.llm.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@prompt-service/src/unstract/prompt_service/core/retrievers/subquestion.py`:
- Around line 43-50: Before calling LLMQuestionGenerator.from_defaults and
SubQuestionQueryEngine.from_defaults, add an explicit guard that checks self.llm
and raises a clear exception (e.g., ValueError) if it is None or falsy so
construction fails fast; update the code path where question_gen and
query_engine are created (LLMQuestionGenerator.from_defaults,
SubQuestionQueryEngine.from_defaults) to perform this check first and
return/raise immediately with a descriptive message referencing self.llm.
ℹ️ Review info
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to Reviews > Disable Cache setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (4)
prompt-service/src/unstract/prompt_service/core/retrievers/base_retriever.pyprompt-service/src/unstract/prompt_service/core/retrievers/fusion.pyprompt-service/src/unstract/prompt_service/core/retrievers/keyword_table.pyprompt-service/src/unstract/prompt_service/core/retrievers/subquestion.py
✅ Files skipped from review due to trivial changes (2)
- prompt-service/src/unstract/prompt_service/core/retrievers/keyword_table.py
- prompt-service/src/unstract/prompt_service/core/retrievers/fusion.py
Contributor
|
@pk-zipstack Let's address the following: Accessing private attributes of LLM in RetrieverLLM
Add public accessors on
|
hari-kuriakose
requested changes
Mar 13, 2026
Contributor
hari-kuriakose
left a comment
hari-kuriakose
left a comment
There was a problem hiding this comment.
@pk-zipstack LGTM overall.
Please check the comments though.
- Add LLMCompat.from_llm() classmethod to encapsulate access to LLM's private attributes within SDK1, avoiding cross-package coupling - Update RetrieverLLM to use the factory method instead of accessing LLM._adapter_id, _adapter_metadata, etc. directly - Add 20 unit tests covering: LLMCompat.from_llm(), RetrieverLLM isinstance checks, chat/complete/achat/acomplete delegation, type conversion, emulated types, and _to_litellm_messages - Add prompt-service tox environment for running tests - Add pytest-asyncio and pytest-md-report to test dependencies Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
Greptile SummaryThis PR introduces a two-layer bridge architecture ( Key changes:
Issues found:
Confidence Score: 3/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant LI as llama-index Component<br/>(KeywordTable/Fusion/Router/SubQuestion)
participant RL as RetrieverLLM<br/>(llama-index LLM subclass)
participant LC as LLMCompat<br/>(plain Python class)
participant SDK as SDK1 LLM<br/>(litellm wrapper)
participant LiteLLM as litellm
Note over LI,RL: RetrieverLLM passes isinstance(llm, LLM) check
LI->>RL: chat(messages) / complete(prompt)
RL->>LC: _compat.chat(messages) / _compat.complete(prompt)
Note over LC: _messages_to_prompt() flattens<br/>all messages with role prefixes
LC->>SDK: _llm_instance.complete(prompt, **kwargs)
SDK->>LiteLLM: litellm.completion(messages, **kwargs)
LiteLLM-->>SDK: raw response
SDK-->>LC: {"response": LLMResponseCompat}
LC-->>RL: emulated ChatResponse / CompletionResponse
Note over RL: converts emulated → real llama-index types
RL-->>LI: real llama-index ChatResponse / CompletionResponse
Prompt To Fix All With AIThis is a comment left during a code review.
Path: unstract/sdk1/src/unstract/sdk1/llm.py
Line: 693-710
Comment:
**Double `CallbackManager` registration for non-public adapters**
`LLMCompat.__init__()` creates an inner `LLM` instance (which already calls `CallbackManager.set_callback(model=inner_llm, ...)` inside `LLM.__init__`) and then immediately calls `CallbackManager.set_callback(model=self, ...)` again for the `LLMCompat` wrapper. Since every actual LiteLLM call goes through `self._llm_instance.complete()` (not through `LLMCompat` methods directly), the callback registered for the `LLMCompat` wrapper in step 2 will never be triggered by a LiteLLM invocation.
If `CallbackManager.set_callback` appends to a global litellm callback list (rather than being model-scoped), the callback for the inner `LLM` fires once per call while the callback for `LLMCompat` is never invoked but still occupies a slot. More critically, if the callback list is global and accumulates entries across repeated `LLMCompat.__init__()` calls, it creates a memory/reporting drift.
Note that `from_llm()` (the preferred construction path for all retriever use-cases) has the same pattern and the same concern. Consider removing the `CallbackManager.set_callback` call from `LLMCompat.__init__()` entirely, since the inner `LLM` already registers its own callback, or add a comment explaining why a second registration for the wrapper is intentional.
How can I resolve this? If you propose a fix, please make it concise.
---
This is a comment left during a code review.
Path: prompt-service/src/unstract/prompt_service/tests/unit/test_retriever_llm.py
Line: 113-214
Comment:
**Missing test coverage for `require_llm()` error path**
`TestBaseRetrieverLlmProperty` covers the `llm` lazy property thoroughly but has no test for `BaseRetriever.require_llm()`. Specifically, the error path — when `_llm is None` — is untested. Given that `require_llm()` is the primary guard against silent OpenAI fallbacks in `fusion.py`, `keyword_table.py`, and `subquestion.py`, it warrants at least two test cases:
```python
def test_require_llm_raises_when_no_llm(self, base_retriever_cls):
"""require_llm() should raise ValueError when no LLM was provided."""
retriever = base_retriever_cls(
vector_db=MagicMock(), prompt="test", doc_id="doc-1", top_k=5
)
with pytest.raises(ValueError, match="requires an LLM"):
retriever.require_llm()
def test_require_llm_returns_retriever_llm_when_llm_provided(
self, base_retriever_cls, mock_sdk1_llm
):
"""require_llm() should return the RetrieverLLM when an LLM is set."""
with patch.object(LLMCompat, "from_llm", return_value=MagicMock()):
retriever = base_retriever_cls(
vector_db=MagicMock(), prompt="test", doc_id="doc-1",
top_k=5, llm=mock_sdk1_llm,
)
result = retriever.require_llm()
assert isinstance(result, RetrieverLLM)
```
How can I resolve this? If you propose a fix, please make it concise.
---
This is a comment left during a code review.
Path: unstract/sdk1/src/unstract/sdk1/llm.py
Line: 735-743
Comment:
**`formatted` parameter silently dropped in `complete()` and `acomplete()`**
Both `complete()` and `acomplete()` accept a `formatted: bool = False` parameter (matching the llama-index interface), but neither forwards it to `self._llm_instance.complete()` / `self._llm_instance.acomplete()`:
```python
def complete(self, prompt: str, formatted: bool = False, **kwargs: Any) -> CompletionResponse:
result = self._llm_instance.complete(prompt, **kwargs) # formatted is not forwarded
```
llama-index passes `formatted=True` when the prompt has already been pre-formatted for the model (e.g., for chat-optimized models). If `LLM.complete()` supports this parameter, silently dropping it could cause double-formatting in future. If `LLM.complete()` does not accept `formatted` at all, at minimum add a comment clarifying the intentional drop to avoid future confusion.
The same applies to `acomplete()` at the async equivalent below.
How can I resolve this? If you propose a fix, please make it concise.Last reviewed commit: "Add require_llm() gu..." |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
prompt-service/src/unstract/prompt_service/core/retrievers/retriever_llm.py
Show resolved
Hide resolved
…mpt-service - LLMCompat.from_llm(), emulated types, and _to_litellm_messages tests belong in sdk1/tests since those classes live in SDK1 - RetrieverLLM tests stay in prompt-service since that class lives there - Add type annotations to SDK1 tests to match existing test conventions Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
for more information, see https://pre-commit.ci
- Declare _compat as Pydantic PrivateAttr in RetrieverLLM for proper v2 lifecycle support (model_copy, serialization) - Forward **kwargs to litellm in LLMCompat chat/complete/achat/acomplete and pop cost_model after adapter.validate() to match LLM.complete() - Forward system_prompt in LLMCompat.from_llm() factory method - Remove dead predict()/apredict() from LLMCompat — RetrieverLLM inherits these from LlamaIndexBaseLLM directly - Add NOTE+TODO documenting missing usage audit for retriever LLM calls - Document --noconftest usage in tox.ini and unit conftest.py Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
for more information, see https://pre-commit.ci
…e level - Add _system_prompt to mock fixture and assertion to match from_llm() - Move litellm.drop_params = True to module-level init instead of repeating it per-call in chat/complete/achat/acomplete Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
prompt-service/src/unstract/prompt_service/core/retrievers/base_retriever.py
Outdated
Show resolved
Hide resolved
Simple, Automerging, and Recursive retrievers never use the LLM for llama-index components. Lazy construction via a property avoids the cost of adapter init and CallbackManager setup for those paths. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
prompt-service/src/unstract/prompt_service/core/retrievers/retriever_llm.py
Show resolved
Hide resolved
…ating from_llm() now bypasses __init__ and stores the caller's LLM directly, avoiding redundant adapter validation and PlatformHelper calls. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
LLMCompat is a bridge layer — it should not invoke litellm directly. All chat/complete/achat/acomplete methods now delegate to LLM.complete() and LLM.acomplete(), which handle litellm invocation, error wrapping, and usage auditing in one place. - Replace litellm.completion() calls with self._llm_instance.complete() - Replace litellm.acompletion() calls with self._llm_instance.acomplete() - Replace _to_litellm_messages() and _get_completion_kwargs() with _messages_to_prompt() which extracts the last user message - Update tests: replace _to_litellm_messages tests with _messages_to_prompt tests covering edge cases Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
SDK1 tests (test_llm_compat.py): - Add TestLLMCompatDelegation verifying chat/complete/achat/acomplete delegate to LLM.complete()/acomplete() instead of calling litellm - Test kwargs forwarding, return type wrapping, metadata, and get_model_name/get_metrics/test_connection delegation - Test all four streaming NotImplementedError paths Prompt-service tests (test_retriever_llm.py): - Add TestBaseRetrieverLlmProperty verifying lazy construction: returns None without LLM, returns RetrieverLLM with LLM, defers creation until first access, caches across accesses, calls from_llm() exactly once, preserves raw _llm reference - Stub VectorDB via sys.modules to avoid triggering the full adapter registration chain during unit tests Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
for more information, see https://pre-commit.ci
- Clarify that LLMCompat is a plain class (no llama-index inheritance), unlike EmbeddingCompat which inherits BaseEmbedding - Note that from_llm() is preferred when an LLM instance already exists - Document single-turn assumption in _messages_to_prompt() Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…o dep
_messages_to_prompt now concatenates all messages with role prefixes
instead of extracting only the last user message. This preserves
task-specific system instructions from llama-index components like
LLMQuestionGenerator ("You are an expert Q&A system...") that were
previously silently dropped.
Also add pytest-asyncio to sdk1 test dependencies — it was only
installed transitively, so async tests could silently pass without
executing if the transitive path changed.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
for more information, see https://pre-commit.ci
prompt-service/src/unstract/prompt_service/core/retrievers/subquestion.py
Show resolved
Hide resolved
Subquestion, Fusion, and KeywordTable retrievers pass self.llm to llama-index components. When no LLM is configured, self.llm returns None and llama-index silently falls back to its default OpenAI LLM, producing a confusing API key error. Add BaseRetriever.require_llm() that fails early with a clear message, and use it in all three retrievers. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
Test ResultsSummary
Runner Tests - Full Report
SDK1 Tests - Full Report
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
LLMCompatclass in SDK1 that emulates the llama-index LLM interface without requiring llama-index as a dependencyRetrieverLLMclass in prompt-service that inherits from llama-index'sLLMbase class and delegates calls toLLMCompatRetrieverLLMwhen passing LLM to llama-index componentsSubQuestionQueryEnginefailing due to missingllama-index-question-gen-openaipackageWhy
unstract.sdk1.llm.LLMuseslitellm.completion()directly and does not inherit fromllama_index.core.llms.llm.LLM'LLM' object has no attribute 'predict'errors at runtime for all retrieval strategies except Simpleresolve_llm()assertsisinstance(llm, LLM)which fails for plain classes — breaking Fusion, Router, and Subquestion retrieversHow
Two-layer architecture:
LLMCompat(unstract/sdk1/src/unstract/sdk1/llm.py): A plain Python class (no llama-index dependency) that emulates the llama-index LLM interface. Uses local dataclass emulations of llama-index types (ChatMessage,ChatResponse,CompletionResponse,LLMMetadata,MessageRole). Follows the existingEmbeddingCompatinitialization pattern — takes raw adapter params and creates an SDK1LLMinternally. Implementschat(),complete(),predict(),achat(),acomplete(),apredict().RetrieverLLM(prompt-service/.../core/retrievers/retriever_llm.py): Inherits fromllama_index.core.llms.llm.LLM(passingresolve_llm()isinstance checks) and delegates all LLM calls to an internalLLMCompatinstance. Converts emulated return types back to real llama-index types.Supporting changes:
llm_helper.py: Helper to convert SDK1LLM→RetrieverLLMinstancebase_retriever.py: Addedllama_index_llmproperty that lazily creates and caches theRetrieverLLMwrapperself.llmwithself.llama_index_llmwhere passed to llama-index components (keyword_table, subquestion, fusion, router)LLMQuestionGeneratortoSubQuestionQueryEngine.from_defaults()to avoid import ofllama-index-question-gen-openaiCan this PR break any existing features. If yes, please list possible items. If no, please explain why. (PS: Admins do not merge the PR without this section filled)
LLMCompatandRetrieverLLMare only used when retrievers pass LLM to llama-index components. The Simple, Automerging, and Recursive retrievers don't pass LLM to llama-index and remain unchanged. The bridge delegates all calls to the existing SDK1 LLM, so no behavior changes.Database Migrations
Env Config
Relevant Docs
Related Issues or PRs
Dependencies Versions
LLMCompatuses emulated types (no llama-index imports in SDK1).RetrieverLLMuses llama-index core classes already present in the prompt-service.Notes on Testing
chunk_size > 0and run extraction — should no longer error with'LLM' object has no attribute 'predict'Screenshots
Checklist
I have read and understood the Contribution Guidelines.