OpenAI v1 support

.github/workflows/ci.yml

@@ -51,6 +51,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']
         steps:
             -   uses: actions/checkout@v2
             -   name: Set up Python ${{ matrix.python-version }}
@@ -72,16 +73,27 @@ jobs:
                 run: |
                     make full
                     poetry run pip install pydantic==${{ matrix.pydantic-version }}
+                    poetry run pip install openai==${{ matrix.openai-version }}
 
-            -   if: matrix.pydantic-version == '2.4.2'
-                name: Static analysis with pyright (ignoring pydantic v1)
+            -   if: matrix.pydantic-version == '2.4.2' && matrix.openai-version == '0.28.1'
+                name: Static analysis with pyright (ignoring pydantic v1 and openai v1)
                 run: |
-                    make type-pydantic-v2
+                    make type-pydantic-v2-openai-v0
 
-            -   if: matrix.pydantic-version == '1.10.9'
-                name: Static analysis with mypy (ignoring pydantic v2)
+            -   if: matrix.pydantic-version == '1.10.9' && matrix.openai-version == '0.28.1'
+                name: Static analysis with mypy (ignoring pydantic v2 and openai v1)
                 run: |
-                    make type-pydantic-v1
+                    make type-pydantic-v1-openai-v0
+
+            -   if: matrix.pydantic-version == '2.4.2' && matrix.openai-version == '1.2.4'
+                name: Static analysis with pyright (ignoring pydantic v1 and openai v0)
+                run: |
+                    make type-pydantic-v2-openai-v1
+
+            -   if: matrix.pydantic-version == '1.10.9' && matrix.openai-version == '1.2.4'
+                name: Static analysis with mypy (ignoring pydantic v2 and openai v0)
+                run: |
+                    make type-pydantic-v1-openai-v1
 
     Pytests:
         runs-on: ubuntu-latest
@@ -92,6 +104,7 @@ jobs:
                 # dependencies: ['dev', 'full']
                 dependencies: ['full']
                 pydantic-version: ['1.10.9', '2.4.2']
+                openai-version: ['0.28.1', '1.2.4']
         steps:
             -   uses: actions/checkout@v2
             -   name: Set up Python ${{ matrix.python-version }}
@@ -103,15 +116,16 @@ jobs:
                 uses: actions/cache@v3
                 with:
                     path: ~/.cache/pypoetry
-                    key: poetry-cache-${{ runner.os }}-${{ steps.setup_python.outputs.python-version }}-${{ env.POETRY_VERSION }}-${{ matrix.pydantic-version }}
+                    key: poetry-cache-${{ runner.os }}-${{ steps.setup_python.outputs.python-version }}-${{ env.POETRY_VERSION }}-${{ matrix.pydantic-version }}-${{ matrix.openai-version }}
 
             -   name: Install Poetry
                 uses: snok/install-poetry@v1
 
             -   name: Install Dependencies
                 run: |
                     make ${{ matrix.dependencies }}
-                    python -m pip install pydantic==${{ matrix.pydantic-version }}
+                    poetry run pip install pydantic==${{ matrix.pydantic-version }}
+                    poetry run pip install openai==${{ matrix.openai-version }}
 
             -   name: Run Pytests
                 run: |

Makefile

@@ -8,13 +8,23 @@ autoformat:
 type:
 	poetry run pyright guardrails/
 
-type-pydantic-v1:
-	echo '{"exclude": ["guardrails/utils/pydantic_utils/v2.py"]}' > pyrightconfig.json
+type-pydantic-v1-openai-v0:
+	echo '{"exclude": ["guardrails/utils/pydantic_utils/v2.py", "guardrails/utils/openai_utils/v1.py"]}' > pyrightconfig.json
 	poetry run pyright guardrails/
 	rm pyrightconfig.json
 
-type-pydantic-v2:
-	echo '{"exclude": ["guardrails/utils/pydantic_utils/v1.py"]}' > pyrightconfig.json
+type-pydantic-v1-openai-v1:
+	echo '{"exclude": ["guardrails/utils/pydantic_utils/v2.py", "guardrails/utils/openai_utils/v0.py"]}' > pyrightconfig.json
+	poetry run pyright guardrails/
+	rm pyrightconfig.json
+
+type-pydantic-v2-openai-v0:
+	echo '{"exclude": ["guardrails/utils/pydantic_utils/v1.py", "guardrails/utils/openai_utils/v1.py"]}' > pyrightconfig.json
+	poetry run pyright guardrails/
+	rm pyrightconfig.json
+
+type-pydantic-v2-openai-v1:
+	echo '{"exclude": ["guardrails/utils/pydantic_utils/v1.py", "guardrails/utils/openai_utils/v0.py"]}' > pyrightconfig.json
 	poetry run pyright guardrails/
 	rm pyrightconfig.json
 

guardrails/applications/text2sql.py

@@ -4,11 +4,10 @@
 from string import Template
 from typing import Callable, Dict, Optional, Type
 
-import openai
-
 from guardrails.document_store import DocumentStoreBase, EphemeralDocumentStore
 from guardrails.embedding import EmbeddingBase, OpenAIEmbedding
 from guardrails.guard import Guard
+from guardrails.utils.openai_utils import get_static_openai_create_func
 from guardrails.utils.sql_utils import create_sql_driver
 from guardrails.vectordb import Faiss, VectorDBBase
 
@@ -71,7 +70,7 @@ def __init__(
         rail_params: Optional[Dict] = None,
         example_formatter: Callable = example_formatter,
         reask_prompt: str = REASK_PROMPT,
-        llm_api: Callable = openai.Completion.create,
+        llm_api: Optional[Callable] = None,
         llm_api_kwargs: Optional[Dict] = None,
         num_relevant_examples: int = 2,
     ):
@@ -88,6 +87,8 @@ def __init__(
             example_formatter: Fn to format examples. Defaults to example_formatter.
             reask_prompt: Prompt to use for reasking. Defaults to REASK_PROMPT.
         """
+        if llm_api is None:
+            llm_api = get_static_openai_create_func()
 
         self.example_formatter = example_formatter
         self.llm_api = llm_api
@@ -185,9 +186,10 @@ def __call__(self, text: str) -> Optional[str]:
                 "Async API is not supported in Text2SQL application. "
                 "Please use a synchronous API."
             )
-
+        if self.llm_api is None:
+            return None
         try:
-            output = self.guard(
+            return self.guard(
                 self.llm_api,
                 prompt_params={
                     "nl_instruction": text,
@@ -201,6 +203,4 @@ def __call__(self, text: str) -> Optional[str]:
                 "generated_sql"
             ]
         except TypeError:
-            output = None
-
-        return output
+            return None

guardrails/embedding.py

@@ -1,10 +1,9 @@
-import os
 from abc import ABC, abstractmethod
 from functools import cached_property
 from itertools import islice
 from typing import Callable, List, Optional
 
-import openai
+from guardrails.utils.openai_utils import OpenAIClient
 
 
 class EmbeddingBase(ABC):
@@ -114,9 +113,9 @@ def output_dim(self) -> int:
 class OpenAIEmbedding(EmbeddingBase):
     def __init__(
         self,
-        model: Optional[str] = "text-embedding-ada-002",
-        encoding_name: Optional[str] = "cl100k_base",
-        max_tokens: Optional[int] = 8191,
+        model: str = "text-embedding-ada-002",
+        encoding_name: str = "cl100k_base",
+        max_tokens: int = 8191,
         api_key: Optional[str] = None,
         api_base: Optional[str] = None,
     ):
@@ -137,15 +136,14 @@ def embed_query(self, query: str) -> List[float]:
         return resp[0]
 
     def _get_embedding(self, texts: List[str]) -> List[List[float]]:
-        api_key = (
-            self.api_key
-            if self.api_key is not None
-            else os.environ.get("OPENAI_API_KEY")
+        client = OpenAIClient(
+            api_key=self.api_key,
+            api_base=self.api_base,
         )
-        resp = openai.Embedding.create(
-            api_key=api_key, model=self._model, input=texts, api_base=self.api_base
+        return client.create_embedding(
+            model=self._model,
+            input=texts,
         )
-        return [r["embedding"] for r in resp["data"]]  # type: ignore
 
     @property
     def output_dim(self) -> int: