refactor(instructor): improve client handling and code modularity (#562)

instructor-ai · Apr 3, 2024 · 1f1cb5e · 1f1cb5e
1 parent 700d35d
commit 1f1cb5e
Show file tree

Hide file tree

Showing 4 changed files with 135 additions and 109 deletions.
diff --git a/instructor/__init__.py b/instructor/__init__.py
@@ -12,14 +12,21 @@
 from .function_calls import OpenAISchema, openai_schema
 from .patch import apatch, patch
 from .process_response import handle_parallel_model
-from .client import Instructor, from_openai, from_anthropic, from_litellm, from_groq
+from .client import (
+    Instructor,
+    AsyncInstructor,
+    from_openai,
+    from_litellm,
+    Provider,
+)
+
 
 __all__ = [
     "Instructor",
     "from_openai",
-    "from_anthropic",
     "from_litellm",
-    "from_groq",
+    "AsyncInstructor",
+    "Provider",
     "OpenAISchema",
     "CitationMixin",
     "IterableModel",
@@ -36,3 +43,19 @@
     "handle_parallel_model",
     "handle_response_model",
 ]
+
+try:
+    import anthropic
+    from .client_anthropic import from_anthropic
+
+    __all__.append("from_anthropic")
+except ImportError:
+    pass
+
+try:
+    import groq
+    from .client_groq import from_groq
+
+    __all__.append("from_groq")
+except ImportError:
+    pass
diff --git a/instructor/client.py b/instructor/client.py
@@ -1,8 +1,6 @@
 import openai
 import inspect
 import instructor
-import anthropic
-import groq
 from .utils import Provider, get_provider
 from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
 from anthropic.types import Message
@@ -29,15 +27,15 @@
 
 
 class Instructor:
-    client: openai.OpenAI | anthropic.Anthropic | None
+    client: Any | None
     create_fn: Any
     mode: instructor.Mode
     default_model: str | None = None
     provider: Provider
 
     def __init__(
         self,
-        client: openai.OpenAI | anthropic.Anthropic | None,
+        client: Any | None,
         create: Callable,
         mode: instructor.Mode = instructor.Mode.TOOLS,
         provider: Provider = Provider.OPENAI,
@@ -151,15 +149,15 @@ def handle_kwargs(self, kwargs: dict):
 
 
 class AsyncInstructor(Instructor):
-    client: openai.AsyncOpenAI | anthropic.AsyncAnthropic | None
+    client: Any | None
     create_fn: Any
     mode: instructor.Mode
     default_model: str | None = None
     provider: Provider
 
     def __init__(
         self,
-        client: openai.AsyncOpenAI | anthropic.AsyncAnthropic | None,
+        client: Any | None,
         create: Callable,
         mode: instructor.Mode = instructor.Mode.TOOLS,
         provider: Provider = Provider.OPENAI,
@@ -282,10 +280,7 @@ def from_openai(
             instructor.Mode.JSON,
             instructor.Mode.JSON_SCHEMA,
         }
-    if provider in {Provider.GROQ}:
-        assert mode in {
-            instructor.Mode.MD_JSON,
-        }
+
     if provider in {Provider.OPENAI}:
         assert mode in {
             instructor.Mode.TOOLS,
@@ -352,99 +347,3 @@ def from_litellm(
             mode=mode,
             **kwargs,
         )
-
-
-@overload
-def from_anthropic(
-    client: anthropic.Anthropic,
-    mode: instructor.Mode = instructor.Mode.ANTHROPIC_JSON,
-    **kwargs,
-) -> Instructor: ...
-
-
-@overload
-def from_anthropic(
-    client: anthropic.AsyncAnthropic,
-    mode: instructor.Mode = instructor.Mode.ANTHROPIC_JSON,
-    **kwargs,
-) -> Instructor: ...
-
-
-def from_anthropic(
-    client: anthropic.Anthropic | anthropic.AsyncAnthropic,
-    mode: instructor.Mode = instructor.Mode.ANTHROPIC_JSON,
-    **kwargs,
-) -> Instructor | AsyncInstructor:
-    assert mode in {
-        instructor.Mode.ANTHROPIC_JSON,
-        instructor.Mode.ANTHROPIC_TOOLS,
-    }, "Mode be one of {instructor.Mode.ANTHROPIC_JSON, instructor.Mode.ANTHROPIC_TOOLS}"
-
-    assert isinstance(
-        client, (anthropic.Anthropic, anthropic.AsyncAnthropic)
-    ), "Client must be an instance of anthropic.Anthropic or anthropic.AsyncAnthropic"
-
-    if isinstance(client, anthropic.Anthropic):
-        return Instructor(
-            client=client,
-            create=instructor.patch(create=client.messages.create, mode=mode),
-            provider=Provider.ANTHROPIC,
-            mode=mode,
-            **kwargs,
-        )
-
-    else:
-        return AsyncInstructor(
-            client=client,
-            create=instructor.patch(create=client.messages.create, mode=mode),
-            provider=Provider.ANTHROPIC,
-            mode=mode,
-            **kwargs,
-        )
-
-
-@overload
-def from_groq(
-    client: groq.Groq,
-    mode: instructor.Mode = instructor.Mode.TOOLS,
-    **kwargs,
-) -> Instructor: ...
-
-
-@overload
-def from_groq(
-    client: groq.Groq,
-    mode: instructor.Mode = instructor.Mode.TOOLS,
-    **kwargs,
-) -> Instructor: ...
-
-
-def from_groq(
-    client: groq.Groq,
-    mode: instructor.Mode = instructor.Mode.TOOLS,
-    **kwargs,
-) -> Instructor:
-    assert mode in {
-        instructor.Mode.JSON,
-        instructor.Mode.TOOLS,
-    }, "Mode be one of {instructor.Mode.JSON, instructor.Mode.TOOLS}"
-
-    assert isinstance(client, (groq.Groq)), "Client must be an instance of groq.GROQ"
-
-    if isinstance(client, groq.Groq):
-        return Instructor(
-            client=client,
-            create=instructor.patch(create=client.chat.completions.create, mode=mode),
-            provider=Provider.GROQ,
-            mode=mode,
-            **kwargs,
-        )
-
-    else:
-        return AsyncInstructor(
-            client=client,
-            create=instructor.patch(create=client.messages.create, mode=mode),
-            provider=Provider.GROQ,
-            mode=mode,
-            **kwargs,
-        )
diff --git a/instructor/client_anthropic.py b/instructor/client_anthropic.py
@@ -0,0 +1,53 @@
+import anthropic
+import instructor
+
+from typing import overload
+
+
+@overload
+def from_anthropic(
+    client: anthropic.Anthropic,
+    mode: instructor.Mode = instructor.Mode.ANTHROPIC_JSON,
+    **kwargs,
+) -> instructor.Instructor: ...
+
+
+@overload
+def from_anthropic(
+    client: anthropic.AsyncAnthropic,
+    mode: instructor.Mode = instructor.Mode.ANTHROPIC_JSON,
+    **kwargs,
+) -> instructor.Instructor: ...
+
+
+def from_anthropic(
+    client: anthropic.Anthropic | anthropic.AsyncAnthropic,
+    mode: instructor.Mode = instructor.Mode.ANTHROPIC_JSON,
+    **kwargs,
+) -> instructor.Instructor | instructor.AsyncInstructor:
+    assert mode in {
+        instructor.Mode.ANTHROPIC_JSON,
+        instructor.Mode.ANTHROPIC_TOOLS,
+    }, "Mode be one of {instructor.Mode.ANTHROPIC_JSON, instructor.Mode.ANTHROPIC_TOOLS}"
+
+    assert isinstance(
+        client, (anthropic.Anthropic, anthropic.AsyncAnthropic)
+    ), "Client must be an instance of anthropic.Anthropic or anthropic.AsyncAnthropic"
+
+    if isinstance(client, anthropic.Anthropic):
+        return instructor.Instructor(
+            client=client,
+            create=instructor.patch(create=client.messages.create, mode=mode),
+            provider=instructor.Provider.ANTHROPIC,
+            mode=mode,
+            **kwargs,
+        )
+
+    else:
+        return instructor.AsyncInstructor(
+            client=client,
+            create=instructor.patch(create=client.messages.create, mode=mode),
+            provider=instructor.Provider.ANTHROPIC,
+            mode=mode,
+            **kwargs,
+        )
diff --git a/instructor/client_groq.py b/instructor/client_groq.py
@@ -0,0 +1,51 @@
+from typing import overload
+
+import groq
+import instructor
+
+
+@overload
+def from_groq(
+    client: groq.Groq,
+    mode: instructor.Mode = instructor.Mode.TOOLS,
+    **kwargs,
+) -> instructor.Instructor: ...
+
+
+@overload
+def from_groq(
+    client: groq.Groq,
+    mode: instructor.Mode = instructor.Mode.TOOLS,
+    **kwargs,
+) -> instructor.Instructor: ...
+
+
+def from_groq(
+    client: groq.Groq,
+    mode: instructor.Mode = instructor.Mode.TOOLS,
+    **kwargs,
+) -> instructor.Instructor:
+    assert mode in {
+        instructor.Mode.JSON,
+        instructor.Mode.TOOLS,
+    }, "Mode be one of {instructor.Mode.JSON, instructor.Mode.TOOLS}"
+
+    assert isinstance(client, (groq.Groq)), "Client must be an instance of groq.GROQ"
+
+    if isinstance(client, groq.Groq):
+        return instructor.Instructor(
+            client=client,
+            create=instructor.patch(create=client.chat.completions.create, mode=mode),
+            provider=instructor.Provider.GROQ,
+            mode=mode,
+            **kwargs,
+        )
+
+    else:
+        return instructor.Instructor(
+            client=client,
+            create=instructor.patch(create=client.messages.create, mode=mode),
+            provider=instructor.Provider.GROQ,
+            mode=mode,
+            **kwargs,
+        )