Skip to content

Commit 6ab89c0

Browse files
arimu1arimu1
committed
fix(console): allow Ollama custom model validation (#1551)
Detect Ollama endpoints during model validation and use a dedicated provider path: accept chat completion responses without a usage field and permit private/local addresses via the existing blockPrivate SSRF toggle instead of the strict OpenAI compatibility check.
1 parent ebd65a2 commit 6ab89c0

3 files changed

Lines changed: 111 additions & 3 deletions

File tree

‎console/backend/toolkit/src/main/java/com/iflytek/astron/console/toolkit/service/model/ModelService.java‎

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
import org.springframework.web.client.HttpServerErrorException;
6262
import org.springframework.web.client.RestTemplate;
6363

64+
import java.net.MalformedURLException;
6465
import java.net.URL;
6566
import java.security.interfaces.RSAPrivateKey;
6667
import java.util.*;
@@ -110,6 +111,8 @@ public class ModelService extends ServiceImpl<ModelMapper, Model> {
110111
private static final String PROVIDER_OPENAI = "openai";
111112
private static final String PROVIDER_ANTHROPIC = "anthropic";
112113
private static final String PROVIDER_GOOGLE = "google";
114+
private static final String PROVIDER_OLLAMA = "ollama";
115+
private static final int OLLAMA_DEFAULT_PORT = 11434;
113116
private static final String ANTHROPIC_VERSION = "2023-06-01";
114117

115118
private static final String CODE_XINGCHEN = "xingchen";
@@ -147,7 +150,7 @@ public String validateModel(ModelValidationRequest request) {
147150
}
148151

149152
// 2) Construct/validate URL + request body/headers
150-
final String provider = normalizeProvider(request.getProvider(), true);
153+
final String provider = resolveValidationProvider(request.getProvider(), request.getEndpoint());
151154
final String url = buildModelApiUrlNew(request.getEndpoint(), provider, request.getDomain());
152155
final Map<String, Object> requestBody =
153156
buildValidationPayload(request.getDomain(), provider);
@@ -258,6 +261,7 @@ private String buildModelApiUrlNew(String baseUrl, String provider, String model
258261
// compatibility
259262
ssrfProperties.setIpBlaklist(ipBlacklist);
260263
ssrfProperties.setIpWhitelist(ipWhitelist);
264+
ssrfProperties.setBlockPrivate(!PROVIDER_OLLAMA.equals(provider));
261265

262266
// 0) Remove userInfo and normalize
263267
String stripped = SsrfValidators.stripUserInfo(baseUrl);
@@ -390,9 +394,43 @@ private boolean isValidModelResponse(String responseBody, String provider) throw
390394
if (PROVIDER_GOOGLE.equals(provider)) {
391395
return root.has("candidates") && root.get("candidates").isArray();
392396
}
397+
if (PROVIDER_OLLAMA.equals(provider)) {
398+
return root.has("choices") && root.get("choices").isArray();
399+
}
393400
return root.has("choices") && root.get("choices").isArray() && root.has("usage");
394401
}
395402

403+
/**
404+
* Resolve the provider used for validation. Ollama endpoints are detected explicitly or from the
405+
* endpoint URL so they use a dedicated validation path instead of the strict OpenAI check.
406+
*/
407+
private String resolveValidationProvider(String provider, String endpoint) {
408+
String normalized = normalizeProvider(provider, true);
409+
if (PROVIDER_OLLAMA.equals(normalized) || looksLikeOllamaEndpoint(endpoint)) {
410+
return PROVIDER_OLLAMA;
411+
}
412+
return normalized;
413+
}
414+
415+
private static boolean looksLikeOllamaEndpoint(String endpoint) {
416+
if (StringUtils.isBlank(endpoint)) {
417+
return false;
418+
}
419+
String trimmed = endpoint.trim();
420+
String lower = trimmed.toLowerCase(Locale.ROOT);
421+
if (lower.contains("ollama")) {
422+
return true;
423+
}
424+
try {
425+
URL url = new URL(trimmed.contains("://") ? trimmed : "http://" + trimmed);
426+
int port = url.getPort();
427+
return port == OLLAMA_DEFAULT_PORT
428+
|| (port == -1 && lower.matches(".*:11434(?:/|$).*"));
429+
} catch (MalformedURLException e) {
430+
return lower.matches(".*:11434(?:/|$).*");
431+
}
432+
}
433+
396434
private void saveOrUpdateModel(ModelValidationRequest request) {
397435
final boolean isNew = (request.getId() == null);
398436
final Long spaceId = SpaceInfoUtil.getSpaceId();

‎console/backend/toolkit/src/main/java/com/iflytek/astron/console/toolkit/util/ssrf/SsrfParamGuard.java‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,13 @@ public void validateUrlParam(String url) {
6464

6565
// 2) IP blacklist (compatible with hostnames and IPs)
6666
List<String> ipBlacklist = props.getIpBlaklist();
67-
if (SsrfValidators.isHostDeniedByIpPolicy(
68-
u.getHost(), ipBlacklist, props.getIpWhitelist(), Dns.SYSTEM)) {
67+
List<String> ipWhitelist = props.getIpWhitelist();
68+
boolean denied = props.isBlockPrivate()
69+
? SsrfValidators.isHostDeniedByIpPolicy(
70+
u.getHost(), ipBlacklist, ipWhitelist, Dns.SYSTEM)
71+
: SsrfValidators.isHostBlockedByIpBlacklist(
72+
u.getHost(), ipBlacklist, ipWhitelist, Dns.SYSTEM);
73+
if (denied) {
6974
throw new BusinessException(ResponseEnum.MODEL_URL_CHECK_FAILED);
7075
}
7176

‎console/backend/toolkit/src/test/java/com/iflytek/astron/console/toolkit/service/model/ModelServiceTest.java‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -828,6 +828,71 @@ void buildModelApiUrlRejectsPrivateIpWhenBlacklistIsEmpty() {
828828
"local-model"));
829829
}
830830

831+
@Test
832+
void buildModelApiUrlAllowsPrivateIpForOllamaEndpoint() {
833+
ConfigInfo emptyConfig = new ConfigInfo();
834+
emptyConfig.setValue("");
835+
when(configInfoMapper.getListByCategory("NETWORK_SEGMENT_BLACK_LIST"))
836+
.thenReturn(List.of(emptyConfig));
837+
when(configInfoMapper.getListByCategory("IP_WHITE_LIST"))
838+
.thenReturn(List.of(emptyConfig));
839+
840+
String url = ReflectionTestUtils.invokeMethod(
841+
modelService,
842+
"buildModelApiUrlNew",
843+
"http://127.0.0.1:11434",
844+
"ollama",
845+
"llama3.2");
846+
847+
assertEquals("http://127.0.0.1:11434/v1/chat/completions", url);
848+
}
849+
850+
@Test
851+
void testValidateModel_ollamaResponseWithoutUsage_success() {
852+
ModelValidationRequest req = new ModelValidationRequest();
853+
req.setId(110L);
854+
req.setApiKeyMasked(false);
855+
req.setEndpoint("http://127.0.0.1:11434");
856+
req.setDomain("llama3.2");
857+
req.setModelName("local-ollama");
858+
req.setUid("u1");
859+
req.setTag(Collections.emptyList());
860+
req.setConfig(Collections.emptyList());
861+
862+
Model dbModel = new Model();
863+
dbModel.setId(110L);
864+
dbModel.setUid("u1");
865+
dbModel.setApiKey("ollama");
866+
dbModel.setIsDeleted(false);
867+
doReturn(dbModel).when(modelService).getById(110L);
868+
869+
when(configInfoMapper.getListByCategory("NETWORK_SEGMENT_BLACK_LIST"))
870+
.thenReturn(Collections.singletonList(new ConfigInfo()));
871+
when(configInfoMapper.getListByCategory("IP_WHITE_LIST"))
872+
.thenReturn(Collections.emptyList());
873+
874+
doReturn(dbModel)
875+
.doReturn(null)
876+
.when(modelService)
877+
.getOne(any(LambdaQueryWrapper.class));
878+
when(mapper.updateById(any(Model.class))).thenReturn(1);
879+
doNothing().when(modelCategoryService).saveAll(any(ModelCategoryReq.class));
880+
881+
String ollamaResp = """
882+
{"choices":[{"message":{"role":"assistant","content":"hi"}}],"model":"llama3.2"}
883+
""";
884+
ResponseEntity<String> httpOk = new ResponseEntity<>(ollamaResp, HttpStatus.OK);
885+
when(restTemplate.exchange(anyString(), eq(HttpMethod.POST), any(HttpEntity.class), eq(String.class)))
886+
.thenReturn(httpOk);
887+
888+
String result = modelService.validateModel(req);
889+
890+
assertEquals("Model validation passed", result);
891+
ArgumentCaptor<Model> modelCaptor = ArgumentCaptor.forClass(Model.class);
892+
verify(mapper).updateById(modelCaptor.capture());
893+
assertEquals("ollama", modelCaptor.getValue().getProvider());
894+
}
895+
831896
/**
832897
* Test {@link ModelService#(ModelDto, String)} to ensure public and owner models are merged, sorted
833898
* and paginated correctly.

0 commit comments

Comments
 (0)