Add UT for llm api: OpenAI, vLLM

superduperdb/ext/llm/base.py

@@ -2,6 +2,7 @@
 import asyncio
 import dataclasses as dc
 import inspect
+import typing
 from functools import wraps
 from logging import WARNING, getLogger
 from typing import Any, Callable, List, Optional, Union
@@ -11,6 +12,9 @@
 from superduperdb.components.model import _Predictor
 from superduperdb.ext.utils import format_prompt
 
+if typing.TYPE_CHECKING:
+    from superduperdb.base.datalayer import Datalayer
+
 # Disable httpx info level logging
 getLogger("httpx").setLevel(WARNING)
 
@@ -51,8 +55,26 @@ class _BaseLLM(Component, _Predictor, metaclass=abc.ABCMeta):
     def __post_init__(self):
         super().__post_init__()
         self.takes_context = True
+        self.identifier = self.identifier.replace("/", "-")
         assert "{input}" in self.prompt_template, "Template must contain {input}"
 
+    def to_call(self, X, *args, **kwargs):
+        raise NotImplementedError
+
+    def post_create(self, db: "Datalayer") -> None:
+        # TODO: Do not make sense to add this logic here,
+        # Need a auto DataType to handle this
+        from superduperdb.backends.ibis.data_backend import IbisDataBackend
+        from superduperdb.backends.ibis.field_types import dtype
+
+        if isinstance(db.databackend, IbisDataBackend) and self.encoder is None:
+            self.encoder = dtype('str')
+
+        # since then the `.add` clause is not necessary
+        output_component = db.databackend.create_model_table_or_collection(self)  # type: ignore[arg-type]
+        if output_component is not None:
+            db.add(output_component)
+
     @abc.abstractmethod
     def init(self):
         ...

superduperdb/ext/llm/vllm.py

@@ -39,7 +39,7 @@ def _generate(self, prompt: str, **kwargs) -> str:
         """
         post_data = self.build_post_data(prompt, **kwargs)
         response = requests.post(self.api_url, json=post_data)
-        return response.json()["text"]
+        return response.json()["text"][0]
 
     def _batch_generate(self, prompts: List[str], **kwargs: Any) -> List[str]:
         """
@@ -53,7 +53,7 @@ async def _async_generate(self, session, semaphore, prompt: str, **kwargs) -> st
             try:
                 async with session.post(self.api_url, json=post_data) as response:
                     response_json = await response.json()
-                    return response_json["text"]
+                    return response_json["text"][0]
             except aiohttp.ClientError as e:
                 logging.error(f"HTTP request failed: {e}. Prompt: {prompt}")
                 return ""