fix: Use global LLM in knowledge retrieval

[erichare]

This pull request adds support for additional embedding providers in the retrieval component, improves API key management by introducing global fallback mechanisms, and expands the test coverage for these new providers and scenarios. The changes make it easier to use Google Generative AI, Ollama, and IBM WatsonX embeddings, and ensure proper error handling when required API keys are missing.

Embedding provider support:

• Added support for Google Generative AI, Ollama, and IBM WatsonX embedding providers in the _build_embeddings method of retrieval.py, including logic to retrieve provider-specific configuration variables and error handling for missing API keys.

API key management improvements:

• Implemented a fallback mechanism to retrieve API keys from globally stored provider variables when not supplied in metadata, improving usability for OpenAI, Google Generative AI, and IBM WatsonX providers.
• Updated error messages for missing API keys to clarify that keys can be provided in component settings or configured globally. [1] [2]

Test coverage enhancements:

• Added unit tests for Google Generative AI, Ollama, and IBM WatsonX embedding providers, including checks for correct instantiation and error handling when API keys are missing.
• Added a test for OpenAI provider to verify fallback to the global API key when none is provided in metadata.

Dependency updates:

• Imported new utility functions get_all_variables_for_provider and get_api_key_for_provider from unified_models to support the new embedding provider logic.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 0db3e477-456d-4f4f-adaa-955100e79448

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch fix-knowledge-global-llm

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[github-actions]

Frontend Unit Test Coverage Report

Coverage Summary

Lines	Statements	Branches	Functions
	23.32% (8189/35114)	16.02% (4417/27565)	16.16% (1193/7378)

Unit Test Results

Tests	Skipped	Failures	Errors	Time
2657	0 💤	0 ❌	0 🔥	45.015s ⏱️

[Cristhianzl]

ISSUE (RECOMMENDED): User data in logs

File: retrieval.py:356

logger.info(f"Performing similarity search with query: {self.search_query}")

This logs the user's search query verbatim. While not strictly PII, search queries can contain personal information (names, emails, addresses). Consider logging only that a search was performed, not the content:

logger.info("Performing similarity search")

Severity: RECOMMENDED (not a blocker, but should be addressed)

---

CRITICAL: DRY Principle

Item	Status	Notes
No duplicate types	PASS
No duplicate functions	WARN	See below
No duplicate constants	PASS

ISSUE (RECOMMENDED): Duplicated session/variable resolution pattern

Files: retrieval.py:151-187 and retrieval.py:189-212

Both _resolve_provider_variables and _resolve_api_key follow the same pattern:

1. Open session_scope()
2. Get variable_service
3. Convert self.user_id to UUID
4. Call variable_service.get_variable()
5. Fallback on exception

The UUID conversion line is duplicated verbatim:

user_id = self.user_id if not isinstance(self.user_id, str) else __import__("uuid").UUID(self.user_id)

Suggestion: Extract the UUID conversion to a small helper property or method. The two methods themselves are different enough (one resolves multiple vars, the other resolves one) to remain separate.

Severity: RECOMMENDED

---

CRITICAL: File Structure Limits

Metric	Value	Limit	Status
Lines of code	~404	~500	PASS
Functions (different responsibilities)	6 methods + 1 module func	5 standalone	WARN
Main classes per file	1 (KnowledgeBaseComponent)	1	PASS

ISSUE (IMPORTANT): Function responsibility count

The file has functions with these prefixes:

• _get_* (1: _get_kb_metadata, + 1 module: _get_knowledge_bases_root_path)
• _resolve_* (2: _resolve_provider_variables, _resolve_api_key)
• _build_* (1: _build_embeddings)
• retrieve_* (1: retrieve_data)
• update_* (1: update_build_config)

That's 5 different prefixes — at the limit. Adding any new responsibility prefix would exceed it. However, these are all methods of a single Component class following the Langflow component pattern (update_build_config + business methods), so this is acceptable as-is.

Severity: PASS (at limit, but justified by component pattern)

---

IMPORTANT: Error Handling

ISSUE (IMPORTANT): Overly broad exception catch

File: retrieval.py:183 and retrieval.py:211

`Exception` already covers `ValueError`, making `ValueError` redundant. More importantly, catching bare `Exception` hides unexpected errors (e.g., `ConnectionError`, `RuntimeError`). The
except (ValueError, Exception):  # noqa: BLE001
``` fallback to `os.environ` should only happen for expected failures.

**Suggestion:**
```python
except (ValueError, KeyError, AttributeError):
    # Expected failures when variable not found or service unavailable
    env_value = os.environ.get(var_key)
    ...

Severity: IMPORTANT

ISSUE (RECOMMENDED): Silent fallback may mask real issues

In _resolve_provider_variables (line 183-186), if the variable service throws an unexpected error (e.g., database connection failure), the code silently falls back to environment variables without logging. Consider adding a logger.debug or logger.warning for the fallback.

---

IMPORTANT: Code Quality

ISSUE (RECOMMENDED): Inline import of uuid

File: retrieval.py:169 and retrieval.py:204

user_id = self.user_id if not isinstance(self.user_id, str) else __import__("uuid").UUID(self.user_id)

Using __import__("uuid") inline is unusual and harder to read. uuid is a stdlib module and should be imported at the top of the file.

Suggestion:

# At top of file
import uuid

# In method
user_id = self.user_id if not isinstance(self.user_id, str) else uuid.UUID(self.user_id)

ISSUE (RECOMMENDED): Inline import of os

File: retrieval.py:157

import os

This is imported inside the method body. Since os is stdlib, it should be at the top of the file for consistency.

---

RECOMMENDED: Comments

ISSUE: Comments explaining "what" instead of "why"

Line	Comment	Issue
226	# Handle various providers	States the obvious — the if/elif chain is self-evident
329	# Resolve API key: user override > metadata (decrypted) > global variable	Good — explains priority logic (KEEP)
336	# Resolve provider-specific variables (e.g. base_url for Ollama, project_id for WatsonX)	Borderline — the method name already says this, but examples are useful (OK)
344-345	# Clear Chroma's singleton client cache...	Good — explains non-obvious workaround (KEEP)
355	# Use the search query to perform a similarity search	Redundant — the code is self-evident

---

TESTING

Coverage Analysis

Aspect	Status	Notes
Happy path tests	PASS	All providers tested with valid inputs
Adversarial tests	PARTIAL	Missing some edge cases (see below)
Error states	PASS	No-key errors tested for OpenAI, Cohere, Google, WatsonX
Test independence	PASS	Each test is self-contained
Mocking	PASS	External deps properly mocked

ISSUE (IMPORTANT): contextlib.suppress(Exception) in test

File: test_retrieval.py:475

with contextlib.suppress(Exception):
    await component.retrieve_data()

This suppresses ALL exceptions, meaning the test passes even if the method throws an unexpected error. This defeats the purpose of testing.

Suggestion: Either mock enough to make the method succeed, or use pytest.raises(SpecificError) to assert the expected failure.

ISSUE (IMPORTANT): Missing adversarial tests

The following scenarios are not tested:

1. _resolve_provider_variables — no direct unit test at all

   • What if provider_vars is empty for Ollama? (no OLLAMA_BASE_URL)
   • What if variable_service returns None?
   • What if self.user_id is a string vs UUID?

2. _resolve_api_key — no direct unit test

   • What if provider is not in provider_variable_map?
   • What if user_id is None?

3. _build_embeddings edge cases:

   • Ollama without base_url (should still work, just without kwarg)
   • WatsonX with api_key from provider_vars fallback (line 282: api_key or all_vars.get("WATSONX_APIKEY"))
   • WatsonX with partial vars (only apikey, no project_id or url)
   • Empty metadata dict (no provider, no model)

4. retrieve_data integration:

   • user_id is None
   • User not found in DB
   • search_query is empty vs populated (different code paths at line 354)

ISSUE (RECOMMENDED): Test naming convention

Current: test_build_embeddings_openai_no_key
Rule says: should_[expected]_when_[condition]
Expected: test_should_raise_error_when_openai_api_key_missing

This is a project-wide convention mismatch (existing tests don't follow it either), so it's not a blocker.

---

Summary

Category	Score	Issues
Security & PII	PASS (1 RECOMMENDED)	Search query logged verbatim
DRY	PASS (1 RECOMMENDED)	UUID conversion duplicated
File Structure	PASS	At the 5-prefix limit but justified
Error Handling	1 IMPORTANT, 1 RECOMMENDED	Bare Exception catch, silent fallback
Code Quality	2 RECOMMENDED	Inline imports (__import__, os)
Comments	2 RECOMMENDED	"What" comments instead of "why"
Testing	2 IMPORTANT, 1 RECOMMENDED	suppress(Exception), missing adversarial tests, naming

Verdict

No CRITICAL blockers. The PR is functional and well-structured. The main items to address before merge:

1. IMPORTANT: Replace bare except (ValueError, Exception) with specific exception types
2. IMPORTANT: Remove contextlib.suppress(Exception) from test — use proper assertion
3. IMPORTANT: Add unit tests for _resolve_api_key and _resolve_provider_variables
4. RECOMMENDED: Move uuid and os imports to top of file
5. RECOMMENDED: Remove/redact search query from log at line 356

[Cristhianzl]

lgtm