Skip to content

Commit 98c7252

Browse files
author
wangjiaju.716
committed
Add buildin tool span.
1 parent d67fc9c commit 98c7252

File tree

3 files changed

+375
-131
lines changed

3 files changed

+375
-131
lines changed

‎veadk/tools/builtin_tools/image_edit.py‎

Lines changed: 117 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,11 @@
1818
from volcenginesdkarkruntime import Ark
1919
from veadk.config import getenv
2020
import base64
21-
21+
from opentelemetry import trace
22+
import traceback
23+
import json
24+
from veadk.version import VERSION
25+
from opentelemetry.trace import Span
2226
from veadk.utils.logger import get_logger
2327

2428
logger = get_logger(__name__)
@@ -100,45 +104,70 @@ async def image_edit(
100104
seed = item.get("seed", -1)
101105

102106
try:
103-
response = client.images.generate(
104-
model=getenv("MODEL_EDIT_NAME"),
105-
image=origin_image,
106-
prompt=prompt,
107-
response_format=response_format,
108-
guidance_scale=guidance_scale,
109-
watermark=watermark,
110-
seed=seed,
111-
)
112-
113-
if response.data and len(response.data) > 0:
114-
for item in response.data:
115-
if response_format == "url":
116-
image = item.url
117-
tool_context.state[f"{image_name}_url"] = image
118-
119-
elif response_format == "b64_json":
120-
image = item.b64_json
121-
image_bytes = base64.b64decode(image)
122-
123-
tool_context.state[f"{image_name}_url"] = (
124-
f"data:image/jpeg;base64,{image}"
125-
)
126-
127-
report_artifact = types.Part.from_bytes(
128-
data=image_bytes, mime_type="image/png"
129-
)
130-
await tool_context.save_artifact(image_name, report_artifact)
131-
logger.debug(f"Image saved as ADK artifact: {image_name}")
132-
133-
success_list.append({image_name: image})
134-
else:
135-
error_details = f"No images returned by Doubao model: {response}"
136-
logger.error(error_details)
137-
error_list.append(image_name)
107+
tracer = trace.get_tracer("gcp.vertex.agent")
108+
with tracer.start_as_current_span("call_llm") as span:
109+
inputs = {
110+
"prompt": prompt,
111+
"image": origin_image,
112+
"response_format": response_format,
113+
"guidance_scale": guidance_scale,
114+
"watermark": watermark,
115+
"seed": seed,
116+
}
117+
input_part = {
118+
"role": "user",
119+
"content": json.dumps(inputs, ensure_ascii=False),
120+
}
121+
response = client.images.generate(
122+
model=getenv("MODEL_EDIT_NAME"), **inputs
123+
)
124+
output_part = None
125+
if response.data and len(response.data) > 0:
126+
for item in response.data:
127+
if response_format == "url":
128+
image = item.url
129+
tool_context.state[f"{image_name}_url"] = image
130+
output_part = {
131+
"message.role": "model",
132+
"message.content": image,
133+
}
134+
elif response_format == "b64_json":
135+
image = item.b64_json
136+
image_bytes = base64.b64decode(image)
137+
138+
tool_context.state[f"{image_name}_url"] = (
139+
f"data:image/jpeg;base64,{image}"
140+
)
141+
142+
report_artifact = types.Part.from_bytes(
143+
data=image_bytes, mime_type="image/png"
144+
)
145+
await tool_context.save_artifact(
146+
image_name, report_artifact
147+
)
148+
logger.debug(f"Image saved as ADK artifact: {image_name}")
149+
150+
success_list.append({image_name: image})
151+
else:
152+
error_details = f"No images returned by Doubao model: {response}"
153+
logger.error(error_details)
154+
error_list.append(image_name)
155+
156+
add_span_attributes(
157+
span,
158+
tool_context,
159+
input_part=input_part,
160+
output_part=output_part,
161+
output_tokens=response.usage.output_tokens,
162+
total_tokens=response.usage.total_tokens,
163+
request_model=getenv("MODEL_EDIT_NAME"),
164+
response_model=getenv("MODEL_EDIT_NAME"),
165+
)
138166

139167
except Exception as e:
140168
error_details = f"No images returned by Doubao model: {e}"
141169
logger.error(error_details)
170+
traceback.print_exc()
142171
error_list.append(image_name)
143172

144173
if len(success_list) == 0:
@@ -153,3 +182,55 @@ async def image_edit(
153182
"success_list": success_list,
154183
"error_list": error_list,
155184
}
185+
186+
187+
def add_span_attributes(
188+
span: Span,
189+
tool_context: ToolContext,
190+
input_part: dict = None,
191+
output_part: dict = None,
192+
input_tokens: int = None,
193+
output_tokens: int = None,
194+
total_tokens: int = None,
195+
request_model: str = None,
196+
response_model: str = None,
197+
):
198+
try:
199+
# common attributes
200+
app_name = tool_context._invocation_context.app_name
201+
user_id = tool_context._invocation_context.user_id
202+
agent_name = tool_context.agent_name
203+
session_id = tool_context._invocation_context.session.id
204+
span.set_attribute("gen_ai.agent.name", agent_name)
205+
span.set_attribute("openinference.instrumentation.veadk", VERSION)
206+
span.set_attribute("gen_ai.app.name", app_name)
207+
span.set_attribute("gen_ai.user.id", user_id)
208+
span.set_attribute("gen_ai.session.id", session_id)
209+
span.set_attribute("agent_name", agent_name)
210+
span.set_attribute("agent.name", agent_name)
211+
span.set_attribute("app_name", app_name)
212+
span.set_attribute("app.name", app_name)
213+
span.set_attribute("user.id", user_id)
214+
span.set_attribute("session.id", session_id)
215+
span.set_attribute("cozeloop.report.source", "veadk")
216+
217+
# llm attributes
218+
span.set_attribute("gen_ai.system", "openai")
219+
span.set_attribute("gen_ai.operation.name", "chat")
220+
if request_model:
221+
span.set_attribute("gen_ai.request.model", request_model)
222+
if response_model:
223+
span.set_attribute("gen_ai.response.model", response_model)
224+
if total_tokens:
225+
span.set_attribute("gen_ai.usage.total_tokens", total_tokens)
226+
if output_tokens:
227+
span.set_attribute("gen_ai.usage.output_tokens", output_tokens)
228+
if input_tokens:
229+
span.set_attribute("gen_ai.usage.input_tokens", input_tokens)
230+
if input_part:
231+
span.add_event("gen_ai.user.message", input_part)
232+
if output_part:
233+
span.add_event("gen_ai.choice", output_part)
234+
235+
except Exception:
236+
traceback.print_exc()

‎veadk/tools/builtin_tools/image_generate.py‎

Lines changed: 116 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,11 @@
1919
from veadk.config import getenv
2020
import base64
2121
from volcenginesdkarkruntime import Ark
22-
22+
from opentelemetry import trace
23+
import traceback
24+
import json
25+
from veadk.version import VERSION
26+
from opentelemetry.trace import Span
2327
from veadk.utils.logger import get_logger
2428

2529
logger = get_logger(__name__)
@@ -101,41 +105,65 @@ async def image_generate(
101105
seed = item.get("seed", -1)
102106

103107
try:
104-
response = client.images.generate(
105-
model=getenv("MODEL_IMAGE_NAME"),
106-
prompt=prompt,
107-
response_format=response_format,
108-
size=size,
109-
guidance_scale=guidance_scale,
110-
watermark=watermark,
111-
seed=seed,
112-
)
113-
114-
if response.data and len(response.data) > 0:
115-
for item in response.data:
116-
if response_format == "url":
117-
image = item.url
118-
tool_context.state[f"{image_name}_url"] = image
119-
120-
elif response_format == "b64_json":
121-
image = item.b64_json
122-
image_bytes = base64.b64decode(image)
123-
124-
tool_context.state[f"{image_name}_url"] = (
125-
f"data:image/jpeg;base64,{image}"
126-
)
127-
128-
report_artifact = types.Part.from_bytes(
129-
data=image_bytes, mime_type="image/png"
130-
)
131-
await tool_context.save_artifact(image_name, report_artifact)
132-
logger.debug(f"Image saved as ADK artifact: {image_name}")
133-
134-
success_list.append({image_name: image})
135-
else:
136-
error_details = f"No images returned by Doubao model: {response}"
137-
logger.error(error_details)
138-
error_list.append(image_name)
108+
tracer = trace.get_tracer("gcp.vertex.agent")
109+
with tracer.start_as_current_span("call_llm") as span:
110+
inputs = {
111+
"prompt": prompt,
112+
"response_format": response_format,
113+
"size": size,
114+
"guidance_scale": guidance_scale,
115+
"watermark": watermark,
116+
"seed": seed,
117+
}
118+
input_part = {
119+
"role": "user",
120+
"content": json.dumps(inputs, ensure_ascii=False),
121+
}
122+
response = client.images.generate(
123+
model=getenv("MODEL_IMAGE_NAME"), **inputs
124+
)
125+
output_part = None
126+
if response.data and len(response.data) > 0:
127+
for item in response.data:
128+
if response_format == "url":
129+
image = item.url
130+
tool_context.state[f"{image_name}_url"] = image
131+
output_part = {
132+
"message.role": "model",
133+
"message.content": image,
134+
}
135+
elif response_format == "b64_json":
136+
image = item.b64_json
137+
image_bytes = base64.b64decode(image)
138+
139+
tool_context.state[f"{image_name}_url"] = (
140+
f"data:image/jpeg;base64,{image}"
141+
)
142+
143+
report_artifact = types.Part.from_bytes(
144+
data=image_bytes, mime_type="image/png"
145+
)
146+
await tool_context.save_artifact(
147+
image_name, report_artifact
148+
)
149+
logger.debug(f"Image saved as ADK artifact: {image_name}")
150+
151+
success_list.append({image_name: image})
152+
else:
153+
error_details = f"No images returned by Doubao model: {response}"
154+
logger.error(error_details)
155+
error_list.append(image_name)
156+
157+
add_span_attributes(
158+
span,
159+
tool_context,
160+
input_part=input_part,
161+
output_part=output_part,
162+
output_tokens=response.usage.output_tokens,
163+
total_tokens=response.usage.total_tokens,
164+
request_model=getenv("MODEL_IMAGE_NAME"),
165+
response_model=getenv("MODEL_IMAGE_NAME"),
166+
)
139167

140168
except Exception as e:
141169
error_details = f"No images returned by Doubao model: {e}"
@@ -154,3 +182,55 @@ async def image_generate(
154182
"success_list": success_list,
155183
"error_list": error_list,
156184
}
185+
186+
187+
def add_span_attributes(
188+
span: Span,
189+
tool_context: ToolContext,
190+
input_part: dict = None,
191+
output_part: dict = None,
192+
input_tokens: int = None,
193+
output_tokens: int = None,
194+
total_tokens: int = None,
195+
request_model: str = None,
196+
response_model: str = None,
197+
):
198+
try:
199+
# common attributes
200+
app_name = tool_context._invocation_context.app_name
201+
user_id = tool_context._invocation_context.user_id
202+
agent_name = tool_context.agent_name
203+
session_id = tool_context._invocation_context.session.id
204+
span.set_attribute("gen_ai.agent.name", agent_name)
205+
span.set_attribute("openinference.instrumentation.veadk", VERSION)
206+
span.set_attribute("gen_ai.app.name", app_name)
207+
span.set_attribute("gen_ai.user.id", user_id)
208+
span.set_attribute("gen_ai.session.id", session_id)
209+
span.set_attribute("agent_name", agent_name)
210+
span.set_attribute("agent.name", agent_name)
211+
span.set_attribute("app_name", app_name)
212+
span.set_attribute("app.name", app_name)
213+
span.set_attribute("user.id", user_id)
214+
span.set_attribute("session.id", session_id)
215+
span.set_attribute("cozeloop.report.source", "veadk")
216+
217+
# llm attributes
218+
span.set_attribute("gen_ai.system", "openai")
219+
span.set_attribute("gen_ai.operation.name", "chat")
220+
if request_model:
221+
span.set_attribute("gen_ai.request.model", request_model)
222+
if response_model:
223+
span.set_attribute("gen_ai.response.model", response_model)
224+
if total_tokens:
225+
span.set_attribute("gen_ai.usage.total_tokens", total_tokens)
226+
if output_tokens:
227+
span.set_attribute("gen_ai.usage.output_tokens", output_tokens)
228+
if input_tokens:
229+
span.set_attribute("gen_ai.usage.input_tokens", input_tokens)
230+
if input_part:
231+
span.add_event("gen_ai.user.message", input_part)
232+
if output_part:
233+
span.add_event("gen_ai.choice", output_part)
234+
235+
except Exception:
236+
traceback.print_exc()

0 commit comments

Comments
 (0)