Skip to content

Commit c32388c

Browse files
sudipnextsudipnext
authored
Merge pull request presenton#576 from presenton/feature/add-litellm-opened-template-gen-vision-supported
feat: add LiteLLM support and enhance model configuration
2 parents b1870b9 + 7af7348 commit c32388c

27 files changed

Lines changed: 508 additions & 92 deletions

‎servers/fastapi/constants/llm.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,4 +8,5 @@
88
DEFAULT_OPENROUTER_MODEL = "openai/gpt-4o"
99
DEFAULT_CEREBRAS_MODEL = "llama-3.3-70b"
1010
DEFAULT_ANTHROPIC_MODEL = "claude-sonnet-4-20250514"
11+
DEFAULT_LITELLM_MODEL = "gpt-4.1"
1112
DEFAULT_CODEX_MODEL = "gpt-5.2"

‎servers/fastapi/enums/llm_provider.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,5 +10,6 @@ class LLMProvider(Enum):
1010
OPENROUTER = "openrouter"
1111
CEREBRAS = "cerebras"
1212
ANTHROPIC = "anthropic"
13+
LITELLM = "litellm"
1314
CUSTOM = "custom"
1415
CODEX = "codex"

‎servers/fastapi/models/user_config.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,11 @@ class UserConfig(BaseModel):
3838
CEREBRAS_MODEL: Optional[str] = None
3939
CEREBRAS_BASE_URL: Optional[str] = None
4040

41+
# LiteLLM (OpenAI-compatible gateway / proxy)
42+
LITELLM_BASE_URL: Optional[str] = None
43+
LITELLM_API_KEY: Optional[str] = None
44+
LITELLM_MODEL: Optional[str] = None
45+
4146
# Anthropic
4247
ANTHROPIC_API_KEY: Optional[str] = None
4348
ANTHROPIC_MODEL: Optional[str] = None

‎servers/fastapi/pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ dependencies = [
2626
"pdfplumber>=0.11.7",
2727
"psycopg[binary]>=3.2.0",
2828
"sqlmodel>=0.0.24",
29-
"llmai==0.2.3",
29+
"llmai==0.2.4",
3030
"jsonschema>=4.26.0",
3131
]
3232

‎servers/fastapi/templates/providers.py‎

Lines changed: 32 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -9,15 +9,12 @@
99
from utils.llm_config import get_llm_config
1010
from utils.llm_provider import get_llm_provider, get_model
1111
from utils.llm_utils import extract_text
12+
from utils.template_vision_errors import (
13+
VISION_LAYOUT_USER_MESSAGE,
14+
is_likely_vision_capability_error,
15+
)
1216

1317
MAX_ATTEMPTS_PER_PROVIDER = 4
14-
SUPPORTED_TEMPLATE_PROVIDERS = (
15-
LLMProvider.OPENAI,
16-
LLMProvider.CODEX,
17-
LLMProvider.GOOGLE,
18-
LLMProvider.ANTHROPIC,
19-
LLMProvider.AZURE,
20-
)
2118

2219

2320
def _exception_message(exc: Exception) -> str:
@@ -32,17 +29,9 @@ def _exception_message(exc: Exception) -> str:
3229
return " ".join(message.split())[:500]
3330

3431

35-
def _unsupported_template_provider_message() -> str:
36-
return (
37-
"Template generation only supports OpenAI, Codex, Google, Anthropic, or Azure OpenAI."
38-
)
39-
40-
41-
def _supported_template_provider_or_raise() -> tuple[LLMProvider, str]:
42-
provider = get_llm_provider()
43-
if provider not in SUPPORTED_TEMPLATE_PROVIDERS:
44-
raise HTTPException(status_code=400, detail=_unsupported_template_provider_message())
45-
return provider, get_model()
32+
def _resolve_template_provider_and_model() -> tuple[LLMProvider, str]:
33+
"""Uses the configured text LLM; slide layout generation requires vision (image parts)."""
34+
return get_llm_provider(), get_model()
4635

4736

4837
def _provider_label(provider: LLMProvider) -> str:
@@ -52,10 +41,22 @@ def _provider_label(provider: LLMProvider) -> str:
5241
return "Codex"
5342
if provider == LLMProvider.GOOGLE:
5443
return "Google"
44+
if provider == LLMProvider.VERTEX:
45+
return "Vertex AI"
5546
if provider == LLMProvider.ANTHROPIC:
5647
return "Anthropic"
5748
if provider == LLMProvider.AZURE:
5849
return "Azure OpenAI"
50+
if provider == LLMProvider.OLLAMA:
51+
return "Ollama"
52+
if provider == LLMProvider.OPENROUTER:
53+
return "OpenRouter"
54+
if provider == LLMProvider.CEREBRAS:
55+
return "Cerebras"
56+
if provider == LLMProvider.CUSTOM:
57+
return "Custom"
58+
if provider == LLMProvider.LITELLM:
59+
return "LiteLLM"
5960
return "Template provider"
6061

6162

@@ -108,6 +109,7 @@ async def _run_template_llm_with_retries(
108109
*,
109110
provider_label: str,
110111
call: Callable[[], Awaitable[str]],
112+
requires_vision: bool = False,
111113
) -> str:
112114
last_exception: Optional[Exception] = None
113115

@@ -118,10 +120,18 @@ async def _run_template_llm_with_retries(
118120
return response_text
119121
raise ValueError("No output from template generation provider")
120122
except HTTPException as exc:
123+
if requires_vision and is_likely_vision_capability_error(exc):
124+
raise HTTPException(
125+
status_code=400, detail=VISION_LAYOUT_USER_MESSAGE
126+
) from exc
121127
if 400 <= exc.status_code < 500:
122128
raise exc
123129
last_exception = exc
124130
except Exception as exc:
131+
if requires_vision and is_likely_vision_capability_error(exc):
132+
raise HTTPException(
133+
status_code=400, detail=VISION_LAYOUT_USER_MESSAGE
134+
) from exc
125135
last_exception = exc
126136

127137
if isinstance(last_exception, HTTPException):
@@ -141,7 +151,7 @@ def _template_provider_label_and_call(
141151
image_bytes: Optional[bytes] = None,
142152
media_type: str = "image/png",
143153
) -> tuple[str, Callable[[], Awaitable[str]]]:
144-
provider, model = _supported_template_provider_or_raise()
154+
provider, model = _resolve_template_provider_and_model()
145155
label = _provider_label(provider)
146156
return (
147157
label,
@@ -168,7 +178,9 @@ async def generate_slide_layout_code(
168178
image_bytes=image_bytes,
169179
media_type=media_type,
170180
)
171-
return await _run_template_llm_with_retries(provider_label=label, call=call)
181+
return await _run_template_llm_with_retries(
182+
provider_label=label, call=call, requires_vision=True
183+
)
172184

173185

174186
async def edit_slide_layout_code(

‎servers/fastapi/tests/unit/test_template_providers.py‎

Lines changed: 74 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -36,32 +36,28 @@ def generate(self, **kwargs):
3636
LLMProvider.OPENAI,
3737
LLMProvider.CODEX,
3838
LLMProvider.GOOGLE,
39+
LLMProvider.VERTEX,
3940
LLMProvider.ANTHROPIC,
4041
LLMProvider.AZURE,
42+
LLMProvider.OLLAMA,
43+
LLMProvider.OPENROUTER,
44+
LLMProvider.CEREBRAS,
45+
LLMProvider.LITELLM,
46+
LLMProvider.CUSTOM,
4147
],
4248
)
43-
def test_supported_template_provider_resolution(monkeypatch, provider: LLMProvider):
49+
def test_resolve_template_provider_and_model(monkeypatch, provider: LLMProvider):
4450
monkeypatch.setattr(providers_module, "get_llm_provider", lambda: provider)
4551
monkeypatch.setattr(providers_module, "get_model", lambda: "resolved-model")
4652

4753
resolved_provider, resolved_model = (
48-
providers_module._supported_template_provider_or_raise()
54+
providers_module._resolve_template_provider_and_model()
4955
)
5056

5157
assert resolved_provider == provider
5258
assert resolved_model == "resolved-model"
5359

5460

55-
def test_supported_template_provider_rejects_unsupported(monkeypatch):
56-
monkeypatch.setattr(providers_module, "get_llm_provider", lambda: LLMProvider.VERTEX)
57-
58-
with pytest.raises(HTTPException) as exc_info:
59-
providers_module._supported_template_provider_or_raise()
60-
61-
assert exc_info.value.status_code == 400
62-
assert "Template generation only supports" in str(exc_info.value.detail)
63-
64-
6561
def test_provider_label_fallback_for_unknown_value():
6662
class _UnknownProvider:
6763
pass
@@ -96,8 +92,14 @@ def test_exception_message_handles_http_exception_string_detail():
9692
LLMProvider.OPENAI,
9793
LLMProvider.CODEX,
9894
LLMProvider.GOOGLE,
95+
LLMProvider.VERTEX,
9996
LLMProvider.ANTHROPIC,
10097
LLMProvider.AZURE,
98+
LLMProvider.OLLAMA,
99+
LLMProvider.OPENROUTER,
100+
LLMProvider.CEREBRAS,
101+
LLMProvider.LITELLM,
102+
LLMProvider.CUSTOM,
101103
],
102104
)
103105
def test_generate_slide_layout_code_uses_llmai_for_all_supported_providers(
@@ -143,8 +145,14 @@ def test_generate_slide_layout_code_uses_llmai_for_all_supported_providers(
143145
LLMProvider.OPENAI,
144146
LLMProvider.CODEX,
145147
LLMProvider.GOOGLE,
148+
LLMProvider.VERTEX,
146149
LLMProvider.ANTHROPIC,
147150
LLMProvider.AZURE,
151+
LLMProvider.OLLAMA,
152+
LLMProvider.OPENROUTER,
153+
LLMProvider.CEREBRAS,
154+
LLMProvider.LITELLM,
155+
LLMProvider.CUSTOM,
148156
],
149157
)
150158
def test_edit_slide_layout_code_uses_llmai_text_only(
@@ -297,3 +305,57 @@ async def never_called():
297305
monkeypatch.setattr(providers_module, "MAX_ATTEMPTS_PER_PROVIDER", original_attempts)
298306
assert exc_info.value.status_code == 500
299307
assert "Failed to generate template output" in str(exc_info.value.detail)
308+
309+
def test_run_template_llm_with_retries_maps_vision_errors_when_requires_vision():
310+
async def vision_fail():
311+
raise RuntimeError("This model does not support image inputs")
312+
313+
with pytest.raises(HTTPException) as exc_info:
314+
asyncio.run(
315+
providers_module._run_template_llm_with_retries(
316+
provider_label="OpenAI",
317+
call=vision_fail,
318+
requires_vision=True,
319+
)
320+
)
321+
322+
assert exc_info.value.status_code == 400
323+
assert "TEMPLATE_VISION_MODEL_REQUIRED" in str(exc_info.value.detail)
324+
325+
326+
def test_run_template_llm_with_retries_does_not_map_vision_without_flag():
327+
async def vision_fail():
328+
raise RuntimeError("This model does not support image inputs")
329+
330+
with pytest.raises(HTTPException) as exc_info:
331+
asyncio.run(
332+
providers_module._run_template_llm_with_retries(
333+
provider_label="OpenAI",
334+
call=vision_fail,
335+
requires_vision=False,
336+
)
337+
)
338+
339+
assert exc_info.value.status_code == 502
340+
341+
342+
def test_generate_slide_layout_code_fail_fast_on_vision_error(monkeypatch):
343+
dummy_client = _DummyClient(
344+
outputs=[RuntimeError("image_url is not supported for this model")]
345+
)
346+
monkeypatch.setattr(providers_module, "get_llm_provider", lambda: LLMProvider.OPENAI)
347+
monkeypatch.setattr(providers_module, "get_model", lambda: "text-only")
348+
monkeypatch.setattr(providers_module, "get_llm_config", lambda: {"provider": "openai"})
349+
monkeypatch.setattr(providers_module, "get_client", lambda config: dummy_client)
350+
351+
with pytest.raises(HTTPException) as exc_info:
352+
asyncio.run(
353+
providers_module.generate_slide_layout_code(
354+
system_prompt="sys",
355+
user_text="user",
356+
image_bytes=b"x",
357+
)
358+
)
359+
360+
assert exc_info.value.status_code == 400
361+
assert len(dummy_client.calls) == 1

‎servers/fastapi/utils/available_models.py‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,27 @@
33
from google import genai
44

55

6+
def normalize_openai_compatible_base_url(url: str) -> str:
7+
"""Ensure base URL targets the OpenAI-compatible /v1 root (LiteLLM, vLLM, etc.)."""
8+
u = (url or "").strip().rstrip("/")
9+
if not u:
10+
return u
11+
if u.endswith("/v1"):
12+
return u
13+
base = u.split("?", 1)[0]
14+
if "/v1" in base:
15+
return u
16+
return f"{u}/v1"
17+
18+
619
async def list_available_openai_compatible_models(url: str, api_key: str) -> list[str]:
7-
client = AsyncOpenAI(api_key=api_key, base_url=url)
20+
url = normalize_openai_compatible_base_url(url)
21+
# Local LiteLLM / OpenAI-compatible proxies often omit auth; SDK rejects a blank key.
22+
effective_key = (api_key or "").strip() or "EMPTY"
23+
client = AsyncOpenAI(api_key=effective_key, base_url=url)
824
models = (await client.models.list()).data
925
if models:
10-
return list(map(lambda x: x.id, models))
26+
return [m.id for m in models if m.id]
1127
return []
1228

1329

@@ -31,4 +47,4 @@ async def list_available_anthropic_models(api_key: str) -> list[str]:
3147

3248
async def list_available_google_models(api_key: str) -> list[str]:
3349
client = genai.Client(api_key=api_key)
34-
return list(map(lambda x: x.name, client.models.list(config={"page_size": 50})))
50+
return [x.name for x in client.models.list(config={"page_size": 50}) if x.name]

‎servers/fastapi/utils/get_env.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,18 @@ def get_cerebras_base_url_env():
150150
return os.getenv("CEREBRAS_BASE_URL")
151151

152152

153+
def get_litellm_base_url_env():
154+
return os.getenv("LITELLM_BASE_URL")
155+
156+
157+
def get_litellm_api_key_env():
158+
return os.getenv("LITELLM_API_KEY")
159+
160+
161+
def get_litellm_model_env():
162+
return os.getenv("LITELLM_MODEL")
163+
164+
153165
def get_custom_llm_api_key_env():
154166
return os.getenv("CUSTOM_LLM_API_KEY")
155167

‎servers/fastapi/utils/llm_config.py‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
ChatGPTClientConfig,
1010
ClientConfig,
1111
GoogleClientConfig,
12+
LiteLLMClientConfig,
1213
OpenAIApiType,
1314
OpenAIClientConfig,
1415
OpenRouterClientConfig,
@@ -33,6 +34,8 @@
3334
get_custom_llm_url_env,
3435
get_disable_thinking_env,
3536
get_google_api_key_env,
37+
get_litellm_api_key_env,
38+
get_litellm_base_url_env,
3639
get_ollama_url_env,
3740
get_openai_api_key_env,
3841
get_openrouter_api_key_env,
@@ -43,6 +46,7 @@
4346
get_vertex_project_env,
4447
get_web_grounding_env,
4548
)
49+
from utils.available_models import normalize_openai_compatible_base_url
4650
from utils.llm_provider import get_llm_provider
4751
from utils.parsers import parse_bool_or_none
4852
from utils.set_env import (
@@ -218,6 +222,20 @@ def get_llm_config() -> ClientConfig:
218222
api_key=api_key,
219223
base_url=base_url or None,
220224
)
225+
case LLMProvider.LITELLM:
226+
base_url = normalize_openai_compatible_base_url(
227+
get_litellm_base_url_env() or ""
228+
)
229+
if not base_url:
230+
raise HTTPException(
231+
status_code=400,
232+
detail="LiteLLM base URL is not set (LITELLM_BASE_URL).",
233+
)
234+
lk = (get_litellm_api_key_env() or "").strip()
235+
return LiteLLMClientConfig(
236+
base_url=base_url,
237+
api_key=lk if lk else None,
238+
)
221239
case LLMProvider.OLLAMA:
222240
return OpenAIClientConfig(
223241
base_url=(get_ollama_url_env() or "http://localhost:11434") + "/v1",
@@ -244,7 +262,7 @@ def get_llm_config() -> ClientConfig:
244262
status_code=400,
245263
detail=(
246264
"LLM Provider must be either openai, google, vertex, azure, "
247-
"openrouter, cerebras, anthropic, ollama, "
265+
"openrouter, cerebras, anthropic, litellm, ollama, "
248266
"custom, or codex"
249267
),
250268
)

0 commit comments

Comments
 (0)