veadk-python 0.2.7__py3-none-any.whl → 0.2.9__py3-none-any.whl

veadk/tools/builtin_tools/video_generate.py

@@ -26,14 +26,15 @@ from volcenginesdkarkruntime.types.content_generation.create_task_content_param
 )
 
 from veadk.config import getenv
+from veadk.consts import DEFAULT_MODEL_AGENT_API_BASE, DEFAULT_VIDEO_MODEL_NAME
 from veadk.utils.logger import get_logger
 from veadk.version import VERSION
 
 logger = get_logger(__name__)
 
 client = Ark(
-    api_key=getenv("MODEL_VIDEO_API_KEY"),
-    base_url=getenv("MODEL_VIDEO_API_BASE"),
+    api_key=getenv("MODEL_AGENT_API_KEY"),
+    base_url=DEFAULT_MODEL_AGENT_API_BASE,
 )
 
 
@@ -42,7 +43,7 @@ async def generate(prompt, first_frame_image=None, last_frame_image=None):
         if first_frame_image is None:
             logger.debug("text generation")
             response = client.content_generation.tasks.create(
-                model=getenv("MODEL_VIDEO_NAME"),
+                model=DEFAULT_VIDEO_MODEL_NAME,
                 content=[
                     {"type": "text", "text": prompt},
                 ],
@@ -50,7 +51,7 @@ async def generate(prompt, first_frame_image=None, last_frame_image=None):
         elif last_frame_image is None:
             logger.debug("first frame generation")
             response = client.content_generation.tasks.create(
-                model=getenv("MODEL_VIDEO_NAME"),
+                model=DEFAULT_VIDEO_MODEL_NAME,
                 content=cast(
                     list[CreateTaskContentParam],  # avoid IDE warning
                     [
@@ -65,7 +66,7 @@ async def generate(prompt, first_frame_image=None, last_frame_image=None):
         else:
             logger.debug("last frame generation")
             response = client.content_generation.tasks.create(
-                model=getenv("MODEL_VIDEO_NAME"),
+                model=DEFAULT_VIDEO_MODEL_NAME,
                 content=[
                     {"type": "text", "text": prompt},
                     {
@@ -194,19 +195,19 @@ async def video_generate(params: list, tool_context: ToolContext) -> Dict:
     batch_size = 10
     success_list = []
     error_list = []
-    tracer = trace.get_tracer("gcp.vertex.agent")
-    with tracer.start_as_current_span("call_llm") as span:
-        input_part = {"role": "user"}
-        output_part = {"message.role": "model"}
-
-        for idx, item in enumerate(params):
-            input_part[f"parts.{idx}.type"] = "text"
-            input_part[f"parts.{idx}.text"] = json.dumps(item, ensure_ascii=False)
-
-        for start_idx in range(0, len(params), batch_size):
-            batch = params[start_idx : start_idx + batch_size]
-            task_dict = {}
+
+    for start_idx in range(0, len(params), batch_size):
+        batch = params[start_idx : start_idx + batch_size]
+        task_dict = {}
+        tracer = trace.get_tracer("gcp.vertex.agent")
+        with tracer.start_as_current_span("call_llm") as span:
+            input_part = {"role": "user"}
+            output_part = {"message.role": "model"}
+            total_tokens = 0
             for idx, item in enumerate(batch):
+                input_part[f"parts.{idx}.type"] = "text"
+                input_part[f"parts.{idx}.text"] = json.dumps(item, ensure_ascii=False)
+
                 video_name = item["video_name"]
                 prompt = item["prompt"]
                 first_frame = item.get("first_frame", None)
@@ -223,7 +224,6 @@ async def video_generate(params: list, tool_context: ToolContext) -> Dict:
                     logger.error(f"Error: {e}")
                     error_list.append(video_name)
 
-            total_tokens = 0
             while True:
                 task_list = list(task_dict.keys())
                 if len(task_list) == 0:
@@ -256,29 +256,29 @@ async def video_generate(params: list, tool_context: ToolContext) -> Dict:
                         )
                 time.sleep(10)
 
-        add_span_attributes(
-            span,
-            tool_context,
-            input_part=input_part,
-            output_part=output_part,
-            output_tokens=total_tokens,
-            total_tokens=total_tokens,
-            request_model=getenv("MODEL_VIDEO_NAME"),
-            response_model=getenv("MODEL_VIDEO_NAME"),
-        )
-
-        if len(success_list) == 0:
-            return {
-                "status": "error",
-                "success_list": success_list,
-                "error_list": error_list,
-            }
-        else:
-            return {
-                "status": "success",
-                "success_list": success_list,
-                "error_list": error_list,
-            }
+            add_span_attributes(
+                span,
+                tool_context,
+                input_part=input_part,
+                output_part=output_part,
+                output_tokens=total_tokens,
+                total_tokens=total_tokens,
+                request_model=DEFAULT_VIDEO_MODEL_NAME,
+                response_model=DEFAULT_VIDEO_MODEL_NAME,
+            )
+
+    if len(success_list) == 0:
+        return {
+            "status": "error",
+            "success_list": success_list,
+            "error_list": error_list,
+        }
+    else:
+        return {
+            "status": "success",
+            "success_list": success_list,
+            "error_list": error_list,
+        }
 
 
 def add_span_attributes(

veadk/tools/builtin_tools/web_scraper.py

@@ -72,5 +72,5 @@ def web_scraper(query: str) -> dict[str, Any]:
         return results_str
 
     except requests.exceptions.RequestException as e:
-        error_message = f"Error: {str(e)}"
+        error_message = f"Error: {str(e)}, response: {response}"
         raise ValueError(error_message)

veadk/tools/builtin_tools/web_search.py

@@ -180,13 +180,13 @@ def web_search(query: str) -> list[str]:
         "WebSearch",
         json.dumps(req),
     )
+
     try:
         results: list = response_body["Result"]["WebResults"]
+        final_results = []
+        for result in results:
+            final_results.append(result["Summary"].strip())
+        return final_results
     except Exception as e:
-        logger.error(f"Web search failed: {e}")
-        return []
-
-    final_results = []
-    for result in results:
-        final_results.append(result["Summary"].strip())
-    return final_results
+        logger.error(f"Web search failed {e}, response body: {response_body}")
+        return [response_body]

veadk/tools/load_knowledgebase_tool.py

@@ -24,6 +24,7 @@ from pydantic import BaseModel, Field
 from typing_extensions import override
 
 from veadk.knowledgebase import KnowledgeBase
+from veadk.knowledgebase.entry import KnowledgebaseEntry
 
 if TYPE_CHECKING:
     from google.adk.models.llm_request import LlmRequest
@@ -32,10 +33,6 @@ if TYPE_CHECKING:
 knowledgebase: KnowledgeBase | None = None
 
 
-class KnowledgebaseEntry(BaseModel):
-    content: str
-
-
 class LoadKnowledgebaseResponse(BaseModel):
     knowledges: list[KnowledgebaseEntry] = Field(default_factory=list)
 
@@ -55,10 +52,7 @@ async def search_knowledgebase(
 ) -> SearchKnowledgebaseResponse:
     """Searches the knowledgebase of the current user."""
     if isinstance(knowledgebase, KnowledgeBase):
-        res = knowledgebase.search(query, app_name=app_name)
-        entry_list = []
-        for r in res:
-            entry_list.append(KnowledgebaseEntry(content=r))
+        entry_list = knowledgebase.search(query)
         return SearchKnowledgebaseResponse(knowledges=entry_list)
     else:
         return SearchKnowledgebaseResponse(knowledges=[])

veadk/tracing/telemetry/attributes/extractors/llm_attributes_extractors.py

@@ -244,8 +244,26 @@ def llm_gen_ai_messages(params: LLMAttributesParams) -> ExtractorResponse:
                                     )
                             if part.inline_data:
                                 if len(content.parts) == 1:
-                                    # TODO(qingliang)
-                                    pass
+                                    part = content.parts[0]
+                                    user_event["gen_ai.user.message"].update(
+                                        {
+                                            "parts.0.type": "image_url",
+                                            "parts.0.image_url.name": (
+                                                part.inline_data.display_name.split(
+                                                    "/"
+                                                )[-1]
+                                                if part.inline_data
+                                                and part.inline_data.display_name
+                                                else "<unknown_image_name>"
+                                            ),
+                                            "parts.0.image_url.url": (
+                                                part.inline_data.display_name
+                                                if part.inline_data
+                                                and part.inline_data.display_name
+                                                else "<unknown_image_url>"
+                                            ),
+                                        }
+                                    )
                                 else:
                                     user_event["gen_ai.user.message"].update(
                                         {
@@ -510,7 +528,7 @@ def llm_gen_ai_request_functions(params: LLMAttributesParams) -> ExtractorRespon
                 f"gen_ai.request.functions.{idx}.name": tool_instance.name,
                 f"gen_ai.request.functions.{idx}.description": tool_instance.description,
                 f"gen_ai.request.functions.{idx}.parameters": str(
-                    tool_instance._get_declaration().parameters.model_dump(  # type: ignore
+                    tool_instance._get_declaration().parameters.model_dump_json(  # type: ignore
                         exclude_none=True
                     )
                     if tool_instance._get_declaration()

veadk/tracing/telemetry/exporters/apmplus_exporter.py

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import time
 from dataclasses import dataclass
 from typing import Any
 
+from google.adk.agents.invocation_context import InvocationContext
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
-from opentelemetry import metrics
+from opentelemetry import metrics, trace
 from opentelemetry import metrics as metrics_api
 from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
@@ -193,7 +195,13 @@ class MeterUploader:
             explicit_bucket_boundaries_advisory=_GEN_AI_SERVER_TIME_PER_OUTPUT_TOKEN_BUCKETS,
         )
 
-    def record(self, llm_request: LlmRequest, llm_response: LlmResponse) -> None:
+    def record(
+        self,
+        invocation_context: InvocationContext,
+        event_id: str,
+        llm_request: LlmRequest,
+        llm_response: LlmResponse,
+    ) -> None:
         attributes = {
             "gen_ai_system": "volcengine",
             "gen_ai_response_model": llm_request.model,
@@ -217,10 +225,18 @@ class MeterUploader:
                 token_attributes = {**attributes, "gen_ai_token_type": "output"}
                 self.token_usage.record(output_token, attributes=token_attributes)
 
-            # TODO: Get llm duration
-            # duration = 5.0
-            # if self.duration_histogram:
-            #     self.duration_histogram.record(duration, attributes=attributes)
+            # Get llm duration
+            span = trace.get_current_span()
+            if span and hasattr(span, "start_time") and self.duration_histogram:
+                # We use span start time as the llm request start time
+                tik = span.start_time  # type: ignore
+                # We use current time as the llm request end time
+                tok = time.time_ns()
+                # Calculate duration in seconds
+                duration = (tok - tik) / 1e9
+                self.duration_histogram.record(
+                    duration, attributes=attributes
+                )  # unit in seconds
 
             # Get model request error
             if llm_response.error_code and self.chat_exception_counter:
@@ -269,6 +285,8 @@ class APMPlusExporter(BaseExporter):
     config: APMPlusExporterConfig = Field(default_factory=APMPlusExporterConfig)
 
     def model_post_init(self, context: Any) -> None:
+        logger.info(f"APMPlusExporter sevice name: {self.config.service_name}")
+
         headers = {
             "x-byteapm-appkey": self.config.app_key,
         }

veadk/tracing/telemetry/exporters/cozeloop_exporter.py

@@ -42,6 +42,8 @@ class CozeloopExporter(BaseExporter):
     config: CozeloopExporterConfig = Field(default_factory=CozeloopExporterConfig)
 
     def model_post_init(self, context: Any) -> None:
+        logger.info(f"CozeloopExporter space ID: {self.config.space_id}")
+
         headers = {
             "cozeloop-workspace-id": self.config.space_id,
             "authorization": f"Bearer {self.config.token}",

veadk/tracing/telemetry/exporters/inmemory_exporter.py

@@ -75,11 +75,22 @@ class _InMemoryExporter(export.SpanExporter):
 class _InMemorySpanProcessor(export.SimpleSpanProcessor):
     def __init__(self, exporter: _InMemoryExporter) -> None:
         super().__init__(exporter)
-        self.spans = []
 
     def on_start(self, span, parent_context) -> None:
-        if span.context:
-            self.spans.append(span)
+        if span.name.startswith("invocation"):
+            span.set_attribute("gen_ai.operation.name", "chain")
+            span.set_attribute("gen_ai.usage.total_tokens", 0)
+
+            ctx = set_value("invocation_span_instance", span, context=parent_context)
+            token = attach(ctx)  # mount context on `invocation` root span in Google ADK
+            setattr(span, "_invocation_token", token)  # for later detach
+
+        if span.name.startswith("agent_run"):
+            span.set_attribute("gen_ai.operation.name", "agent")
+
+            ctx = set_value("agent_run_span_instance", span, context=parent_context)
+            token = attach(ctx)
+            setattr(span, "_agent_run_token", token)  # for later detach
 
     def on_end(self, span: ReadableSpan) -> None:
         if span.context:
@@ -92,8 +103,14 @@ class _InMemorySpanProcessor(export.SimpleSpanProcessor):
             except Exception:
                 logger.exception("Exception while exporting Span.")
             detach(token)
-            if span in self.spans:
-                self.spans.remove(span)
+
+            token = getattr(span, "_invocation_token", None)
+            if token:
+                detach(token)
+
+            token = getattr(span, "_agent_run_token", None)
+            if token:
+                detach(token)
 
 
 class InMemoryExporter(BaseExporter):
@@ -106,6 +123,3 @@ class InMemoryExporter(BaseExporter):
 
         self._exporter = _InMemoryExporter()
         self.processor = _InMemorySpanProcessor(self._exporter)
-
-
-_INMEMORY_EXPORTER_INSTANCE = InMemoryExporter()

veadk/tracing/telemetry/exporters/tls_exporter.py

@@ -44,6 +44,8 @@ class TLSExporter(BaseExporter):
     config: TLSExporterConfig = Field(default_factory=TLSExporterConfig)
 
     def model_post_init(self, context: Any) -> None:
+        logger.info(f"TLSExporter topic ID: {self.config.topic_id}")
+
         headers = {
             "x-tls-otel-tracetopic": self.config.topic_id,
             "x-tls-otel-ak": self.config.access_key,

veadk/tracing/telemetry/opentelemetry_tracer.py

@@ -19,10 +19,9 @@ import time
 from typing import Any
 
 from opentelemetry import trace as trace_api
-from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk import trace as trace_sdk
 from opentelemetry.sdk.resources import Resource
-from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace import TracerProvider, SpanLimits
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 from typing_extensions import override
@@ -30,12 +29,10 @@ from typing_extensions import override
 from veadk.tracing.base_tracer import BaseTracer
 from veadk.tracing.telemetry.exporters.apmplus_exporter import APMPlusExporter
 from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter
-from veadk.tracing.telemetry.exporters.inmemory_exporter import (
-    _INMEMORY_EXPORTER_INSTANCE,
-    InMemoryExporter,
-)
+from veadk.tracing.telemetry.exporters.inmemory_exporter import InMemoryExporter
 from veadk.utils.logger import get_logger
 from veadk.utils.patches import patch_google_adk_telemetry
+from veadk.utils.misc import get_temp_dir
 
 logger = get_logger(__name__)
 
@@ -82,13 +79,19 @@ class OpentelemetryTracer(BaseModel, BaseTracer):
 
     def _init_global_tracer_provider(self) -> None:
         # set provider anyway, then get global provider
-        trace_api.set_tracer_provider(trace_sdk.TracerProvider())
+        trace_api.set_tracer_provider(
+            trace_sdk.TracerProvider(
+                span_limits=SpanLimits(
+                    max_attributes=4096,
+                )
+            )
+        )
         global_tracer_provider: TracerProvider = trace_api.get_tracer_provider()  # type: ignore
 
         span_processors = global_tracer_provider._active_span_processor._span_processors
         have_apmplus_exporter = any(
             isinstance(p, (BatchSpanProcessor, SimpleSpanProcessor))
-            and isinstance(p.span_exporter, OTLPSpanExporter)
+            and hasattr(p.span_exporter, "_endpoint")
             and "apmplus" in p.span_exporter._endpoint
             for p in span_processors
         )
@@ -119,7 +122,7 @@ class OpentelemetryTracer(BaseModel, BaseTracer):
                     f"Add span processor for exporter `{exporter.__class__.__name__}` to OpentelemetryTracer failed."
                 )
 
-        self._inmemory_exporter = _INMEMORY_EXPORTER_INSTANCE
+        self._inmemory_exporter = InMemoryExporter()
         if self._inmemory_exporter.processor:
             # make sure the in memory exporter processor is added at index 0
             # because we use this to record all spans
@@ -159,7 +162,7 @@ class OpentelemetryTracer(BaseModel, BaseTracer):
         self,
         user_id: str = "unknown_user_id",
         session_id: str = "unknown_session_id",
-        path: str = "/tmp",
+        path: str = get_temp_dir(),
     ) -> str:
         def _build_trace_file_path(path: str, user_id: str, session_id: str) -> str:
             return f"{path}/{self.name}_{user_id}_{session_id}_{self.trace_id}.json"

veadk/tracing/telemetry/telemetry.py

@@ -20,7 +20,8 @@ from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
 from google.adk.tools import BaseTool
 from opentelemetry import trace
-from opentelemetry.sdk.trace import _Span
+from opentelemetry.context import get_value
+from opentelemetry.sdk.trace import Span, _Span
 
 from veadk.tracing.telemetry.attributes.attributes import ATTRIBUTES
 from veadk.tracing.telemetry.attributes.extractors.types import (
@@ -28,16 +29,14 @@ from veadk.tracing.telemetry.attributes.extractors.types import (
     LLMAttributesParams,
     ToolAttributesParams,
 )
-from veadk.tracing.telemetry.exporters.inmemory_exporter import (
-    _INMEMORY_EXPORTER_INSTANCE,
-)
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
 
 
-def upload_metrics(
+def _upload_metrics(
     invocation_context: InvocationContext,
+    event_id: str,
     llm_request: LlmRequest,
     llm_response: LlmResponse,
 ) -> None:
@@ -48,11 +47,13 @@ def upload_metrics(
         for tracer in tracers:
             for exporter in getattr(tracer, "exporters", []):
                 if getattr(exporter, "meter_uploader", None):
-                    exporter.meter_uploader.record(llm_request, llm_response)
+                    exporter.meter_uploader.record(
+                        invocation_context, event_id, llm_request, llm_response
+                    )
 
 
 def _set_agent_input_attribute(
-    span: _Span, invocation_context: InvocationContext
+    span: Span, invocation_context: InvocationContext
 ) -> None:
     # We only save the original user input as the agent input
     # hence once the `agent.input` has been set, we don't overwrite it
@@ -106,7 +107,7 @@ def _set_agent_input_attribute(
                 )
 
 
-def _set_agent_output_attribute(span: _Span, llm_response: LlmResponse) -> None:
+def _set_agent_output_attribute(span: Span, llm_response: LlmResponse) -> None:
     content = llm_response.content
     if content and content.parts:
         for idx, part in enumerate(content.parts):
@@ -126,67 +127,64 @@ def set_common_attributes_on_model_span(
     current_span: _Span,
     **kwargs,
 ) -> None:
-    if current_span.context:
-        current_span_id = current_span.context.trace_id
-    else:
-        logger.warning(
-            "Current span context is missing, failed to get `trace_id` to set common attributes."
-        )
-        return
-
+    common_attributes = ATTRIBUTES.get("common", {})
     try:
-        spans = _INMEMORY_EXPORTER_INSTANCE.processor.spans  # type: ignore
-
-        spans_in_current_trace = [
-            span
-            for span in spans
-            if span.context and span.context.trace_id == current_span_id
-        ]
-
-        common_attributes = ATTRIBUTES.get("common", {})
-        for span in spans_in_current_trace:
-            if span.is_recording():
-                if span.name.startswith("invocation"):
-                    span.set_attribute("gen_ai.operation.name", "chain")
-                    _set_agent_input_attribute(span, invocation_context)
-                    _set_agent_output_attribute(span, llm_response)
-                elif span.name.startswith("agent_run"):
-                    span.set_attribute("gen_ai.operation.name", "agent")
-                    _set_agent_input_attribute(span, invocation_context)
-                    _set_agent_output_attribute(span, llm_response)
-                for attr_name, attr_extractor in common_attributes.items():
-                    value = attr_extractor(**kwargs)
-                    span.set_attribute(attr_name, value)
+        invocation_span: Span = get_value("invocation_span_instance")  # type: ignore
+        agent_run_span: Span = get_value("agent_run_span_instance")  # type: ignore
+
+        if invocation_span and invocation_span.name.startswith("invocation"):
+            _set_agent_input_attribute(invocation_span, invocation_context)
+            _set_agent_output_attribute(invocation_span, llm_response)
+            for attr_name, attr_extractor in common_attributes.items():
+                value = attr_extractor(**kwargs)
+                invocation_span.set_attribute(attr_name, value)
+
+            # Calculate the token usage for the whole invocation span
+            current_step_token_usage = (
+                llm_response.usage_metadata.total_token_count
+                if llm_response.usage_metadata
+                and llm_response.usage_metadata.total_token_count
+                else 0
+            )
+            prev_total_token_usage = (
+                invocation_span.attributes["gen_ai.usage.total_tokens"]
+                if invocation_span.attributes
+                else 0
+            )
+            accumulated_total_token_usage = (
+                current_step_token_usage + int(prev_total_token_usage)  # type: ignore
+            )  # we can ignore this warning, cause we manually set the attribute to int before
+            invocation_span.set_attribute(
+                "gen_ai.usage.total_tokens", accumulated_total_token_usage
+            )
+
+        if agent_run_span and agent_run_span.name.startswith("agent_run"):
+            _set_agent_input_attribute(agent_run_span, invocation_context)
+            _set_agent_output_attribute(agent_run_span, llm_response)
+            for attr_name, attr_extractor in common_attributes.items():
+                value = attr_extractor(**kwargs)
+                agent_run_span.set_attribute(attr_name, value)
+
+        for attr_name, attr_extractor in common_attributes.items():
+            value = attr_extractor(**kwargs)
+            current_span.set_attribute(attr_name, value)
     except Exception as e:
         logger.error(f"Failed to set common attributes for spans: {e}")
 
 
 def set_common_attributes_on_tool_span(current_span: _Span) -> None:
-    # find parent span (generally a llm span)
-    if not current_span.context:
-        logger.warning(
-            f"Get tool span's context failed. Skip setting common attributes for span {current_span.name}"
-        )
-        return
-
-    if not current_span.parent:
-        logger.warning(
-            f"Get tool span's parent failed. Skip setting common attributes for span {current_span.name}"
-        )
-        return
-
-    parent_span_id = current_span.parent.span_id
-    for span in _INMEMORY_EXPORTER_INSTANCE.processor.spans:  # type: ignore
-        if span.context.span_id == parent_span_id:
-            common_attributes = ATTRIBUTES.get("common", {})
-            for attr_name in common_attributes.keys():
-                if hasattr(span.attributes, attr_name):
-                    current_span.set_attribute(attr_name, span.attributes[attr_name])
-                else:
-                    logger.error(f"Parent span does not have attribute {attr_name}")
+    common_attributes = ATTRIBUTES.get("common", {})
 
+    invocation_span: Span = get_value("invocation_span_instance")  # type: ignore
 
-def trace_send_data(): ...
+    for attr_name in common_attributes.keys():
+        if (
+            invocation_span
+            and invocation_span.name.startswith("invocation")
+            and invocation_span.attributes
+            and attr_name in invocation_span.attributes
+        ):
+            current_span.set_attribute(attr_name, invocation_span.attributes[attr_name])
 
 
 def trace_tool_call(
@@ -212,7 +210,7 @@ def trace_call_llm(
     llm_request: LlmRequest,
     llm_response: LlmResponse,
 ) -> None:
-    span = trace.get_current_span()
+    span: Span = trace.get_current_span()  # type: ignore
 
     from veadk.agent import Agent
 
@@ -234,6 +232,7 @@ def trace_call_llm(
             span.context.trace_state.get("call_type", "")
             if (
                 hasattr(span, "context")
+                and span.context
                 and hasattr(span.context, "trace_state")
                 and hasattr(span.context.trace_state, "get")
             )
@@ -253,4 +252,8 @@ def trace_call_llm(
         response: ExtractorResponse = attr_extractor(params)
         ExtractorResponse.update_span(span, attr_name, response)
 
-    upload_metrics(invocation_context, llm_request, llm_response)
+    _upload_metrics(invocation_context, event_id, llm_request, llm_response)
+
+
+# Do not modify this function
+def trace_send_data(): ...

veadk/utils/misc.py

@@ -148,3 +148,18 @@ def set_envs(config_yaml_path: str) -> tuple[dict, dict]:
         os.environ[k] = str(v)
 
     return config_dict, veadk_environments
+
+
+def get_temp_dir():
+    """
+    Return the corresponding temporary directory based on the operating system
+    - For Windows systems, return the system's default temporary directory
+    - For other systems (macOS, Linux, etc.), return the /tmp directory
+    """
+    # First determine if it is a Windows system
+    if sys.platform.startswith("win"):
+        # Windows systems use the temporary directory from environment variables
+        return os.environ.get("TEMP", os.environ.get("TMP", r"C:\WINDOWS\TEMP"))
+    else:
+        # Non-Windows systems (macOS, Linux, etc.) uniformly return /tmp
+        return "/tmp"

veadk/version.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-VERSION = "0.2.7"
+VERSION = "0.2.9"