Repository: explosion/spacy-llm (language: python)

Issue:
Dolly/HF backend
## Description
Easy to configure:
```
[components.llm.backend]
@llm_backends = "spacy.DollyHF.v1"
model = "databricks/dolly-v2-3b"
```

cf also https://huggingface.co/databricks/dolly-v2-3b

### Types of change
feature

## Checklist
- [x] I confirm that I have the right to submit this contribution under the project's MIT license.
- [x] I ran the tests, and all new and existing tests passed.
- [x] My changes don't require a change to the documentation, or if they do, I've added all required information.



Discussion / hints:

Fix patch:
diff --git a/spacy_llm/backends/__init__.py b/spacy_llm/backends/__init__.py
index cd8f76c1..9ac666e0 100644
--- a/spacy_llm/backends/__init__.py
+++ b/spacy_llm/backends/__init__.py
@@ -1,5 +1,6 @@
 from .langchain import backend_langchain, query_langchain
 from .minichain import backend_minichain, query_minichain
+from .rest import backend_rest
 
 from dotenv import load_dotenv
 
@@ -10,4 +11,5 @@
     "query_langchain",
     "backend_minichain",
     "backend_langchain",
+    "backend_rest",
 ]
diff --git a/spacy_llm/backends/rest/__init__.py b/spacy_llm/backends/rest/__init__.py
new file mode 100644
index 00000000..ba0848cc
--- /dev/null
+++ b/spacy_llm/backends/rest/__init__.py
@@ -0,0 +1,3 @@
+from .registry import backend_rest
+
+__all__ = ["backend_rest"]
diff --git a/spacy_llm/backends/rest/backend.py b/spacy_llm/backends/rest/backend.py
new file mode 100644
index 00000000..de10f9b9
--- /dev/null
+++ b/spacy_llm/backends/rest/backend.py
@@ -0,0 +1,76 @@
+import os
+from typing import Iterable, Dict, Any, List
+
+import requests  # type: ignore
+import srsly
+
+
+class Backend:
+    """Queries LLMs via their REST APIs."""
+
+    _SUPPORTED_APIS = ("OpenAI",)
+
+    def __init__(self, api: str, config: Dict[Any, Any], strict: bool):
+        """Initializes new Backend instance.
+        api (str): Name of LLM API.
+        config (Dict[Any, Any]): Config passed on to LLM API.
+        strict (bool): If True, ValueError is raised if the LLM API returns a malformed response (i. e. any kind of JSON
+            or other response object that does not conform to the expectation of how a well-formed response object from
+            this API should look like). If False, the API error responses are returned by __call__(), but no error will
+            be raised.
+            Note that only response object structure will be checked, not the prompt response text per se.
+        """
+        if api not in Backend._SUPPORTED_APIS:
+            raise ValueError(
+                f"{api} is not one of the supported APIs ({Backend._SUPPORTED_APIS})."
+            )
+        self._api = api
+        self._config = config
+        self._strict = strict
+        self._calls = {"OpenAI": self._call_openai}
+
+    def __call__(self, prompts: Iterable[str]) -> Iterable[str]:
+        """Executes prompts on specified API.
+        prompts (Iterable[str]): Prompts to execute.
+        RETURNS (Iterable[str]): API responses.
+        """
+        return self._calls[self._api](prompts)
+
+    def _call_openai(self, prompts: Iterable[str]) -> Iterable[str]:
+        """
+        Calls OpenAI API. Note that this currently only implements the completions endpoint
+        (https://api.openai.com/v1/completions), not the chat endpoint (https://api.openai.com/v1/chat/completions).
+        See https://platform.openai.com/docs/api-reference/.
+        prompts (Iterable[str]): Prompts to execute.
+        RETURNS (Iterable[str]): API responses.
+        """
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f'Bearer {os.getenv("OPENAI_API_KEY", "")}',
+        }
+        api_responses: List[str] = []
+
+        prompts = list(prompts)
+        json_data = {"prompt": prompts, **self._config}
+        responses = requests.post(
+            self._config.pop("url")
+            if "url" in self._config
+            else "https://api.openai.com/v1/completions",
+            headers=headers,
+            json=json_data,
+        ).json()
+
+        if "error" in responses:
+            if self._strict:
+                raise ValueError(f"API call failed: {responses}.")
+            else:
+                return [srsly.json_dumps(responses)] * len(prompts)
+        assert len(responses["choices"]) == len(prompts)
+
+        for prompt, response in zip(prompts, responses["choices"]):
+            if "text" in response:
+                api_responses.append(response["text"])
+            else:
+                api_responses.append(srsly.json_dumps(response))
+
+        return api_responses
diff --git a/spacy_llm/backends/rest/registry.py b/spacy_llm/backends/rest/registry.py
new file mode 100644
index 00000000..bcc39fd9
--- /dev/null
+++ b/spacy_llm/backends/rest/registry.py
@@ -0,0 +1,30 @@
+from typing import Callable, Iterable, Dict, Any
+
+import spacy
+from spacy.util import SimpleFrozenDict
+from .backend import Backend
+
+
+@spacy.registry.llm_backends("spacy.REST.v1")
+def backend_rest(
+    api: str,
+    config: Dict[Any, Any] = SimpleFrozenDict(),
+    strict: bool = True,
+) -> Callable[[Iterable[str]], Iterable[str]]:
+    """Returns Callable using minimal REST backend to prompt specified API.
+    api (str): Name of any API. Currently supported: "OpenAI".
+    config (Dict[Any, Any]): LLM config arguments passed on to the initialization of the Backend instance.
+    strict (bool): If True, ValueError is raised if the LLM API returns a malformed response (i. e. any kind of JSON
+        or other response object that does not conform to the expectation of how a well-formed response object from
+        this API should look like). If False, the API error responses are returned by __call__(), but no error will
+        be raised.
+    RETURNS (Callable[[Iterable[str]], Iterable[str]]]): Callable using the querying the specified API using a
+        Backend instance.
+    """
+
+    backend = Backend(api=api, config=config, strict=strict)
+
+    def _query(prompts: Iterable[str]) -> Iterable[str]:
+        return backend(prompts)
+
+    return _query
diff --git a/spacy_llm/pipeline/llm.py b/spacy_llm/pipeline/llm.py
index cf49ff16..3b0417bf 100644
--- a/spacy_llm/pipeline/llm.py
+++ b/spacy_llm/pipeline/llm.py
@@ -22,10 +22,10 @@
     default_config={
         "task": None,
         "backend": {
-            "@llm_backends": "spacy.MiniChain.v1",
+            "@llm_backends": "spacy.REST.v1",
             "api": "OpenAI",
-            "config": {},
-            "query": {"@llm_queries": "spacy.RunMiniChain.v1"},
+            "config": {"model": "text-davinci-003"},
+            "strict": True,
         },
     },
 )


Test patch:
diff --git a/spacy_llm/tests/backend/__init__.py b/spacy_llm/tests/backend/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/spacy_llm/tests/backend/test_rest.py b/spacy_llm/tests/backend/test_rest.py
new file mode 100644
index 00000000..bc4f993b
--- /dev/null
+++ b/spacy_llm/tests/backend/test_rest.py
@@ -0,0 +1,50 @@
+import pytest
+import spacy
+import srsly
+
+PIPE_CFG = {
+    "backend": {
+        "@llm_backends": "spacy.REST.v1",
+        "api": "OpenAI",
+        "config": {"temperature": 0.3, "model": "text-davinci-003"},
+    },
+    "task": {"@llm_tasks": "spacy.NoOp.v1"},
+}
+
+
+def test_initialization():
+    """Test initialization and simple run"""
+    nlp = spacy.blank("en")
+    nlp.add_pipe("llm", config=PIPE_CFG)
+    nlp("This is a test.")
+
+
+@pytest.mark.parametrize("strict", (False, True))
+def test_rest_backend_error_handling(strict: bool):
+    """Test error handling for default/minimal REST backend.
+    strict (bool): Whether to use strict mode.
+    """
+    nlp = spacy.blank("en")
+    nlp.add_pipe(
+        "llm",
+        config={
+            "task": {"@llm_tasks": "spacy.NoOp.v1"},
+            "backend": {"config": {"model": "x-text-davinci-003"}, "strict": strict},
+        },
+    )
+
+    if strict:
+        with pytest.raises(
+            ValueError,
+            match="API call failed: {'error': {'message': 'The model `x-text-davinci-003` does not exist', 'type': "
+            "'invalid_request_error', 'param': None, 'code': None}}.",
+        ):
+            nlp("this is a test")
+    else:
+        response = nlp.get_pipe("llm")._backend(["this is a test"])
+        assert len(response) == 1
+        response = srsly.json_loads(response[0])
+        assert (
+            response["error"]["message"]
+            == "The model `x-text-davinci-003` does not exist"
+        )
diff --git a/spacy_llm/tests/pipeline/test_llm.py b/spacy_llm/tests/pipeline/test_llm.py
index 9d5e748c..23a2d5e9 100644
--- a/spacy_llm/tests/pipeline/test_llm.py
+++ b/spacy_llm/tests/pipeline/test_llm.py
@@ -82,30 +82,33 @@ def test_integrations(config: Dict[str, Any]):
     nlp("This is a test.")
 
 
-def test_type_checking() -> None:
-    """Tests type checking for consistency between functions."""
+def test_type_checking_valid() -> None:
+    """Test type checking for consistency between functions."""
+    # Ensure default config doesn't raise warnings.
+    nlp = spacy.blank("en")
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        nlp.add_pipe("llm", config={"task": {"@llm_tasks": "spacy.NoOp.v1"}})
+
+
+def test_type_checking_invalid() -> None:
+    """Test type checking for consistency between functions."""
 
     @spacy.registry.llm_tasks("spacy.TestIncorrect.v1")
     def noop_task_incorrect() -> Tuple[
         Callable[[Iterable[Doc]], Iterable[int]],
-        Callable[[Iterable[Doc], Iterable[int]], Iterable[Doc]],
+        Callable[[Iterable[Doc], Iterable[float]], Iterable[Doc]],
     ]:
         def template(docs: Iterable[Doc]) -> Iterable[int]:
             return [0] * len(list(docs))
 
         def parse(
-            docs: Iterable[Doc], prompt_responses: Iterable[int]
+            docs: Iterable[Doc], prompt_responses: Iterable[float]
         ) -> Iterable[Doc]:
             return docs
 
         return template, parse
 
-    # Ensure default config doesn't raise warnings.
-    nlp = spacy.blank("en")
-    with warnings.catch_warnings():
-        warnings.simplefilter("error")
-        nlp.add_pipe("llm", config={"task": {"@llm_tasks": "spacy.NoOp.v1"}})
-
     nlp = spacy.blank("en")
     with pytest.warns(UserWarning) as record:
         nlp.add_pipe(
@@ -121,5 +124,5 @@ def parse(
     assert (
         str(record[1].message)
         == "Type returned from `backend` (`typing.Iterable[str]`) doesn't match type "
-        "expected by `parse` (`typing.Iterable[int]`)."
+        "expected by `parse` (`typing.Iterable[float]`)."
     )