Skip to content

georgii.zorabov new parsing strategy tools #228

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.

Sign up for GitHub

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

Open
georgiizorabov wants to merge 1 commit into
base: new-parsing-strategy
Choose a base branch
from
Open

georgii.zorabov new parsing strategy tools #228

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
53 changes: 21 additions & 32 deletions ...rc/jvmMain/kotlin/org/jetbrains/ai/tracy/anthropic/adapters/AnthropicLLMTracingAdapter.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ package org.jetbrains.ai.tracy.anthropic.adapters
import org.jetbrains.ai.tracy.anthropic.model.*
import org.jetbrains.ai.tracy.core.adapters.LLMTracingAdapter
import org.jetbrains.ai.tracy.core.adapters.LLMTracingAdapter.Companion.PayloadType
import org.jetbrains.ai.tracy.core.adapters.LLMTracingAdapter.Companion.populateUnmappedAttributes
import org.jetbrains.ai.tracy.core.adapters.media.*
import org.jetbrains.ai.tracy.core.http.protocol.TracyHttpRequest
import org.jetbrains.ai.tracy.core.http.protocol.TracyHttpResponse
Expand All @@ -17,7 +16,6 @@ import org.jetbrains.ai.tracy.core.http.protocol.asJson
import org.jetbrains.ai.tracy.core.model.*
import org.jetbrains.ai.tracy.core.policy.ContentKind
import org.jetbrains.ai.tracy.core.policy.contentTracingAllowed
import org.jetbrains.ai.tracy.core.policy.orRedactedOutput
import io.opentelemetry.api.trace.Span
import io.opentelemetry.semconv.incubating.GenAiIncubatingAttributes.*
import kotlinx.serialization.json.*
Expand Down Expand Up @@ -91,9 +89,9 @@ class AnthropicLLMTracingAdapter : LLMTracingAdapter(genAISystem = GenAiSystemIn

val tools = req.tools?.map { tool ->
TracedTool(
type = "function",
name = tool.name,
description = tool.description,
type = tool.type,
parameters = tool.inputSchema?.toString(),
)
} ?: emptyList()
Expand Down Expand Up @@ -123,38 +121,29 @@ class AnthropicLLMTracingAdapter : LLMTracingAdapter(genAISystem = GenAiSystemIn
}

private fun normalizeResponse(resp: AnthropicMessagesResponse): TracedResponse {
val completions = resp.content?.map { block ->
when (block.type) {
"text" -> TracedCompletion(
type = block.type,
content = block.text,
)
var textContent: String? = null
val toolCalls = mutableListOf<TracedToolCall>()

"tool_use" -> {
// Anthropic uses flat tool attributes without a sub-index:
// gen_ai.completion.$index.tool.call.id
// gen_ai.completion.$index.tool.call.type
// gen_ai.completion.$index.tool.name
// gen_ai.completion.$index.tool.arguments
val toolExtras = buildMap {
block.id?.let { put("tool.call.id", it) }
block.type?.let { put("tool.call.type", it) }
block.name?.let { put("tool.name", it.orRedactedOutput()) }
block.input?.let { put("tool.arguments", it.toString().orRedactedOutput()) }
}
TracedCompletion(
type = block.type,
extraAttributes = toolExtras,
resp.content?.forEach { block ->
when (block.type) {
"text" -> textContent = listOfNotNull(textContent, block.text).joinToString("\n")
"tool_use" -> toolCalls.add(
TracedToolCall(
id = block.id,
type = "function",
name = block.name,
arguments = block.input?.toString(),
)
}

else -> TracedCompletion(
type = block.type,
// pass the whole block as content for unknown types
content = block.toString(),
)
}
} ?: emptyList()
}

val completion = TracedCompletion(
role = resp.role ?: "assistant",
content = textContent,
finishReason = resp.stopReason,
toolCalls = toolCalls,
)

val usage = resp.usage?.let { u ->
val extras = buildMap<String, Any?> {
Expand All @@ -177,7 +166,7 @@ class AnthropicLLMTracingAdapter : LLMTracingAdapter(genAISystem = GenAiSystemIn
return TracedResponse(
id = resp.id,
model = resp.model,
completions = completions,
completions = listOf(completion),
finishReasons = listOfNotNull(resp.stopReason),
usage = usage,
extraAttributes = extras,
Expand Down
34 changes: 13 additions & 21 deletions ...ing/anthropic/src/jvmTest/kotlin/org/jetbrains/ai/tracy/anthropic/AnthropicTracingTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,12 +89,9 @@ class AnthropicTracingTest : BaseAnthropicTracingTest() {
val finishReasons = trace.attributes[AttributeKey.stringArrayKey("gen_ai.response.finish_reasons")]
Assumptions.assumeTrue { finishReasons?.contains("tool_use") == true }

// when completion is null, then the tool call index will be 0,
// otherwise 1 (i.e., coming after the normal assistant response)
val toolCallIndex = if (completion == null) 0 else 1

val toolName = trace.attributes[AttributeKey.stringKey("gen_ai.completion.$toolCallIndex.tool.name")]
val toolArgs = trace.attributes[AttributeKey.stringKey("gen_ai.completion.$toolCallIndex.tool.arguments")]
// Tool call attributes at gen_ai.completion.0.tool.0.*
val toolName = trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.tool.0.name")]
val toolArgs = trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.tool.0.arguments")]

if (!policy.captureOutputs) {
assertEquals("REDACTED", toolName, "Tool name content should be redacted")
Expand Down Expand Up @@ -148,24 +145,19 @@ class AnthropicTracingTest : BaseAnthropicTracingTest() {

// Check tool definitions in the request
assertEquals("hi", trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.name")])
assertEquals("custom", trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.type")])
assertEquals("function", trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.type")])
assertFalse(trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.description")].isNullOrEmpty())
assertFalse(trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.parameters")].isNullOrEmpty())

// assert tool use requests when LLM finished with a tool call
// assert tool use in response — merged completion at index 0
if (trace.attributes[GEN_AI_RESPONSE_FINISH_REASONS]?.contains("tool_use") == true) {
// expect any of the indices to capture AI's tool call request
val index = listOf(0, 1).firstOrNull {
trace.attributes[AttributeKey.stringKey("gen_ai.completion.$it.tool.call.id")]?.isNotEmpty() == true
}

assertEquals("hi", trace.attributes[AttributeKey.stringKey("gen_ai.completion.$index.tool.name")])
assertEquals("hi", trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.tool.0.name")])
assertEquals(
"tool_use",
trace.attributes[AttributeKey.stringKey("gen_ai.completion.$index.tool.call.type")]
"function",
trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.tool.0.call.type")]
)
assertFalse(trace.attributes[AttributeKey.stringKey("gen_ai.completion.$index.tool.call.id")].isNullOrEmpty())
assertFalse(trace.attributes[AttributeKey.stringKey("gen_ai.completion.$index.tool.arguments")].isNullOrEmpty())
assertFalse(trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.tool.0.call.id")].isNullOrEmpty())
assertFalse(trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.tool.0.arguments")].isNullOrEmpty())
}
}

Expand Down Expand Up @@ -361,9 +353,9 @@ class AnthropicTracingTest : BaseAnthropicTracingTest() {

assertTrue(trace.attributes[AttributeKey.stringKey("gen_ai.response.model")]?.commonPrefixWith(model.asString()) == "claude-haiku-4-5")

val type = trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.type")]
assertNotNull(type)
assertTrue(type.isNotEmpty())
val role = trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.role")]
assertNotNull(role)
assertEquals("assistant", role)

val text = trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.content")]
assertNotNull(text)
Expand Down
31 changes: 4 additions & 27 deletions tracing/core/src/jvmMain/kotlin/org/jetbrains/ai/tracy/core/model/SpanSerializer.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import kotlinx.serialization.json.*
* - Generation config attributes (model, temperature, etc.)
* - System prompt (text or structured blocks)
* - Indexed messages with redaction by [ContentKind]
* - Indexed tool definitions (flat and Gemini nested functions)
* - Indexed tool definitions
* - Indexed completions with tool calls
* - Usage statistics
* - Media content extraction
Expand Down Expand Up @@ -64,18 +64,6 @@ object SpanSerializer {
tool.description?.let { span.setAttribute("gen_ai.tool.$index.description", it.orRedactedInput()) }
tool.parameters?.let { span.setAttribute("gen_ai.tool.$index.parameters", it.orRedactedInput()) }
tool.strict?.let { span.setAttribute("gen_ai.tool.$index.strict", it) }

// Gemini nested function declarations
tool.functions?.forEachIndexed { fnIdx, fn ->
fn.type?.let { span.setAttribute("gen_ai.tool.$index.function.$fnIdx.type", it) }
fn.name?.let { span.setAttribute("gen_ai.tool.$index.function.$fnIdx.name", it.orRedactedInput()) }
fn.description?.let {
span.setAttribute("gen_ai.tool.$index.function.$fnIdx.description", it.orRedactedInput())
}
fn.parameters?.let {
span.setAttribute("gen_ai.tool.$index.function.$fnIdx.parameters", it.orRedactedInput())
}
}
}

// Media content
Expand Down Expand Up @@ -183,20 +171,9 @@ object SpanSerializer {
for (tool in request.tools) {
addJsonObject {
tool.type?.let { put("type", it) }
if (tool.functions == null) {
putJsonObject("function") {
putFunctionFields(tool.name, tool.description, tool.parameters)
tool.strict?.let { put("strict", it) }
}
} else {
putJsonArray("function_declarations") {
for (fn in tool.functions) {
addJsonObject {
fn.type?.let { put("type", it) }
putFunctionFields(fn.name, fn.description, fn.parameters)
}
}
}
putJsonObject("function") {
putFunctionFields(tool.name, tool.description, tool.parameters)
tool.strict?.let { put("strict", it) }
}
}
}
Expand Down
16 changes: 1 addition & 15 deletions tracing/core/src/jvmMain/kotlin/org/jetbrains/ai/tracy/core/model/TracedRequest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,26 +70,12 @@ data class SystemBlock(val type: String?, val content: String?)

/**
* A tool definition from the request.
* All providers normalize to this flat structure with type="function".
*/
data class TracedTool(
val type: String? = null,
val name: String? = null,
val description: String? = null,
val parameters: String? = null,
val strict: String? = null,
/**
* Gemini: nested function declarations within a single tool object.
* OpenAI/Anthropic leave this null (each tool has a single function).
*/
val functions: List<TracedToolFunction>? = null,
)

/**
* A function declaration nested within a Gemini tool.
*/
data class TracedToolFunction(
val type: String? = null,
val name: String? = null,
val description: String? = null,
val parameters: String? = null,
)
5 changes: 0 additions & 5 deletions ...mini/src/jvmMain/kotlin/org/jetbrains/ai/tracy/gemini/adapters/GeminiLLMTracingAdapter.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,11 +43,6 @@ import io.opentelemetry.semconv.incubating.GenAiIncubatingAttributes.*
*/
class GeminiLLMTracingAdapter : LLMTracingAdapter(genAISystem = GenAiSystemIncubatingValues.GEMINI) {
override fun getRequestBodyAttributes(span: Span, request: TracyHttpRequest) {
val (model, operation) = request.url.modelAndOperation()

model?.let { span.setAttribute(GEN_AI_REQUEST_MODEL, model) }
operation?.let { span.setAttribute(GEN_AI_OPERATION_NAME, operation) }

val handler = selectHandler(request.url)
handler.handleRequestAttributes(span, request)
}
Expand Down
35 changes: 20 additions & 15 deletions ...jvmMain/kotlin/org/jetbrains/ai/tracy/gemini/adapters/handlers/GeminiContentGenHandler.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ class GeminiContentGenHandler(

TracedCompletion(
role = candidate.content?.role,
content = textContent ?: parts?.toString(),
content = textContent ?: if (toolCalls.isEmpty()) parts?.toString() else null,
finishReason = candidate.finishReason,
toolCalls = toolCalls,
)
Expand Down Expand Up @@ -174,18 +174,18 @@ class GeminiContentGenHandler(
val tools = body["tools"]
if (tools !is JsonArray) return emptyList()

return tools.jsonArray.map { tool ->
TracedTool(
functions = tool.jsonObject["functionDeclarations"]?.jsonArray?.map { fn ->
TracedToolFunction(
type = fn.jsonObject["parametersJsonSchema"]?.jsonObject
?.get("type")?.jsonPrimitive?.content,
name = fn.jsonObject["name"]?.jsonPrimitive?.contentOrNull,
description = fn.jsonObject["description"]?.jsonPrimitive?.contentOrNull,
parameters = fn.jsonObject["parameters"]?.toString(),
)
},
)
return tools.jsonArray.flatMap { tool ->
val fns = tool.jsonObject["functionDeclarations"]?.jsonArray
?: return@flatMap emptyList()
fns.map { fn ->
TracedTool(
type = "function",
name = fn.jsonObject["name"]?.jsonPrimitive?.contentOrNull,
description = fn.jsonObject["description"]?.jsonPrimitive?.contentOrNull,
parameters = (fn.jsonObject["parameters"]
?: fn.jsonObject["parametersJsonSchema"])?.toString(),
)
}
}
}

Expand All @@ -208,7 +208,7 @@ class GeminiContentGenHandler(
if (parts == null || parts.size != 1) return null
val item = parts.first()
if (item !is JsonObject) return null
if ("text" in item.keys) return item["text"]?.toString()
if (item.keys.size == 1 && "text" in item.keys) return item["text"]?.jsonPrimitive?.content
return null
}

Expand All @@ -217,16 +217,21 @@ class GeminiContentGenHandler(
*/
private fun extractToolCalls(parts: List<JsonElement>?): List<TracedToolCall> {
if (parts == null) return emptyList()
var callIndex = 0
return buildList {
for (part in parts) {
if (part !is JsonObject) continue
val functionCall = part["functionCall"]?.jsonObject ?: continue
val name = functionCall["name"]?.jsonPrimitive?.content
add(
TracedToolCall(
name = functionCall["name"]?.jsonPrimitive?.content,
id = "call_${name ?: callIndex}", // Gemini doesn't provide tool call IDs
type = "function",
name = name,
arguments = functionCall["args"]?.toString(),
)
)
callIndex++
}
}
}
Expand Down
24 changes: 13 additions & 11 deletions tracing/gemini/src/jvmTest/kotlin/org/jetbrains/ai/tracy/gemini/GeminiTracingTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,9 +71,9 @@ class GeminiTracingTest : BaseGeminiTracingTest() {

// input side
val prompt = trace.attributes[AttributeKey.stringKey("gen_ai.prompt.0.content")]
val name = trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.function.0.name")]
val description = trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.function.0.description")]
val parameters = trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.function.0.parameters")]
val name = trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.name")]
val description = trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.description")]
val parameters = trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.parameters")]

if (!policy.captureInputs) {
assertEquals("REDACTED", prompt, "User prompt should be redacted")
Expand Down Expand Up @@ -199,16 +199,17 @@ class GeminiTracingTest : BaseGeminiTracingTest() {
val trace = traces.first()

// assert request
assertEquals("hi", trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.function.0.name")])
assertEquals("hi", trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.name")])
assertEquals(
"Say hi to the user",
trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.function.0.description")]
trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.description")]
)
assertEquals("object", trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.function.0.type")])
assertEquals("function", trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.type")])

// assert response
assertEquals("hi", trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.tool.0.name")])
assertFalse(trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.tool.0.arguments")].isNullOrEmpty())
assertEquals("function", trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.tool.0.call.type")])
}

@Test
Expand Down Expand Up @@ -281,16 +282,17 @@ class GeminiTracingTest : BaseGeminiTracingTest() {
val trace = traces.first()

// assert request
assertEquals("hi", trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.function.0.name")])
assertEquals("hi", trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.name")])
assertEquals(
"Say hi to the user",
trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.function.0.description")]
trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.description")]
)
assertEquals("object", trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.function.0.type")])
assertEquals("function", trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.type")])

// assert response
assertEquals("hi", trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.tool.0.name")])
assertFalse(trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.tool.0.arguments")].isNullOrEmpty())
assertEquals("function", trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.tool.0.call.type")])
}

@Test
Expand Down Expand Up @@ -353,8 +355,8 @@ class GeminiTracingTest : BaseGeminiTracingTest() {
val trace = traces.first()

// Assert both tools are declared in the request
assertEquals("hi", trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.function.0.name")])
assertEquals("goodbye", trace.attributes[AttributeKey.stringKey("gen_ai.tool.1.function.0.name")])
assertEquals("hi", trace.attributes[AttributeKey.stringKey("gen_ai.tool.0.name")])
assertEquals("goodbye", trace.attributes[AttributeKey.stringKey("gen_ai.tool.1.name")])

// If function calls were made, assert both appear in completion metadata
if ((trace.attributes[AttributeKey.stringKey("gen_ai.completion.0.finish_reason")]
Expand Down
Loading