custom trace connect ptaas trace (#31)

1. support custom trace connect ptaas trace
2. fix async ptaas httpclient, modify to use AsyncClient

Author: JqRrt

CHANGLOG.md

@@ -1,3 +1,8 @@
+## [0.1.16] - 2025-09-24
+### Added
+- support custom trace connect ptaas trace
+- fix async ptaas httpclient, modify to use AsyncClient
+
 ## [0.1.15] - 2025-09-17
 ### Added
 - modify cachetools version to >=5.5.2,<7.0.0

cozeloop/_client.py

@@ -181,12 +181,14 @@ def __init__(
             jwt_oauth_private_key=jwt_oauth_private_key,
             jwt_oauth_public_key_id=jwt_oauth_public_key_id
         )
+
         http_client = httpclient.Client(
             api_base_url=api_base_url,
             http_client=inner_client,
             auth=auth,
             timeout=timeout,
-            upload_timeout=upload_timeout
+            upload_timeout=upload_timeout,
+            header_injector=self._create_default_header_injector(),
         )
         finish_pro = default_finish_event_processor
         if trace_finish_event_processor:
@@ -219,6 +221,17 @@ def combined_processor(event_info: FinishEventInfo):
             prompt_trace=prompt_trace
         )
 
+    def _create_default_header_injector(self) -> Callable[[], Dict[str, str]]:
+        def default_header_injector() -> Dict[str, str]:
+            try:
+                span = self.get_span_from_context()
+                if span and hasattr(span, 'to_header'):
+                    return span.to_header()
+            except Exception:
+                pass
+            return {}
+        return default_header_injector
+
     def _get_from_env(self, val: str, env_key: str) -> str:
         if val:
             return val

cozeloop/internal/httpclient/client.py

@@ -3,7 +3,7 @@
 
 import logging
 import os
-from typing import Optional, Dict, Union, IO, Type, Tuple, Any
+from typing import Optional, Dict, Union, IO, Type, Tuple, Any, Callable
 
 import httpx
 import pydantic
@@ -29,16 +29,17 @@ def __init__(
             auth: Auth,
             timeout: int = consts.DEFAULT_TIMEOUT,
             upload_timeout: int = consts.DEFAULT_UPLOAD_TIMEOUT,
+            header_injector: Optional[Callable[[], Dict[str, str]]] = None,
     ):
         self.api_base_url = api_base_url
         self.http_client = http_client
         self.auth = auth
         self.timeout = timeout
         self.upload_timeout = upload_timeout
+        self.header_injector = header_injector
 
     def _build_url(self, path: str) -> str:
         return f"{self.api_base_url}{path}"
-
     def _set_headers(self, headers: Optional[Dict[str, str]] = None) -> Dict[str, str]:
         res = user_agent_header()
         if headers:
@@ -52,6 +53,14 @@ def _set_headers(self, headers: Optional[Dict[str, str]] = None) -> Dict[str, st
         if ppe_env:
             res["x-use-ppe"] = "1"
 
+        if self.header_injector:
+            try:
+                injected_headers = self.header_injector()
+                if injected_headers:
+                    res.update(injected_headers)
+            except Exception as e:
+                logger.debug(f"Header injection failed: {e}")
+
         return res
 
     def request(

cozeloop/internal/version.py

@@ -1,4 +1,4 @@
 # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
 # SPDX-License-Identifier: MIT
 
-VERSION = 'v0.1.11'
+VERSION = 'v0.1.16'

examples/prompt/ptaas/ptaas.py

@@ -14,6 +14,8 @@
 import asyncio
 import os
 
+from anyio import sleep
+
 from cozeloop import new_client, Client
 from cozeloop.entities.prompt import Message, Role, ExecuteResult
 
@@ -141,7 +143,8 @@ async def async_stream_example(client: Client) -> None:
 async def main():
     """Main function"""
     client = setup_client()
-    
+
+    root_span = client.start_span("root", "custom")
     try:
         # Sync non-stream call
         sync_non_stream_example(client)
@@ -157,6 +160,7 @@ async def main():
 
     finally:
         # Close client
+        root_span.finish()
         if hasattr(client, 'close'):
             client.close()
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cozeloop"
-version = "0.1.15"
+version = "0.1.16"
 description = "coze loop sdk"
 authors = ["JiangQi715 <[email protected]>"]
 license = "MIT"