Fix plugins

Author: blythed

.github/workflows/ci_code.yml

@@ -64,7 +64,6 @@ jobs:
         python -m pip install plugins/mongodb
         python -m pip install plugins/openai
         python -m pip install plugins/ibis
-        python -m pip install plugins/sqlalchemy
 
 
     - name: Lint and type-check

.github/workflows/ci_templates.yaml

@@ -136,7 +136,6 @@ jobs:
         run: |
           export SUPERDUPER_TEMPLATE=${{ matrix.template }}
           export SUPERDUPER_DATA_BACKEND='mongomock://test_db'
-          cd superduper/templates && ln -s ../../templates/* . && cd ../../
           pytest test/integration/template/test_template.py -s
         env:
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}

plugins/anthropic/superduper_anthropic/model.py

@@ -1,12 +1,13 @@
 import dataclasses as dc
+import os
 import typing as t
 
 import anthropic
 from anthropic import APIConnectionError, APIError, APIStatusError, APITimeoutError
 from superduper.base.query_dataset import QueryDataset
 from superduper.components.model import APIBaseModel
 from superduper.misc.retry import Retry
-from superduper.misc.utils import format_prompt, get_key
+from superduper.misc.utils import format_prompt
 
 retry = Retry(
     exception_types=(APIConnectionError, APIError, APIStatusError, APITimeoutError)
@@ -34,7 +35,7 @@ def init(self, db=None):
         :param db: The database to use.
         """
         self.client = anthropic.Anthropic(
-            api_key=get_key(KEY_NAME), **self.client_kwargs
+            api_key=os.environ[KEY_NAME], **self.client_kwargs
         )
         super().init(db=db)
 

plugins/cohere/plugin_test/test_model_cohere.py

@@ -21,7 +21,6 @@ def test_embed_one():
     embed = CohereEmbed(identifier='embed-english-v2.0')
     resp = embed.predict('Hello world')
 
-    assert len(resp) == embed.shape[0]
     assert isinstance(resp, list)
     assert all(isinstance(x, float) for x in resp)
 
@@ -35,7 +34,6 @@ def test_embed_batch():
     resp = embed.predict_batches(['Hello', 'world'])
 
     assert len(resp) == 2
-    assert len(resp[0]) == embed.shape[0]
     assert isinstance(resp[0], list)
     assert all(isinstance(x, float) for x in resp[0])
 

plugins/cohere/superduper_cohere/model.py

@@ -1,4 +1,5 @@
 import dataclasses as dc
+import os
 import typing as t
 
 import cohere
@@ -7,7 +8,7 @@
 from superduper.base.query_dataset import QueryDataset
 from superduper.components.model import APIBaseModel
 from superduper.misc.retry import Retry
-from superduper.misc.utils import format_prompt, get_key
+from superduper.misc.utils import format_prompt
 
 retry = Retry(exception_types=(CohereAPIError, CohereConnectionError))
 
@@ -42,30 +43,21 @@ class CohereEmbed(Cohere):
 
     """
 
-    shapes: t.ClassVar[t.Dict] = {'embed-english-v2.0': (4096,)}
-    shape: t.Optional[t.Sequence[int]] = None
     batch_size: int = 100
-    signature: str = 'singleton'
-
-    def postinit(self):
-        """Post-initialization method."""
-        if self.shape is None:
-            self.shape = self.shapes[self.identifier]
-        return super().postinit()
 
     @retry
     def predict(self, X: str):
         """Predict the embedding of a single text.
 
         :param X: The text to predict the embedding of.
         """
-        client = cohere.Client(get_key(KEY_NAME), **self.client_kwargs)
+        client = cohere.Client(os.environ[KEY_NAME], **self.client_kwargs)
         e = client.embed(texts=[X], model=self.identifier, **self.predict_kwargs)
         return e.embeddings[0]
 
     @retry
     def _predict_a_batch(self, texts: t.List[str]):
-        client = cohere.Client(get_key(KEY_NAME), **self.client_kwargs)
+        client = cohere.Client(os.environ[KEY_NAME], **self.client_kwargs)
         out = client.embed(texts=texts, model=self.identifier, **self.predict_kwargs)
         return [r for r in out.embeddings]
 
@@ -111,7 +103,7 @@ def predict(self, prompt: str, context: t.Optional[t.List[str]] = None):
         """
         if context is not None:
             prompt = format_prompt(prompt, self.prompt, context=context)
-        client = cohere.Client(get_key(KEY_NAME), **self.client_kwargs)
+        client = cohere.Client(os.environ[KEY_NAME], **self.client_kwargs)
         resp = client.generate(
             prompt=prompt, model=self.identifier, **self.predict_kwargs
         )

plugins/ibis/plugin_test/config.yaml

@@ -3,4 +3,4 @@ auto_schema: false
 force_apply: true
 json_native: false
 datatype_presets:
-  vector: superduper.components.datatype.Array
+  vector: superduper.base.datatype.Array

plugins/ibis/pyproject.toml

@@ -29,13 +29,12 @@ dependencies = [
   "ibis-framework[sqlite]>=9.0.1,<10.0.0",
   "click",
   'pandas',
-  "sqlalchemy"
+  "sqlalchemy>=1.4.0",
 ]
 
 [project.optional-dependencies]
 test = [
     # Annotation plugin dependencies will be installed in CI
-    # :CI: plugins/sqlalchemy
 ]
 
 [project.urls]

plugins/jina/superduper_jina/client.py

@@ -1,11 +1,11 @@
+import os
 from typing import List, Optional
 
 import aiohttp
 import requests
 from aiohttp import ClientConnectionError, ClientResponseError
 from requests.exceptions import HTTPError
 from superduper.misc.retry import Retry
-from superduper.misc.utils import get_key
 
 JINA_API_URL: str = "https://api.jina.ai/v1/embeddings"
 KEY_NAME = 'JINA_API_KEY'
@@ -34,7 +34,7 @@ def __init__(
         # if the user does not provide the API key,
         # check if it is set in the environment variable
         if api_key is None:
-            api_key = get_key(KEY_NAME)
+            api_key = os.environ[KEY_NAME]
 
         self.model_name = model_name
         self._session = requests.Session()

plugins/openai/plugin_test/test_model_openai.py

@@ -11,8 +11,6 @@
     OpenAIAudioTranslation,
     OpenAIChatCompletion,
     OpenAIEmbedding,
-    OpenAIImageCreation,
-    OpenAIImageEdit,
     _available_models,
 )
 
@@ -122,76 +120,6 @@ def test_batch_chat():
     assert isinstance(resp[0], str)
 
 
-@vcr.use_cassette()
-def test_create_url():
-    e = OpenAIImageCreation(
-        identifier='dall-e',
-        prompt='a close up, studio photographic portrait of a {context}',
-        response_format='url',
-    )
-    resp = e.predict('cat')
-
-    # PNG 8-byte signature
-    assert resp[0:16] == PNG_BYTE_SIGNATURE
-
-
-@vcr.use_cassette()
-def test_create_url_batch():
-    e = OpenAIImageCreation(
-        identifier='dall-e',
-        prompt='a close up, studio photographic portrait of a',
-        response_format='url',
-    )
-    resp = e.predict_batches(['cat', 'dog'])
-
-    for img in resp:
-        # PNG 8-byte signature
-        assert img[0:16] == PNG_BYTE_SIGNATURE
-
-
-@vcr.use_cassette()
-def test_edit_url():
-    e = OpenAIImageEdit(
-        identifier='dall-e',
-        prompt='A celebration party at the launch of {context}',
-        response_format='url',
-    )
-    with open('test/material/data/rickroll.png', 'rb') as f:
-        buffer = io.BytesIO(f.read())
-    resp = e.predict(buffer, context=['superduper'])
-    buffer.close()
-
-    # PNG 8-byte signature
-    assert resp[0:16] == PNG_BYTE_SIGNATURE
-
-
-@vcr.use_cassette()
-def test_edit_url_batch():
-    e = OpenAIImageEdit(
-        identifier='dall-e',
-        prompt='A celebration party at the launch of superduper',
-        response_format='url',
-    )
-    with open('test/material/data/rickroll.png', 'rb') as f:
-        buffer_one = io.BytesIO(f.read())
-    with open('test/material/data/rickroll.png', 'rb') as f:
-        buffer_two = io.BytesIO(f.read())
-
-    resp = e.predict_batches(
-        [
-            ((buffer_one,), {}),
-            ((buffer_two,), {}),
-        ]
-    )
-
-    buffer_one.close()
-    buffer_two.close()
-
-    for img in resp:
-        # PNG 8-byte signature
-        assert img[0:16] == PNG_BYTE_SIGNATURE
-
-
 @vcr.use_cassette()
 def test_transcribe():
     with open('test/material/data/test.wav', 'rb') as f: