test_connection on Flow Model always fails for embedding models (chat ping via litellm.completion)
- Dominant language
- Python
- Stars
- 69
- Forks
- 28
- PR merge metrics
- No merged PRs in 30d
Description
## Description
The **Test Connection** button on the Flow Model doctype always fails for embedding models, because `FlowModel.test_connection` unconditionally sends a chat ping via `litellm.completion()`:
```python
kwargs = {
"model": self.model_id,
"api_key": api_key,
"messages": [{"role": "user", "content": "ping"}],
"max_tokens": 1,
"timeout": 15,
}
...
litellm.completion(**kwargs)
```
For an embedding model this routes to the provider's chat/completion endpoint (e.g. Gemini's `generateContent`), which embedding models don't support — so the test fails even when the model, credentials, and routing are all correct. The actual knowledge pipeline (`flow/knowledge/embedder.py`) uses `litellm.embedding()` and works fine with the same record.
This is misleading during setup: the error looks like a model/credential problem and sent us down the wrong debugging path (we suspected a stale LiteLLM model registry) before reading the source.
## Steps to reproduce
1. Create a Flow Provider for `gemini` with a valid API key.
2. Create a Flow Model with Model ID `gemini/gemini-embedding-2` (or `gemini/gemini-embedding-001` — any embedding model reproduces it).
3. Click **Test Connection**.
## Actual result
```
litellm.NotFoundError: GeminiException - {
"error": {
"code": 404,
"message": "models/gemini-embedding-2 is not found for API version v1beta, or is not supported for generateContent. Call ModelService.ListModels to see the list of available models and their supported methods.",
"status": "NOT_FOUND"
}
}
```
Meanwhile, selecting the same Flow Model as the embedding model in Flow Knowledge Settings works: `_sync_embedding_dimension` → `probe_dimension` → `litellm.embedding()` succeeds and correctly detects 3072 dimensions.
## Expected result
Test Connection should succeed for a correctly configured embedding model, or at minimum indicate that the model is embedding-only rather than surfacing a provider 404 from the wrong endpoint.
## Suggested fix
Detect the model mode and pick the matching call, falling back to a completion ping when the mode is unknown:
```python
mode = None
try:
mode = litellm.get_model_info(self.model_id).get("mode")
except Exception:
pass # unmapped model — assume chat
if mode == "embedding":
litellm.embedding(model=self.model_id, input=["ping"], api_key=api_key, timeout=15, **extra)
else:
litellm.completion(**kwargs)
```
Since `get_model_info` can't classify models absent from LiteLLM's registry, an alternative (or complement) would be an explicit "Model Type: Chat / Embedding" field on Flow Model — which would also let the UI hide chat-only fields like Params for embedding models.
## Environment
- Flow: `develop` branch (pre-alpha)
- Frappe: 16.30.0
- Hosted on Frappe Cloud (Private Bench)
- Provider: Gemini API (Google AI Studio key)
Contributor guide
No contributing guide indexed for this repository
Research direction
Start at the FlowModel.test_connection entry point and compare its litellm.completion() call with flow/knowledge/embedder.py, where _sync_embedding_dimension and probe_dimension use litellm.embedding(). Determine how the model mode is identified, then make Test Connection use the appropriate path while retaining a chat fallback for unknown models. Done means a valid embedding model can be tested successfully without a chat-endpoint error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- api, backend
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100