diff --git a/CHANGLOG.md b/CHANGLOG.md index d927009..5961ee3 100644 --- a/CHANGLOG.md +++ b/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 diff --git a/cozeloop/_client.py b/cozeloop/_client.py index 805f48d..8327147 100644 --- a/cozeloop/_client.py +++ b/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 diff --git a/cozeloop/internal/httpclient/client.py b/cozeloop/internal/httpclient/client.py index 9965886..dee44bb 100644 --- a/cozeloop/internal/httpclient/client.py +++ b/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( diff --git a/cozeloop/internal/version.py b/cozeloop/internal/version.py index 5f90db9..2223429 100644 --- a/cozeloop/internal/version.py +++ b/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' \ No newline at end of file +VERSION = 'v0.1.16' \ No newline at end of file diff --git a/examples/prompt/ptaas/ptaas.py b/examples/prompt/ptaas/ptaas.py index 8e8b166..1d595d3 100755 --- a/examples/prompt/ptaas/ptaas.py +++ b/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() diff --git a/pyproject.toml b/pyproject.toml index a3e590c..df48804 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "cozeloop" -version = "0.1.15" +version = "0.1.16" description = "coze loop sdk" authors = ["JiangQi715 "] license = "MIT"