Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implements list sources operation for chroma directly through apis. B… #29

Merged
merged 1 commit into from
Nov 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
<modelVersion>4.0.0</modelVersion>
<groupId>com.mule.mulechain</groupId>
<artifactId>mulechain-vectors</artifactId>
<version>0.1.88-SNAPSHOT</version>
<version>0.1.89-SNAPSHOT</version>
<packaging>mule-extension</packaging>
<name>MAC Vectors</name>

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ public class EmbeddingOperationValidator {
// Constants.VECTOR_STORE_ELASTICSEARCH, // Needs to be tested
// Constants.VECTOR_STORE_OPENSEARCH, // Needs to be tested
Constants.VECTOR_STORE_MILVUS,
// Constants.VECTOR_STORE_CHROMA, // Needs to be tested
Constants.VECTOR_STORE_CHROMA,
// Constants.VECTOR_STORE_PINECONE, // Do not support GTE with strings.
Constants.VECTOR_STORE_AI_SEARCH // Needs to be tested
)));
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package org.mule.extension.mulechain.vectors.internal.store.chroma;

import org.json.JSONArray;
import org.json.JSONObject;
import org.mule.extension.mulechain.vectors.internal.config.Configuration;
import org.mule.extension.mulechain.vectors.internal.constant.Constants;
Expand All @@ -10,6 +11,7 @@

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.util.HashMap;
Expand All @@ -36,7 +38,7 @@ public ChromaStore(String storeName, Configuration configuration, QueryParameter
super(storeName, configuration, queryParams, modelParams);

JSONObject config = JsonUtils.readConfigFile(configuration.getConfigFilePath());
JSONObject vectorStoreConfig = config.getJSONObject(Constants.VECTOR_STORE_ELASTICSEARCH);
JSONObject vectorStoreConfig = config.getJSONObject(Constants.VECTOR_STORE_CHROMA);
this.url = vectorStoreConfig.getString("CHROMA_URL");
}

Expand All @@ -62,6 +64,14 @@ public JSONObject listSources() {

while(offset < segmentCount) {

JSONArray metadataObjects = getMetadataObjects(collectionId, offset, queryParams.embeddingPageSize());
for(int i = 0; i< metadataObjects.length(); i++) {

JSONObject metadataObject = metadataObjects.getJSONObject(i);
JSONObject sourceObject = getSourceObject(metadataObject);
addOrUpdateSourceObjectIntoSourceObjectMap(sourceObjectMap, sourceObject);
}
offset = offset + metadataObjects.length();
}

} catch (Exception e) {
Expand All @@ -76,6 +86,66 @@ public JSONObject listSources() {
return jsonObject;
}

private JSONArray getMetadataObjects(String collectionId, long offset, long limit) {

JSONArray metadataObjects = new JSONArray();
try {

String urlString = url + "/api/v1/collections/" + collectionId + "/get";
URL url = new URL(urlString);

// Open connection and configure HTTP request
HttpURLConnection connection = (HttpURLConnection) url.openConnection();
connection.setRequestMethod("POST");
connection.setRequestProperty("Content-Type", "application/json");
connection.setDoOutput(true); // Enable output for the connection

JSONObject jsonRequest = new JSONObject();
jsonRequest.put("limit", limit);
jsonRequest.put("offset", offset);

JSONArray jsonInclude = new JSONArray();
jsonInclude.put("metadatas");

jsonRequest.put("include", jsonInclude);

// Write JSON body to the request output stream
try (OutputStream os = connection.getOutputStream()) {
byte[] input = jsonRequest.toString().getBytes("utf-8");
os.write(input, 0, input.length);
}

// Check the response code and handle accordingly
if (connection.getResponseCode() == 200) {

BufferedReader in = new BufferedReader(new InputStreamReader(connection.getInputStream()));
StringBuilder responseBuilder = new StringBuilder();
String line;

// Read response line by line
while ((line = in.readLine()) != null) {
responseBuilder.append(line);
}
in.close();

// Parse JSON response
JSONObject jsonResponse = new JSONObject(responseBuilder.toString());
metadataObjects = jsonResponse.getJSONArray("metadatas");

} else {

// Log any error responses from the server
LOGGER.error("Error: " + connection.getResponseCode() + " " + connection.getResponseMessage());
}

} catch (Exception e) {

// Handle any exceptions that occur during the process
LOGGER.error("Error getting collection segments", e);
}
return metadataObjects;
}

/**
* Retrieves the total number of segments in the specified collection.
*
Expand Down