get attr in all openai patchees

GreptimeTeam · May 7, 2024 · f615c5e · f615c5e
1 parent b4979f5
commit f615c5e
Show file tree

Hide file tree

Showing 9 changed files with 52 additions and 47 deletions.
diff --git a/src/greptimeai/patchee/openai_patchee/chat_completion.py b/src/greptimeai/patchee/openai_patchee/chat_completion.py
@@ -4,32 +4,33 @@
 from openai import AsyncOpenAI, OpenAI
 
 from greptimeai.patchee import Patchee
+from greptimeai.utils.attr import get_attr, get_optional_attr
 
 from . import _SPAN_NAME_COMPLETION, OpenaiPatchees
 
 
 class ChatCompletionPatchees(OpenaiPatchees):
     def __init__(self, client: Union[OpenAI, AsyncOpenAI, None] = None):
         chat_completion_create = Patchee(
-            obj=client.chat.completions if client else openai.chat.completions,
+            obj=get_optional_attr([client, openai], ["chat", "completions"]),
             method_name="create",
             span_name=_SPAN_NAME_COMPLETION,
             event_name="chat.completions.create",
         )
 
         chat_raw_completion_create = Patchee(
-            obj=client.chat.with_raw_response.completions
-            if client
-            else openai.chat.with_raw_response.completions,
+            obj=get_optional_attr(
+                [client, openai], ["chat", "with_raw_response", "completions"]
+            ),
             method_name="create",
             span_name=_SPAN_NAME_COMPLETION,
             event_name="chat.with_raw_response.completions.create",
         )
 
         chat_completion_raw_create = Patchee(
-            obj=client.chat.completions.with_raw_response
-            if client
-            else openai.chat.completions.with_raw_response,
+            obj=get_optional_attr(
+                [client, openai], ["chat", "completions", "with_raw_response"]
+            ),
             method_name="create",
             span_name=_SPAN_NAME_COMPLETION,
             event_name="chat.completions.with_raw_response.create",
@@ -43,7 +44,7 @@ def __init__(self, client: Union[OpenAI, AsyncOpenAI, None] = None):
 
         if client:
             raw_chat_completion_create = Patchee(
-                obj=client.with_raw_response.chat.completions,
+                obj=get_attr(client, ["with_raw_response", "chat", "completions"]),
                 method_name="create",
                 span_name=_SPAN_NAME_COMPLETION,
                 event_name="with_raw_response.chat.completions.create",

diff --git a/src/greptimeai/patchee/openai_patchee/completion.py b/src/greptimeai/patchee/openai_patchee/completion.py
@@ -4,23 +4,24 @@
 from openai import AsyncOpenAI, OpenAI
 
 from greptimeai.patchee import Patchee
+from greptimeai.utils.attr import get_attr, get_optional_attr
 
 from . import _SPAN_NAME_COMPLETION, OpenaiPatchees
 
 
 class CompletionPatchees(OpenaiPatchees):
     def __init__(self, client: Union[OpenAI, AsyncOpenAI, None] = None):
         completion_create = Patchee(
-            obj=client.completions if client else openai.completions,
+            obj=get_optional_attr([client, openai], ["completions"]),
             method_name="create",
             span_name=_SPAN_NAME_COMPLETION,
             event_name="completions.create",
         )
 
         completion_raw_create = Patchee(
-            obj=client.completions.with_raw_response
-            if client
-            else openai.completions.with_raw_response,
+            obj=get_optional_attr(
+                [client, openai], ["completions", "with_raw_response"]
+            ),
             method_name="create",
             span_name=_SPAN_NAME_COMPLETION,
             event_name="completions.with_raw_response.create",
@@ -33,7 +34,7 @@ def __init__(self, client: Union[OpenAI, AsyncOpenAI, None] = None):
 
         if client:
             raw_completion_create = Patchee(
-                obj=client.with_raw_response.completions,
+                obj=get_attr(client, ["with_raw_response", "completions"]),
                 method_name="create",
                 span_name=_SPAN_NAME_COMPLETION,
                 event_name="with_raw_response.completions.create",

diff --git a/src/greptimeai/patchee/openai_patchee/embedding.py b/src/greptimeai/patchee/openai_patchee/embedding.py
@@ -4,23 +4,24 @@
 from openai import AsyncOpenAI, OpenAI
 
 from greptimeai.patchee import Patchee
+from greptimeai.utils.attr import get_attr, get_optional_attr
 
 from . import _SPAN_NAME_EMBEDDING, OpenaiPatchees
 
 
 class EmbeddingPatchees(OpenaiPatchees):
     def __init__(self, client: Union[OpenAI, AsyncOpenAI, None] = None):
         embeddings_create = Patchee(
-            obj=client.embeddings if client else openai.embeddings,
+            obj=get_optional_attr([client, openai], ["embeddings"]),
             method_name="create",
             span_name=_SPAN_NAME_EMBEDDING,
             event_name="embeddings.create",
         )
 
         embeddings_raw_create = Patchee(
-            obj=client.embeddings.with_raw_response
-            if client
-            else openai.embeddings.with_raw_response,
+            obj=get_optional_attr(
+                [client, openai], ["embeddings", "with_raw_response"]
+            ),
             method_name="create",
             span_name=_SPAN_NAME_EMBEDDING,
             event_name="embeddings.with_raw_response.create",
@@ -33,7 +34,7 @@ def __init__(self, client: Union[OpenAI, AsyncOpenAI, None] = None):
 
         if client:
             raw_embeddings_create = Patchee(
-                obj=client.with_raw_response.embeddings,
+                obj=get_attr(client, ["with_raw_response", "embeddings"]),
                 method_name="create",
                 span_name=_SPAN_NAME_EMBEDDING,
                 event_name="with_raw_response.embeddings.create",

diff --git a/src/greptimeai/patchee/openai_patchee/file.py b/src/greptimeai/patchee/openai_patchee/file.py
@@ -4,23 +4,22 @@
 from openai import AsyncOpenAI, OpenAI
 
 from greptimeai.patchee import Patchee
+from greptimeai.utils.attr import get_attr, get_optional_attr
 
 from . import _SPAN_NAME_FILE, OpenaiPatchees
 
 
 class _FilePatchees:
     def __init__(self, client: Union[OpenAI, AsyncOpenAI, None], method_name: str):
         files = Patchee(
-            obj=client.files if client else openai.files,
+            obj=get_optional_attr([client, openai], ["files"]),
             method_name=method_name,
             span_name=_SPAN_NAME_FILE,
             event_name=f"files.{method_name}",
         )
 
         files_raw = Patchee(
-            obj=client.files.with_raw_response
-            if client
-            else openai.files.with_raw_response,
+            obj=get_optional_attr([client, openai], ["files", "with_raw_response"]),
             method_name=method_name,
             span_name=_SPAN_NAME_FILE,
             event_name=f"files.with_raw_response.{method_name}",
@@ -30,7 +29,7 @@ def __init__(self, client: Union[OpenAI, AsyncOpenAI, None], method_name: str):
 
         if client:
             raw_files = Patchee(
-                obj=client.with_raw_response.files,
+                obj=get_attr(client, ["with_raw_response", "files"]),
                 method_name=method_name,
                 span_name=_SPAN_NAME_FILE,
                 event_name=f"with_raw_response.files.{method_name}",

diff --git a/src/greptimeai/patchee/openai_patchee/fine_tuning.py b/src/greptimeai/patchee/openai_patchee/fine_tuning.py
@@ -4,32 +4,33 @@
 from openai import AsyncOpenAI, OpenAI
 
 from greptimeai.patchee import Patchee
+from greptimeai.utils.attr import get_attr, get_optional_attr
 
 from . import _SPAN_NAME_FINE_TUNNING, OpenaiPatchees
 
 
 class _FineTuningPatchees:
     def __init__(self, client: Union[OpenAI, AsyncOpenAI, None], method_name: str):
         fine_tuning_jobs = Patchee(
-            obj=client.fine_tuning.jobs if client else openai.fine_tuning.jobs,
+            obj=get_optional_attr([client, openai], ["fine_tuning", "jobs"]),
             method_name=method_name,
             span_name=_SPAN_NAME_FINE_TUNNING,
             event_name=f"fine_tuning.jobs.{method_name}",
         )
 
         fine_tuning_raw_jobs = Patchee(
-            obj=client.fine_tuning.with_raw_response.jobs
-            if client
-            else openai.fine_tuning.with_raw_response.jobs,
+            obj=get_optional_attr(
+                [client, openai], ["fine_tuning", "with_raw_response", "jobs"]
+            ),
             method_name=method_name,
             span_name=_SPAN_NAME_FINE_TUNNING,
             event_name=f"fine_tuning.with_raw_response.jobs.{method_name}",
         )
 
         fine_tuning_jobs_raw = Patchee(
-            obj=client.fine_tuning.jobs.with_raw_response
-            if client
-            else openai.fine_tuning.jobs.with_raw_response,
+            obj=get_optional_attr(
+                [client, openai], ["fine_tuning", "jobs", "with_raw_response"]
+            ),
             method_name=method_name,
             span_name=_SPAN_NAME_FINE_TUNNING,
             event_name=f"fine_tuning.jobs.with_raw_response.{method_name}",
@@ -43,7 +44,7 @@ def __init__(self, client: Union[OpenAI, AsyncOpenAI, None], method_name: str):
 
         if client:
             raw_fine_tuning = Patchee(
-                obj=client.with_raw_response.fine_tuning.jobs,
+                obj=get_attr(client, ["with_raw_response", "fine_tuning", "jobs"]),
                 method_name=method_name,
                 span_name=_SPAN_NAME_FINE_TUNNING,
                 event_name=f"with_raw_response.fine_tuning.jobs.{method_name}",

diff --git a/src/greptimeai/patchee/openai_patchee/image.py b/src/greptimeai/patchee/openai_patchee/image.py
@@ -4,23 +4,22 @@
 from openai import AsyncOpenAI, OpenAI
 
 from greptimeai.patchee import Patchee
+from greptimeai.utils.attr import get_attr, get_optional_attr
 
 from . import _SPAN_NAME_IMAGE, OpenaiPatchees
 
 
 class _ImagePatchees:
     def __init__(self, client: Union[OpenAI, AsyncOpenAI, None], method_name: str):
         images = Patchee(
-            obj=client.images if client else openai.images,
+            obj=get_optional_attr([client, openai], ["images"]),
             method_name=method_name,
             span_name=_SPAN_NAME_IMAGE,
             event_name=f"images.{method_name}",
         )
 
         images_raw = Patchee(
-            obj=client.images.with_raw_response
-            if client
-            else openai.images.with_raw_response,
+            obj=get_optional_attr([client, openai], ["images", "with_raw_response"]),
             method_name=method_name,
             span_name=_SPAN_NAME_IMAGE,
             event_name=f"images.with_raw_response.{method_name}",
@@ -30,7 +29,7 @@ def __init__(self, client: Union[OpenAI, AsyncOpenAI, None], method_name: str):
 
         if client:
             raw_images = Patchee(
-                obj=client.with_raw_response.images,
+                obj=get_attr(client, ["with_raw_response", "images"]),
                 method_name=method_name,
                 span_name=_SPAN_NAME_IMAGE,
                 event_name=f"with_raw_response.images.{method_name}",

diff --git a/src/greptimeai/patchee/openai_patchee/model.py b/src/greptimeai/patchee/openai_patchee/model.py
@@ -4,23 +4,22 @@
 from openai import AsyncOpenAI, OpenAI
 
 from greptimeai.patchee import Patchee
+from greptimeai.utils.attr import get_attr, get_optional_attr
 
 from . import _SPAN_NAME_MODEL, OpenaiPatchees
 
 
 class _ModelPatchees:
     def __init__(self, client: Union[OpenAI, AsyncOpenAI, None], method_name: str):
         models = Patchee(
-            obj=client.models if client else openai.models,
+            obj=get_optional_attr([client, openai], ["models"]),
             method_name=method_name,
             span_name=_SPAN_NAME_MODEL,
             event_name=f"models.{method_name}",
         )
 
         models_raw = Patchee(
-            obj=client.models.with_raw_response
-            if client
-            else openai.models.with_raw_response,
+            obj=get_optional_attr([client, openai], ["models", "with_raw_response"]),
             method_name=method_name,
             span_name=_SPAN_NAME_MODEL,
             event_name=f"models.with_raw_response.{method_name}",
@@ -30,7 +29,7 @@ def __init__(self, client: Union[OpenAI, AsyncOpenAI, None], method_name: str):
 
         if client:
             raw_models = Patchee(
-                obj=client.with_raw_response.models,
+                obj=get_attr(client, ["with_raw_response", "models"]),
                 method_name=method_name,
                 span_name=_SPAN_NAME_MODEL,
                 event_name=f"with_raw_response.models.{method_name}",

diff --git a/src/greptimeai/patchee/openai_patchee/moderation.py b/src/greptimeai/patchee/openai_patchee/moderation.py
@@ -4,23 +4,24 @@
 from openai import AsyncOpenAI, OpenAI
 
 from greptimeai.patchee import Patchee
+from greptimeai.utils.attr import get_attr, get_optional_attr
 
 from . import _SPAN_NAME_MODERATION, OpenaiPatchees
 
 
 class ModerationPatchees(OpenaiPatchees):
     def __init__(self, client: Union[OpenAI, AsyncOpenAI, None] = None):
         moderations_create = Patchee(
-            obj=client.moderations if client else openai.moderations,
+            obj=get_optional_attr([client, openai], ["moderations"]),
             method_name="create",
             span_name=_SPAN_NAME_MODERATION,
             event_name="moderations.create",
         )
 
         moderations_raw_create = Patchee(
-            obj=client.moderations.with_raw_response
-            if client
-            else openai.moderations.with_raw_response,
+            obj=get_optional_attr(
+                [client, openai], ["moderations", "with_raw_response"]
+            ),
             method_name="create",
             span_name=_SPAN_NAME_MODERATION,
             event_name="moderations.with_raw_response.create",
@@ -30,7 +31,7 @@ def __init__(self, client: Union[OpenAI, AsyncOpenAI, None] = None):
 
         if client:
             raw_moderations_create = Patchee(
-                obj=client.with_raw_response.moderations,
+                obj=get_attr(client, ["with_raw_response", "moderations"]),
                 method_name="create",
                 span_name=_SPAN_NAME_MODERATION,
                 event_name="with_raw_response.moderations.create",

diff --git a/src/greptimeai/patchee/openai_patchee/retry.py b/src/greptimeai/patchee/openai_patchee/retry.py
@@ -10,8 +10,11 @@
 
 class RetryPatchees(OpenaiPatchees):
     def __init__(self, client: Union[OpenAI, AsyncOpenAI, None] = None):
+        obj = client
+        if obj is None and hasattr(openai, "_client"):
+            obj = openai._client
         retry = Patchee(
-            obj=client or openai._client,
+            obj=obj,
             method_name="_retry_request",
             span_name="",  # retry is only event, so span_name won't be used
             event_name="retry_request",