frappe / frappe/flow_client

test_connection on Flow Model always fails for embedding models (chat ping via litellm.completion)

Open
#96 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.