Skip to content

Commit

Permalink
feat(vector-stores): Infeer queryRootUrl in VertexAIMatchingEngine (#133
Browse files Browse the repository at this point in the history 
)
  • Loading branch information
@davidmigloz
davidmigloz authored Aug 19, 2023
1 parent 2c508dc commit c535336
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 31 deletions.
68 changes: 48 additions & 20 deletions packages/langchain_google/lib/src/vector_stores/matching_engine.dart
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,10 @@ import 'package:vertex_ai/vertex_ai.dart';
/// Vertex AI Matching Engine documentation:
/// https://cloud.google.com/vertex-ai/docs/matching-engine/overview
///
/// Currently it only supports Batch Updates, it doesn't support Streaming
/// Updates. Batch Updates take around 1h to be applied to the index. See:
/// https://cloud.google.com/vertex-ai/docs/matching-engine/update-rebuild-index#update_index_content_with_batch_updates
///
/// ### Set up your Google Cloud Platform project
///
/// 1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager).
Expand Down Expand Up @@ -65,7 +69,6 @@ import 'package:vertex_ai/vertex_ai.dart';
/// authHttpClient: authClient,
/// project: 'your-project-id',
/// location: 'europe-west1',
/// queryRootUrl: 'https://your-query-root-url.vdb.vertexai.goog/',
/// indexId: 'your-index-id',
/// gcsBucketName: 'your-gcs-bucket-name',
/// embeddings: embeddings,
Expand Down Expand Up @@ -101,8 +104,7 @@ class VertexAIMatchingEngine extends VectorStore {
/// `https://$location-aiplatform.googleapis.com/`.
/// - [queryRootUrl] The root URL of the Vertex AI Matching Engine index
/// endpoint. For example: `https://your-query-root-url.vdb.vertexai.goog/`.
/// You can find this URL when you create the index endpoint. See:
/// https://cloud.google.com/vertex-ai/docs/matching-engine/deploy-index-public#get_the_index_domain_name
/// Only needed if you have multiple endpoints for the same index.
/// - [indexId] The ID of the index to use. You can find this ID when you
/// create the index or in the Vertex AI console.
/// - [indexEndpointId] The ID of the IndexEndpoint to use. Only needed if
Expand All @@ -120,10 +122,10 @@ class VertexAIMatchingEngine extends VectorStore {
/// added to the index without deleting the previous ones.
VertexAIMatchingEngine({
required final AuthClient authHttpClient,
required final String project,
required final String location,
required this.project,
required this.location,
final String? rootUrl,
required final String queryRootUrl,
final String? queryRootUrl,
required this.indexId,
final String? indexEndpointId,
final String? deployedIndexId,
Expand All @@ -132,31 +134,37 @@ class VertexAIMatchingEngine extends VectorStore {
this.gcsIndexesFolder = 'indexes',
this.completeOverwriteWhenAdding = false,
required super.embeddings,
}) : _queryClient = VertexAIMatchingEngineClient(
authHttpClient: authHttpClient,
project: project,
location: location,
rootUrl: queryRootUrl,
),
}) : _authHttpClient = authHttpClient,
_managementClient = VertexAIMatchingEngineClient(
authHttpClient: authHttpClient,
project: project,
location: location,
rootUrl: rootUrl ?? 'https://$location-aiplatform.googleapis.com/',
),
_storageClient = Storage(authHttpClient, project),
_queryRootUrl = queryRootUrl,
_indexEndpointId = indexEndpointId,
_deployedIndexId = deployedIndexId;

/// An authenticated HTTP client.
final AuthClient _authHttpClient;

/// A client for querying the Vertex AI Matching Engine index.
final VertexAIMatchingEngineClient _queryClient;
VertexAIMatchingEngineClient? _queryClient;

/// A client for managing the Vertex AI Matching Engine index.
final VertexAIMatchingEngineClient _managementClient;

/// A client for interacting with Google Cloud Storage.
final Storage _storageClient;

/// The ID of the Google Cloud project to use.
final String project;

/// The Google Cloud location to use. Vertex AI and Cloud Storage should have
/// the same location.
final String location;

/// The id of the index to use.
final String indexId;

Expand All @@ -176,6 +184,9 @@ class VertexAIMatchingEngine extends VectorStore {
/// deleting the previous ones.
final bool completeOverwriteWhenAdding;

/// The root URL of the Vertex AI Matching Engine index endpoint.
String? _queryRootUrl;

/// The ID of the index endpoint to use.
String? _indexEndpointId;

Expand Down Expand Up @@ -260,8 +271,10 @@ class VertexAIMatchingEngine extends VectorStore {
required final List<double> embedding,
final int k = 4,
}) async {
final (indexEndpointId, deployedIndexId) = await _getIndexIds();
final queryRes = await _queryClient.indexEndpoints.findNeighbors(
final (queryRootUrl, indexEndpointId, deployedIndexId) =
await _getIndexIds();
final queryClient = _getQueryClient(queryRootUrl);
final queryRes = await queryClient.indexEndpoints.findNeighbors(
indexEndpointId: indexEndpointId,
deployedIndexId: deployedIndexId,
queries: [
Expand Down Expand Up @@ -292,10 +305,12 @@ class VertexAIMatchingEngine extends VectorStore {
}

/// Returns the index endpoint and deployed index ids.
Future<(String indexEndpointId, String deployedIndexId)>
Future<(String queryRootUrl, String indexEndpointId, String deployedIndexId)>
_getIndexIds() async {
if (_indexEndpointId != null && _deployedIndexId != null) {
return (_indexEndpointId!, _deployedIndexId!);
if (_queryRootUrl != null &&
_indexEndpointId != null &&
_deployedIndexId != null) {
return (_queryRootUrl!, _indexEndpointId!, _deployedIndexId!);
}

final index = await _managementClient.indexes.get(id: indexId);
Expand All @@ -310,11 +325,24 @@ class VertexAIMatchingEngine extends VectorStore {
'that you want to use.',
);
}

final deployedIndex = index.deployedIndexes.first;
final indexEndpoint = await _managementClient.indexEndpoints
.get(id: deployedIndex.indexEndpointId);

_queryRootUrl = 'http://${indexEndpoint.publicEndpointDomainName}/';
_indexEndpointId = deployedIndex.indexEndpointId;
_deployedIndexId = deployedIndex.deployedIndexId;
return (_indexEndpointId!, _deployedIndexId!);
return (_queryRootUrl!, _indexEndpointId!, _deployedIndexId!);
}

/// Returns a client for querying the Vertex AI Matching Engine index.
VertexAIMatchingEngineClient _getQueryClient(final String queryRootUrl) {
return _queryClient ??= VertexAIMatchingEngineClient(
authHttpClient: _authHttpClient,
project: project,
location: location,
rootUrl: queryRootUrl,
);
}

/// Returns the document with the given id from the Storage bucket.
Expand Down
18 changes: 7 additions & 11 deletions packages/langchain_google/test/vector_stores/matching_engine.dart
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,16 +19,13 @@ void main() async {
authHttpClient: authHttpClient,
project: Platform.environment['VERTEX_AI_PROJECT_ID']!,
location: 'europe-west1',
queryRootUrl:
'https://1451028425.europe-west1-706285145183.vdb.vertexai.goog/',
indexId: '5086059315115065344',
gcsBucketName: 'vertex-ai',
gcsIndexesFolder: 'index',
indexId: '6394355006866194432',
gcsBucketName: 'public_knowledge_base_index',
embeddings: embeddings,
);

group('VertexAIMatchingEngine tests', skip: true, () {
test('Test VertexAIMatchingEngine add new vectors', () async {
group('VertexAIMatchingEngine tests', () {
test('Test VertexAIMatchingEngine add new vectors', skip: true, () async {
final res = await vectorStore.addDocuments(
documents: const [
Document(
Expand All @@ -49,14 +46,13 @@ void main() async {

test('Test VertexAIMatchingEngine query', () async {
final res = await vectorStore.similaritySearch(
query: 'Can I add tags to my index?',
query: 'Can I pay by credit card?',
k: 1,
);
expect(res.length, 1);
expect(
res.first.pageContent,
'You can also add optional tags to your index to help '
'with diversifying results or filtering pre index query',
res.first.id,
'faq_621656c96b5ff317d867d019',
);
});
});
Expand Down

0 comments on commit c535336

Please sign in to comment.