Skip to content

Commit

Permalink
feat: specify span name based on matching controller
Browse files Browse the repository at this point in the history
  • Loading branch information
@alekseyvdovenko
alekseyvdovenko committed Oct 8, 2024
1 parent df98208 commit 7a8e5a4
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 18 deletions.
41 changes: 23 additions & 18 deletions src/main/java/com/epam/aidial/core/controller/ControllerSelector.java
View file
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(path, matcher);
return matcher;
}
return null;
}

private String resourcePath(String url) {
Expand Down
41 changes: 41 additions & 0 deletions src/main/java/com/epam/aidial/core/util/SpanUtil.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
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.Matcher;

@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 void updateName(String path, Matcher matcher) {
String spanName = path;
if (Span.current() instanceof ReadableSpan span) {
String method = span.getAttribute(ATTR_HTTP_METHOD);
if (method != null) {
spanName = method + " " + spanName;
}
}
if (matcher.groupCount() > 0) {
for (String group : GROUPS) {
spanName = replace(spanName, group, matcher);
}
}
Span.current().updateName(spanName);
}

private static String replace(String spanName, String group, Matcher matcher) {
try {
return spanName.replaceAll(matcher.group(group), "{%s}".formatted(group));
} catch (IllegalArgumentException ignored) {
return spanName;
}
}
}

0 comments on commit 7a8e5a4

Please sign in to comment.