import asyncio import inspect from typing import ( Any, Awaitable, Callable, Dict, Iterator, List, Optional, Union, cast, ) from guardrails.prompt import Prompt, Instructions from guardrails.errors import UserFacingException from guardrails.classes.llm.llm_response import LLMResponse from guardrails.classes.llm.prompt_callable import ( CALLABLE_FAILURE_SUFFIX, PromptCallableBase, PromptCallableException, ) from guardrails.types.inputs import MessageHistory import warnings from guardrails.utils.safe_get import safe_get from guardrails.telemetry import trace_llm_call, trace_operation from guardrails.utils.prompt_utils import messages_to_prompt_string ### # Synchronous wrappers ### def nonchat_prompt(prompt: str, instructions: Optional[str] = None) -> str: """Prepare final prompt for nonchat engine.""" if instructions: prompt = "\n\n".join([instructions, prompt]) return prompt def chat_prompt( prompt: Optional[str], instructions: Optional[str] = None, messages: Optional[List[Dict]] = None, ) -> List[Dict[str, str]]: """Prepare final prompt for chat engine.""" if messages: return messages if prompt is None: raise PromptCallableException( "You must pass in either `text` or `messages` to `guard.__call__`." ) if not instructions: instructions = "You are a helpful assistant." return [ {"role": "system", "content": instructions}, {"role": "user", "content": prompt}, ] def litellm_messages( prompt: Optional[str], instructions: Optional[str] = None, messages: Optional[List[Dict]] = None, ) -> List[Dict[str, str]]: """Prepare messages for LiteLLM.""" if messages: return messages if prompt is None: raise PromptCallableException( "Either `text` or `messages` required for `guard.__call__`." ) if instructions: prompt = "\n\n".join([instructions, prompt]) return [{"role": "user", "content": prompt}] class ManifestCallable(PromptCallableBase): def _invoke_llm( self, text: str, client: Any, instructions: Optional[str] = None, *args, **kwargs, ) -> LLMResponse: """Wrapper for manifest client. To use manifest for guardrailse, do ``` client = Manifest(client_name=..., client_connection=...) raw_llm_response, validated_response, *rest = guard( client, prompt_params={...}, ... ``` """ try: import manifest # noqa: F401 # type: ignore except ImportError: raise PromptCallableException( "The `manifest` package is not installed. " "Install with `poetry add manifest-ml`" ) client = cast(manifest.Manifest, client) prompt = nonchat_prompt(prompt=text, instructions=instructions) trace_operation( input_mime_type="application/json", input_value={ **kwargs, "prompt": prompt, "args": args, }, ) trace_llm_call( input_messages=chat_prompt(text, instructions), invocation_parameters={ **kwargs, "prompt": prompt, }, ) manifest_response = client.run(prompt, *args, **kwargs) trace_operation( output_mime_type="application/json", output_value=manifest_response ) trace_llm_call( output_messages=[{"role": "assistant", "content": manifest_response}] ) return LLMResponse( output=manifest_response, ) class LiteLLMCallable(PromptCallableBase): def _invoke_llm( self, text: Optional[str] = None, model: str = "gpt-3.5-turbo", messages: Optional[List[Dict]] = None, *args, **kwargs, ) -> LLMResponse: """Wrapper for Lite LLM completions. To use Lite LLM for guardrails, do ``` from litellm import completion raw_llm_response, validated_response = guard( completion, model="gpt-3.5-turbo", prompt_params={...}, temperature=..., ... ) ``` """ try: from litellm import completion # type: ignore except ImportError as e: raise PromptCallableException( "The `litellm` package is not installed. " "Install with `pip install litellm`" ) from e if messages is not None: messages = litellm_messages(prompt=text, messages=messages) kwargs["messages"] = messages trace_operation( input_mime_type="application/json", input_value={ **kwargs, "model": model, "args": args, }, ) function_calling_tools = [ tool.get("function") for tool in kwargs.get("tools", []) if isinstance(tool, Dict) and tool.get("type") == "function" ] trace_llm_call( input_messages=kwargs.get("messages"), invocation_parameters={ **kwargs, "model": model, }, function_call=kwargs.get( "function_call", safe_get(function_calling_tools, 0) ), ) # these are gr only and should not be getting passed to llms kwargs.pop("reask_messages", None) response = completion( model=model, *args, **kwargs, ) if kwargs.get("stream", False): # If stream is defined and set to True, # the callable returns a generator object llm_response = cast(Iterator[str], response) return LLMResponse( output="", # FIXME: Why is this different from the async streaming implementation? streamOutput=llm_response, ) trace_operation(output_mime_type="application/json", output_value=response) if response.choices[0].message.content is not None: # type: ignore output = response.choices[0].message.content # type: ignore else: try: output = response.choices[0].message.function_call.arguments # type: ignore except AttributeError: try: choice = response.choices[0] # type: ignore output = choice.message.tool_calls[-1].function.arguments # type: ignore except AttributeError as ae_tools: raise ValueError( "No message content or function" " call arguments returned from OpenAI" ) from ae_tools completion_tokens = response.usage.completion_tokens # type: ignore prompt_tokens = response.usage.prompt_tokens # type: ignore total_tokens = None if completion_tokens or prompt_tokens: total_tokens = (completion_tokens or 0) + (prompt_tokens or 0) trace_llm_call( output_messages=[choice.message for choice in response.choices], # type: ignore token_count_completion=completion_tokens, # type: ignore token_count_prompt=prompt_tokens, # type: ignore token_count_total=total_tokens, # type: ignore ) return LLMResponse( output=output, # type: ignore prompt_token_count=prompt_tokens, # type: ignore response_token_count=completion_tokens, # type: ignore ) class HuggingFaceModelCallable(PromptCallableBase): def _invoke_llm( self, model_generate: Any, *args, messages: Union[ list[dict[str, Union[str, Prompt, Instructions]]], MessageHistory ], **kwargs, ) -> LLMResponse: try: import transformers # noqa: F401 # type: ignore except ImportError: raise PromptCallableException( "The `transformers` package is not installed. " "Install with `pip install transformers`" ) try: import torch except ImportError: raise PromptCallableException( "The `torch` package is not installed. Install with `pip install torch`" ) prompt = messages_to_prompt_string(messages) tokenizer = kwargs.pop("tokenizer") if not tokenizer: raise UserFacingException( ValueError( "'tokenizer' must be provided in order to use Hugging Face models!" ) ) torch_device = "cuda" if torch.cuda.is_available() else "cpu" return_tensors = kwargs.pop("return_tensors", "pt") skip_special_tokens = kwargs.pop("skip_special_tokens", True) input_ids = kwargs.pop("input_ids", None) input_values = kwargs.pop("input_values", None) input_features = kwargs.pop("input_features", None) pixel_values = kwargs.pop("pixel_values", None) model_inputs = kwargs.pop("model_inputs", {}) if ( input_ids is None and input_values is None and input_features is None and pixel_values is None and not model_inputs ): model_inputs = tokenizer(prompt, return_tensors=return_tensors).to( torch_device ) else: model_inputs["input_ids"] = input_ids model_inputs["input_values"] = input_values model_inputs["input_features"] = input_features model_inputs["pixel_values"] = pixel_values do_sample = kwargs.pop("do_sample", None) temperature = kwargs.pop("temperature", None) if not do_sample and temperature == 0: temperature = None model_inputs["do_sample"] = do_sample model_inputs["temperature"] = temperature trace_operation( input_mime_type="application/json", input_value={ **model_inputs, **kwargs, }, ) trace_llm_call( input_messages=messages, invocation_parameters={ **model_inputs, **kwargs, }, ) output = model_generate( **model_inputs, **kwargs, ) trace_operation(output_mime_type="application/json", output_value=output) # NOTE: This is currently restricted to single outputs # Should we choose to support multiple return sequences, # We would need to either validate all of them # and choose the one with the least failures, # or accept a selection function decoded_output = tokenizer.decode( output[0], skip_special_tokens=skip_special_tokens ) trace_llm_call( output_messages=[{"role": "assistant", "content": decoded_output}] ) return LLMResponse(output=decoded_output) class HuggingFacePipelineCallable(PromptCallableBase): def _invoke_llm( self, pipeline: Any, *args, messages: Union[ list[dict[str, Union[str, Prompt, Instructions]]], MessageHistory ], **kwargs, ) -> LLMResponse: try: import transformers # noqa: F401 # type: ignore except ImportError: raise PromptCallableException( "The `transformers` package is not installed. " "Install with `pip install transformers`" ) try: import torch # noqa: F401 # type: ignore except ImportError: raise PromptCallableException( "The `torch` package is not installed. Install with `pip install torch`" ) content_key = kwargs.pop("content_key", "generated_text") temperature = kwargs.pop("temperature", None) if temperature == 0: temperature = None prompt = messages_to_prompt_string(messages) trace_operation( input_mime_type="application/json", input_value={ **kwargs, "prompt": prompt, "temperature": temperature, "args": args, }, ) trace_llm_call( input_messages=chat_prompt(prompt, kwargs.get("instructions")), invocation_parameters={ **kwargs, "prompt": prompt, "temperature": temperature, }, ) output = pipeline( prompt, temperature=temperature, *args, **kwargs, ) trace_operation(output_mime_type="application/json", output_value=output) # NOTE: This is currently restricted to single outputs # Should we choose to support multiple return sequences, # We would need to either validate all of them # and choose the one with the least failures, # or accept a selection function content = safe_get(output[0], content_key) trace_llm_call(output_messages=[{"role": "assistant", "content": content}]) return LLMResponse(output=content) 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.varkw: raise ValueError("Custom LLM callables must accept **kwargs!") if not llm_api_args.kwonlyargs or "messages" not in llm_api_args.kwonlyargs: warnings.warn( "We recommend including 'messages'" " as keyword-only arguments for custom LLM callables." " Doing so ensures these arguments are not unintentionally" " passed through to other calls via **kwargs.", UserWarning, ) self.llm_api = llm_api super().__init__(*args, **kwargs) def _invoke_llm(self, *args, **kwargs) -> LLMResponse: """Wrapper for arbitrary callable. To use an arbitrary callable for guardrails, do ``` raw_llm_response, validated_response, *rest = guard( my_callable, prompt_params={...}, ... ) ``` """ trace_operation( input_mime_type="application/json", input_value={ **kwargs, "args": args, }, ) trace_llm_call( input_messages=chat_prompt( kwargs.get("prompt", ""), kwargs.get("instructions") ), invocation_parameters={ **kwargs, }, ) # Get the response from the callable # The LLM response should either be a # string or an generator object of strings llm_response = self.llm_api(*args, **kwargs) # type: ignore # Check if kwargs stream is passed in if kwargs.get("stream", False): # If stream is defined and set to True, # the callable returns a generator object llm_response = cast(Iterator[str], llm_response) return LLMResponse( output="", # FIXME: Why is this different from the async streaming implementation? streamOutput=llm_response, ) trace_operation(output_mime_type="application/json", output_value=llm_response) trace_llm_call(output_messages=[{"role": "assistant", "content": llm_response}]) # Else, the callable returns a string llm_response = cast(str, llm_response) return LLMResponse( output=llm_response, ) def get_llm_ask( llm_api: Optional[Callable] = None, *args, **kwargs, ) -> Optional[PromptCallableBase]: if "temperature" not in kwargs: model = kwargs.get("model", "") if not ( isinstance(model, str) and (model.startswith("gpt-5") or model.startswith("openai/gpt-5")) ): warnings.warn( "The default value of 0 for temperature is deprecated " "and will be removed in guardrails-ai v0.8.x and higher.", DeprecationWarning, ) kwargs.update({"temperature": 0}) try: from litellm import completion if llm_api == completion or (llm_api is None and kwargs.get("model")): return LiteLLMCallable(*args, **kwargs) except ImportError: pass if llm_api is not None: llm_self = getattr(llm_api, "__self__", None) if ( llm_self is not None and hasattr(llm_self, "__class__") and getattr(llm_self.__class__, "__name__", None) == "GuardrailsEngine" and getattr(llm_api, "__name__", None) == "engine_api" ): return ArbitraryCallable(*args, llm_api=llm_api, **kwargs) try: import manifest # noqa: F401 # type: ignore if isinstance(llm_api, manifest.Manifest): return ManifestCallable(*args, client=llm_api, **kwargs) except ImportError: pass try: from transformers import ( # noqa: F401 # type: ignore FlaxPreTrainedModel, GenerationMixin, PreTrainedModel, TFPreTrainedModel, ) api_self = getattr(llm_api, "__self__", None) if ( isinstance(api_self, PreTrainedModel) or isinstance(api_self, TFPreTrainedModel) or isinstance(api_self, FlaxPreTrainedModel) ): if ( hasattr(llm_api, "__func__") and llm_api.__func__ == GenerationMixin.generate # type: ignore ): return HuggingFaceModelCallable(*args, model_generate=llm_api, **kwargs) raise ValueError("Only text generation models are supported at this time.") except ImportError: pass try: from transformers import Pipeline # noqa: F401 # type: ignore if isinstance(llm_api, Pipeline): # Couldn't find a constant for this if llm_api.task == "text-generation": return HuggingFacePipelineCallable(*args, pipeline=llm_api, **kwargs) raise ValueError( "Only text generation pipelines are supported at this time." ) except ImportError: pass # Let the user pass in an arbitrary callable. if llm_api is not None: return ArbitraryCallable(*args, llm_api=llm_api, **kwargs) ### # Async wrappers ### class AsyncPromptCallableBase(PromptCallableBase): async def invoke_llm( self, *args, **kwargs, ) -> LLMResponse: raise NotImplementedError async def __call__(self, *args, **kwargs) -> LLMResponse: try: result = await self.invoke_llm( *self.init_args, *args, **self.init_kwargs, **kwargs ) except Exception as e: raise PromptCallableException( "The callable `fn` passed to `Guard(fn, ...)` failed" f" with the following error: `{e}`. {CALLABLE_FAILURE_SUFFIX}" ) if not isinstance(result, LLMResponse): raise PromptCallableException( "The callable `fn` passed to `Guard(fn, ...)` returned" f" a non-string value: {result}. {CALLABLE_FAILURE_SUFFIX}" ) return result class AsyncLiteLLMCallable(AsyncPromptCallableBase): async def invoke_llm( self, text: Optional[str] = None, instructions: Optional[str] = None, messages: Optional[List[Dict]] = None, *args, **kwargs, ): """Wrapper for Lite LLM completions. To use Lite LLM for guardrails, do ``` from litellm import completion raw_llm_response, validated_response = guard( completion, model="gpt-3.5-turbo", prompt_params={...}, temperature=..., ... ) ``` """ try: from litellm import acompletion, CustomStreamWrapper # type: ignore except ImportError as e: raise PromptCallableException( "The `litellm` package is not installed. " "Install with `pip install litellm`" ) from e if text is not None or instructions is not None or messages is not None: messages = litellm_messages( prompt=text, instructions=instructions, messages=messages, ) kwargs["messages"] = messages trace_operation( input_mime_type="application/json", input_value={ **kwargs, "args": args, }, ) function_calling_tools = [ tool.get("function") for tool in kwargs.get("tools", []) if isinstance(tool, Dict) and tool.get("type") == "function" ] trace_llm_call( input_messages=kwargs.get("messages"), invocation_parameters={**kwargs}, function_call=kwargs.get( "function_call", safe_get(function_calling_tools, 0) ), ) # these are gr only and should not be getting passed to llms kwargs.pop("reask_messages", None) response = await acompletion( *args, **kwargs, ) if kwargs.get("stream", False): completion_stream = cast(CustomStreamWrapper, response) # If stream is defined and set to True, # the callable returns a generator object # response = cast(AsyncIterator[str], response) return LLMResponse( output="", asyncStreamOutput=completion_stream, ) trace_operation(output_mime_type="application/json", output_value=response) if response.choices[0].message.content is not None: # type: ignore output = response.choices[0].message.content # type: ignore else: try: output = response.choices[0].message.function_call.arguments # type: ignore except AttributeError: try: choice = response.choices[0] # type: ignore output = choice.message.tool_calls[-1].function.arguments # type: ignore except AttributeError as ae_tools: raise ValueError( "No message content or function" " call arguments returned from OpenAI" ) from ae_tools completion_tokens = response.usage.completion_tokens # type: ignore prompt_tokens = response.usage.prompt_tokens # type: ignore total_tokens = None if completion_tokens or prompt_tokens: total_tokens = (completion_tokens or 0) + (prompt_tokens or 0) trace_llm_call( output_messages=[choice.message for choice in response.choices], # type: ignore token_count_completion=completion_tokens, # type: ignore token_count_prompt=prompt_tokens, # type: ignore token_count_total=total_tokens, # type: ignore ) return LLMResponse( output=output, # type: ignore prompt_token_count=prompt_tokens, # type: ignore response_token_count=completion_tokens, # type: ignore ) class AsyncManifestCallable(AsyncPromptCallableBase): async def invoke_llm( self, text: str, client: Any, instructions: Optional[str] = None, *args, **kwargs, ): """Async wrapper for manifest client. To use manifest for guardrails, do ``` client = Manifest(client_name=..., client_connection=...) raw_llm_response, validated_response, *rest = guard( client, prompt_params={...}, ... ``` """ try: import manifest # noqa: F401 # type: ignore except ImportError: raise PromptCallableException( "The `manifest` package is not installed. " "Install with `poetry add manifest-ml`" ) prompts = [nonchat_prompt(prompt=text, instructions=instructions)] trace_operation( input_mime_type="application/json", input_value={ **kwargs, "prompts": prompts, "args": args, }, ) trace_llm_call( input_messages=chat_prompt(text, instructions), invocation_parameters={ **kwargs, "prompts": prompts, }, ) client = cast(manifest.Manifest, client) manifest_response = await client.arun_batch( prompts=prompts, *args, **kwargs, ) if kwargs.get("stream", False): raise NotImplementedError( "Manifest async streaming is not yet supported by manifest." ) trace_operation( output_mime_type="application/json", output_value=manifest_response ) trace_llm_call( output_messages=[{"role": "assistant", "content": manifest_response[0]}] ) return LLMResponse( output=manifest_response[0], ) class AsyncArbitraryCallable(AsyncPromptCallableBase): def __init__(self, llm_api: Callable, *args, **kwargs): llm_api_args = inspect.getfullargspec(llm_api) if not llm_api_args.varkw: raise ValueError("Custom LLM callables must accept **kwargs!") if not llm_api_args.kwonlyargs or "messages" not in llm_api_args.kwonlyargs: warnings.warn( "We recommend including 'messages'" " as keyword-only arguments for custom LLM callables." " Doing so ensures these arguments are not unintentionally" " passed through to other calls via **kwargs.", UserWarning, ) self.llm_api = llm_api super().__init__(*args, **kwargs) async def invoke_llm(self, *args, **kwargs) -> LLMResponse: """Wrapper for arbitrary callable. To use an arbitrary callable for guardrails, do ``` raw_llm_response, validated_response, *rest = guard( my_callable, prompt_params={...}, ... ) ``` """ trace_operation( input_mime_type="application/json", input_value={ **kwargs, "args": args, }, ) trace_llm_call( input_messages=chat_prompt( kwargs.get("prompt", ""), kwargs.get("instructions") ), invocation_parameters={ **kwargs, }, ) output = await self.llm_api(*args, **kwargs) if kwargs.get("stream", False): # If stream is defined and set to True, # the callable returns a generator object return LLMResponse( output="", # FIXME: Why is this different from the synchronous streaming implementation? ## noqa: E501 # This shouldn't be necessary: https://docs.litellm.ai/docs/completion/stream#async-streaming asyncStreamOutput=output.completion_stream, ) trace_operation(output_mime_type="application/json", output_value=output) trace_llm_call(output_messages=[{"role": "assistant", "content": output}]) return LLMResponse( output=output, ) def get_async_llm_ask( llm_api: Callable[..., Awaitable[Any]], *args, **kwargs ) -> AsyncPromptCallableBase: try: import litellm if llm_api == litellm.acompletion or (llm_api is None and kwargs.get("model")): return AsyncLiteLLMCallable(*args, **kwargs) except ImportError: pass try: import manifest # noqa: F401 # type: ignore if isinstance(llm_api, manifest.Manifest): return AsyncManifestCallable(*args, client=llm_api, **kwargs) except ImportError: pass 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[..., Awaitable[Any]]]] = None, *args, **kwargs, ) -> bool: if not llm_api: return True # TODO: Support other models; requires server-side updates model = get_llm_ask(llm_api, *args, **kwargs) if asyncio.iscoroutinefunction(llm_api): model = get_async_llm_ask(llm_api, *args, **kwargs) return isinstance(model, LiteLLMCallable) or isinstance(model, AsyncLiteLLMCallable)