modify get_default_client

Author: JqRrt

cozeloop/_client.py

@@ -406,14 +406,21 @@ def set_default_client(client: Client):
 
 def get_default_client() -> Client:
     global _default_client
-    if _default_client is None:
+    temp_client = None
+    with _client_lock:
+        temp_client = _default_client
+    if temp_client is None:
         with _client_lock:
-            if _default_client is None:
-                try:
-                    _default_client = new_client()
-                    atexit.register(_graceful_shutdown)
-                except Exception as e:
-                    new_exception = e
+            temp_client = _default_client
+        if temp_client is None:
+            try:
+                temp_client = new_client()
+                with _client_lock:
+                    _default_client = temp_client
+                atexit.register(_graceful_shutdown)
+            except Exception as e:
+                new_exception = e
+                with _client_lock:
                     _default_client = _NoopClient(new_exception)
     return _default_client
 

cozeloop/decorator/decorator.py

@@ -4,8 +4,6 @@
 from typing import Optional, Callable, Any, overload, Dict, Generic, Iterator, TypeVar, List, cast, AsyncIterator
 from functools import wraps
 
-from langchain_core.runnables import RunnableLambda, RunnableConfig
-
 from cozeloop import Client, Span, start_span
 from cozeloop.decorator.utils import is_async_func, is_gen_func, is_async_gen_func, is_class_func
 
@@ -333,6 +331,7 @@ async def scorer_leader(state: MyState) -> dict | str:
         """
 
         def decorator(func: Callable):
+            from langchain_core.runnables import RunnableLambda, RunnableConfig
 
             @wraps(func)
             def sync_wrapper(*args: Any, **kwargs: Any):

cozeloop/internal/trace/span.py

@@ -398,6 +398,9 @@ def set_deployment_env(self, deployment_env: str) -> None:
         self.set_tags({DEPLOYMENT_ENV: deployment_env})
 
     def set_finish_time(self, finish_time: datetime) -> None:
+        """
+        Set the finish time of the span. DO NOT use this method unless you know what you are doing.
+        """
         self.finish_time = finish_time
 
 

cozeloop/span.py

@@ -195,7 +195,7 @@ def set_deployment_env(self, deployment_env: str) -> None:
     @abstractmethod
     def set_finish_time(self, finish_time: datetime) -> None:
         """
-        Set the finish time of the span.
+        Set the finish time of the span. DO NOT use this method unless you know what you are doing.
         """