update initpy

openlit · Mar 16, 2024 · 0915604 · 0915604
1 parent 8ccb598
commit 0915604
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/src/dokumetry/__init__.py b/src/dokumetry/__init__.py
@@ -52,21 +52,19 @@ def init(llm, doku_url, api_key, environment="default", application_name="defaul
 
     # pylint: disable=no-else-return, line-too-long
     if hasattr(llm, 'moderations') and callable(llm.chat.completions.create) and ('.openai.azure.com/' not in str(llm.base_url)):
+        print("OpenAI")
         if isinstance(llm, OpenAI):
             init_openai(llm, doku_url, api_key, environment, application_name, skip_resp)
         elif isinstance(llm, AsyncOpenAI):
             init_async_openai(llm, doku_url, api_key, environment, application_name, skip_resp)
         return
     if hasattr(llm, 'moderations') and callable(llm.chat.completions.create) and ('.openai.azure.com/' in str(llm.base_url)):
         if isinstance(llm, AzureOpenAI):
-            init_openai(llm, doku_url, api_key, environment, application_name, skip_resp)
+            init_azure_openai(llm, doku_url, api_key, environment, application_name, skip_resp)
         elif isinstance(llm, AsyncAzureOpenAI):
-            init_async_openai(llm, doku_url, api_key, environment, application_name, skip_resp)
-        return
-    elif hasattr(llm, 'generate') and callable(llm.generate):
-        init_cohere(llm, doku_url, api_key, environment, application_name, skip_resp)
+            init_async_azure_openai(llm, doku_url, api_key, environment, application_name, skip_resp)
         return
-    elif isinstance(llm, MistralClient):
+    if isinstance(llm, MistralClient):
         init_mistral(llm, doku_url, api_key, environment, application_name, skip_resp)
         return
     elif isinstance(llm, MistralAsyncClient):
@@ -78,3 +76,6 @@ def init(llm, doku_url, api_key, environment="default", application_name="defaul
     elif isinstance(llm, Anthropic):
         init_anthropic(llm, doku_url, api_key, environment, application_name, skip_resp)
         return
+    elif hasattr(llm, 'generate') and callable(llm.generate):
+        init_cohere(llm, doku_url, api_key, environment, application_name, skip_resp)
+        return