Types!!! (#372)

Co-authored-by: Luke Van Seters <lukevanseters@gmail.com>
Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

instructor/dsl/iterable.py

@@ -110,7 +110,7 @@ def IterableModel(
     subtask_class: Type[BaseModel],
     name: Optional[str] = None,
     description: Optional[str] = None,
-):
+) -> Type[BaseModel]:
     """
     Dynamically create a IterableModel OpenAISchema that can be used to segment multiple
     tasks given a base class. This creates class that can be used to create a toolkit
@@ -187,5 +187,7 @@ def from_streaming_response(cls, completion) -> Generator[User]:
         if description is None
         else description
     )
-
+    assert issubclass(
+        new_cls, OpenAISchema
+    ), "The new class should be a subclass of OpenAISchema"
     return new_cls

instructor/dsl/partialjson.py

@@ -125,7 +125,7 @@ def parse_string(self, s, e):
         s = s[end + 1 :]
         return json.loads(str_val), s
 
-    def parse_number(self, s, e):
+    def parse_number(self, s):
         i = 0
         while i < len(s) and s[i] in "0123456789.-":
             i += 1
@@ -139,7 +139,7 @@ def parse_number(self, s, e):
                 if "." in num_str or "e" in num_str or "E" in num_str
                 else int(num_str)
             )
-        except ValueError:
+        except ValueError as e:
             raise e
         return num, s
 

instructor/function_calls.py

@@ -1,11 +1,13 @@
+from typing import Type, TypeVar, Self
 from docstring_parser import parse
 from functools import wraps
 from pydantic import BaseModel, create_model
 from instructor.exceptions import IncompleteOutputException
-
 import enum
 import warnings
 
+T = TypeVar("T")
+
 
 class Mode(enum.Enum):
     """The mode to use for patching the client"""
@@ -118,11 +120,11 @@ def openai_schema(cls):
     @classmethod
     def from_response(
         cls,
-        completion,
-        validation_context=None,
+        completion: T,
+        validation_context: dict = None,
         strict: bool = None,
         mode: Mode = Mode.TOOLS,
-    ):
+    ) -> Self:
         """Execute the function from the response of an openai chat completion
 
         Parameters:
@@ -176,9 +178,11 @@ def from_response(
     async def from_response_async(
         cls,
         completion,
-        validation_context=None,
+        validation_context: dict = None,
         strict: bool = None,
         mode: Mode = Mode.TOOLS,
+        stream_multitask: bool = False,
+        stream_partial: bool = False,
     ):
         """Execute the function from the response of an openai chat completion
 
@@ -230,7 +234,7 @@ async def from_response_async(
             raise ValueError(f"Invalid patch mode: {mode}")
 
 
-def openai_schema(cls) -> OpenAISchema:
+def openai_schema(cls: Type[BaseModel]) -> OpenAISchema:
     if not issubclass(cls, BaseModel):
         raise TypeError("Class must be a subclass of pydantic.BaseModel")
 
@@ -239,4 +243,4 @@ def openai_schema(cls) -> OpenAISchema:
             cls.__name__,
             __base__=(cls, OpenAISchema),
         )
-    )  # type: ignore
+    )