enforce method signature for custom llm callables

Author: CalebCourier

guardrails/formatters/json_formatter.py

@@ -1,5 +1,5 @@
 import json
-from typing import Optional, Union
+from typing import Dict, List, Optional, Union
 
 from guardrails.formatters.base_formatter import BaseFormatter
 from guardrails.llm_providers import (
@@ -99,32 +99,48 @@ def wrap_callable(self, llm_callable) -> ArbitraryCallable:
 
         if isinstance(llm_callable, HuggingFacePipelineCallable):
             model = llm_callable.init_kwargs["pipeline"]
-            return ArbitraryCallable(
-                lambda p: json.dumps(
+
+            def fn(
+                prompt: str,
+                *args,
+                instructions: Optional[str] = None,
+                msg_history: Optional[List[Dict[str, str]]] = None,
+                **kwargs,
+            ) -> str:
+                return json.dumps(
                     Jsonformer(
                         model=model.model,
                         tokenizer=model.tokenizer,
                         json_schema=self.output_schema,
-                        prompt=p,
+                        prompt=prompt,
                     )()
                 )
-            )
+
+            return ArbitraryCallable(fn)
         elif isinstance(llm_callable, HuggingFaceModelCallable):
             # This will not work because 'model_generate' is the .gen method.
             # model = self.api.init_kwargs["model_generate"]
             # Use the __self__ to grab the base mode for passing into JF.
             model = llm_callable.init_kwargs["model_generate"].__self__
             tokenizer = llm_callable.init_kwargs["tokenizer"]
-            return ArbitraryCallable(
-                lambda p: json.dumps(
+
+            def fn(
+                prompt: str,
+                *args,
+                instructions: Optional[str] = None,
+                msg_history: Optional[List[Dict[str, str]]] = None,
+                **kwargs,
+            ) -> str:
+                return json.dumps(
                     Jsonformer(
                         model=model,
                         tokenizer=tokenizer,
                         json_schema=self.output_schema,
-                        prompt=p,
+                        prompt=prompt,
                     )()
                 )
-            )
+
+            return ArbitraryCallable(fn)
         else:
             raise ValueError(
                 "JsonFormatter can only be used with HuggingFace*Callable."

guardrails/llm_providers.py

@@ -1,5 +1,6 @@
 import asyncio
 
+import inspect
 from typing import (
     Any,
     Awaitable,
@@ -711,6 +712,26 @@ def _invoke_llm(self, prompt: str, pipeline: Any, *args, **kwargs) -> LLMRespons
 
 class ArbitraryCallable(PromptCallableBase):
     def __init__(self, llm_api: Optional[Callable] = None, *args, **kwargs):
+        llm_api_args = inspect.getfullargspec(llm_api)
+        if not llm_api_args.args:
+            raise ValueError(
+                "Custom LLM callables must accept"
+                " at least one positional argument for prompt!"
+            )
+        if not llm_api_args.varkw:
+            raise ValueError("Custom LLM callables must accept **kwargs!")
+        if (
+            not llm_api_args.kwonlyargs
+            or "instructions" not in llm_api_args.kwonlyargs
+            or "msg_history" not in llm_api_args.kwonlyargs
+        ):
+            warnings.warn(
+                "We recommend including 'instructions' and 'msg_history'"
+                " as keyword-only arguments for custom LLM callables."
+                " Doing so ensures these arguments are not uninentionally"
+                " passed through to other calls via **kwargs.",
+                UserWarning,
+            )
         self.llm_api = llm_api
         super().__init__(*args, **kwargs)
 
@@ -1190,6 +1211,26 @@ async def invoke_llm(
 
 class AsyncArbitraryCallable(AsyncPromptCallableBase):
     def __init__(self, llm_api: Callable, *args, **kwargs):
+        llm_api_args = inspect.getfullargspec(llm_api)
+        if not llm_api_args.args:
+            raise ValueError(
+                "Custom LLM callables must accept"
+                " at least one positional argument for prompt!"
+            )
+        if not llm_api_args.varkw:
+            raise ValueError("Custom LLM callables must accept **kwargs!")
+        if (
+            not llm_api_args.kwonlyargs
+            or "instructions" not in llm_api_args.kwonlyargs
+            or "msg_history" not in llm_api_args.kwonlyargs
+        ):
+            warnings.warn(
+                "We recommend including 'instructions' and 'msg_history'"
+                " as keyword-only arguments for custom LLM callables."
+                " Doing so ensures these arguments are not uninentionally"
+                " passed through to other calls via **kwargs.",
+                UserWarning,
+            )
         self.llm_api = llm_api
         super().__init__(*args, **kwargs)
 
@@ -1241,7 +1282,7 @@ async def invoke_llm(self, *args, **kwargs) -> LLMResponse:
 
 
 def get_async_llm_ask(
-    llm_api: Callable[[Any], Awaitable[Any]], *args, **kwargs
+    llm_api: Callable[..., Awaitable[Any]], *args, **kwargs
 ) -> AsyncPromptCallableBase:
     try:
         import litellm
@@ -1268,11 +1309,12 @@ def get_async_llm_ask(
     except ImportError:
         pass
 
-    return AsyncArbitraryCallable(*args, llm_api=llm_api, **kwargs)
+    if llm_api is not None:
+        return AsyncArbitraryCallable(*args, llm_api=llm_api, **kwargs)
 
 
 def model_is_supported_server_side(
-    llm_api: Optional[Union[Callable, Callable[[Any], Awaitable[Any]]]] = None,
+    llm_api: Optional[Union[Callable, Callable[..., Awaitable[Any]]]] = None,
     *args,
     **kwargs,
 ) -> bool:
@@ -1292,7 +1334,7 @@ def model_is_supported_server_side(
 
 # CONTINUOUS FIXME: Update with newly supported LLMs
 def get_llm_api_enum(
-    llm_api: Callable[[Any], Awaitable[Any]], *args, **kwargs
+    llm_api: Callable[..., Awaitable[Any]], *args, **kwargs
 ) -> Optional[LLMResource]:
     # TODO: Distinguish between v1 and v2
     model = get_llm_ask(llm_api, *args, **kwargs)

tests/integration_tests/test_assets/custom_llm.py

@@ -0,0 +1,21 @@
+from typing import Dict, List, Optional
+
+
+def mock_llm(
+    prompt: Optional[str] = None,
+    *args,
+    instructions: Optional[str] = None,
+    msg_history: Optional[List[Dict[str, str]]] = None,
+    **kwargs,
+) -> str:
+    return ""
+
+
+async def mock_async_llm(
+    prompt: Optional[str] = None,
+    *args,
+    instructions: Optional[str] = None,
+    msg_history: Optional[List[Dict[str, str]]] = None,
+    **kwargs,
+) -> str:
+    return ""

tests/integration_tests/test_async.py

@@ -3,6 +3,7 @@
 from guardrails import AsyncGuard, Prompt
 from guardrails.utils import docs_utils
 from guardrails.classes.llm.llm_response import LLMResponse
+from tests.integration_tests.test_assets.custom_llm import mock_async_llm
 from tests.integration_tests.test_assets.fixtures import (  # noqa
     fixture_llm_output,
     fixture_rail_spec,
@@ -12,10 +13,6 @@
 from .mock_llm_outputs import entity_extraction
 
 
-async def mock_llm(*args, **kwargs):
-    return ""
-
-
 @pytest.mark.asyncio
 async def test_entity_extraction_with_reask(mocker):
     """Test that the entity extraction works with re-asking."""
@@ -45,7 +42,7 @@ async def test_entity_extraction_with_reask(mocker):
     preprocess_prompt_spy = mocker.spy(async_runner, "preprocess_prompt")
 
     final_output = await guard(
-        llm_api=mock_llm,
+        llm_api=mock_async_llm,
         prompt_params={"document": content[:6000]},
         num_reasks=1,
     )
@@ -104,7 +101,7 @@ async def test_entity_extraction_with_noop(mocker):
     content = docs_utils.read_pdf("docs/examples/data/chase_card_agreement.pdf")
     guard = AsyncGuard.from_rail_string(entity_extraction.RAIL_SPEC_WITH_NOOP)
     final_output = await guard(
-        llm_api=mock_llm,
+        llm_api=mock_async_llm,
         prompt_params={"document": content[:6000]},
         num_reasks=1,
     )
@@ -151,7 +148,7 @@ async def test_entity_extraction_with_noop_pydantic(mocker):
         prompt=entity_extraction.PYDANTIC_PROMPT,
     )
     final_output = await guard(
-        llm_api=mock_llm,
+        llm_api=mock_async_llm,
         prompt_params={"document": content[:6000]},
         num_reasks=1,
     )
@@ -192,7 +189,7 @@ async def test_entity_extraction_with_filter(mocker):
     content = docs_utils.read_pdf("docs/examples/data/chase_card_agreement.pdf")
     guard = AsyncGuard.from_rail_string(entity_extraction.RAIL_SPEC_WITH_FILTER)
     final_output = await guard(
-        llm_api=mock_llm,
+        llm_api=mock_async_llm,
         prompt_params={"document": content[:6000]},
         num_reasks=1,
     )
@@ -232,7 +229,7 @@ async def test_entity_extraction_with_fix(mocker):
     content = docs_utils.read_pdf("docs/examples/data/chase_card_agreement.pdf")
     guard = AsyncGuard.from_rail_string(entity_extraction.RAIL_SPEC_WITH_FIX)
     final_output = await guard(
-        llm_api=mock_llm,
+        llm_api=mock_async_llm,
         prompt_params={"document": content[:6000]},
         num_reasks=1,
     )
@@ -269,7 +266,7 @@ async def test_entity_extraction_with_refrain(mocker):
     content = docs_utils.read_pdf("docs/examples/data/chase_card_agreement.pdf")
     guard = AsyncGuard.from_rail_string(entity_extraction.RAIL_SPEC_WITH_REFRAIN)
     final_output = await guard(
-        llm_api=mock_llm,
+        llm_api=mock_async_llm,
         prompt_params={"document": content[:6000]},
         num_reasks=1,
     )
@@ -295,7 +292,7 @@ async def test_rail_spec_output_parse(rail_spec, llm_output, validated_output):
     guard = AsyncGuard.from_rail_string(rail_spec)
     output = await guard.parse(
         llm_output,
-        llm_api=mock_llm,
+        llm_api=mock_async_llm,
     )
     assert output.validated_output == validated_output
 
@@ -334,7 +331,7 @@ async def test_string_rail_spec_output_parse(
     guard: AsyncGuard = AsyncGuard.from_rail_string(string_rail_spec)
     output = await guard.parse(
         string_llm_output,
-        llm_api=mock_llm,
+        llm_api=mock_async_llm,
         num_reasks=0,
     )
     assert output.validated_output == validated_string_output