Merge pull request #783 from limnick/main

allow passing of kwargs to typed predictor decorators

Author: okhat

dspy/functional/functional.py

@@ -14,18 +14,32 @@
 from dspy.signatures.signature import ensure_signature, make_signature
 
 
-def predictor(func) -> dspy.Module:
-    """Decorator that creates a predictor module based on the provided function."""
-    signature = _func_to_signature(func)
-    *_, output_key = signature.output_fields.keys()
-    return _StripOutput(TypedPredictor(signature), output_key)
-
-
-def cot(func) -> dspy.Module:
-    """Decorator that creates a chain of thought module based on the provided function."""
-    signature = _func_to_signature(func)
-    *_, output_key = signature.output_fields.keys()
-    return _StripOutput(TypedChainOfThought(signature), output_key)
+def predictor(*args, **kwargs):
+    def _predictor(func) -> dspy.Module:
+        """Decorator that creates a predictor module based on the provided function."""
+        signature = _func_to_signature(func)
+        *_, output_key = signature.output_fields.keys()
+        return _StripOutput(TypedPredictor(signature, **kwargs), output_key)
+
+    # if we have only a single callable argument, the decorator was invoked with no key word arguments
+    #  so we just return the wrapped function 
+    if len(args) == 1 and callable(args[0]) and len(kwargs) == 0:
+        return _predictor(args[0])
+    return _predictor
+
+
+def cot(*args, **kwargs):
+    def _cot(func) -> dspy.Module:
+        """Decorator that creates a chain of thought module based on the provided function."""
+        signature = _func_to_signature(func)
+        *_, output_key = signature.output_fields.keys()
+        return _StripOutput(TypedChainOfThought(signature, **kwargs), output_key)
+
+    # if we have only a single callable argument, the decorator was invoked with no key word arguments
+    #  so we just return the wrapped function 
+    if len(args) == 1 and callable(args[0]) and len(kwargs) == 0:
+        return _cot(args[0])
+    return _cot
 
 
 class _StripOutput(dspy.Module):