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"]