281 lines
8.0 KiB
Python
281 lines
8.0 KiB
Python
from enum import Enum
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
AsyncIterator,
|
|
Dict,
|
|
Iterator,
|
|
List,
|
|
Optional,
|
|
Union,
|
|
)
|
|
|
|
import httpx
|
|
from typing_extensions import Unpack
|
|
|
|
from replicate import identifier
|
|
from replicate.exceptions import ReplicateError
|
|
from replicate.helpers import transform_output
|
|
|
|
try:
|
|
from pydantic import v1 as pydantic # type: ignore
|
|
except ImportError:
|
|
import pydantic # type: ignore
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from replicate.client import Client
|
|
from replicate.identifier import ModelVersionIdentifier
|
|
from replicate.model import Model
|
|
from replicate.prediction import Predictions
|
|
from replicate.version import Version
|
|
|
|
|
|
class ServerSentEvent(pydantic.BaseModel): # type: ignore
|
|
"""
|
|
A server-sent event.
|
|
"""
|
|
|
|
class EventType(Enum):
|
|
"""
|
|
A server-sent event type.
|
|
"""
|
|
|
|
OUTPUT = "output"
|
|
LOGS = "logs"
|
|
ERROR = "error"
|
|
DONE = "done"
|
|
|
|
event: EventType
|
|
data: str
|
|
id: str
|
|
retry: Optional[int]
|
|
|
|
def __str__(self) -> str:
|
|
if self.event == ServerSentEvent.EventType.OUTPUT:
|
|
return self.data
|
|
|
|
return ""
|
|
|
|
|
|
class EventSource:
|
|
"""
|
|
A server-sent event source.
|
|
"""
|
|
|
|
client: "Client"
|
|
response: "httpx.Response"
|
|
use_file_output: bool
|
|
|
|
def __init__(
|
|
self,
|
|
client: "Client",
|
|
response: "httpx.Response",
|
|
*,
|
|
use_file_output: Optional[bool] = True,
|
|
) -> None:
|
|
self.client = client
|
|
self.response = response
|
|
self.use_file_output = use_file_output or True
|
|
content_type, _, _ = response.headers["content-type"].partition(";")
|
|
if content_type != "text/event-stream":
|
|
raise ValueError(
|
|
"Expected response Content-Type to be 'text/event-stream', "
|
|
f"got {content_type!r}"
|
|
)
|
|
|
|
class Decoder:
|
|
"""
|
|
A decoder for server-sent events.
|
|
"""
|
|
|
|
event: Optional["ServerSentEvent.EventType"]
|
|
data: List[str]
|
|
last_event_id: Optional[str]
|
|
retry: Optional[int]
|
|
|
|
def __init__(self) -> None:
|
|
self.event = None
|
|
self.data = []
|
|
self.last_event_id = None
|
|
self.retry = None
|
|
|
|
def decode(self, line: str) -> Optional[ServerSentEvent]:
|
|
"""
|
|
Decode a line and return a server-sent event if applicable.
|
|
"""
|
|
|
|
if not line:
|
|
if (
|
|
not any([self.event, self.data, self.last_event_id, self.retry])
|
|
or self.event is None
|
|
or self.last_event_id is None
|
|
):
|
|
return None
|
|
|
|
sse = ServerSentEvent(
|
|
event=self.event,
|
|
data="\n".join(self.data),
|
|
id=self.last_event_id,
|
|
retry=self.retry,
|
|
)
|
|
|
|
self.event = None
|
|
self.data = []
|
|
self.retry = None
|
|
|
|
return sse
|
|
|
|
if line.startswith(":"):
|
|
return None
|
|
|
|
fieldname, _, value = line.partition(":")
|
|
value = value[1:] if value.startswith(" ") else value
|
|
|
|
if fieldname == "event":
|
|
if event := ServerSentEvent.EventType(value):
|
|
self.event = event
|
|
elif fieldname == "data":
|
|
self.data.append(value)
|
|
elif fieldname == "id":
|
|
if "\0" not in value:
|
|
self.last_event_id = value
|
|
elif fieldname == "retry":
|
|
try:
|
|
self.retry = int(value)
|
|
except (TypeError, ValueError):
|
|
pass
|
|
|
|
return None
|
|
|
|
def __iter__(self) -> Iterator[ServerSentEvent]:
|
|
decoder = EventSource.Decoder()
|
|
|
|
for line in self.response.iter_lines():
|
|
line = line.rstrip("\n")
|
|
sse = decoder.decode(line)
|
|
if sse is not None:
|
|
if sse.event == ServerSentEvent.EventType.ERROR:
|
|
raise RuntimeError(sse.data)
|
|
|
|
if (
|
|
self.use_file_output
|
|
and sse.event == ServerSentEvent.EventType.OUTPUT
|
|
):
|
|
sse.data = transform_output(sse.data, client=self.client)
|
|
|
|
yield sse
|
|
|
|
if sse.event == ServerSentEvent.EventType.DONE:
|
|
return
|
|
|
|
async def __aiter__(self) -> AsyncIterator[ServerSentEvent]:
|
|
decoder = EventSource.Decoder()
|
|
async for line in self.response.aiter_lines():
|
|
line = line.rstrip("\n")
|
|
sse = decoder.decode(line)
|
|
if sse is not None:
|
|
if sse.event == ServerSentEvent.EventType.ERROR:
|
|
raise RuntimeError(sse.data)
|
|
|
|
if (
|
|
self.use_file_output
|
|
and sse.event == ServerSentEvent.EventType.OUTPUT
|
|
):
|
|
sse.data = transform_output(sse.data, client=self.client)
|
|
|
|
yield sse
|
|
|
|
if sse.event == ServerSentEvent.EventType.DONE:
|
|
return
|
|
|
|
|
|
def stream(
|
|
client: "Client",
|
|
ref: Union["Model", "Version", "ModelVersionIdentifier", str],
|
|
input: Optional[Dict[str, Any]] = None,
|
|
*,
|
|
use_file_output: Optional[bool] = True,
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> Iterator[ServerSentEvent]:
|
|
"""
|
|
Run a model and stream its output.
|
|
"""
|
|
|
|
params = params or {}
|
|
params["stream"] = True
|
|
|
|
version, owner, name, version_id = identifier._resolve(ref)
|
|
|
|
if version or version_id:
|
|
prediction = client.predictions.create(
|
|
version=(version or version_id), input=input or {}, **params
|
|
)
|
|
elif owner and name:
|
|
prediction = client.models.predictions.create(
|
|
model=(owner, name), input=input or {}, **params
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid argument: {ref}. Expected model, version, or reference in the format owner/name or owner/name:version"
|
|
)
|
|
|
|
url = prediction.urls and prediction.urls.get("stream", None)
|
|
if not url or not isinstance(url, str):
|
|
raise ReplicateError("Model does not support streaming")
|
|
|
|
headers = {}
|
|
headers["Accept"] = "text/event-stream"
|
|
headers["Cache-Control"] = "no-store"
|
|
|
|
with client._client.stream("GET", url, headers=headers) as response:
|
|
yield from EventSource(client, response, use_file_output=use_file_output)
|
|
|
|
|
|
async def async_stream(
|
|
client: "Client",
|
|
ref: Union["Model", "Version", "ModelVersionIdentifier", str],
|
|
input: Optional[Dict[str, Any]] = None,
|
|
*,
|
|
use_file_output: Optional[bool] = True,
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> AsyncIterator[ServerSentEvent]:
|
|
"""
|
|
Run a model and stream its output asynchronously.
|
|
"""
|
|
|
|
params = params or {}
|
|
params["stream"] = True
|
|
|
|
version, owner, name, version_id = identifier._resolve(ref)
|
|
|
|
if version or version_id:
|
|
prediction = await client.predictions.async_create(
|
|
version=(version or version_id), input=input or {}, **params
|
|
)
|
|
elif owner and name:
|
|
prediction = await client.models.predictions.async_create(
|
|
model=(owner, name), input=input or {}, **params
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid argument: {ref}. Expected model, version, or reference in the format owner/name or owner/name:version"
|
|
)
|
|
|
|
url = prediction.urls and prediction.urls.get("stream", None)
|
|
if not url or not isinstance(url, str):
|
|
raise ReplicateError("Model does not support streaming")
|
|
|
|
headers = {}
|
|
headers["Accept"] = "text/event-stream"
|
|
headers["Cache-Control"] = "no-store"
|
|
|
|
async with client._async_client.stream("GET", url, headers=headers) as response:
|
|
async for event in EventSource(
|
|
client, response, use_file_output=use_file_output
|
|
):
|
|
yield event
|
|
|
|
|
|
__all__ = ["ServerSentEvent"]
|