-
Notifications
You must be signed in to change notification settings - Fork 276
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: add ai server proxy to builder server (#13467)
- Loading branch information
Showing
11 changed files
with
281 additions
and
6 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
67 changes: 67 additions & 0 deletions
67
...er-server/src/main/kotlin/io/airbyte/connector_builder/exceptions/AssistProxyException.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,67 @@ | ||
| @file:Suppress("ktlint:standard:package-name") | ||
|
|
||
| package io.airbyte.connector_builder.exceptions | ||
|
|
||
| import com.fasterxml.jackson.databind.JsonNode | ||
| import io.airbyte.commons.server.errors.KnownException | ||
|
|
||
| class AssistProxyException(private var responseCode: Int, jsonBody: JsonNode) : | ||
| KnownException(getStringFromResponse(jsonBody), getThrowableFromResponse(jsonBody)) { | ||
| override fun getHttpCode(): Int { | ||
| return responseCode | ||
| } | ||
| } | ||
|
|
||
| fun getStringFromResponse(jsonBody: JsonNode): String { | ||
| if (jsonBody.has("message")) { | ||
| return jsonBody.get("message").asText() | ||
| } | ||
| return "Unknown AI Assist error" | ||
| } | ||
|
|
||
| fun getThrowableFromResponse(jsonBody: JsonNode): Throwable? { | ||
| if (jsonBody.has("exceptionStack")) { | ||
| val message = getStringFromResponse(jsonBody) | ||
| val givenStack = jsonBody.get("exceptionStack") | ||
| val givenClassName = jsonBody.get("exceptionClassName")?.asText() ?: "Python" | ||
| val stackTrace = convertToStackTrace(givenStack, givenClassName) ?: return null | ||
|
|
||
| val throwable = Throwable(message) | ||
| throwable.stackTrace = stackTrace | ||
| return throwable | ||
| } | ||
| return null | ||
| } | ||
|
|
||
| fun convertToStackTrace( | ||
| exceptionStack: JsonNode, | ||
| exceptionClassName: String, | ||
| ): Array<StackTraceElement>? { | ||
| if (!exceptionStack.isArray) return null | ||
|
|
||
| // exceptionStack is an array of strings from python | ||
| return exceptionStack.mapIndexed { index, stackLine -> | ||
| val stackTraceParts = stackLine.asText().split(":") | ||
| val (fileName, lineNumber, functionName) = parseStackTraceParts(stackTraceParts, index) | ||
| StackTraceElement(exceptionClassName, functionName, fileName, lineNumber) | ||
| }.toTypedArray() | ||
| } | ||
|
|
||
| private fun parseStackTraceParts( | ||
| parts: List<String>, | ||
| index: Int, | ||
| ): Triple<String, Int, String> { | ||
| return when (parts.size) { | ||
| 3 -> Triple(parts[0], parseLineNumber(parts[1], index), parts[2]) | ||
| 2 -> Triple(parts[0], parseLineNumber(parts[1], index), "unknown_function") | ||
| 1 -> Triple("unknown_file.py", index + 1, parts[0]) | ||
| else -> Triple("unknown_file.py", index + 1, "unknown_function") | ||
| } | ||
| } | ||
|
|
||
| private fun parseLineNumber( | ||
| lineNumber: String, | ||
| index: Int, | ||
| ): Int { | ||
| return lineNumber.toIntOrNull() ?: (index + 1) | ||
| } |
34 changes: 34 additions & 0 deletions
34
...uilder-server/src/main/kotlin/io/airbyte/connector_builder/handlers/AssistProxyHandler.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,34 @@ | ||
| /* | ||
| * Copyright (c) 2020-2024 Airbyte, Inc., all rights reserved. | ||
| */ | ||
| @file:Suppress("ktlint:standard:package-name") | ||
|
|
||
| package io.airbyte.connector_builder.handlers | ||
|
|
||
| import io.airbyte.commons.json.Jsons | ||
| import io.airbyte.connector_builder.exceptions.ConnectorBuilderException | ||
| import io.airbyte.connector_builder.requester.assist.AssistConfiguration | ||
| import io.airbyte.connector_builder.requester.assist.AssistProxy | ||
| import jakarta.inject.Inject | ||
| import jakarta.inject.Singleton | ||
|
|
||
| /** | ||
| * Proxy to the Assist API. | ||
| */ | ||
| @Singleton | ||
| class AssistProxyHandler | ||
| @Inject | ||
| constructor(private val proxyConfig: AssistConfiguration) { | ||
| /** | ||
| * Call the Assistant to get connector data | ||
| */ | ||
| @Throws(ConnectorBuilderException::class) | ||
| fun process(requestBody: Map<String, Object>): Map<String, Object> { | ||
| val path = "/v1/process" | ||
| val proxy = AssistProxy(this.proxyConfig) | ||
|
|
||
| val jsonBody = Jsons.jsonNode(requestBody) | ||
| val result = proxy.post(path, jsonBody) | ||
| return Jsons.`object`(result, Map::class.java) as Map<String, Object> | ||
| } | ||
| } |
17 changes: 17 additions & 0 deletions
17
...rver/src/main/kotlin/io/airbyte/connector_builder/requester/assist/AssistConfiguration.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,17 @@ | ||
| /* | ||
| * Copyright (c) 2020-2024 Airbyte, Inc., all rights reserved. | ||
| */ | ||
| @file:Suppress("ktlint:standard:package-name") | ||
|
|
||
| package io.airbyte.connector_builder.requester.assist | ||
|
|
||
| import java.io.IOException | ||
| import java.net.HttpURLConnection | ||
|
|
||
| /** | ||
| * Proxy to the Assist Service. Blocks until the job completes. | ||
| */ | ||
| interface AssistConfiguration { | ||
| @Throws(IOException::class) | ||
| fun getConnection(path: String): HttpURLConnection | ||
| } |
39 changes: 39 additions & 0 deletions
39
.../src/main/kotlin/io/airbyte/connector_builder/requester/assist/AssistConfigurationImpl.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,39 @@ | ||
| /* | ||
| * Copyright (c) 2020-2024 Airbyte, Inc., all rights reserved. | ||
| */ | ||
| @file:Suppress("ktlint:standard:package-name") | ||
|
|
||
| package io.airbyte.connector_builder.requester.assist | ||
|
|
||
| import io.micronaut.context.annotation.Value | ||
| import jakarta.inject.Singleton | ||
| import org.jooq.tools.StringUtils | ||
| import java.io.IOException | ||
| import java.net.HttpURLConnection | ||
| import java.net.MalformedURLException | ||
| import java.net.ProtocolException | ||
| import java.net.URL | ||
|
|
||
| /** | ||
| * Construct and send requests to the CDK's Connector Builder handler. | ||
| */ | ||
| @Singleton | ||
| class AssistConfigurationImpl( | ||
| @Value("\${airbyte.connector-builder-server.ai-assist.url-base}") private val targetApiBaseUrl: String, | ||
| ) : AssistConfiguration { | ||
| @Throws(IOException::class) | ||
| override fun getConnection(path: String): HttpURLConnection { | ||
| if (StringUtils.isBlank(targetApiBaseUrl)) { | ||
| throw RuntimeException("Assist Service URL is not set.") | ||
| } | ||
| try { | ||
| val url = URL("$targetApiBaseUrl$path") | ||
| val connection = url.openConnection() as HttpURLConnection | ||
| return connection | ||
| } catch (e: ProtocolException) { | ||
| throw RuntimeException(e) | ||
| } catch (e: MalformedURLException) { | ||
| throw RuntimeException(e) | ||
| } | ||
| } | ||
| } |
66 changes: 66 additions & 0 deletions
66
...ilder-server/src/main/kotlin/io/airbyte/connector_builder/requester/assist/AssistProxy.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,66 @@ | ||
| /* | ||
| * Copyright (c) 2020-2024 Airbyte, Inc., all rights reserved. | ||
| */ | ||
| @file:Suppress("ktlint:standard:package-name") | ||
|
|
||
| package io.airbyte.connector_builder.requester.assist | ||
|
|
||
| import com.fasterxml.jackson.databind.JsonNode | ||
| import com.fasterxml.jackson.databind.ObjectMapper | ||
| import io.airbyte.connector_builder.exceptions.AssistProxyException | ||
| import io.airbyte.connector_builder.exceptions.ConnectorBuilderException | ||
| import java.io.IOException | ||
| import java.io.InputStreamReader | ||
|
|
||
| class AssistProxy(private val proxyConfig: AssistConfiguration) { | ||
| fun post( | ||
| path: String, | ||
| jsonBody: JsonNode?, | ||
| ): JsonNode { | ||
| val connection = proxyConfig.getConnection(path) | ||
| connection.apply { | ||
| requestMethod = "POST" | ||
| setRequestProperty("Content-Type", "application/json") | ||
| doOutput = true | ||
| } | ||
|
|
||
| connection.outputStream.use { outputStream -> | ||
| objectMapper.writeValue(outputStream, jsonBody) | ||
| } | ||
| val responseCode: Int | ||
| val jsonResponse: JsonNode | ||
|
|
||
| try { | ||
| responseCode = connection.responseCode | ||
| val inputStream = | ||
| if (responseCode in 200..299) { | ||
| connection.inputStream | ||
| } else { | ||
| connection.errorStream | ||
| } | ||
|
|
||
| jsonResponse = | ||
| inputStream.use { inputStream -> | ||
| InputStreamReader(inputStream, "utf-8").use { reader -> | ||
| reader.readText().let { | ||
| objectMapper.readTree(it) | ||
| } | ||
| } | ||
| } | ||
| } catch (e: IOException) { | ||
| throw ConnectorBuilderException("AI Assist processing error", e) | ||
| } finally { | ||
| connection.disconnect() | ||
| } | ||
|
|
||
| if (responseCode !in 200..299) { | ||
| throw AssistProxyException(responseCode, jsonResponse) | ||
| } | ||
|
|
||
| return jsonResponse | ||
| } | ||
|
|
||
| companion object { | ||
| private val objectMapper = ObjectMapper() | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.