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

feat: added span name processor #510

Merged
merged 5 commits into from
Oct 9, 2024
Merged
Show file tree
Hide file tree
Changes from 2 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
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import com.epam.aidial.core.ProxyContext;
import com.epam.aidial.core.config.Deployment;
import com.epam.aidial.core.config.Features;
import com.epam.aidial.core.util.SpanUtil;
import com.epam.aidial.core.util.UrlUtil;
import io.vertx.core.http.HttpMethod;
import lombok.experimental.UtilityClass;
Expand All @@ -16,47 +17,47 @@
@UtilityClass
public class ControllerSelector {

private static final Pattern PATTERN_POST_DEPLOYMENT = Pattern.compile("^/+openai/deployments/(.+?)/(completions|chat/completions|embeddings)$");
private static final Pattern PATTERN_DEPLOYMENT = Pattern.compile("^/+openai/deployments/(.+?)$");
private static final Pattern PATTERN_POST_DEPLOYMENT = Pattern.compile("^/+openai/deployments/(?<id>.+?)/(completions|chat/completions|embeddings)$");
private static final Pattern PATTERN_DEPLOYMENT = Pattern.compile("^/+openai/deployments/(?<id>.+?)$");
private static final Pattern PATTERN_DEPLOYMENTS = Pattern.compile("^/+openai/deployments$");

private static final Pattern PATTERN_MODEL = Pattern.compile("^/+openai/models/(.+?)$");
private static final Pattern PATTERN_MODEL = Pattern.compile("^/+openai/models/(?<id>.+?)$");
private static final Pattern PATTERN_MODELS = Pattern.compile("^/+openai/models$");

private static final Pattern PATTERN_ADDON = Pattern.compile("^/+openai/addons/(.+?)$");
private static final Pattern PATTERN_ADDON = Pattern.compile("^/+openai/addons/(?<id>.+?)$");
private static final Pattern PATTERN_ADDONS = Pattern.compile("^/+openai/addons$");

private static final Pattern PATTERN_ASSISTANT = Pattern.compile("^/+openai/assistants/(.+?)$");
private static final Pattern PATTERN_ASSISTANT = Pattern.compile("^/+openai/assistants/(?<id>.+?)$");
private static final Pattern PATTERN_ASSISTANTS = Pattern.compile("^/+openai/assistants$");

private static final Pattern PATTERN_APPLICATION = Pattern.compile("^/+openai/applications/(.+?)$");
private static final Pattern PATTERN_APPLICATION = Pattern.compile("^/+openai/applications/(?<id>.+?)$");
private static final Pattern PATTERN_APPLICATIONS = Pattern.compile("^/+openai/applications$");


private static final Pattern PATTERN_BUCKET = Pattern.compile("^/v1/bucket$");

private static final Pattern PATTERN_FILES = Pattern.compile("^/v1/files/[a-zA-Z0-9]+/.*");
private static final Pattern PATTERN_FILES_METADATA = Pattern.compile("^/v1/metadata/files/[a-zA-Z0-9]+/.*");
private static final Pattern PATTERN_FILES = Pattern.compile("^/v1/files/(?<bucket>[a-zA-Z0-9]+)/(?<path>.*)");
private static final Pattern PATTERN_FILES_METADATA = Pattern.compile("^/v1/metadata/files/(?<bucket>[a-zA-Z0-9]+)/(?<path>.*)");

private static final Pattern PATTERN_RESOURCE = Pattern.compile("^/v1/(conversations|prompts|applications)/[a-zA-Z0-9]+/.*");
private static final Pattern PATTERN_RESOURCE_METADATA = Pattern.compile("^/v1/metadata/(conversations|prompts|applications)/[a-zA-Z0-9]+/.*");
private static final Pattern PATTERN_RESOURCE = Pattern.compile("^/v1/(conversations|prompts|applications)/(?<bucket>[a-zA-Z0-9]+)/(?<path>.*)");
private static final Pattern PATTERN_RESOURCE_METADATA = Pattern.compile("^/v1/metadata/(conversations|prompts|applications)/(?<bucket>[a-zA-Z0-9]+)/(?<path>.*)");

// deployment feature patterns
private static final Pattern PATTERN_RATE_RESPONSE = Pattern.compile("^/+v1/(.+?)/rate$");
private static final Pattern PATTERN_TOKENIZE = Pattern.compile("^/+v1/deployments/(.+?)/tokenize$");
private static final Pattern PATTERN_TRUNCATE_PROMPT = Pattern.compile("^/+v1/deployments/(.+?)/truncate_prompt$");
private static final Pattern PATTERN_CONFIGURATION = Pattern.compile("^/+v1/deployments/(.+?)/configuration$");
private static final Pattern PATTERN_RATE_RESPONSE = Pattern.compile("^/+v1/(?<id>.+?)/rate$");
private static final Pattern PATTERN_TOKENIZE = Pattern.compile("^/+v1/deployments/(?<id>.+?)/tokenize$");
private static final Pattern PATTERN_TRUNCATE_PROMPT = Pattern.compile("^/+v1/deployments/(?<id>.+?)/truncate_prompt$");
private static final Pattern PATTERN_CONFIGURATION = Pattern.compile("^/+v1/deployments/(?<id>.+?)/configuration$");

private static final Pattern SHARE_RESOURCE_OPERATIONS = Pattern.compile("^/v1/ops/resource/share/(create|list|discard|revoke|copy)$");
private static final Pattern INVITATIONS = Pattern.compile("^/v1/invitations$");
private static final Pattern INVITATION = Pattern.compile("^/v1/invitations/([a-zA-Z0-9]+)$");
private static final Pattern INVITATION = Pattern.compile("^/v1/invitations/(?<id>[a-zA-Z0-9]+)$");
private static final Pattern PUBLICATIONS = Pattern.compile("^/v1/ops/publication/(list|get|create|delete|approve|reject)$");
private static final Pattern PUBLISHED_RESOURCES = Pattern.compile("^/v1/ops/publication/resource/list$");
private static final Pattern PUBLICATION_RULES = Pattern.compile("^/v1/ops/publication/rule/list$");

private static final Pattern RESOURCE_OPERATIONS = Pattern.compile("^/v1/ops/resource/(move|subscribe)$");

private static final Pattern DEPLOYMENT_LIMITS = Pattern.compile("^/v1/deployments/(.+?)/limits$");
private static final Pattern DEPLOYMENT_LIMITS = Pattern.compile("^/v1/deployments/(?<id>.+?)/limits$");

private static final Pattern NOTIFICATIONS = Pattern.compile("^/v1/ops/notification/(list|delete)$");

Expand Down Expand Up @@ -295,7 +296,7 @@ private static Controller selectPost(Proxy proxy, ProxyContext context, String p
String operation = match.group(1);
ResourceOperationController controller = new ResourceOperationController(proxy, context);

return switch (operation) {
return switch (operation) {
case "move" -> controller::move;
case "subscribe" -> controller::subscribe;
default -> null;
Expand Down Expand Up @@ -364,7 +365,11 @@ private static Controller selectPut(Proxy proxy, ProxyContext context, String pa

private Matcher match(Pattern pattern, String path) {
Matcher matcher = pattern.matcher(path);
return matcher.find() ? matcher : null;
if (matcher.find()) {
SpanUtil.updateName(pattern, path);
return matcher;
}
return null;
}

private String resourcePath(String url) {
Expand Down
62 changes: 62 additions & 0 deletions src/main/java/com/epam/aidial/core/util/RegexUtil.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
package com.epam.aidial.core.util;

import lombok.experimental.UtilityClass;

import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

@UtilityClass
public class RegexUtil {

public String replaceNamedGroups(Pattern pattern, String input, List<String> groups) {
alekseyvdovenko marked this conversation as resolved.
Show resolved Hide resolved
if (groups == null || groups.isEmpty()) {
return input;
}
var regexGroups = collectGroups(pattern, input, groups);
if (regexGroups.isEmpty()) {
return input;
}
StringBuilder nameBuilder = new StringBuilder();
for (int i = 0; i < input.length(); ) {
Optional<RegexGroup> groupOpt = find(regexGroups, i);
if (groupOpt.isEmpty()) {
nameBuilder.append(input.charAt(i));
i++;
} else {
RegexGroup group = groupOpt.get();
nameBuilder.append("{").append(group.group()).append("}");
i = group.end();
}
}
return nameBuilder.toString();
}

private List<RegexGroup> collectGroups(Pattern pattern, String input, List<String> groups) {
List<RegexGroup> regexGroups = new ArrayList<>();
Matcher matcher = pattern.matcher(input);
if (matcher.matches() && matcher.groupCount() > 0) {
for (String group : groups) {
try {
int start = matcher.start(group);
int end = matcher.end(group);
regexGroups.add(new RegexGroup(group, start, end));
} catch (IllegalStateException | IllegalArgumentException ignored) {
// Ignore
alekseyvdovenko marked this conversation as resolved.
Show resolved Hide resolved
}
}
}
return regexGroups;
}

private Optional<RegexGroup> find(List<RegexGroup> groups, int position) {
return groups.stream()
.filter(g -> g.start() <= position && position < g.end())
.findFirst();
}

private record RegexGroup(String group, int start, int end) {
}
}
28 changes: 28 additions & 0 deletions src/main/java/com/epam/aidial/core/util/SpanUtil.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
package com.epam.aidial.core.util;

import io.opencensus.contrib.http.util.HttpTraceAttributeConstants;
import io.opentelemetry.api.common.AttributeKey;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.sdk.trace.ReadableSpan;
import lombok.experimental.UtilityClass;

import java.util.List;
import java.util.regex.Pattern;

@UtilityClass
public class SpanUtil {

private static final AttributeKey<String> ATTR_HTTP_METHOD = AttributeKey.stringKey(HttpTraceAttributeConstants.HTTP_METHOD);
private static final List<String> GROUPS = List.of("id", "bucket", "path");

public static void updateName(Pattern pathPattern, String path) {
alekseyvdovenko marked this conversation as resolved.
Show resolved Hide resolved
String spanName = RegexUtil.replaceNamedGroups(pathPattern, path, GROUPS);
if (Span.current() instanceof ReadableSpan span) {
String method = span.getAttribute(ATTR_HTTP_METHOD);
if (method != null) {
spanName = method + " " + spanName;
}
}
Span.current().updateName(spanName);
}
}
86 changes: 86 additions & 0 deletions src/test/java/com/epam/aidial/core/util/RegexUtilTest.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
package com.epam.aidial.core.util;

import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;

import java.util.List;
import java.util.regex.Pattern;

import static org.junit.jupiter.api.Assertions.assertEquals;

class RegexUtilTest {

@ParameterizedTest
@MethodSource("datasource")
void getName(String pathPattern, String path, String expectedName) {
var pattern = Pattern.compile(pathPattern);
var groups = List.of("id", "bucket", "path");
assertEquals(expectedName, RegexUtil.replaceNamedGroups(pattern, path, groups));
}

public static List<Arguments> datasource() {
return List.of(
Arguments.of(
"^/+openai/deployments/(?<id>.+?)/(completions|chat/completions|embeddings)$",
"/openai/deployments/gpt/chat/completions",
"/openai/deployments/{id}/chat/completions"
),
Arguments.of(
"^/+openai/deployments/(?<id>.+?)/(completions|chat/completions|embeddings)$",
"/openai/deployments/l/embeddings",
"/openai/deployments/{id}/embeddings"
),
Arguments.of(
"^/+openai/deployments/(?<id>.+?)$",
"/openai/deployments/gpt",
"/openai/deployments/{id}"
),
Arguments.of(
"^/+openai/deployments/(?<id>.+?)$",
"/openai/deployments/l",
"/openai/deployments/{id}"
),
Arguments.of(
"^/v1/bucket$",
"/v1/bucket",
"/v1/bucket"
),
Arguments.of(
"^/v1/files/(?<bucket>[a-zA-Z0-9]+)/(?<path>.*)",
"/v1/files/GHyJjv7CfrGRiv6RNajWsde7ET6bGTrbD45JatdSfsPK/path/to/file.pdf",
"/v1/files/{bucket}/{path}"
),
Arguments.of(
"^/v1/files/(?<bucket>[a-zA-Z0-9]+)/(?<path>.*)",
"/v1/files/f/i/l/e.pdf",
"/v1/files/{bucket}/{path}"
),
Arguments.of(
"^/v1/files/(?<bucket>[a-zA-Z0-9]+)/(?<path>.*)",
"/v1/files/f/f",
"/v1/files/{bucket}/{path}"
),
Arguments.of(
"^/v1/ops/resource/share/(create|list|discard|revoke|copy)$",
"/v1/ops/resource/share/list",
"/v1/ops/resource/share/list"
),
Arguments.of(
"^/v1/invitations/(?<id>[a-zA-Z0-9]+)$",
"/v1/invitations/123abc",
"/v1/invitations/{id}"
),
Arguments.of(
"^/api/(?<somethingElse>.+?)$",
"/api/123",
"/api/123"
),
Arguments.of(
"^/v1/(?<somethingElse>.+?)$",
"/api/123",
"/api/123"
)
);
}
}
alekseyvdovenko marked this conversation as resolved.
Show resolved Hide resolved