Skip to content

Commit 0afb4c4

Browse files
willschippwillschipp
authored
Merge pull request #20 from embabel/18-add-google-vertex-as-a-source
Added google models API as a source
2 parents 47d0474 + 4e7407d commit 0afb4c4

File tree

11 files changed

+761
-2
lines changed

11 files changed

+761
-2
lines changed

‎README.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,10 @@ The Database is loaded from the Project [Llm LeaderBoard](https://github.com/Jon
4141

4242
`GET /api/v1/models` returns an array of JSON objects representing a `ModelMetadata` object
4343
`GET /api/v1/models/search/findByName?name={model name}` returns a list of matching `ModelMetadata` including providers and costs for each matching model
44+
`GET /api/v1/models/search/findByNameContains?name={model name}` returns a list of matching `ModelMetadata` where the name contians the 'name' passed
45+
`GET /api/v1/models/search/findByNameAndProvider?name={model name}&provider={provider name}` returns a list of matching `ModelMetadata`
46+
`GET /api/v1/models/search/findByTask?task={task name}` returns a list of matching `ModelMetadata`
47+
`GET /api/v1/models/count` returns the number of models in the repository
4448
`GET /api/v1/models/lastUpdate` returns a timestamp for when the repository was last refreshed
4549

4650
Repository maintenance is via an Agent approach. The server provides an MCP Server compliant toolset as well as a direct, manual mechanism to trigger the Agent. The Agent will validate if the repository needs refreshing.

‎TESTING.md

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,13 @@
22

33
Projects comprise of Unit Tests `*Test.java` and Integation Tests `*ITest.java`. Integration tests are excluded by default.
44

5-
To run integration tests `mvn clean test -Dtest=*ITest`
5+
To run integration tests `mvn clean test -Dtest=*ITest`
6+
7+
8+
## Local Running
9+
10+
1. install ollama `curl -fsSL https://ollama.com/install.sh | sh`
11+
2. start ollama `ollama serve > /dev/null 2>&1 &`
12+
3. pull model `ollama pull llama3.1:8b`
13+
4. build `mvn clean install`
14+
5. run `java -jar embabel-database-server/target/embabel-database-server-${VERISON}.jar`

‎embabel-database-agent/src/main/java/com/embabel/database/agent/service/GoogleMetadataDiscoveryService.java

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/*
2+
* Copyright 2024-2025 Embabel Software, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.embabel.database.agent.service;
17+
18+
import java.util.ArrayList;
19+
import java.util.List;
20+
21+
import org.springframework.beans.factory.annotation.Value;
22+
import org.springframework.web.reactive.function.client.WebClient;
23+
24+
import com.embabel.common.ai.model.ModelMetadata;
25+
import com.embabel.database.agent.util.ModelMetadataParser;
26+
27+
public class GoogleMetadataDiscoveryService implements ModelMetadataDiscoveryService {
28+
29+
@Value("${google.models.base-url:https://generativelanguage.googleapis.com}")
30+
private String googleModelsBaseUrl;
31+
32+
@Value("${google.models.url:/v1/models?key=}")
33+
private String googleModelsUrl;
34+
35+
@Value("${google.models.api-key}")
36+
private String apiKey;
37+
38+
private ModelMetadataParser googleModelMetadataParser;
39+
40+
private WebClient webClient;
41+
42+
public GoogleMetadataDiscoveryService() {}
43+
44+
public GoogleMetadataDiscoveryService(ModelMetadataParser googleModelMetadataParser) {
45+
this.googleModelMetadataParser = googleModelMetadataParser;
46+
}
47+
48+
@Override
49+
public List<ModelMetadata> retrieveModelMetadata() {
50+
// init
51+
List<ModelMetadata> listModelMetadata = new ArrayList<>();
52+
//retrieve the objects from the endpoint
53+
//build the final urls
54+
String url = googleModelsUrl + apiKey;
55+
//test
56+
if (webClient == null) {
57+
webClient = WebClient.builder().baseUrl(googleModelsBaseUrl).build();
58+
} //end if
59+
String json = webClient.get()
60+
.uri(url)
61+
.retrieve()
62+
.bodyToMono(String.class)
63+
.block();
64+
//convert to n accessible object
65+
listModelMetadata = googleModelMetadataParser.parse(json);
66+
//return
67+
return listModelMetadata;
68+
}
69+
70+
}

‎embabel-database-agent/src/main/java/com/embabel/database/agent/util/GoogleParser.java

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
/*
2+
* Copyright 2024-2025 Embabel Software, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.embabel.database.agent.util;
17+
18+
import java.nio.file.Path;
19+
import java.time.LocalDate;
20+
import java.util.ArrayList;
21+
import java.util.List;
22+
import java.util.Map;
23+
24+
import org.apache.commons.logging.Log;
25+
import org.apache.commons.logging.LogFactory;
26+
import org.springframework.beans.factory.annotation.Autowired;
27+
28+
import com.embabel.common.ai.model.ModelMetadata;
29+
import com.embabel.common.ai.model.PerTokenPricingModel;
30+
import com.embabel.common.ai.model.PricingModel;
31+
import com.embabel.database.core.repository.LlmModelMetadata;
32+
import com.fasterxml.jackson.core.type.TypeReference;
33+
import com.fasterxml.jackson.databind.ObjectMapper;
34+
35+
public class GoogleParser implements ModelMetadataParser {
36+
37+
private static Log logger = LogFactory.getLog(GoogleParser.class);
38+
39+
private static final String MODEL_LIST_KEY = "models";
40+
41+
@Autowired
42+
TaskParser googleTaskParser;
43+
44+
@Autowired
45+
ObjectMapper objectMapper;
46+
47+
@SuppressWarnings("unchecked")
48+
@Override
49+
public List<ModelMetadata> parse(String json) {
50+
//init
51+
List<ModelMetadata> listModelMetadata = new ArrayList<>();
52+
//get the raw list
53+
try {
54+
Map<String,Object> modelsHighLevel = objectMapper.readValue(json,new TypeReference<Map<String,Object>>(){});
55+
//get the list
56+
List<Map<String,Object>> models = (List<Map<String, Object>>) modelsHighLevel.get(MODEL_LIST_KEY);
57+
//loop
58+
for (Map<String,Object> model : models) {
59+
//get components
60+
String providerName = "google";
61+
String modelName = model.get("name").toString();
62+
//date parse
63+
LocalDate knowledgeCutoffDate = LocalDate.of(1970, 1, 1); //TODO need to fix updated at
64+
//price
65+
PricingModel pricingModel = new PerTokenPricingModel(0.0, 0.0);
66+
//tokens
67+
Long paramCount = model.get("inputTokenLimit") != null ? Long.parseLong(model.get("inputTokenLimit").toString()) : 0l;
68+
//get the task
69+
String task = googleTaskParser.getTask(model);
70+
//build the metadatamodel
71+
ModelMetadata modelMetadata = LlmModelMetadata.Companion.create(modelName,providerName,knowledgeCutoffDate,pricingModel,0l,task,this.getClass().getSimpleName());
72+
//add
73+
listModelMetadata.add(modelMetadata);
74+
} //end for
75+
} catch (Exception e) {
76+
logger.error("error parsing",e);
77+
}
78+
return listModelMetadata;
79+
}
80+
81+
@Override
82+
public List<ModelMetadata> parse(Path path) {
83+
logger.error("Not implemented for " + this.getClass().getName());
84+
return null;
85+
}
86+
87+
@Override
88+
public List<ModelMetadata> parse(List<?> list) {
89+
logger.error("Not implemented for " + this.getClass().getName());
90+
return null;
91+
}
92+
93+
}

‎embabel-database-agent/src/main/java/com/embabel/database/agent/util/GoogleTaskParser.java

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
/*
2+
* Copyright 2024-2025 Embabel Software, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.embabel.database.agent.util;
17+
18+
import java.util.HashMap;
19+
import java.util.List;
20+
import java.util.Map;
21+
22+
import org.springframework.beans.factory.annotation.Autowired;
23+
24+
import com.fasterxml.jackson.databind.ObjectMapper;
25+
26+
public class GoogleTaskParser implements TaskParser {
27+
28+
private List<Map<String,Object>> tasks;
29+
30+
@Autowired
31+
ObjectMapper objectMapper;
32+
33+
@Override
34+
public String getTask(Map<String, Object> attributes) {
35+
//init
36+
String modelTask = null;
37+
//load up the options
38+
if (tasks == null || tasks.isEmpty()) {
39+
tasks = this.getTasks(objectMapper, RESOURCE_LOCATION);
40+
} //end if
41+
//google response structure is
42+
/*
43+
"supportedGenerationMethods": [
44+
"generateContent",
45+
"countTokens",
46+
"createCachedContent"
47+
]
48+
*/
49+
Map<String,Object> matches = new HashMap<>();
50+
//convert that into the boolean map
51+
// List<String> methods = (List<String>) attributes.get("supportedGenerationMethods");
52+
// for (String method : methods) {
53+
if (attributes.get("name").toString().contains("embedding")) {
54+
//embedding model
55+
return "text-to-embedding"; //TODO make a bit more elegant and flexible
56+
} else if (attributes.get("name").toString().contains("image")) {
57+
//return image
58+
return "text-to-image-text"; //TODO make a bit more flexible
59+
} else {
60+
return "audio-video-image-text-to-text"; //TODO make a bit more flexible
61+
} //end if
62+
}
63+
64+
}

‎embabel-database-agent/src/main/resources/data/task_types.json

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,17 @@
3232
"outputVideo": false,
3333
"outputImage": true
3434
},
35+
{
36+
"Classification": "text-to-image-text",
37+
"inputText": true,
38+
"inputAudio": false,
39+
"inputVideo": false,
40+
"inputImage": false,
41+
"outputText": true,
42+
"outputAudio": false,
43+
"outputVideo": false,
44+
"outputImage": true
45+
},
3546
{
3647
"Classification": "text-to-video",
3748
"inputText": true,
@@ -108,5 +119,27 @@
108119
"outputAudio": false,
109120
"outputVideo": false,
110121
"outputImage": false
122+
},
123+
{
124+
"Classification": "audio-video-image-text-to-text",
125+
"inputText": true,
126+
"inputAudio": true,
127+
"inputVideo": true,
128+
"inputImage": true,
129+
"outputText": true,
130+
"outputAudio": false,
131+
"outputVideo": false,
132+
"outputImage": false
133+
},
134+
{
135+
"Classification": "text-to-embedding",
136+
"inputText": true,
137+
"inputAudio": false,
138+
"inputVideo": false,
139+
"inputImage": false,
140+
"outputText": false,
141+
"outputAudio": false,
142+
"outputVideo": false,
143+
"outputImage": false
111144
}
112145
]

‎embabel-database-agent/src/test/java/com/embabel/database/agent/util/GoogleParserTest.java

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
/*
2+
* Copyright 2024-2025 Embabel Software, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.embabel.database.agent.util;
17+
18+
import static org.junit.jupiter.api.Assertions.assertNotNull;
19+
import static org.junit.jupiter.api.Assertions.assertTrue;
20+
21+
import java.io.InputStream;
22+
import java.nio.charset.StandardCharsets;
23+
import java.util.List;
24+
import java.util.Scanner;
25+
26+
import org.junit.jupiter.api.Test;
27+
import org.springframework.test.util.ReflectionTestUtils;
28+
29+
import com.embabel.common.ai.model.ModelMetadata;
30+
import com.fasterxml.jackson.databind.ObjectMapper;
31+
32+
public class GoogleParserTest {
33+
34+
@Test
35+
void testJsonStringParse() throws Exception {
36+
//load up the json string
37+
String json;
38+
String json_path = "json/google.response.json";
39+
//load from test resources
40+
try (InputStream is = getClass().getClassLoader().getResourceAsStream(json_path)) {
41+
if (is == null) {
42+
throw new IllegalArgumentException("test resource not found");
43+
}//end if
44+
try (Scanner scanner = new Scanner(is, StandardCharsets.UTF_8.name())) {
45+
json = scanner.useDelimiter("\\A").next();
46+
}
47+
}
48+
//setup the injection parts
49+
ObjectMapper objectMapper = new ObjectMapper();
50+
GoogleParser googleParser = new GoogleParser();
51+
GoogleTaskParser googleTaskParser = new GoogleTaskParser();
52+
ReflectionTestUtils.setField(googleParser, "objectMapper", objectMapper);
53+
ReflectionTestUtils.setField(googleParser, "googleTaskParser", googleTaskParser);
54+
//execute
55+
List<ModelMetadata> models = googleParser.parse(json);
56+
//check
57+
assertNotNull(models);
58+
assertTrue(models.size() > 0);
59+
60+
}
61+
}

0 commit comments

Comments
 (0)