Improve mlchain_context

mlchain/__init__.py

@@ -1,5 +1,15 @@
 from os import environ
 
+# Mlchain Context
+from contextvars import ContextVar
+from typing import Any, Dict
+
+_request_scope_context_storage: ContextVar[Dict[Any, Any]] = ContextVar(
+    "mlchain_context"
+)
+from .context import mlchain_context
+
+# Gevent fix
 if "DISABLE_GEVENT_FIX" not in environ:
     # Fix gevent
     try:
@@ -18,28 +28,18 @@
         # Handle target environment that doesn't support HTTPS verification
         ssl._create_default_https_context = _create_unverified_https_context
 
-# Mlchain Context
-from contextvars import ContextVar
-from typing import Any, Dict
-
-_request_scope_context_storage: ContextVar[Dict[Any, Any]] = ContextVar(
-    "mlchain_context"
-)
 
 # Parameters of MLchain
-__version__ = "0.3.1"
+__version__ = "0.3.2"
 
 HOST = "https://www.api.mlchain.ml"
 WEB_HOST = HOST
 API_ADDRESS = HOST
 MODEL_ID = None
 
-from os import environ
-
 environ['OBJC_DISABLE_INITIALIZE_FORK_SAFETY'] = 'YES'
 
 from mlchain.base.log import logger
-from .context import mlchain_context
 
 from .base.exceptions import *
 from .config import mlconfig

mlchain/server/flask_server.py

@@ -74,6 +74,7 @@ def __get_msgpack_blosc_response(self, output, status=200):
 
     def init_context(self): 
         uid = str(uuid4())
+        mlchain_context.set_mlchain_context_id(uid)
         mlchain_context['MLCHAIN_CONTEXT_ID'] = uid
         return uid 
 

mlchain/server/starlette_server.py

@@ -79,6 +79,7 @@ def __get_msgpack_blosc_response(self, output, status=200):
 
     def init_context(self): 
         uid = str(uuid4())
+        mlchain_context.set_mlchain_context_id(uid)
         mlchain_context['MLCHAIN_CONTEXT_ID'] = uid
         return uid 
 

mlchain/server/view.py

@@ -52,14 +52,16 @@ def init_context_with_headers(self, headers, context_id:str = None):
         mlchain_context.set(context)
 
         if mlchain_context.MLCHAIN_CONTEXT_ID is None: 
+            mlchain_context.set_mlchain_context_id(context_id)
             mlchain_context['MLCHAIN_CONTEXT_ID'] = context_id 
         else: 
-            context_id = mlchain_context['MLCHAIN_CONTEXT_ID']
+            context_id = mlchain_context.MLCHAIN_CONTEXT_ID
 
         return context_id
 
     def init_context(self): 
         uid = str(uuid4())
+        mlchain_context.set_mlchain_context_id(uid)
         mlchain_context['MLCHAIN_CONTEXT_ID'] = uid
         return uid 
 

setup.py

@@ -5,7 +5,7 @@
 import os
 from setuptools import setup, find_packages
 
-__version__ = "0.3.1"
+__version__ = "0.3.2"
 
 project = "mlchain"