modify header_injector

Author: JqRrt

cozeloop/_client.py

@@ -74,7 +74,6 @@ def new_client(
         tag_truncate_conf: Optional[TagTruncateConf] = None,
         api_base_path: Optional[APIBasePath] = None,
         trace_queue_conf: Optional[QueueConf] = None,
-        header_injector: Optional[Callable[[], Dict[str, str]]] = None,
 ) -> Client:
     cache_key = _generate_cache_key(  # all args are used to generate cache key
         api_base_url,
@@ -94,7 +93,6 @@ def new_client(
         tag_truncate_conf,
         api_base_path,
         trace_queue_conf,
-        header_injector,
     )
 
     with _cache_lock:
@@ -120,7 +118,6 @@ def new_client(
             tag_truncate_conf=tag_truncate_conf,
             api_base_path=api_base_path,
             trace_queue_conf=trace_queue_conf,
-            header_injector=header_injector,
         )
         _client_cache[cache_key] = client
         return client
@@ -152,7 +149,6 @@ def __init__(
             tag_truncate_conf: Optional[TagTruncateConf] = None,
             api_base_path: Optional[APIBasePath] = None,
             trace_queue_conf: Optional[QueueConf] = None,
-            header_injector: Optional[Callable[[], Dict[str, str]]] = None,
     ):
         workspace_id = self._get_from_env(workspace_id, ENV_WORKSPACE_ID)
         api_base_url = self._get_from_env(api_base_url, ENV_API_BASE_URL)
@@ -185,18 +181,14 @@ def __init__(
             jwt_oauth_private_key=jwt_oauth_private_key,
             jwt_oauth_public_key_id=jwt_oauth_public_key_id
         )
-        
-        # 创建默认的header注入函数
-        if header_injector is None:
-            header_injector = self._create_default_header_injector()
-        
+
         http_client = httpclient.Client(
             api_base_url=api_base_url,
             http_client=inner_client,
             auth=auth,
             timeout=timeout,
             upload_timeout=upload_timeout,
-            header_injector=header_injector,
+            header_injector=self._create_default_header_injector(),
         )
         finish_pro = default_finish_event_processor
         if trace_finish_event_processor:
@@ -230,14 +222,12 @@ def combined_processor(event_info: FinishEventInfo):
         )
 
     def _create_default_header_injector(self) -> Callable[[], Dict[str, str]]:
-        """创建默认的header注入函数，包含span header注入逻辑"""
         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

cozeloop/internal/httpclient/client.py

@@ -42,17 +42,6 @@ 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()
-        
-        # 调用传入的header注入函数
-        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}")
-        
         if headers:
             res.update(headers)
         res[consts.AUTHORIZE_HEADER] = f"Bearer {self.auth.token}"
@@ -64,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(

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()