Merge pull request #706 from guardrails-ai/feat/llm-inputs-lite2

Feat/llm inputs lite2

Author: zsimjee

.github/workflows/ci.yml

@@ -50,7 +50,7 @@ jobs:
       matrix:
         python-version: ["3.8", "3.9", "3.10", "3.11"]
         pydantic-version: ["1.10.9", "2.4.2"]
-        openai-version: ["0.28.1", "1.2.4"]
+        openai-version: ["1.2.4"]
     steps:
       - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
@@ -87,7 +87,7 @@ jobs:
         # dependencies: ['dev', 'full']
         dependencies: ["full"]
         pydantic-version: ["1.10.9", "2.4.2"]
-        openai-version: ["0.28.1", "1.2.4"]
+        openai-version: ["1.2.4"]
     steps:
       - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}

guardrails/llm_providers.py

@@ -395,6 +395,16 @@ def _invoke_llm(
             *args,
             **kwargs,
         )
+
+        if kwargs.get("stream", False):
+            # If stream is defined and set to True,
+            # the callable returns a generator object
+            llm_response = cast(Iterable[str], response)
+            return LLMResponse(
+                output="",
+                stream_output=llm_response,
+            )
+
         return LLMResponse(
             output=response.choices[0].message.content,  # type: ignore
             prompt_token_count=response.usage.prompt_tokens,  # type: ignore
@@ -782,6 +792,53 @@ async def invoke_llm(
         )
 
 
+class AsyncLiteLLMCallable(AsyncPromptCallableBase):
+    async def invoke_llm(
+        self,
+        text: str,
+        instructions: Optional[str] = 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  # type: ignore
+        except ImportError as e:
+            raise PromptCallableException(
+                "The `litellm` package is not installed. "
+                "Install with `pip install litellm`"
+            ) from e
+
+        response = await acompletion(
+            messages=litellm_messages(
+                prompt=text,
+                instructions=instructions,
+            ),
+            *args,
+            **kwargs,
+        )
+
+        return LLMResponse(
+            output=response.choices[0].message.content,  # type: ignore
+            prompt_token_count=response.usage.prompt_tokens,  # type: ignore
+            response_token_count=response.usage.completion_tokens,  # type: ignore
+        )
+
+
 class AsyncManifestCallable(AsyncPromptCallableBase):
     async def invoke_llm(
         self,
@@ -860,6 +917,14 @@ def get_async_llm_ask(
     except ImportError:
         pass
 
+    try:
+        import litellm
+
+        if llm_api == litellm.acompletion:
+            return AsyncLiteLLMCallable(*args, **kwargs)
+    except ImportError:
+        pass
+
     return AsyncArbitraryCallable(*args, llm_api=llm_api, **kwargs)
 
 

guardrails/run/stream_runner.py

@@ -5,6 +5,7 @@
 from guardrails.classes.validation_outcome import ValidationOutcome
 from guardrails.datatypes import verify_metadata_requirements
 from guardrails.llm_providers import (
+    LiteLLMCallable,
     OpenAICallable,
     OpenAIChatCallable,
     PromptCallableBase,
@@ -227,6 +228,11 @@ def get_chunk_text(self, chunk: Any, api: Union[PromptCallableBase, None]) -> st
                 content = chunk.choices[0].delta.content
                 if not finished and content:
                     chunk_text = content
+        elif isinstance(api, LiteLLMCallable):
+            finished = chunk.choices[0].finish_reason
+            content = chunk.choices[0].delta.content
+            if not finished and content:
+                chunk_text = content
         else:
             try:
                 chunk_text = chunk