From fb77d4ff23fed673fee05afa4acb09fd522db79a Mon Sep 17 00:00:00 2001 From: Joel Date: Mon, 2 Oct 2023 21:57:54 +0200 Subject: [PATCH] Implement semantic search API for unified resources (#31881) * Implement semantic search API for unified resources * use require.eventually to complete test faster * reset mod and sum * Update lib/auth/assist/assistv1/test/service_test.go Co-authored-by: Jakub Nyckowski * use correct proto numbering * reset e * make grpc --------- Co-authored-by: Jakub Nyckowski --- api/gen/proto/go/assist/v1/assist.pb.go | 356 +++++++++++++----- api/gen/proto/go/assist/v1/assist_grpc.pb.go | 39 ++ api/proto/teleport/assist/v1/assist.proto | 20 + lib/ai/embeddingprocessor_test.go | 80 ++-- lib/ai/mock_embedder.go | 52 +++ lib/auth/assist/assistv1/service.go | 108 +++++- lib/auth/assist/assistv1/test/service_test.go | 147 +++++++- lib/auth/auth_with_roles.go | 107 +----- lib/auth/grpcserver.go | 2 +- lib/services/unified_resource.go | 122 +++++- 10 files changed, 756 insertions(+), 277 deletions(-) create mode 100644 lib/ai/mock_embedder.go diff --git a/api/gen/proto/go/assist/v1/assist.pb.go b/api/gen/proto/go/assist/v1/assist.pb.go index 127c2e31b37ff..8b0f63f2a31d7 100644 --- a/api/gen/proto/go/assist/v1/assist.pb.go +++ b/api/gen/proto/go/assist/v1/assist.pb.go @@ -21,6 +21,7 @@ package assist import ( + proto "github.com/gravitational/teleport/api/client/proto" protoreflect "google.golang.org/protobuf/reflect/protoreflect" protoimpl "google.golang.org/protobuf/runtime/protoimpl" emptypb "google.golang.org/protobuf/types/known/emptypb" @@ -958,6 +959,122 @@ func (x *GetAssistantEmbeddingsResponse) GetEmbeddings() []*EmbeddedDocument { return nil } +// SearchUnifiedResourcesRequest is a request to search for one or more resource kinds using similiarity search. +type SearchUnifiedResourcesRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + // query is the query used for similarity search. + Query string `protobuf:"bytes,1,opt,name=query,proto3" json:"query,omitempty"` + // limit is the number of embeddings to return (also known as k). + Limit int32 `protobuf:"varint,2,opt,name=limit,proto3" json:"limit,omitempty"` + // kinds is the kind of embeddings to return (ex, node). Returns all supported kinds if empty. + Kinds []string `protobuf:"bytes,3,rep,name=kinds,proto3" json:"kinds,omitempty"` +} + +func (x *SearchUnifiedResourcesRequest) Reset() { + *x = SearchUnifiedResourcesRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_teleport_assist_v1_assist_proto_msgTypes[16] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *SearchUnifiedResourcesRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SearchUnifiedResourcesRequest) ProtoMessage() {} + +func (x *SearchUnifiedResourcesRequest) ProtoReflect() protoreflect.Message { + mi := &file_teleport_assist_v1_assist_proto_msgTypes[16] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SearchUnifiedResourcesRequest.ProtoReflect.Descriptor instead. +func (*SearchUnifiedResourcesRequest) Descriptor() ([]byte, []int) { + return file_teleport_assist_v1_assist_proto_rawDescGZIP(), []int{16} +} + +func (x *SearchUnifiedResourcesRequest) GetQuery() string { + if x != nil { + return x.Query + } + return "" +} + +func (x *SearchUnifiedResourcesRequest) GetLimit() int32 { + if x != nil { + return x.Limit + } + return 0 +} + +func (x *SearchUnifiedResourcesRequest) GetKinds() []string { + if x != nil { + return x.Kinds + } + return nil +} + +// SearchUnifiedResourcesResponse is a response from the assistant service with a similarity-ordered list of resources. +type SearchUnifiedResourcesResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + // resources is the list of resources. + Resources []*proto.PaginatedResource `protobuf:"bytes,1,rep,name=resources,proto3" json:"resources,omitempty"` +} + +func (x *SearchUnifiedResourcesResponse) Reset() { + *x = SearchUnifiedResourcesResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_teleport_assist_v1_assist_proto_msgTypes[17] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *SearchUnifiedResourcesResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SearchUnifiedResourcesResponse) ProtoMessage() {} + +func (x *SearchUnifiedResourcesResponse) ProtoReflect() protoreflect.Message { + mi := &file_teleport_assist_v1_assist_proto_msgTypes[17] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SearchUnifiedResourcesResponse.ProtoReflect.Descriptor instead. +func (*SearchUnifiedResourcesResponse) Descriptor() ([]byte, []int) { + return file_teleport_assist_v1_assist_proto_rawDescGZIP(), []int{17} +} + +func (x *SearchUnifiedResourcesResponse) GetResources() []*proto.PaginatedResource { + if x != nil { + return x.Resources + } + return nil +} + var File_teleport_assist_v1_assist_proto protoreflect.FileDescriptor var file_teleport_assist_v1_assist_proto_rawDesc = []byte{ @@ -968,6 +1085,9 @@ var file_teleport_assist_v1_assist_proto_rawDesc = []byte{ 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2f, 0x65, 0x6d, 0x70, 0x74, 0x79, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x1a, 0x1f, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2f, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x2e, 0x70, 0x72, + 0x6f, 0x74, 0x6f, 0x1a, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2f, 0x6c, 0x65, + 0x67, 0x61, 0x63, 0x79, 0x2f, 0x63, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x2f, 0x70, 0x72, 0x6f, 0x74, + 0x6f, 0x2f, 0x61, 0x75, 0x74, 0x68, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0x62, 0x0a, 0x1b, 0x47, 0x65, 0x74, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x27, 0x0a, 0x0f, 0x63, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, @@ -1070,75 +1190,95 @@ var file_teleport_assist_v1_assist_proto_rawDesc = []byte{ 0x24, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x45, 0x6d, 0x62, 0x65, 0x64, 0x64, 0x65, 0x64, 0x44, 0x6f, 0x63, 0x75, 0x6d, 0x65, 0x6e, 0x74, 0x52, 0x0a, 0x65, 0x6d, 0x62, 0x65, 0x64, 0x64, 0x69, 0x6e, 0x67, - 0x73, 0x32, 0xdd, 0x06, 0x0a, 0x0d, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x53, 0x65, 0x72, 0x76, - 0x69, 0x63, 0x65, 0x12, 0x8e, 0x01, 0x0a, 0x1b, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x41, 0x73, + 0x73, 0x22, 0x61, 0x0a, 0x1d, 0x53, 0x65, 0x61, 0x72, 0x63, 0x68, 0x55, 0x6e, 0x69, 0x66, 0x69, + 0x65, 0x64, 0x52, 0x65, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, + 0x73, 0x74, 0x12, 0x14, 0x0a, 0x05, 0x71, 0x75, 0x65, 0x72, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x05, 0x71, 0x75, 0x65, 0x72, 0x79, 0x12, 0x14, 0x0a, 0x05, 0x6c, 0x69, 0x6d, 0x69, + 0x74, 0x18, 0x02, 0x20, 0x01, 0x28, 0x05, 0x52, 0x05, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x12, 0x14, + 0x0a, 0x05, 0x6b, 0x69, 0x6e, 0x64, 0x73, 0x18, 0x03, 0x20, 0x03, 0x28, 0x09, 0x52, 0x05, 0x6b, + 0x69, 0x6e, 0x64, 0x73, 0x22, 0x58, 0x0a, 0x1e, 0x53, 0x65, 0x61, 0x72, 0x63, 0x68, 0x55, 0x6e, + 0x69, 0x66, 0x69, 0x65, 0x64, 0x52, 0x65, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x73, 0x52, 0x65, + 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x36, 0x0a, 0x09, 0x72, 0x65, 0x73, 0x6f, 0x75, 0x72, + 0x63, 0x65, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x18, 0x2e, 0x70, 0x72, 0x6f, 0x74, + 0x6f, 0x2e, 0x50, 0x61, 0x67, 0x69, 0x6e, 0x61, 0x74, 0x65, 0x64, 0x52, 0x65, 0x73, 0x6f, 0x75, + 0x72, 0x63, 0x65, 0x52, 0x09, 0x72, 0x65, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x73, 0x32, 0xde, + 0x07, 0x0a, 0x0d, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, + 0x12, 0x8e, 0x01, 0x0a, 0x1b, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x41, 0x73, 0x73, 0x69, 0x73, + 0x74, 0x61, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, + 0x12, 0x36, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, + 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x41, 0x73, 0x73, 0x69, + 0x73, 0x74, 0x61, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, + 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x37, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, + 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x72, + 0x65, 0x61, 0x74, 0x65, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x43, 0x6f, 0x6e, + 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, + 0x65, 0x12, 0x88, 0x01, 0x0a, 0x19, 0x47, 0x65, 0x74, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, + 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x12, + 0x34, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, + 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, + 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x52, 0x65, + 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x35, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, + 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, - 0x69, 0x6f, 0x6e, 0x12, 0x36, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, - 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x41, - 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, - 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x37, 0x2e, 0x74, 0x65, + 0x69, 0x6f, 0x6e, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x6d, 0x0a, 0x1b, + 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x43, + 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x36, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, - 0x2e, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, - 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, - 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x88, 0x01, 0x0a, 0x19, 0x47, 0x65, 0x74, 0x41, 0x73, 0x73, 0x69, - 0x73, 0x74, 0x61, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, - 0x6e, 0x73, 0x12, 0x34, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, - 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x73, 0x73, 0x69, 0x73, - 0x74, 0x61, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, - 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x35, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, - 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x47, 0x65, - 0x74, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, - 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, - 0x6d, 0x0a, 0x1b, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, - 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x36, + 0x2e, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, + 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, + 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, + 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x12, 0x79, 0x0a, 0x14, 0x47, + 0x65, 0x74, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x4d, 0x65, 0x73, 0x73, 0x61, + 0x67, 0x65, 0x73, 0x12, 0x2f, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, + 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x73, 0x73, 0x69, + 0x73, 0x74, 0x61, 0x6e, 0x74, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x73, 0x52, 0x65, 0x71, + 0x75, 0x65, 0x73, 0x74, 0x1a, 0x30, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, + 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x73, 0x73, + 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x73, 0x52, 0x65, + 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x63, 0x0a, 0x16, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, + 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, + 0x12, 0x31, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, + 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x41, 0x73, 0x73, 0x69, + 0x73, 0x74, 0x61, 0x6e, 0x74, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, + 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, + 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x12, 0x75, 0x0a, 0x1f, 0x55, + 0x70, 0x64, 0x61, 0x74, 0x65, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x43, 0x6f, + 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x6e, 0x66, 0x6f, 0x12, 0x3a, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, - 0x2e, 0x76, 0x31, 0x2e, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, - 0x61, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, - 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x12, 0x79, - 0x0a, 0x14, 0x47, 0x65, 0x74, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x4d, 0x65, - 0x73, 0x73, 0x61, 0x67, 0x65, 0x73, 0x12, 0x2f, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, - 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x47, 0x65, 0x74, 0x41, - 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x73, - 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x30, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, + 0x2e, 0x76, 0x31, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, + 0x61, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x49, + 0x6e, 0x66, 0x6f, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, 0x67, 0x6f, 0x6f, + 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x45, 0x6d, 0x70, + 0x74, 0x79, 0x12, 0x6a, 0x0a, 0x0f, 0x49, 0x73, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x45, 0x6e, + 0x61, 0x62, 0x6c, 0x65, 0x64, 0x12, 0x2a, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, + 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x49, 0x73, 0x41, 0x73, 0x73, + 0x69, 0x73, 0x74, 0x45, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x1a, 0x2b, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, + 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x49, 0x73, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x45, + 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x7f, + 0x0a, 0x16, 0x53, 0x65, 0x61, 0x72, 0x63, 0x68, 0x55, 0x6e, 0x69, 0x66, 0x69, 0x65, 0x64, 0x52, + 0x65, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x73, 0x12, 0x31, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, + 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x53, 0x65, + 0x61, 0x72, 0x63, 0x68, 0x55, 0x6e, 0x69, 0x66, 0x69, 0x65, 0x64, 0x52, 0x65, 0x73, 0x6f, 0x75, + 0x72, 0x63, 0x65, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x32, 0x2e, 0x74, 0x65, + 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, + 0x2e, 0x53, 0x65, 0x61, 0x72, 0x63, 0x68, 0x55, 0x6e, 0x69, 0x66, 0x69, 0x65, 0x64, 0x52, 0x65, + 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x32, + 0x99, 0x01, 0x0a, 0x16, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x45, 0x6d, 0x62, 0x65, 0x64, 0x64, + 0x69, 0x6e, 0x67, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x7f, 0x0a, 0x16, 0x47, 0x65, + 0x74, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x45, 0x6d, 0x62, 0x65, 0x64, 0x64, + 0x69, 0x6e, 0x67, 0x73, 0x12, 0x31, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, + 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x73, 0x73, + 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x45, 0x6d, 0x62, 0x65, 0x64, 0x64, 0x69, 0x6e, 0x67, 0x73, + 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x32, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x47, 0x65, 0x74, - 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, - 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x63, 0x0a, 0x16, 0x43, 0x72, 0x65, - 0x61, 0x74, 0x65, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x4d, 0x65, 0x73, 0x73, - 0x61, 0x67, 0x65, 0x12, 0x31, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, - 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x41, - 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x52, - 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x12, 0x75, - 0x0a, 0x1f, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, - 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x6e, 0x66, - 0x6f, 0x12, 0x3a, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, - 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x41, 0x73, 0x73, - 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x61, 0x74, 0x69, - 0x6f, 0x6e, 0x49, 0x6e, 0x66, 0x6f, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, - 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, - 0x45, 0x6d, 0x70, 0x74, 0x79, 0x12, 0x6a, 0x0a, 0x0f, 0x49, 0x73, 0x41, 0x73, 0x73, 0x69, 0x73, - 0x74, 0x45, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x12, 0x2a, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, - 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x49, 0x73, - 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x45, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x52, 0x65, 0x71, - 0x75, 0x65, 0x73, 0x74, 0x1a, 0x2b, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x2e, - 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x49, 0x73, 0x41, 0x73, 0x73, 0x69, - 0x73, 0x74, 0x45, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, - 0x65, 0x32, 0x99, 0x01, 0x0a, 0x16, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x45, 0x6d, 0x62, 0x65, - 0x64, 0x64, 0x69, 0x6e, 0x67, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x7f, 0x0a, 0x16, - 0x47, 0x65, 0x74, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x45, 0x6d, 0x62, 0x65, - 0x64, 0x64, 0x69, 0x6e, 0x67, 0x73, 0x12, 0x31, 0x2e, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, - 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x47, 0x65, 0x74, 0x41, - 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x45, 0x6d, 0x62, 0x65, 0x64, 0x64, 0x69, 0x6e, - 0x67, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x32, 0x2e, 0x74, 0x65, 0x6c, 0x65, - 0x70, 0x6f, 0x72, 0x74, 0x2e, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2e, 0x76, 0x31, 0x2e, 0x47, - 0x65, 0x74, 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x45, 0x6d, 0x62, 0x65, 0x64, - 0x64, 0x69, 0x6e, 0x67, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x42, 0x45, 0x5a, - 0x43, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x67, 0x72, 0x61, 0x76, - 0x69, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x61, 0x6c, 0x2f, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, - 0x72, 0x74, 0x2f, 0x61, 0x70, 0x69, 0x2f, 0x67, 0x65, 0x6e, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, - 0x2f, 0x67, 0x6f, 0x2f, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2f, 0x76, 0x31, 0x3b, 0x61, 0x73, - 0x73, 0x69, 0x73, 0x74, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, + 0x41, 0x73, 0x73, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x45, 0x6d, 0x62, 0x65, 0x64, 0x64, 0x69, + 0x6e, 0x67, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x42, 0x45, 0x5a, 0x43, 0x67, + 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x67, 0x72, 0x61, 0x76, 0x69, 0x74, + 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x61, 0x6c, 0x2f, 0x74, 0x65, 0x6c, 0x65, 0x70, 0x6f, 0x72, 0x74, + 0x2f, 0x61, 0x70, 0x69, 0x2f, 0x67, 0x65, 0x6e, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2f, 0x67, + 0x6f, 0x2f, 0x61, 0x73, 0x73, 0x69, 0x73, 0x74, 0x2f, 0x76, 0x31, 0x3b, 0x61, 0x73, 0x73, 0x69, + 0x73, 0x74, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, } var ( @@ -1153,7 +1293,7 @@ func file_teleport_assist_v1_assist_proto_rawDescGZIP() []byte { return file_teleport_assist_v1_assist_proto_rawDescData } -var file_teleport_assist_v1_assist_proto_msgTypes = make([]protoimpl.MessageInfo, 16) +var file_teleport_assist_v1_assist_proto_msgTypes = make([]protoimpl.MessageInfo, 18) var file_teleport_assist_v1_assist_proto_goTypes = []interface{}{ (*GetAssistantMessagesRequest)(nil), // 0: teleport.assist.v1.GetAssistantMessagesRequest (*AssistantMessage)(nil), // 1: teleport.assist.v1.AssistantMessage @@ -1171,38 +1311,44 @@ var file_teleport_assist_v1_assist_proto_goTypes = []interface{}{ (*GetAssistantEmbeddingsRequest)(nil), // 13: teleport.assist.v1.GetAssistantEmbeddingsRequest (*EmbeddedDocument)(nil), // 14: teleport.assist.v1.EmbeddedDocument (*GetAssistantEmbeddingsResponse)(nil), // 15: teleport.assist.v1.GetAssistantEmbeddingsResponse - (*timestamppb.Timestamp)(nil), // 16: google.protobuf.Timestamp - (*emptypb.Empty)(nil), // 17: google.protobuf.Empty + (*SearchUnifiedResourcesRequest)(nil), // 16: teleport.assist.v1.SearchUnifiedResourcesRequest + (*SearchUnifiedResourcesResponse)(nil), // 17: teleport.assist.v1.SearchUnifiedResourcesResponse + (*timestamppb.Timestamp)(nil), // 18: google.protobuf.Timestamp + (*proto.PaginatedResource)(nil), // 19: proto.PaginatedResource + (*emptypb.Empty)(nil), // 20: google.protobuf.Empty } var file_teleport_assist_v1_assist_proto_depIdxs = []int32{ - 16, // 0: teleport.assist.v1.AssistantMessage.created_time:type_name -> google.protobuf.Timestamp + 18, // 0: teleport.assist.v1.AssistantMessage.created_time:type_name -> google.protobuf.Timestamp 1, // 1: teleport.assist.v1.CreateAssistantMessageRequest.message:type_name -> teleport.assist.v1.AssistantMessage 1, // 2: teleport.assist.v1.GetAssistantMessagesResponse.messages:type_name -> teleport.assist.v1.AssistantMessage - 16, // 3: teleport.assist.v1.ConversationInfo.created_time:type_name -> google.protobuf.Timestamp + 18, // 3: teleport.assist.v1.ConversationInfo.created_time:type_name -> google.protobuf.Timestamp 5, // 4: teleport.assist.v1.GetAssistantConversationsResponse.conversations:type_name -> teleport.assist.v1.ConversationInfo - 16, // 5: teleport.assist.v1.CreateAssistantConversationRequest.created_time:type_name -> google.protobuf.Timestamp + 18, // 5: teleport.assist.v1.CreateAssistantConversationRequest.created_time:type_name -> google.protobuf.Timestamp 14, // 6: teleport.assist.v1.GetAssistantEmbeddingsResponse.embeddings:type_name -> teleport.assist.v1.EmbeddedDocument - 7, // 7: teleport.assist.v1.AssistService.CreateAssistantConversation:input_type -> teleport.assist.v1.CreateAssistantConversationRequest - 4, // 8: teleport.assist.v1.AssistService.GetAssistantConversations:input_type -> teleport.assist.v1.GetAssistantConversationsRequest - 12, // 9: teleport.assist.v1.AssistService.DeleteAssistantConversation:input_type -> teleport.assist.v1.DeleteAssistantConversationRequest - 0, // 10: teleport.assist.v1.AssistService.GetAssistantMessages:input_type -> teleport.assist.v1.GetAssistantMessagesRequest - 2, // 11: teleport.assist.v1.AssistService.CreateAssistantMessage:input_type -> teleport.assist.v1.CreateAssistantMessageRequest - 9, // 12: teleport.assist.v1.AssistService.UpdateAssistantConversationInfo:input_type -> teleport.assist.v1.UpdateAssistantConversationInfoRequest - 10, // 13: teleport.assist.v1.AssistService.IsAssistEnabled:input_type -> teleport.assist.v1.IsAssistEnabledRequest - 13, // 14: teleport.assist.v1.AssistEmbeddingService.GetAssistantEmbeddings:input_type -> teleport.assist.v1.GetAssistantEmbeddingsRequest - 8, // 15: teleport.assist.v1.AssistService.CreateAssistantConversation:output_type -> teleport.assist.v1.CreateAssistantConversationResponse - 6, // 16: teleport.assist.v1.AssistService.GetAssistantConversations:output_type -> teleport.assist.v1.GetAssistantConversationsResponse - 17, // 17: teleport.assist.v1.AssistService.DeleteAssistantConversation:output_type -> google.protobuf.Empty - 3, // 18: teleport.assist.v1.AssistService.GetAssistantMessages:output_type -> teleport.assist.v1.GetAssistantMessagesResponse - 17, // 19: teleport.assist.v1.AssistService.CreateAssistantMessage:output_type -> google.protobuf.Empty - 17, // 20: teleport.assist.v1.AssistService.UpdateAssistantConversationInfo:output_type -> google.protobuf.Empty - 11, // 21: teleport.assist.v1.AssistService.IsAssistEnabled:output_type -> teleport.assist.v1.IsAssistEnabledResponse - 15, // 22: teleport.assist.v1.AssistEmbeddingService.GetAssistantEmbeddings:output_type -> teleport.assist.v1.GetAssistantEmbeddingsResponse - 15, // [15:23] is the sub-list for method output_type - 7, // [7:15] is the sub-list for method input_type - 7, // [7:7] is the sub-list for extension type_name - 7, // [7:7] is the sub-list for extension extendee - 0, // [0:7] is the sub-list for field type_name + 19, // 7: teleport.assist.v1.SearchUnifiedResourcesResponse.resources:type_name -> proto.PaginatedResource + 7, // 8: teleport.assist.v1.AssistService.CreateAssistantConversation:input_type -> teleport.assist.v1.CreateAssistantConversationRequest + 4, // 9: teleport.assist.v1.AssistService.GetAssistantConversations:input_type -> teleport.assist.v1.GetAssistantConversationsRequest + 12, // 10: teleport.assist.v1.AssistService.DeleteAssistantConversation:input_type -> teleport.assist.v1.DeleteAssistantConversationRequest + 0, // 11: teleport.assist.v1.AssistService.GetAssistantMessages:input_type -> teleport.assist.v1.GetAssistantMessagesRequest + 2, // 12: teleport.assist.v1.AssistService.CreateAssistantMessage:input_type -> teleport.assist.v1.CreateAssistantMessageRequest + 9, // 13: teleport.assist.v1.AssistService.UpdateAssistantConversationInfo:input_type -> teleport.assist.v1.UpdateAssistantConversationInfoRequest + 10, // 14: teleport.assist.v1.AssistService.IsAssistEnabled:input_type -> teleport.assist.v1.IsAssistEnabledRequest + 16, // 15: teleport.assist.v1.AssistService.SearchUnifiedResources:input_type -> teleport.assist.v1.SearchUnifiedResourcesRequest + 13, // 16: teleport.assist.v1.AssistEmbeddingService.GetAssistantEmbeddings:input_type -> teleport.assist.v1.GetAssistantEmbeddingsRequest + 8, // 17: teleport.assist.v1.AssistService.CreateAssistantConversation:output_type -> teleport.assist.v1.CreateAssistantConversationResponse + 6, // 18: teleport.assist.v1.AssistService.GetAssistantConversations:output_type -> teleport.assist.v1.GetAssistantConversationsResponse + 20, // 19: teleport.assist.v1.AssistService.DeleteAssistantConversation:output_type -> google.protobuf.Empty + 3, // 20: teleport.assist.v1.AssistService.GetAssistantMessages:output_type -> teleport.assist.v1.GetAssistantMessagesResponse + 20, // 21: teleport.assist.v1.AssistService.CreateAssistantMessage:output_type -> google.protobuf.Empty + 20, // 22: teleport.assist.v1.AssistService.UpdateAssistantConversationInfo:output_type -> google.protobuf.Empty + 11, // 23: teleport.assist.v1.AssistService.IsAssistEnabled:output_type -> teleport.assist.v1.IsAssistEnabledResponse + 17, // 24: teleport.assist.v1.AssistService.SearchUnifiedResources:output_type -> teleport.assist.v1.SearchUnifiedResourcesResponse + 15, // 25: teleport.assist.v1.AssistEmbeddingService.GetAssistantEmbeddings:output_type -> teleport.assist.v1.GetAssistantEmbeddingsResponse + 17, // [17:26] is the sub-list for method output_type + 8, // [8:17] is the sub-list for method input_type + 8, // [8:8] is the sub-list for extension type_name + 8, // [8:8] is the sub-list for extension extendee + 0, // [0:8] is the sub-list for field type_name } func init() { file_teleport_assist_v1_assist_proto_init() } @@ -1403,6 +1549,30 @@ func file_teleport_assist_v1_assist_proto_init() { return nil } } + file_teleport_assist_v1_assist_proto_msgTypes[16].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*SearchUnifiedResourcesRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_teleport_assist_v1_assist_proto_msgTypes[17].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*SearchUnifiedResourcesResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } } type x struct{} out := protoimpl.TypeBuilder{ @@ -1410,7 +1580,7 @@ func file_teleport_assist_v1_assist_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: file_teleport_assist_v1_assist_proto_rawDesc, NumEnums: 0, - NumMessages: 16, + NumMessages: 18, NumExtensions: 0, NumServices: 2, }, diff --git a/api/gen/proto/go/assist/v1/assist_grpc.pb.go b/api/gen/proto/go/assist/v1/assist_grpc.pb.go index cb399e7316958..f86a586be0689 100644 --- a/api/gen/proto/go/assist/v1/assist_grpc.pb.go +++ b/api/gen/proto/go/assist/v1/assist_grpc.pb.go @@ -41,6 +41,7 @@ const ( AssistService_CreateAssistantMessage_FullMethodName = "/teleport.assist.v1.AssistService/CreateAssistantMessage" AssistService_UpdateAssistantConversationInfo_FullMethodName = "/teleport.assist.v1.AssistService/UpdateAssistantConversationInfo" AssistService_IsAssistEnabled_FullMethodName = "/teleport.assist.v1.AssistService/IsAssistEnabled" + AssistService_SearchUnifiedResources_FullMethodName = "/teleport.assist.v1.AssistService/SearchUnifiedResources" ) // AssistServiceClient is the client API for AssistService service. @@ -61,6 +62,8 @@ type AssistServiceClient interface { UpdateAssistantConversationInfo(ctx context.Context, in *UpdateAssistantConversationInfoRequest, opts ...grpc.CallOption) (*emptypb.Empty, error) // IsAssistEnabled returns true if the assist is enabled or not on the auth level. IsAssistEnabled(ctx context.Context, in *IsAssistEnabledRequest, opts ...grpc.CallOption) (*IsAssistEnabledResponse, error) + // SearchUnifiedResources returns a similarity-ordered list of resources from the unified resource cache. + SearchUnifiedResources(ctx context.Context, in *SearchUnifiedResourcesRequest, opts ...grpc.CallOption) (*SearchUnifiedResourcesResponse, error) } type assistServiceClient struct { @@ -134,6 +137,15 @@ func (c *assistServiceClient) IsAssistEnabled(ctx context.Context, in *IsAssistE return out, nil } +func (c *assistServiceClient) SearchUnifiedResources(ctx context.Context, in *SearchUnifiedResourcesRequest, opts ...grpc.CallOption) (*SearchUnifiedResourcesResponse, error) { + out := new(SearchUnifiedResourcesResponse) + err := c.cc.Invoke(ctx, AssistService_SearchUnifiedResources_FullMethodName, in, out, opts...) + if err != nil { + return nil, err + } + return out, nil +} + // AssistServiceServer is the server API for AssistService service. // All implementations must embed UnimplementedAssistServiceServer // for forward compatibility @@ -152,6 +164,8 @@ type AssistServiceServer interface { UpdateAssistantConversationInfo(context.Context, *UpdateAssistantConversationInfoRequest) (*emptypb.Empty, error) // IsAssistEnabled returns true if the assist is enabled or not on the auth level. IsAssistEnabled(context.Context, *IsAssistEnabledRequest) (*IsAssistEnabledResponse, error) + // SearchUnifiedResources returns a similarity-ordered list of resources from the unified resource cache. + SearchUnifiedResources(context.Context, *SearchUnifiedResourcesRequest) (*SearchUnifiedResourcesResponse, error) mustEmbedUnimplementedAssistServiceServer() } @@ -180,6 +194,9 @@ func (UnimplementedAssistServiceServer) UpdateAssistantConversationInfo(context. func (UnimplementedAssistServiceServer) IsAssistEnabled(context.Context, *IsAssistEnabledRequest) (*IsAssistEnabledResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method IsAssistEnabled not implemented") } +func (UnimplementedAssistServiceServer) SearchUnifiedResources(context.Context, *SearchUnifiedResourcesRequest) (*SearchUnifiedResourcesResponse, error) { + return nil, status.Errorf(codes.Unimplemented, "method SearchUnifiedResources not implemented") +} func (UnimplementedAssistServiceServer) mustEmbedUnimplementedAssistServiceServer() {} // UnsafeAssistServiceServer may be embedded to opt out of forward compatibility for this service. @@ -319,6 +336,24 @@ func _AssistService_IsAssistEnabled_Handler(srv interface{}, ctx context.Context return interceptor(ctx, in, info, handler) } +func _AssistService_SearchUnifiedResources_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(SearchUnifiedResourcesRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(AssistServiceServer).SearchUnifiedResources(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: AssistService_SearchUnifiedResources_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(AssistServiceServer).SearchUnifiedResources(ctx, req.(*SearchUnifiedResourcesRequest)) + } + return interceptor(ctx, in, info, handler) +} + // AssistService_ServiceDesc is the grpc.ServiceDesc for AssistService service. // It's only intended for direct use with grpc.RegisterService, // and not to be introspected or modified (even as a copy) @@ -354,6 +389,10 @@ var AssistService_ServiceDesc = grpc.ServiceDesc{ MethodName: "IsAssistEnabled", Handler: _AssistService_IsAssistEnabled_Handler, }, + { + MethodName: "SearchUnifiedResources", + Handler: _AssistService_SearchUnifiedResources_Handler, + }, }, Streams: []grpc.StreamDesc{}, Metadata: "teleport/assist/v1/assist.proto", diff --git a/api/proto/teleport/assist/v1/assist.proto b/api/proto/teleport/assist/v1/assist.proto index 6157198d791c0..c5b6631a030de 100644 --- a/api/proto/teleport/assist/v1/assist.proto +++ b/api/proto/teleport/assist/v1/assist.proto @@ -18,6 +18,7 @@ package teleport.assist.v1; import "google/protobuf/empty.proto"; import "google/protobuf/timestamp.proto"; +import "teleport/legacy/client/proto/authservice.proto"; option go_package = "github.com/gravitational/teleport/api/gen/proto/go/assist/v1;assist"; @@ -150,6 +151,22 @@ message GetAssistantEmbeddingsResponse { repeated EmbeddedDocument embeddings = 1; } +// SearchUnifiedResourcesRequest is a request to search for one or more resource kinds using similiarity search. +message SearchUnifiedResourcesRequest { + // query is the query used for similarity search. + string query = 1; + // limit is the number of embeddings to return (also known as k). + int32 limit = 2; + // kinds is the kind of embeddings to return (ex, node). Returns all supported kinds if empty. + repeated string kinds = 3; +} + +// SearchUnifiedResourcesResponse is a response from the assistant service with a similarity-ordered list of resources. +message SearchUnifiedResourcesResponse { + // resources is the list of resources. + repeated proto.PaginatedResource resources = 1; +} + // AssistService is a service that provides an ability to communicate with the Teleport Assist. service AssistService { // CreateNewConversation creates a new conversation and returns the UUID of it. @@ -172,6 +189,9 @@ service AssistService { // IsAssistEnabled returns true if the assist is enabled or not on the auth level. rpc IsAssistEnabled(IsAssistEnabledRequest) returns (IsAssistEnabledResponse); + + // SearchUnifiedResources returns a similarity-ordered list of resources from the unified resource cache. + rpc SearchUnifiedResources(SearchUnifiedResourcesRequest) returns (SearchUnifiedResourcesResponse); } // AssistEmbeddingService is a service that provides an ability to communicate with the Assist Embedding service. diff --git a/lib/ai/embeddingprocessor_test.go b/lib/ai/embeddingprocessor_test.go index 2ce32b1b6647d..bbef40f0b2756 100644 --- a/lib/ai/embeddingprocessor_test.go +++ b/lib/ai/embeddingprocessor_test.go @@ -41,53 +41,6 @@ import ( "github.com/gravitational/teleport/lib/utils" ) -// MockEmbedder returns embeddings based on the sha256 hash function. Those -// embeddings have no semantic meaning but ensure different embedded content -// provides different embeddings. -type MockEmbedder struct { - timesCalled map[string]int -} - -func (m *MockEmbedder) ComputeEmbeddings(_ context.Context, input []string) ([]embedding.Vector64, error) { - result := make([]embedding.Vector64, len(input)) - for i, text := range input { - name := strings.Split(text, "\n")[0] - m.timesCalled[name]++ - hash := sha256.Sum256([]byte(text)) - vector := make(embedding.Vector64, len(hash)) - for j, x := range hash { - vector[j] = 1 / float64(int(x)+1) - } - result[i] = vector - } - return result, nil -} - -type mockResourceGetter struct { - services.Presence - services.AccessLists -} - -func (m *mockResourceGetter) GetDatabaseServers(_ context.Context, _ string, _ ...services.MarshalOption) ([]types.DatabaseServer, error) { - return nil, nil -} - -func (m *mockResourceGetter) GetKubernetesServers(_ context.Context) ([]types.KubeServer, error) { - return nil, nil -} - -func (m *mockResourceGetter) GetApplicationServers(_ context.Context, _ string) ([]types.AppServer, error) { - return nil, nil -} - -func (m *mockResourceGetter) GetWindowsDesktops(_ context.Context, _ types.WindowsDesktopFilter) ([]types.WindowsDesktop, error) { - return nil, nil -} - -func (m *mockResourceGetter) ListSAMLIdPServiceProviders(_ context.Context, _ int, _ string) ([]types.SAMLIdPServiceProvider, string, error) { - return nil, "", nil -} - func TestNodeEmbeddingGeneration(t *testing.T) { t.Parallel() @@ -104,8 +57,8 @@ func TestNodeEmbeddingGeneration(t *testing.T) { }) require.NoError(t, err) - embedder := MockEmbedder{ - timesCalled: make(map[string]int), + embedder := ai.MockEmbedder{ + TimesCalled: make(map[string]int), } events := local.NewEventsService(bk) accessLists, err := local.NewAccessListService(bk, clock) @@ -163,7 +116,7 @@ func TestNodeEmbeddingGeneration(t *testing.T) { nodesAcquired, embeddings.GetAllEmbeddings(ctx)) - for k, v := range embedder.timesCalled { + for k, v := range embedder.TimesCalled { require.Equal(t, 1, v, "expected %v to be computed once, was %d", k, v) } @@ -184,7 +137,7 @@ func TestNodeEmbeddingGeneration(t *testing.T) { return len(items) == numInitialNodes+1 }, 7*time.Second, 200*time.Millisecond) - for k, v := range embedder.timesCalled { + for k, v := range embedder.TimesCalled { expected := 1 if strings.Contains(k, "node1") { expected = 2 @@ -331,3 +284,28 @@ func Test_batchReducer_Add(t *testing.T) { }) } } + +type mockResourceGetter struct { + services.Presence + services.AccessLists +} + +func (m *mockResourceGetter) GetDatabaseServers(_ context.Context, _ string, _ ...services.MarshalOption) ([]types.DatabaseServer, error) { + return nil, nil +} + +func (m *mockResourceGetter) GetKubernetesServers(_ context.Context) ([]types.KubeServer, error) { + return nil, nil +} + +func (m *mockResourceGetter) GetApplicationServers(_ context.Context, _ string) ([]types.AppServer, error) { + return nil, nil +} + +func (m *mockResourceGetter) GetWindowsDesktops(_ context.Context, _ types.WindowsDesktopFilter) ([]types.WindowsDesktop, error) { + return nil, nil +} + +func (m *mockResourceGetter) ListSAMLIdPServiceProviders(_ context.Context, _ int, _ string) ([]types.SAMLIdPServiceProvider, string, error) { + return nil, "", nil +} diff --git a/lib/ai/mock_embedder.go b/lib/ai/mock_embedder.go new file mode 100644 index 0000000000000..615fa690bdb0a --- /dev/null +++ b/lib/ai/mock_embedder.go @@ -0,0 +1,52 @@ +/* + * Copyright 2023 Gravitational, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package ai + +import ( + "context" + "crypto/sha256" + "strings" + "sync" + + "github.com/gravitational/teleport/lib/ai/embedding" +) + +// MockEmbedder returns embeddings based on the sha256 hash function. Those +// embeddings have no semantic meaning but ensure different embedded content +// provides different embeddings. +type MockEmbedder struct { + mu sync.Mutex + TimesCalled map[string]int +} + +func (m *MockEmbedder) ComputeEmbeddings(_ context.Context, input []string) ([]embedding.Vector64, error) { + m.mu.Lock() + defer m.mu.Unlock() + + result := make([]embedding.Vector64, len(input)) + for i, text := range input { + name := strings.Split(text, "\n")[0] + m.TimesCalled[name]++ + hash := sha256.Sum256([]byte(text)) + vector := make(embedding.Vector64, len(hash)) + for j, x := range hash { + vector[j] = 1 / float64(int(x)+1) + } + result[i] = vector + } + return result, nil +} diff --git a/lib/auth/assist/assistv1/service.go b/lib/auth/assist/assistv1/service.go index 335d89dcf32c1..6a578345fb7f4 100644 --- a/lib/auth/assist/assistv1/service.go +++ b/lib/auth/assist/assistv1/service.go @@ -23,6 +23,7 @@ import ( "github.com/gravitational/trace" "github.com/sirupsen/logrus" + "golang.org/x/exp/slices" "google.golang.org/protobuf/types/known/emptypb" "github.com/gravitational/teleport/api/defaults" @@ -34,6 +35,12 @@ import ( "github.com/gravitational/teleport/lib/services" ) +const ( + // maxSearchLimit is the maximum number of search results to return. + // We have a hard cap due the simplistic design of our retriever which has quadratic complexity. + maxSearchLimit = 100 +) + // ServiceConfig holds configuration options for // the assist gRPC service. type ServiceConfig struct { @@ -250,9 +257,37 @@ func (a *Service) GetAssistantEmbeddings(ctx context.Context, msg *assist.GetAss // Use default values for the id and content, as we only care about the embeddings. queryEmbeddings := embeddinglib.NewEmbedding(msg.Kind, "", embeddings[0], [32]byte{}) - accessChecker := accessCheckerForKind(ctx, a, authCtx, msg.Kind) + accessChecker := makeAccessChecker(ctx, a, authCtx, msg.Kind) documents := a.embeddings.GetRelevant(queryEmbeddings, int(msg.Limit), accessChecker) - return assembleEmbeddingResponseForKind(ctx, a, msg.Kind, documents) + return assembleEmbeddingResponse(ctx, a, documents) +} + +// SearchUnifiedResources returns a similarity-ordered list of resources from the unified resource cache +func (a *Service) SearchUnifiedResources(ctx context.Context, msg *assist.SearchUnifiedResourcesRequest) (*assist.SearchUnifiedResourcesResponse, error) { + if a.embedder == nil { + return nil, trace.BadParameter("assist is not configured in auth server") + } + + // Call the openAI API to get the embeddings for the query. + embeddings, err := a.embedder.ComputeEmbeddings(ctx, []string{msg.Query}) + if err != nil { + return nil, trace.Wrap(err) + } + if len(embeddings) == 0 { + return nil, trace.NotFound("OpenAI embeddings returned no results") + } + + authCtx, err := a.authorizer.Authorize(ctx) + if err != nil { + return nil, authz.ConvertAuthorizerError(ctx, a.log, err) + } + + // Use default values for the id and content, as we only care about the embeddings. + queryEmbeddings := embeddinglib.NewEmbedding("", "", embeddings[0], [32]byte{}) + limit := max(msg.Limit, maxSearchLimit) + accessChecker := makeAccessChecker(ctx, a, authCtx, msg.Kinds...) + documents := a.embeddings.GetRelevant(queryEmbeddings, int(limit), accessChecker) + return assembleSearchResponse(ctx, a, documents) } // userHasAccess returns true if the user should have access to the resource. @@ -260,13 +295,68 @@ func userHasAccess(authCtx *authz.Context, req interface{ GetUsername() string } return !authz.IsCurrentUser(*authCtx, req.GetUsername()) && !authz.HasBuiltinRole(*authCtx, string(types.RoleAdmin)) } -func assembleEmbeddingResponseForKind(ctx context.Context, a *Service, kind string, documents []*ai.Document) (*assist.GetAssistantEmbeddingsResponse, error) { +func assembleSearchResponse(ctx context.Context, a *Service, documents []*ai.Document) (*assist.SearchUnifiedResourcesResponse, error) { + resources := make([]types.ResourceWithLabels, 0, len(documents)) + + for _, doc := range documents { + var resource types.ResourceWithLabels + var err error + + switch doc.EmbeddedKind { + case types.KindNode: + resource, err = a.resourceGetter.GetNode(ctx, defaults.Namespace, doc.GetEmbeddedID()) + case types.KindKubernetesCluster: + resource, err = a.resourceGetter.GetKubernetesCluster(ctx, doc.GetEmbeddedID()) + case types.KindApp: + resource, err = a.resourceGetter.GetApp(ctx, doc.GetEmbeddedID()) + case types.KindDatabase: + resource, err = a.resourceGetter.GetDatabase(ctx, doc.GetEmbeddedID()) + case types.KindWindowsDesktop: + desktops, err := a.resourceGetter.GetWindowsDesktops(ctx, types.WindowsDesktopFilter{ + Name: doc.GetEmbeddedID(), + }) + if err != nil { + return nil, trace.Wrap(err) + } + + for _, d := range desktops { + if d.GetName() == doc.GetEmbeddedID() { + resource = d + break + } + } + + if resource == nil { + return nil, trace.NotFound("windows desktop %q not found", doc.GetEmbeddedID()) + } + default: + return nil, trace.BadParameter("resource kind %v is not supported", doc.EmbeddedKind) + } + + if err != nil { + return nil, trace.Wrap(err) + } + + resources = append(resources, resource) + } + + paginated, err := services.MakePaginatedResources(types.KindUnifiedResource, resources) + if err != nil { + return nil, trace.Wrap(err) + } + + return &assist.SearchUnifiedResourcesResponse{ + Resources: paginated, + }, nil +} + +func assembleEmbeddingResponse(ctx context.Context, a *Service, documents []*ai.Document) (*assist.GetAssistantEmbeddingsResponse, error) { protoDocs := make([]*assist.EmbeddedDocument, 0, len(documents)) for _, doc := range documents { var content []byte - switch kind { + switch doc.EmbeddedKind { case types.KindNode: node, err := a.resourceGetter.GetNode(ctx, defaults.Namespace, doc.GetEmbeddedID()) if err != nil { @@ -345,23 +435,22 @@ func assembleEmbeddingResponseForKind(ctx context.Context, a *Service, kind stri }, nil } -func accessCheckerForKind(ctx context.Context, a *Service, authCtx *authz.Context, kind string) func(id string, embedding *embeddinglib.Embedding) bool { +func makeAccessChecker(ctx context.Context, a *Service, authCtx *authz.Context, kinds ...string) func(id string, embedding *embeddinglib.Embedding) bool { return func(id string, embedding *embeddinglib.Embedding) bool { - if embedding.EmbeddedKind != kind { + if !slices.Contains(kinds, embedding.EmbeddedKind) && len(kinds) > 0 { return false } var resource services.AccessCheckable var err error - switch kind { + switch embedding.EmbeddedKind { case types.KindNode: resource, err = a.resourceGetter.GetNode(ctx, defaults.Namespace, embedding.GetEmbeddedID()) if err != nil { a.log.Tracef("failed to get node %q: %v", embedding.GetName(), err) return false } - case types.KindKubernetesCluster: resource, err = a.resourceGetter.GetKubernetesCluster(ctx, embedding.GetEmbeddedID()) if err != nil { @@ -400,6 +489,9 @@ func accessCheckerForKind(ctx context.Context, a *Service, authCtx *authz.Contex a.log.Tracef("failed to find windows desktop %q: %v", embedding.GetName(), err) return false } + default: + a.log.Tracef("resource kind %v is not supported", embedding.EmbeddedKind) + return false } return authCtx.Checker.CheckAccess(resource, services.AccessState{MFAVerified: true}) == nil diff --git a/lib/auth/assist/assistv1/test/service_test.go b/lib/auth/assist/assistv1/test/service_test.go index 3fda580cd204b..05cbccad7c319 100644 --- a/lib/auth/assist/assistv1/test/service_test.go +++ b/lib/auth/assist/assistv1/test/service_test.go @@ -19,18 +19,25 @@ package assistv1_test import ( "context" "testing" + "time" + "github.com/gravitational/trace" + "github.com/jonboulle/clockwork" + log "github.com/sirupsen/logrus" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "google.golang.org/protobuf/types/known/timestamppb" + "github.com/gravitational/teleport" assistpb "github.com/gravitational/teleport/api/gen/proto/go/assist/v1" "github.com/gravitational/teleport/api/types" + "github.com/gravitational/teleport/api/utils/retryutils" "github.com/gravitational/teleport/lib/ai" "github.com/gravitational/teleport/lib/assist" "github.com/gravitational/teleport/lib/auth/assist/assistv1" "github.com/gravitational/teleport/lib/authz" "github.com/gravitational/teleport/lib/backend/memory" + "github.com/gravitational/teleport/lib/defaults" "github.com/gravitational/teleport/lib/services" "github.com/gravitational/teleport/lib/services/local" "github.com/gravitational/teleport/lib/tlsca" @@ -268,6 +275,42 @@ func TestService_InsertAssistantMessage(t *testing.T) { } } +func TestService_SearchUnifiedResources(t *testing.T) { + t.Parallel() + + tests := []struct { + username string + req *assistpb.SearchUnifiedResourcesRequest + returnedLen int + }{ + { + username: defaultUser, + req: &assistpb.SearchUnifiedResourcesRequest{ + Kinds: []string{types.KindNode}, + }, + returnedLen: 2, + }, + { + username: noAccessUser, + req: &assistpb.SearchUnifiedResourcesRequest{ + Kinds: []string{types.KindNode}, + }, + returnedLen: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.username, func(t *testing.T) { + ctxs, svc := initSvc(t) + require.Eventually(t, func() bool { + resp, err := svc.SearchUnifiedResources(ctxs[tt.username], tt.req) + require.NoError(t, err) + return tt.returnedLen == len(resp.GetResources()) + }, 5*time.Second, 100*time.Millisecond) + }) + } +} + func initSvc(t *testing.T) (map[string]context.Context, *assistv1.Service) { ctx := context.Background() backend, err := memory.New(memory.Config{}) @@ -278,6 +321,7 @@ func initSvc(t *testing.T) (map[string]context.Context, *assistv1.Service) { trustSvc := local.NewCAService(backend) roleSvc := local.NewAccessService(backend) userSvc := local.NewIdentityService(backend) + presenceSvc := local.NewPresenceService(backend) require.NoError(t, clusterConfigSvc.SetAuthPreference(ctx, types.DefaultAuthPreference())) require.NoError(t, clusterConfigSvc.SetClusterAuditConfig(ctx, types.DefaultClusterAuditConfig())) @@ -289,13 +333,27 @@ func initSvc(t *testing.T) (map[string]context.Context, *assistv1.Service) { services.Trust services.RoleGetter services.UserGetter + services.Presence }{ ClusterConfiguration: clusterConfigSvc, Trust: trustSvc, RoleGetter: roleSvc, UserGetter: userSvc, + Presence: presenceSvc, } + n1, err := types.NewServer("node-1", types.KindNode, types.ServerSpecV2{}) + require.NoError(t, err) + n2, err := types.NewServer("node-2", types.KindNode, types.ServerSpecV2{}) + require.NoError(t, err) + _, err = presenceSvc.UpsertNode(ctx, n1) + require.NoError(t, err) + _, err = presenceSvc.UpsertNode(ctx, n2) + require.NoError(t, err) + + accesslistSvc, err := local.NewAccessListService(backend, clockwork.NewFakeClock()) + require.NoError(t, err) + accessService := local.NewAccessService(backend) eventService := local.NewEventsService(backend) lockWatcher, err := services.NewLockWatcher(ctx, services.LockWatcherConfig{ @@ -318,11 +376,17 @@ func initSvc(t *testing.T) (map[string]context.Context, *assistv1.Service) { role, err := types.NewRole("allow-rules", types.RoleSpecV6{ Allow: types.RoleConditions{ + Namespaces: []string{}, + NodeLabels: types.Labels{types.Wildcard: []string{types.Wildcard}}, Rules: []types.Rule{ { Resources: []string{types.KindAssistant}, Verbs: []string{types.VerbList, types.VerbRead, types.VerbUpdate, types.VerbCreate, types.VerbDelete}, }, + { + Resources: []string{types.KindNode}, + Verbs: []string{types.VerbList, types.VerbRead}, + }, }, }, }) @@ -358,36 +422,99 @@ func initSvc(t *testing.T) (map[string]context.Context, *assistv1.Service) { ctxs[user.GetName()] = ctx } + embedder := ai.MockEmbedder{ + TimesCalled: make(map[string]int), + } + + embeddings := &ai.SimpleRetriever{} + embeddingSrv := local.NewEmbeddingsService(backend) svc, err := assistv1.NewService(&assistv1.ServiceConfig{ - Backend: local.NewAssistService(backend), - Authorizer: authorizer, - Embeddings: &ai.SimpleRetriever{}, - ResourceGetter: &resourceGetterFake{}, + Backend: local.NewAssistService(backend), + Authorizer: authorizer, + Embeddings: embeddings, + ResourceGetter: &resourceGetterAllImpl{ + PresenceService: presenceSvc, + AccessLists: accesslistSvc, + }, + Embedder: &embedder, + }) + require.NoError(t, err) + + unifiedResourcesCache, err := services.NewUnifiedResourceCache(ctx, services.UnifiedResourceCacheConfig{ + ResourceWatcherConfig: services.ResourceWatcherConfig{ + QueueSize: defaults.UnifiedResourcesQueueSize, + Component: teleport.ComponentUnifiedResource, + Client: eventService, + MaxStaleness: time.Second, + }, + ResourceGetter: &resourceGetterAllImpl{ + PresenceService: presenceSvc, + AccessLists: accesslistSvc, + }, }) require.NoError(t, err) + log.Debugf("Starting embedding watcher") + embeddingProcessor := ai.NewEmbeddingProcessor(&ai.EmbeddingProcessorConfig{ + AIClient: &embedder, + EmbeddingsRetriever: embeddings, + EmbeddingSrv: embeddingSrv, + NodeSrv: unifiedResourcesCache, + Jitter: retryutils.NewFullJitter(), + Log: log.NewEntry(log.StandardLogger()), + }) + log.Debugf("Starting embedding processor") + + embeddingProcessorCtx, embeddingProcessorCancel := context.WithCancel(context.Background()) + go embeddingProcessor.Run(embeddingProcessorCtx, time.Millisecond*100, time.Millisecond*100) + t.Cleanup(embeddingProcessorCancel) return ctxs, svc } -type resourceGetterFake struct { +type resourceGetterAllImpl struct { + *local.PresenceService + services.AccessLists } -func (g *resourceGetterFake) GetNode(ctx context.Context, namespace, name string) (types.Server, error) { +func (g *resourceGetterAllImpl) GetKubernetesCluster(ctx context.Context, name string) (types.KubeCluster, error) { + kubeServers, err := g.PresenceService.GetKubernetesServers(ctx) + if err != nil { + return nil, err + } + + for _, kubeServer := range kubeServers { + if kubeServer.GetName() == name { + return kubeServer.GetCluster(), nil + } + } + + return nil, trace.NotFound("cluster not found") +} + +func (g *resourceGetterAllImpl) GetApp(ctx context.Context, name string) (types.Application, error) { + return nil, nil +} + +func (g *resourceGetterAllImpl) GetDatabase(ctx context.Context, name string) (types.Database, error) { return nil, nil } -func (g *resourceGetterFake) GetKubernetesCluster(ctx context.Context, name string) (types.KubeCluster, error) { +func (g *resourceGetterAllImpl) GetWindowsDesktops(ctx context.Context, _ types.WindowsDesktopFilter) ([]types.WindowsDesktop, error) { return nil, nil } -func (g *resourceGetterFake) GetApp(ctx context.Context, name string) (types.Application, error) { +func (m *resourceGetterAllImpl) GetDatabaseServers(_ context.Context, _ string, _ ...services.MarshalOption) ([]types.DatabaseServer, error) { return nil, nil } -func (g *resourceGetterFake) GetDatabase(ctx context.Context, name string) (types.Database, error) { +func (m *resourceGetterAllImpl) GetKubernetesServers(_ context.Context) ([]types.KubeServer, error) { return nil, nil } -func (g *resourceGetterFake) GetWindowsDesktops(ctx context.Context, _ types.WindowsDesktopFilter) ([]types.WindowsDesktop, error) { +func (m *resourceGetterAllImpl) GetApplicationServers(_ context.Context, _ string) ([]types.AppServer, error) { return nil, nil } + +func (m *resourceGetterAllImpl) ListSAMLIdPServiceProviders(_ context.Context, _ int, _ string) ([]types.SAMLIdPServiceProvider, string, error) { + return nil, "", nil +} diff --git a/lib/auth/auth_with_roles.go b/lib/auth/auth_with_roles.go index 871ad4aea800c..4b0c03827d289 100644 --- a/lib/auth/auth_with_roles.go +++ b/lib/auth/auth_with_roles.go @@ -1501,111 +1501,6 @@ func (a *ServerWithRoles) GetNode(ctx context.Context, namespace, name string) ( return node, nil } -func (s *ServerWithRoles) MakePaginatedResources(requestType string, resources []types.ResourceWithLabels) ([]*proto.PaginatedResource, error) { - paginatedResources := make([]*proto.PaginatedResource, 0, len(resources)) - for _, resource := range resources { - var protoResource *proto.PaginatedResource - resourceKind := requestType - if requestType == types.KindUnifiedResource { - resourceKind = resource.GetKind() - } - switch resourceKind { - case types.KindDatabaseServer: - database, ok := resource.(*types.DatabaseServerV3) - if !ok { - return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) - } - - protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_DatabaseServer{DatabaseServer: database}} - case types.KindDatabaseService: - databaseService, ok := resource.(*types.DatabaseServiceV1) - if !ok { - return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) - } - - protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_DatabaseService{DatabaseService: databaseService}} - case types.KindAppServer: - app, ok := resource.(*types.AppServerV3) - if !ok { - return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) - } - - protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_AppServer{AppServer: app}} - case types.KindNode: - srv, ok := resource.(*types.ServerV2) - if !ok { - return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) - } - - protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_Node{Node: srv}} - case types.KindKubeServer: - srv, ok := resource.(*types.KubernetesServerV3) - if !ok { - return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) - } - - protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_KubernetesServer{KubernetesServer: srv}} - case types.KindWindowsDesktop: - desktop, ok := resource.(*types.WindowsDesktopV3) - if !ok { - return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) - } - - protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_WindowsDesktop{WindowsDesktop: desktop}} - case types.KindWindowsDesktopService: - desktopService, ok := resource.(*types.WindowsDesktopServiceV3) - if !ok { - return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) - } - - protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_WindowsDesktopService{WindowsDesktopService: desktopService}} - case types.KindKubernetesCluster: - cluster, ok := resource.(*types.KubernetesClusterV3) - if !ok { - return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) - } - - protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_KubeCluster{KubeCluster: cluster}} - case types.KindUserGroup: - userGroup, ok := resource.(*types.UserGroupV1) - if !ok { - return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) - } - - protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_UserGroup{UserGroup: userGroup}} - case types.KindSAMLIdPServiceProvider, types.KindAppOrSAMLIdPServiceProvider: - switch appOrSP := resource.(type) { - case *types.AppServerV3: - protoResource = &proto.PaginatedResource{ - Resource: &proto.PaginatedResource_AppServerOrSAMLIdPServiceProvider{ - AppServerOrSAMLIdPServiceProvider: &types.AppServerOrSAMLIdPServiceProviderV1{ - Resource: &types.AppServerOrSAMLIdPServiceProviderV1_AppServer{ - AppServer: appOrSP, - }, - }, - }} - case *types.SAMLIdPServiceProviderV1: - protoResource = &proto.PaginatedResource{ - Resource: &proto.PaginatedResource_AppServerOrSAMLIdPServiceProvider{ - AppServerOrSAMLIdPServiceProvider: &types.AppServerOrSAMLIdPServiceProviderV1{ - Resource: &types.AppServerOrSAMLIdPServiceProviderV1_SAMLIdPServiceProvider{ - SAMLIdPServiceProvider: appOrSP, - }, - }, - }} - default: - return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) - } - - default: - return nil, trace.NotImplemented("resource type %s doesn't support pagination", resource.GetKind()) - } - - paginatedResources = append(paginatedResources, protoResource) - } - return paginatedResources, nil -} - // ListUnifiedResources returns a paginated list of unified resources filtered by user access. func (a *ServerWithRoles) ListUnifiedResources(ctx context.Context, req *proto.ListUnifiedResourcesRequest) (*proto.ListUnifiedResourcesResponse, error) { // Fetch full list of resources in the backend. @@ -1726,7 +1621,7 @@ func (a *ServerWithRoles) ListUnifiedResources(ctx context.Context, req *proto.L return nil, trace.Wrap(err) } - paginatedResources, err := a.MakePaginatedResources(types.KindUnifiedResource, resp.Resources) + paginatedResources, err := services.MakePaginatedResources(types.KindUnifiedResource, resp.Resources) if err != nil { return nil, trace.Wrap(err, "making paginated unified resources") } diff --git a/lib/auth/grpcserver.go b/lib/auth/grpcserver.go index 0153aa3ca04e2..b6ea3ef1c3de5 100644 --- a/lib/auth/grpcserver.go +++ b/lib/auth/grpcserver.go @@ -4404,7 +4404,7 @@ func (g *GRPCServer) ListResources(ctx context.Context, req *authpb.ListResource return nil, trace.Wrap(err) } - paginatedResources, err := auth.MakePaginatedResources(req.ResourceType, resp.Resources) + paginatedResources, err := services.MakePaginatedResources(req.ResourceType, resp.Resources) if err != nil { return nil, trace.Wrap(err, "making paginated resources") } diff --git a/lib/services/unified_resource.go b/lib/services/unified_resource.go index ea74624f66526..8c44e49b20db2 100644 --- a/lib/services/unified_resource.go +++ b/lib/services/unified_resource.go @@ -26,6 +26,7 @@ import ( log "github.com/sirupsen/logrus" "github.com/gravitational/teleport" + "github.com/gravitational/teleport/api/client/proto" apidefaults "github.com/gravitational/teleport/api/defaults" "github.com/gravitational/teleport/api/types" "github.com/gravitational/teleport/api/types/accesslist" @@ -33,6 +34,9 @@ import ( "github.com/gravitational/teleport/lib/utils" ) +// UnifiedResourceKinds is a list of all kinds that are stored in the unified resource cache. +var UnifiedResourceKinds []string = []string{types.KindNode, types.KindKubeServer, types.KindDatabaseServer, types.KindAppServer, types.KindSAMLIdPServiceProvider, types.KindWindowsDesktop, types.KindAccessList} + // UnifiedResourceCacheConfig is used to configure a UnifiedResourceCache type UnifiedResourceCacheConfig struct { // BTreeDegree is a degree of B-Tree, 2 for example, will create a @@ -485,15 +489,12 @@ func (c *UnifiedResourceCache) processEventAndUpdateCurrent(ctx context.Context, // resourceKinds returns a list of resources to be watched. func (c *UnifiedResourceCache) resourceKinds() []types.WatchKind { - return []types.WatchKind{ - {Kind: types.KindNode}, - {Kind: types.KindDatabaseServer}, - {Kind: types.KindAppServer}, - {Kind: types.KindSAMLIdPServiceProvider}, - {Kind: types.KindWindowsDesktop}, - {Kind: types.KindKubeServer}, - {Kind: types.KindAccessList}, + watchKinds := make([]types.WatchKind, 0, len(UnifiedResourceKinds)) + for _, kind := range UnifiedResourceKinds { + watchKinds = append(watchKinds, types.WatchKind{Kind: kind}) } + + return watchKinds } func (c *UnifiedResourceCache) defineCollectorAsInitialized() { @@ -543,3 +544,108 @@ type item struct { const ( prefix = "unified_resource" ) + +// MakePaginatedResources converts a list of resources into a list of paginated proto representations. +func MakePaginatedResources(requestType string, resources []types.ResourceWithLabels) ([]*proto.PaginatedResource, error) { + paginatedResources := make([]*proto.PaginatedResource, 0, len(resources)) + for _, resource := range resources { + var protoResource *proto.PaginatedResource + resourceKind := requestType + if requestType == types.KindUnifiedResource { + resourceKind = resource.GetKind() + } + switch resourceKind { + case types.KindDatabaseServer: + database, ok := resource.(*types.DatabaseServerV3) + if !ok { + return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) + } + + protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_DatabaseServer{DatabaseServer: database}} + case types.KindDatabaseService: + databaseService, ok := resource.(*types.DatabaseServiceV1) + if !ok { + return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) + } + + protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_DatabaseService{DatabaseService: databaseService}} + case types.KindAppServer: + app, ok := resource.(*types.AppServerV3) + if !ok { + return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) + } + + protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_AppServer{AppServer: app}} + case types.KindNode: + srv, ok := resource.(*types.ServerV2) + if !ok { + return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) + } + + protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_Node{Node: srv}} + case types.KindKubeServer: + srv, ok := resource.(*types.KubernetesServerV3) + if !ok { + return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) + } + + protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_KubernetesServer{KubernetesServer: srv}} + case types.KindWindowsDesktop: + desktop, ok := resource.(*types.WindowsDesktopV3) + if !ok { + return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) + } + + protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_WindowsDesktop{WindowsDesktop: desktop}} + case types.KindWindowsDesktopService: + desktopService, ok := resource.(*types.WindowsDesktopServiceV3) + if !ok { + return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) + } + + protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_WindowsDesktopService{WindowsDesktopService: desktopService}} + case types.KindKubernetesCluster: + cluster, ok := resource.(*types.KubernetesClusterV3) + if !ok { + return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) + } + + protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_KubeCluster{KubeCluster: cluster}} + case types.KindUserGroup: + userGroup, ok := resource.(*types.UserGroupV1) + if !ok { + return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) + } + + protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_UserGroup{UserGroup: userGroup}} + case types.KindSAMLIdPServiceProvider, types.KindAppOrSAMLIdPServiceProvider: + switch appOrSP := resource.(type) { + case *types.AppServerV3: + protoResource = &proto.PaginatedResource{ + Resource: &proto.PaginatedResource_AppServerOrSAMLIdPServiceProvider{ + AppServerOrSAMLIdPServiceProvider: &types.AppServerOrSAMLIdPServiceProviderV1{ + Resource: &types.AppServerOrSAMLIdPServiceProviderV1_AppServer{ + AppServer: appOrSP, + }, + }, + }} + case *types.SAMLIdPServiceProviderV1: + protoResource = &proto.PaginatedResource{ + Resource: &proto.PaginatedResource_AppServerOrSAMLIdPServiceProvider{ + AppServerOrSAMLIdPServiceProvider: &types.AppServerOrSAMLIdPServiceProviderV1{ + Resource: &types.AppServerOrSAMLIdPServiceProviderV1_SAMLIdPServiceProvider{ + SAMLIdPServiceProvider: appOrSP, + }, + }, + }} + default: + return nil, trace.BadParameter("%s has invalid type %T", resourceKind, resource) + } + default: + return nil, trace.NotImplemented("resource type %s doesn't support pagination", resource.GetKind()) + } + + paginatedResources = append(paginatedResources, protoResource) + } + return paginatedResources, nil +}