Enhance type hints of adapters (#7896)

* enhance type hints of adapters

Author: TomeHirata

dspy/adapters/base.py

@@ -1,13 +1,14 @@
 from abc import ABC, abstractmethod
-
-from litellm import ContextWindowExceededError
+from typing import Type, Any, Optional, TYPE_CHECKING
 
 from dspy.adapters.types import History
-from dspy.utils.callback import with_callbacks
-
+from dspy.utils.callback import BaseCallback, with_callbacks
+from dspy.signatures.signature import Signature
+if TYPE_CHECKING:
+    from dspy.clients.lm import LM
 
 class Adapter(ABC):
-    def __init__(self, callbacks=None):
+    def __init__(self, callbacks: Optional[list[BaseCallback]] = None):
         self.callbacks = callbacks or []
 
     def __init_subclass__(cls, **kwargs) -> None:
@@ -17,60 +18,53 @@ def __init_subclass__(cls, **kwargs) -> None:
         cls.format = with_callbacks(cls.format)
         cls.parse = with_callbacks(cls.parse)
 
-    def __call__(self, lm, lm_kwargs, signature, demos, inputs):
+    def __call__(self, lm: "LM", lm_kwargs: dict[str, Any], signature: Type[Signature], demos: list[dict[str, Any]], inputs: dict[str, Any]) -> list[dict[str, Any]]:
         inputs_ = self.format(signature, demos, inputs)
         inputs_ = dict(prompt=inputs_) if isinstance(inputs_, str) else dict(messages=inputs_)
 
         outputs = lm(**inputs_, **lm_kwargs)
         values = []
 
-        try:
-            for output in outputs:
-                output_logprobs = None
-
-                if isinstance(output, dict):
-                    output, output_logprobs = output["text"], output["logprobs"]
+        for output in outputs:
+            output_logprobs = None
 
-                value = self.parse(signature, output)
+            if isinstance(output, dict):
+                output, output_logprobs = output["text"], output["logprobs"]
 
-                if set(value.keys()) != set(signature.output_fields.keys()):
-                    raise ValueError(
-                        "Parsed output fields do not match signature output fields. "
-                        f"Expected: {set(signature.output_fields.keys())}, Got: {set(value.keys())}"
-                    )
+            value = self.parse(signature, output)
 
-                if output_logprobs is not None:
-                    value["logprobs"] = output_logprobs
+            if set(value.keys()) != set(signature.output_fields.keys()):
+                raise ValueError(
+                    "Parsed output fields do not match signature output fields. "
+                    f"Expected: {set(signature.output_fields.keys())}, Got: {set(value.keys())}"
+                )
 
-                values.append(value)
+            if output_logprobs is not None:
+                value["logprobs"] = output_logprobs
 
-            return values
+            values.append(value)
 
-        except Exception as e:
-            if isinstance(e, ContextWindowExceededError):
-                # On context window exceeded error, we don't want to retry with a different adapter.
-                raise e
-            from dspy.adapters.json_adapter import JSONAdapter
+        return values
 
-            if not isinstance(self, JSONAdapter):
-                return JSONAdapter()(lm, lm_kwargs, signature, demos, inputs)
-            raise e
 
     @abstractmethod
-    def format(self, signature, demos, inputs):
+    def format(self, signature: Type[Signature], demos: list[dict[str, Any]], inputs: dict[str, Any]) -> list[dict[str, Any]]:
         raise NotImplementedError
 
     @abstractmethod
-    def parse(self, signature, completion):
+    def parse(self, signature: Type[Signature], completion: str) -> dict[str, Any]:
         raise NotImplementedError
-
-    def format_finetune_data(self, signature, demos, inputs, outputs):
+    
+    def format_fields(self, signature: Type[Signature], values: dict[str, Any], role: str) -> str:
+        raise NotImplementedError
+    
+    def format_finetune_data(self, signature: Type[Signature], demos: list[dict[str, Any]], inputs: dict[str, Any], outputs: dict[str, Any]) -> dict[str, list[Any]]:
         raise NotImplementedError
 
-    def format_turn(self, signature, values, role, incomplete=False, is_conversation_history=False):
-        pass
+    def format_turn(self, signature: Type[Signature], values, role: str, incomplete: bool = False, is_conversation_history: bool = False) -> dict[str, Any]:
+        raise NotImplementedError
 
-    def format_conversation_history(self, signature, inputs):
+    def format_conversation_history(self, signature: Type[Signature], inputs: dict[str, Any]) -> list[dict[str, Any]]:
         history_field_name = None
         for name, field in signature.input_fields.items():
             if field.annotation == History:

dspy/adapters/chat_adapter.py

@@ -4,11 +4,11 @@
 import re
 import textwrap
 from collections.abc import Mapping
-from itertools import chain
-from typing import Any, Dict, Literal, NamedTuple
+from typing import Any, Dict, Literal, NamedTuple, Optional, Type
 
 import pydantic
 from pydantic.fields import FieldInfo
+from litellm import ContextWindowExceededError
 
 from dspy.adapters.base import Adapter
 from dspy.adapters.types.image import try_expand_image_tags
@@ -17,6 +17,9 @@
 from dspy.signatures.field import OutputField
 from dspy.signatures.signature import Signature, SignatureMeta
 from dspy.signatures.utils import get_dspy_field_type
+from dspy.adapters.json_adapter import JSONAdapter
+from dspy.clients.lm import LM
+from dspy.utils.callback import BaseCallback
 
 field_header_pattern = re.compile(r"\[\[ ## (\w+) ## \]\]")
 
@@ -31,7 +34,20 @@ class FieldInfoWithName(NamedTuple):
 
 
 class ChatAdapter(Adapter):
-    def format(self, signature: Signature, demos: list[dict[str, Any]], inputs: dict[str, Any]) -> list[dict[str, Any]]:
+    def __init__(self, callbacks: Optional[list[BaseCallback]] = None):
+        super().__init__(callbacks)
+
+    def __call__(self, lm: LM, lm_kwargs: dict[str, Any], signature: Type[Signature], demos: list[dict[str, Any]], inputs: dict[str, Any]) -> list[dict[str, Any]]:
+        try:
+            return super().__call__(lm, lm_kwargs, signature, demos, inputs)
+        except Exception as e:
+            if isinstance(e, ContextWindowExceededError):
+                # On context window exceeded error, we don't want to retry with a different adapter.
+                raise e
+            # fallback to JSONAdapter
+            return JSONAdapter()(lm, lm_kwargs, signature, demos, inputs)
+    
+    def format(self, signature: Type[Signature], demos: list[dict[str, Any]], inputs: dict[str, Any]) -> list[dict[str, Any]]:
         messages: list[dict[str, Any]] = []
 
         # Extract demos where some of the output_fields are not filled in.
@@ -64,7 +80,7 @@ def format(self, signature: Signature, demos: list[dict[str, Any]], inputs: dict
         messages = try_expand_image_tags(messages)
         return messages
 
-    def parse(self, signature, completion):
+    def parse(self, signature: Type[Signature], completion: str) -> dict[str, Any]:
         sections = [(None, [])]
 
         for line in completion.splitlines():
@@ -92,7 +108,7 @@ def parse(self, signature, completion):
         return fields
 
     # TODO(PR): Looks ok?
-    def format_finetune_data(self, signature, demos, inputs, outputs):
+    def format_finetune_data(self, signature: Type[Signature], demos: list[dict[str, Any]], inputs: dict[str, Any], outputs: dict[str, Any]) -> dict[str, list[Any]]:
         # Get system + user messages
         messages = self.format(signature, demos, inputs)
 
@@ -105,7 +121,7 @@ def format_finetune_data(self, signature, demos, inputs, outputs):
         # Wrap the messages in a dictionary with a "messages" key
         return dict(messages=messages)
 
-    def format_fields(self, signature, values, role):
+    def format_fields(self, signature: Type[Signature], values: dict[str, Any], role: str) -> str:
         fields_with_values = {
             FieldInfoWithName(name=field_name, info=field_info): values.get(
                 field_name, "Not supplied for this particular example."
@@ -115,7 +131,7 @@ def format_fields(self, signature, values, role):
         }
         return format_fields(fields_with_values)
 
-    def format_turn(self, signature, values, role, incomplete=False, is_conversation_history=False):
+    def format_turn(self, signature: Type[Signature], values: dict[str, Any], role: str, incomplete: bool = False, is_conversation_history: bool = False) -> dict[str, Any]:
         return format_turn(signature, values, role, incomplete, is_conversation_history)
 
 
@@ -139,7 +155,7 @@ def format_fields(fields_with_values: Dict[FieldInfoWithName, Any]) -> str:
     return "\n\n".join(output).strip()
 
 
-def format_turn(signature, values, role, incomplete=False, is_conversation_history=False):
+def format_turn(signature: Type[Signature], values: dict[str, Any], role: str, incomplete=False, is_conversation_history=False):
     """
     Constructs a new message ("turn") to append to a chat thread. The message is carefully formatted
     so that it can instruct an LLM to generate responses conforming to the specified DSPy signature.
@@ -212,11 +228,6 @@ def type_info(v):
     return {"role": role, "content": joined_messages}
 
 
-def flatten_messages(messages):
-    """Flatten nested message lists."""
-    return list(chain.from_iterable(item if isinstance(item, list) else [item] for item in messages))
-
-
 def enumerate_fields(fields: dict) -> str:
     parts = []
     for idx, (k, v) in enumerate(fields.items()):

dspy/adapters/json_adapter.py

@@ -4,7 +4,7 @@
 import logging
 import textwrap
 from copy import deepcopy
-from typing import Any, Dict, KeysView, Literal, NamedTuple
+from typing import Any, Dict, KeysView, Literal, NamedTuple, Type
 
 import json_repair
 import litellm
@@ -16,7 +16,8 @@
 from dspy.adapters.types.image import Image
 from dspy.adapters.types.history import History
 from dspy.adapters.utils import format_field_value, get_annotation_name, parse_value, serialize_for_json
-from dspy.signatures.signature import SignatureMeta
+from dspy.clients.lm import LM
+from dspy.signatures.signature import SignatureMeta, Signature
 from dspy.signatures.utils import get_dspy_field_type
 
 logger = logging.getLogger(__name__)
@@ -26,12 +27,11 @@ class FieldInfoWithName(NamedTuple):
     name: str
     info: FieldInfo
 
-
 class JSONAdapter(Adapter):
     def __init__(self):
         pass
 
-    def __call__(self, lm, lm_kwargs, signature, demos, inputs):
+    def __call__(self, lm: LM, lm_kwargs: dict[str, Any], signature: Type[Signature], demos: list[dict[str, Any]], inputs: dict[str, Any]) -> list[dict[str, Any]]:
         inputs = self.format(signature, demos, inputs)
         inputs = dict(prompt=inputs) if isinstance(inputs, str) else dict(messages=inputs)
 
@@ -66,7 +66,7 @@ def __call__(self, lm, lm_kwargs, signature, demos, inputs):
 
         return values
 
-    def format(self, signature, demos, inputs):
+    def format(self, signature: Type[Signature], demos: list[dict[str, Any]], inputs: dict[str, Any]) -> list[dict[str, Any]]:
         messages = []
 
         # Extract demos where some of the output_fields are not filled in.
@@ -94,7 +94,7 @@ def format(self, signature, demos, inputs):
 
         return messages
 
-    def parse(self, signature, completion):
+    def parse(self, signature: Type[Signature], completion: str) -> dict[str, Any]:
         fields = json_repair.loads(completion)
         fields = {k: v for k, v in fields.items() if k in signature.output_fields}
 
@@ -108,10 +108,7 @@ def parse(self, signature, completion):
 
         return fields
 
-    def format_turn(self, signature, values, role, incomplete=False, is_conversation_history=False):
-        return format_turn(signature, values, role, incomplete, is_conversation_history)
-
-    def format_fields(self, signature, values, role):
+    def format_fields(self, signature: Type[Signature], values: dict[str, Any], role: str) -> str:
         fields_with_values = {
             FieldInfoWithName(name=field_name, info=field_info): values.get(
                 field_name, "Not supplied for this particular example."
@@ -121,6 +118,13 @@ def format_fields(self, signature, values, role):
         }
 
         return format_fields(role=role, fields_with_values=fields_with_values)
+    
+    def format_turn(self, signature: Type[Signature], values, role: str, incomplete: bool = False, is_conversation_history: bool = False) -> dict[str, Any]:
+        return format_turn(signature, values, role, incomplete, is_conversation_history)
+    
+    def format_finetune_data(self, signature: Type[Signature], demos: list[dict[str, Any]], inputs: dict[str, Any], outputs: dict[str, Any]) -> dict[str, list[Any]]:
+        # TODO: implement format_finetune_data method in JSONAdapter
+        raise NotImplementedError
 
 
 def _format_field_value(field_info: FieldInfo, value: Any) -> str:

dspy/clients/lm.py

@@ -6,7 +6,7 @@
 import uuid
 from datetime import datetime
 from hashlib import sha256
-from typing import Any, Dict, List, Literal, Optional, cast
+from typing import Any, Dict, List, Literal, Optional, cast, TYPE_CHECKING
 
 import litellm
 import pydantic
@@ -17,11 +17,12 @@
 from litellm import RetryPolicy
 
 import dspy
-from dspy.adapters.base import Adapter
 from dspy.clients.openai import OpenAIProvider
 from dspy.clients.provider import Provider, TrainingJob
 from dspy.clients.utils_finetune import TrainDataFormat
 from dspy.utils.callback import BaseCallback, with_callbacks
+if TYPE_CHECKING:
+    from dspy.adapters.base import Adapter
 
 from .base_lm import BaseLM
 
@@ -219,7 +220,7 @@ def infer_provider(self) -> Provider:
         # providers in this file. Is this okay?
         return Provider()
 
-    def infer_adapter(self) -> Adapter:
+    def infer_adapter(self) -> "Adapter":
         import dspy
 
         if dspy.settings.adapter:

dspy/predict/predict.py

@@ -9,6 +9,8 @@
 from dspy.primitives.program import Module
 from dspy.signatures.signature import ensure_signature
 from dspy.utils.callback import with_callbacks
+from dspy.dsp.utils import settings
+from dspy.adapters.chat_adapter import ChatAdapter
 
 
 class Predict(Module, Parameter):
@@ -71,16 +73,14 @@ def __call__(self, **kwargs):
         return self.forward(**kwargs)
 
     def forward(self, **kwargs):
-        import dspy
-
         # Extract the three privileged keyword arguments.
         assert "new_signature" not in kwargs, "new_signature is no longer a valid keyword argument."
         signature = ensure_signature(kwargs.pop("signature", self.signature))
         demos = kwargs.pop("demos", self.demos)
         config = dict(**self.config, **kwargs.pop("config", {}))
 
         # Get the right LM to use.
-        lm = kwargs.pop("lm", self.lm) or dspy.settings.lm
+        lm = kwargs.pop("lm", self.lm) or settings.lm
         assert isinstance(lm, BaseLM), "No LM is loaded."
 
         # If temperature is 0.0 but its n > 1, set temperature to 0.7.
@@ -96,9 +96,7 @@ def forward(self, **kwargs):
             missing = [k for k in signature.input_fields if k not in kwargs]
             print(f"WARNING: Not all input fields were provided to module. Present: {present}. Missing: {missing}.")
 
-        import dspy
-
-        adapter = dspy.settings.adapter or dspy.ChatAdapter()
+        adapter = settings.adapter or ChatAdapter()
         completions = adapter(
             lm,
             lm_kwargs=config,
@@ -109,8 +107,8 @@ def forward(self, **kwargs):
 
         pred = Prediction.from_completions(completions, signature=signature)
 
-        if kwargs.pop("_trace", True) and dspy.settings.trace is not None:
-            trace = dspy.settings.trace
+        if kwargs.pop("_trace", True) and settings.trace is not None:
+            trace = settings.trace
             trace.append((self, {**kwargs}, pred))
 
         return pred